| #!/usr/bin/env python3 |
| |
| # Copyright 2021 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Wrapper script to automatically lock devices for crosperf.""" |
| |
| import os |
| import sys |
| import argparse |
| import subprocess |
| import contextlib |
| import json |
| from typing import Optional, Any |
| import dataclasses |
| |
| # Have to do sys.path hackery because crosperf relies on PYTHONPATH |
| # modifications. |
| PARENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.append(PARENT_DIR) |
| |
| |
| def main(sys_args: list[str]) -> Optional[str]: |
| """Run crosperf_autolock. Returns error msg or None""" |
| args, leftover_args = parse_args(sys_args) |
| fleet_params = [ |
| CrosfleetParams(board=args.board, |
| pool=args.pool, |
| lease_time=args.lease_time) |
| for _ in range(args.num_leases) |
| ] |
| if not fleet_params: |
| return ('No board names identified. If you want to use' |
| ' a known host, just use crosperf directly.') |
| try: |
| _run_crosperf(fleet_params, args.dut_lock_timeout, leftover_args) |
| except BoardLockError as e: |
| _eprint('ERROR:', e) |
| _eprint('May need to login to crosfleet? Run "crosfleet login"') |
| _eprint('The leases may also be successful later on. ' |
| 'Check with "crosfleet dut leases"') |
| return 'crosperf_autolock failed' |
| except BoardReleaseError as e: |
| _eprint('ERROR:', e) |
| _eprint('May need to re-run "crosfleet dut abandon"') |
| return 'crosperf_autolock failed' |
| return None |
| |
| |
| def parse_args(args: list[str]) -> tuple[Any, list]: |
| """Parse the CLI arguments.""" |
| parser = argparse.ArgumentParser( |
| 'crosperf_autolock', |
| description='Wrapper around crosperf' |
| ' to autolock DUTs from crosfleet.', |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| parser.add_argument('--board', |
| type=str, |
| help='Space or comma separated list of boards to lock', |
| required=True, |
| default=argparse.SUPPRESS) |
| parser.add_argument('--num-leases', |
| type=int, |
| help='Number of boards to lock.', |
| metavar='NUM', |
| default=1) |
| parser.add_argument('--pool', |
| type=str, |
| help='Pool to pull from.', |
| default='DUT_POOL_QUOTA') |
| parser.add_argument('--dut-lock-timeout', |
| type=float, |
| metavar='SEC', |
| help='Number of seconds we want to try to lease a board' |
| ' from crosfleet. This option does NOT change the' |
| ' lease length.', |
| default=600) |
| parser.add_argument('--lease-time', |
| type=int, |
| metavar='MIN', |
| help='Number of minutes to lock the board. Max is 1440.', |
| default=1440) |
| parser.epilog = ( |
| 'For more detailed flags, you have to read the args taken by the' |
| ' crosperf executable. Args are passed transparently to crosperf.') |
| return parser.parse_known_args(args) |
| |
| |
| class BoardLockError(Exception): |
| """Error to indicate failure to lock a board.""" |
| |
| def __init__(self, msg: str): |
| self.msg = 'BoardLockError: ' + msg |
| super().__init__(self.msg) |
| |
| |
| class BoardReleaseError(Exception): |
| """Error to indicate failure to release a board.""" |
| |
| def __init__(self, msg: str): |
| self.msg = 'BoardReleaseError: ' + msg |
| super().__init__(self.msg) |
| |
| |
| @dataclasses.dataclass(frozen=True) |
| class CrosfleetParams: |
| """Dataclass to hold all crosfleet parameterizations.""" |
| board: str |
| pool: str |
| lease_time: int |
| |
| |
| def _eprint(*msg, **kwargs): |
| print(*msg, file=sys.stderr, **kwargs) |
| |
| |
| def _run_crosperf(crosfleet_params: list[CrosfleetParams], lock_timeout: float, |
| leftover_args: list[str]): |
| """Autolock devices and run crosperf with leftover arguments. |
| |
| Raises: |
| BoardLockError: When board was unable to be locked. |
| BoardReleaseError: When board was unable to be released. |
| """ |
| if not crosfleet_params: |
| raise ValueError('No crosfleet params given; cannot call crosfleet.') |
| |
| # We'll assume all the boards are the same type, which seems to be the case |
| # in experiments that actually get used. |
| passed_board_arg = crosfleet_params[0].board |
| with contextlib.ExitStack() as stack: |
| dut_hostnames = [] |
| for param in crosfleet_params: |
| print( |
| f'Sent lock request for {param.board} for {param.lease_time} minutes' |
| '\nIf this fails, you may need to run "crosfleet dut abandon <...>"') |
| # May raise BoardLockError, abandoning previous DUTs. |
| dut_hostname = stack.enter_context( |
| crosfleet_machine_ctx( |
| param.board, |
| param.lease_time, |
| lock_timeout, |
| {'label-pool': param.pool}, |
| )) |
| if dut_hostname: |
| print(f'Locked {param.board} machine: {dut_hostname}') |
| dut_hostnames.append(dut_hostname) |
| |
| # We import crosperf late, because this import is extremely slow. |
| # We don't want the user to wait several seconds just to get |
| # help info. |
| import crosperf |
| for dut_hostname in dut_hostnames: |
| crosperf.Main([ |
| sys.argv[0], |
| '--no_lock', |
| 'True', |
| '--remote', |
| dut_hostname, |
| '--board', |
| passed_board_arg, |
| ] + leftover_args) |
| |
| |
| @contextlib.contextmanager |
| def crosfleet_machine_ctx(board: str, |
| lease_minutes: int, |
| lock_timeout: float, |
| dims: dict[str, Any], |
| abandon_timeout: float = 120.0) -> Any: |
| """Acquire dut from crosfleet, and release once it leaves the context. |
| |
| Args: |
| board: Board type to lease. |
| lease_minutes: Length of lease, in minutes. |
| lock_timeout: How long to wait for a lock until quitting. |
| dims: Dictionary of dimension arguments to pass to crosfleet's '-dims' |
| abandon_timeout (optional): How long to wait for releasing until quitting. |
| |
| Yields: |
| A string representing the crosfleet DUT hostname. |
| |
| Raises: |
| BoardLockError: When board was unable to be locked. |
| BoardReleaseError: When board was unable to be released. |
| """ |
| # This lock may raise an exception, but if it does, we can't release |
| # the DUT anyways as we won't have the dut_hostname. |
| dut_hostname = crosfleet_autolock(board, lease_minutes, dims, lock_timeout) |
| try: |
| yield dut_hostname |
| finally: |
| if dut_hostname: |
| crosfleet_release(dut_hostname, abandon_timeout) |
| |
| |
| def crosfleet_autolock(board: str, lease_minutes: int, dims: dict[str, Any], |
| timeout_sec: float) -> str: |
| """Lock a device using crosfleet, paramaterized by the board type. |
| |
| Args: |
| board: Board of the DUT we want to lock. |
| lease_minutes: Number of minutes we're trying to lease the DUT for. |
| dims: Dictionary of dimension arguments to pass to crosfleet's '-dims' |
| timeout_sec: Number of seconds to try to lease the DUT. Default 120s. |
| |
| Returns: |
| The hostname of the board, or empty string if it couldn't be parsed. |
| |
| Raises: |
| BoardLockError: When board was unable to be locked. |
| """ |
| crosfleet_cmd_args = [ |
| 'crosfleet', |
| 'dut', |
| 'lease', |
| '-json', |
| '-reason="crosperf autolock"', |
| f'-board={board}', |
| f'-minutes={lease_minutes}', |
| ] |
| if dims: |
| dims_arg = ','.join('{}={}'.format(k, v) for k, v in dims.items()) |
| crosfleet_cmd_args.extend(['-dims', f'{dims_arg}']) |
| |
| try: |
| output = subprocess.check_output(crosfleet_cmd_args, |
| timeout=timeout_sec, |
| encoding='utf-8') |
| except subprocess.CalledProcessError as e: |
| raise BoardLockError( |
| f'crosfleet dut lease failed with exit code: {e.returncode}') |
| except subprocess.TimeoutExpired as e: |
| raise BoardLockError(f'crosfleet dut lease timed out after {timeout_sec}s;' |
| ' please abandon the dut manually.') |
| |
| try: |
| json_obj = json.loads(output) |
| dut_hostname = json_obj['DUT']['Hostname'] |
| if not isinstance(dut_hostname, str): |
| raise TypeError('dut_hostname was not a string') |
| except (json.JSONDecodeError, IndexError, KeyError, TypeError) as e: |
| raise BoardLockError( |
| f'crosfleet dut lease output was parsed incorrectly: {e!r};' |
| f' observed output was {output}') |
| return _maybe_append_suffix(dut_hostname) |
| |
| |
| def crosfleet_release(dut_hostname: str, timeout_sec: float = 120.0): |
| """Release a crosfleet device. |
| |
| Consider using the context managed crosfleet_machine_context |
| |
| Args: |
| dut_hostname: Name of the device we want to release. |
| timeout_sec: Number of seconds to try to release the DUT. Default is 120s. |
| |
| Raises: |
| BoardReleaseError: Potentially failed to abandon the lease. |
| """ |
| crosfleet_cmd_args = [ |
| 'crosfleet', |
| 'dut', |
| 'abandon', |
| dut_hostname, |
| ] |
| exit_code = subprocess.call(crosfleet_cmd_args, timeout=timeout_sec) |
| if exit_code != 0: |
| raise BoardReleaseError( |
| f'"crosfleet dut abandon" had exit code {exit_code}') |
| |
| |
| def _maybe_append_suffix(hostname: str) -> str: |
| if hostname.endswith('.cros') or '.cros.' in hostname: |
| return hostname |
| return hostname + '.cros' |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main(sys.argv[1:])) |