diff options
Diffstat (limited to 'crosperf/crosperf_autolock.py')
| -rwxr-xr-x | crosperf/crosperf_autolock.py | 508 |
1 files changed, 270 insertions, 238 deletions
diff --git a/crosperf/crosperf_autolock.py b/crosperf/crosperf_autolock.py index b593fa9c..011f01e3 100755 --- a/crosperf/crosperf_autolock.py +++ b/crosperf/crosperf_autolock.py @@ -1,19 +1,20 @@ #!/usr/bin/env python3 -# Copyright 2021 The Chromium OS Authors. All rights reserved. +# Copyright 2021 The ChromiumOS Authors # 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 +import json +import os +import subprocess +import sys +from typing import Any, Dict, List, Optional, Tuple + # Have to do sys.path hackery because crosperf relies on PYTHONPATH # modifications. @@ -21,261 +22,292 @@ 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) +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.""" + """Error to indicate failure to lock a board.""" - def __init__(self, msg: str): - self.msg = 'BoardLockError: ' + msg - super().__init__(self.msg) + def __init__(self, msg: str): + self.msg = "BoardLockError: " + msg + super().__init__(self.msg) class BoardReleaseError(Exception): - """Error to indicate failure to release a board.""" + """Error to indicate failure to release a board.""" - def __init__(self, msg: str): - self.msg = 'BoardReleaseError: ' + msg - super().__init__(self.msg) + 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 + """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) + 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_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: 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(f"{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}') + """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 hostname.endswith(".cros") or ".cros." in hostname: + return hostname + return hostname + ".cros" -if __name__ == '__main__': - sys.exit(main(sys.argv[1:])) +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) |