Add crosperf_autolock.py wrapper

This crosperf wrapper allows the user to lock devices by board
type. It will query the swarm using the crosfleet API to find
and lock an appropriate board.

The remaining arguments are passed to crosperf, so the command
will act like crosperf otherwise, once the remote has been set.

One limitation is that this wrapper does not conduct any SSH
port forwarding (which would be required for flashing
devices from a cloudtop, for example).

BUG=b:187535367
TEST=crosperf_autolock.py --board chell chell_exp.exp

Change-Id: If8dc4630c42be63d6427716654e7dd7fa4132df5
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/3221950
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>
Commit-Queue: Jordan R Abrahams <ajordanr@google.com>
Tested-by: Jordan R Abrahams <ajordanr@google.com>

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:]))