• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2018 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Yet another domain specific client for Swarming."""
6
7from __future__ import absolute_import
8from __future__ import division
9from __future__ import print_function
10
11import json
12import os
13import urllib
14
15from lucifer import autotest
16from skylab_staging import errors
17
18# This is hard-coded everywhere -- on builders requesting skylab suite as well
19# as in commands to be generated by users requesting one off suites.
20_SKYLAB_RUN_SUITE_PATH = '/usr/local/autotest/bin/run_suite_skylab'
21_SWARMING_POOL_SKYLAB_BOTS = 'ChromeOSSkylab'
22_SWARMING_POOL_SKYLAB_SUITE_BOTS = 'ChromeOSSkylab-suite'
23# Test push creates all suites at the highest allowed non-admin task priority.
24# This ensures that test push tasks are prioritized over any user created tasks
25# in the staging lab.
26_TEST_PUSH_SUITE_PRIORITY = 50
27
28
29
30class Client(object):
31  """A domain specific client for Swarming service."""
32
33  def __init__(self, cli_path, host, service_account_json=None):
34    self._cli_path = cli_path
35    self._host = host
36    self._service_account_json = service_account_json
37
38  def num_ready_duts(self, board, pool):
39    """Count the number of DUTs in the given board, pool in dut_state ready.
40
41    @param board: The board autotest label of the DUTs.
42    @param pool: The pool autotest label of the DUTs.
43    @returns number of DUTs in dut_state ready.
44    """
45    qargs = [
46        ('dimensions', 'pool:%s' % _SWARMING_POOL_SKYLAB_BOTS),
47        ('dimensions', 'label-board:%s' % board),
48        ('dimensions', 'label-pool:%s' % pool),
49        ('dimensions', 'dut_state:ready'),
50    ]
51    result = self.query('bots/count', qargs)
52    if not result:
53      return 0
54    return int(result['count']) - (int(result['busy'])
55                                   + int(result['dead'])
56                                   + int(result['quarantined'])
57                                   + int(result['maintenance']))
58
59  def query(self, path, qargs):
60    """Run a Swarming 'query' call.
61
62    @param path: Path of the query RPC call.
63    @qargs: Arguments for the RPC call.
64    @returns: json response from the Swarming call.
65    """
66    cros_build_lib = autotest.chromite_load('cros_build_lib')
67    cmdarg = path
68    if qargs:
69      cmdarg += "?%s" % urllib.urlencode(qargs)
70
71    cmd = self._base_cmd('query') + [cmdarg]
72
73    result = cros_build_lib.RunCommand(cmd, capture_output=True)
74    return json.loads(result.output)
75
76  def trigger_suite(self, board, pool, build, suite_name, timeout_s):
77    """Trigger an autotest suite. Use wait_for_suite to wait for results.
78
79    @param board: The board autotest label of the DUTs.
80    @param pool: The pool autotest label of the DUTs.
81    @param build: The build to test, e.g. link-paladin/R70-10915.0.0-rc1.
82    @param suite_name: The name of the suite to run, e.g. provision.
83    @param timeout_s: Timeout for the suite, in seconds.
84    @returns: The task ID of the kicked off suite.
85    """
86    raw_cmd = self._suite_cmd_common(board, pool, build, suite_name, timeout_s)
87    raw_cmd += ['--create_and_return']
88    return self._run(board, pool, build, suite_name, timeout_s, raw_cmd)
89
90  def wait_for_suite(self, task_id, board, pool, build, suite_name, timeout_s):
91    """Wait for a suite previously kicked off via trigger_suite().
92
93    @param task_id: Task ID of the suite, as returned by trigger_suite().
94    @param board: The board autotest label of the DUTs.
95    @param pool: The pool autotest label of the DUTs.
96    @param build: The build to test, e.g. link-paladin/R70-10915.0.0-rc1.
97    @param suite_name: The name of the suite to run, e.g. provision.
98    @param timeout_s: Timeout for the suite, in seconds.
99    @returns: The task ID of the kicked off suite.
100    """
101    raw_cmd = self._suite_cmd_common(board, pool, build, suite_name, timeout_s)
102    raw_cmd += ['--suite_id', task_id]
103    return self._run(board, pool, build, suite_name, timeout_s, raw_cmd)
104
105  def _suite_cmd_common(self, board, pool, build, suite_name, timeout_s):
106    return [
107        _SKYLAB_RUN_SUITE_PATH,
108        '--board', board,
109        '--build', build,
110        '--max_retries', '5',
111        '--pool', _old_style_pool_label(pool),
112        '--priority', str(_TEST_PUSH_SUITE_PRIORITY),
113        '--suite_args', json.dumps({'num_required': 1}),
114        '--suite_name', suite_name,
115        '--test_retry',
116        '--timeout_mins', str(int(timeout_s / 60)),
117    ]
118
119  def _run(self, board, pool, build, suite_name, timeout_s, raw_cmd):
120    timeout_s = str(int(timeout_s))
121    task_name = '%s-%s' % (build, suite_name)
122    # This is a subset of the tags used by builders when creating suites.
123    # These tags are used by the result reporting pipeline in various ways.
124    tags = {
125        'board': board,
126        'build': build,
127        # Required for proper rendering of MILO UI.
128        'luci_project': 'chromeos',
129        'skylab': 'run_suite',
130        'skylab': 'staging',
131        'suite': suite_name,
132        'task_name': task_name,
133    }
134    osutils = autotest.chromite_load('osutils')
135    with osutils.TempDir() as tempdir:
136      summary_file = os.path.join(tempdir, 'summary.json')
137      cmd = self._base_cmd('run') + [
138          '--dimension', 'pool',  _SWARMING_POOL_SKYLAB_SUITE_BOTS,
139          '--expiration', timeout_s,
140          '--io-timeout', timeout_s,
141          '--hard-timeout', timeout_s,
142          '--print-status-update',
143          '--priority', str(_TEST_PUSH_SUITE_PRIORITY),
144          '--raw-cmd',
145          '--task-name', task_name,
146          '--task-summary-json', summary_file,
147          '--timeout', timeout_s,
148      ]
149      for key, val in tags.iteritems():
150        cmd += ['--tags', '%s:%s' % (key, val)]
151      cmd += ['--'] + raw_cmd
152
153      cros_build_lib = autotest.chromite_load('cros_build_lib')
154      cros_build_lib.RunCommand(cmd, error_code_ok=True)
155      return _extract_run_id(summary_file)
156
157
158  def _base_cmd(self, subcommand):
159    cmd = [
160        self._cli_path, subcommand,
161        '--swarming', self._host,
162    ]
163    if self._service_account_json is not None:
164      cmd += ['--auth-service-account-json', self._service_account_json]
165    return cmd
166
167
168  def task_url(self, task_id):
169    """Generate the task url based on task id."""
170    return '%s/user/task/%s' % (self._host, task_id)
171
172
173def _extract_run_id(path):
174  if not os.path.isfile(path):
175    raise errors.TestPushError('No task summary at %s' % path)
176  with open(path) as f:
177    summary = json.load(f)
178  if not summary.get('shards') or len(summary['shards']) != 1:
179    raise errors.TestPushError('Corrupted task summary at %s' % path)
180  run_id = summary['shards'][0].get('run_id')
181  if not run_id:
182    raise errors.TestPushError('No run_id in task summary at %s' % path)
183  return run_id
184
185
186def _old_style_pool_label(label):
187  _POOL_LABEL_PREFIX = 'dut_pool_'
188  label = label.lower()
189  if label.startswith(_POOL_LABEL_PREFIX):
190    return label[len(_POOL_LABEL_PREFIX):]
191  return label
192