282 lines
9.3 KiB
Python
Executable File
282 lines
9.3 KiB
Python
Executable File
#!/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:]))
|