diff options
Diffstat (limited to 'crosperf/crosperf_autolock.py')
| -rwxr-xr-x | crosperf/crosperf_autolock.py | 281 |
1 files changed, 281 insertions, 0 deletions
diff --git a/crosperf/crosperf_autolock.py b/crosperf/crosperf_autolock.py new file mode 100755 index 00000000..b593fa9c --- /dev/null +++ b/crosperf/crosperf_autolock.py @@ -0,0 +1,281 @@ +#!/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:])) |