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