cros_helper.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2017 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.
"""Helper script to manipulate chromeos DUT or query info."""
from __future__ import print_function
import asyncio
import argparse
import json
import logging
import random
import time

from bisect_kit import cli
from bisect_kit import common
from bisect_kit import cros_lab_util
from bisect_kit import cros_util
from bisect_kit import errors

DEFAULT_DUT_POOL = 'DUT_POOL_QUOTA'
logger = logging.getLogger(__name__)


def cmd_version_info(opts):
  info = cros_util.version_info(opts.board, opts.version)
  if opts.name:
    if opts.name not in info:
      logger.error('unknown name=%s', opts.name)
    print(info[opts.name])
  else:
    print(json.dumps(info, sort_keys=True, indent=4))


def cmd_query_dut_board(opts):
  assert cros_util.is_dut(opts.dut)
  print(cros_util.query_dut_board(opts.dut))


def cmd_reboot(opts):
  assert cros_util.is_dut(opts.dut)
  cros_util.reboot(
      opts.dut, force_reboot_callback=cros_lab_util.reboot_via_servo)


def _get_label_by_prefix(info, prefix):
  for label in info['Labels']:
    if label.startswith(prefix + ':'):
      return label
  return None


def cmd_lease_dut(opts):
  if opts.duration is not None and opts.duration < 60:
    raise errors.ArgumentError('--duration', 'must be at least 60 seconds')
  host = cros_lab_util.dut_host_name(opts.dut)
  logger.info('trying to lease %s', host)
  if cros_lab_util.skylab_lease_dut(host, opts.duration):
    logger.info('leased %s', host)
  else:
    raise Exception('unable to lease %s' % host)


def cmd_release_dut(opts):
  host = cros_lab_util.dut_host_name(opts.dut)
  cros_lab_util.skylab_release_dut(host)
  logger.info('%s released', host)


def verify_dimensions_by_lab(dimensions):
  result = []
  bots_dimensions = cros_lab_util.swarming_bots_dimensions()
  for dimension in dimensions:
    key, value = dimension.split(':', 1)
    if value in bots_dimensions.get(key, []):
      result.append(dimension)
    else:
      logger.warning('dimension=%s is unknown in the lab, typo? ignored',
                     dimension)
  return result


def select_available_bots_randomly(dimensions, variants, num=1, is_busy=None):
  bots = []
  for variant in variants:
    # There might be thousand bots available, set 'limit' to reduce swarming
    # API cost. This is not uniform random, but should be good enough.
    bots += cros_lab_util.swarming_bots_list(
        dimensions + [variant], is_busy=is_busy, limit=10)
  if not bots:
    return None
  return random.sample(bots, min(num, len(bots)))


def filter_dimensions_by_board(boards_with_prebuilt, dimensions):
  result = []
  for dimension in dimensions:
    bots = cros_lab_util.swarming_bots_list([dimension], is_busy=None, limit=1)
    if not bots:
      continue
    board = bots[0]['dimensions']['label-board'][0]
    if board not in boards_with_prebuilt:
      logger.warning(
          'dimension=%s (board=%s) does not have corresponding '
          'prebuilt image, ignore', dimension, board)
      continue
    result.append(dimension)
  return result


async def lease_dut_parallelly(duration, bots, timeout=None):
  tasks = []
  hosts = []
  for bot in bots:
    host = bot['dimensions']['dut_name'][0]
    hosts.append(host)
    tasks.append(
        asyncio.create_task(cros_lab_util.async_lease(host, duration=duration)))

  try:
    logger.info('trying to lease %d DUTs: %s', len(hosts), hosts)
    for coro in asyncio.as_completed(tasks, timeout=timeout):
      host = await coro
      if host:
        logger.info('leased %s', host)
        # Unfinished lease tasks will be cancelled when asyncio.run is
        # finishing.
        return host
    return None
  except asyncio.TimeoutError:
    return None


def do_allocate_dut(opts):
  """Helper of cmd_allocate_dut.

  Returns:
    (todo, host, board_to_build)
      todo: 'ready' or 'wait'
      host: leased host name
      board_to_build: board name for building image
  """
  if not opts.dut_name and not opts.pool:
    raise errors.ArgumentError('--pool',
                               'need to be specified if not --dut_name')
  if opts.version_hint:
    for v in opts.version_hint.split(','):
      if cros_util.is_cros_version(v) or cros_util.is_cros_snapshot_version(v):
        continue
      raise errors.ArgumentError(
          '--version_hint',
          'should be Chrome OS version numbers, separated by comma')
  if opts.duration is not None and opts.duration < 60:
    raise errors.ArgumentError('--duration', 'must be at least 60 seconds')

  t0 = time.time()
  dimensions = ['dut_state:ready']
  if not opts.dut_name:
    dimensions.append('label-pool:' + opts.pool)

  variants = []
  if opts.board:
    for board in opts.board.split(','):
      variants.append('label-board:' +
                      cros_lab_util.normalize_board_name(board))
  if opts.model:
    for model in opts.model.split(','):
      variants.append('label-model:' + model)
  if opts.sku:
    for sku in opts.sku.split(','):
      variants.append('label-hwid_sku:' + cros_lab_util.normalize_sku_name(sku))
  if opts.dut_name:
    for dut_name in opts.dut_name.split(','):
      variants.append('dut_name:' + dut_name)

  variants = verify_dimensions_by_lab(variants)
  variants = sorted(set(variants))  # dedup
  if not variants:
    raise errors.NoDutAvailable(
        'Invalid constraints: %s;%s;%s;%s' %
        (opts.board, opts.model, opts.sku, opts.dut_name))

  # Filter variants by prebuilt images.
  if opts.version_hint:
    if not opts.builder_hint:
      opts.builder_hint = opts.board
    if not opts.builder_hint:
      raise errors.ArgumentError('--builder_hint',
                                 'must be specified along with --version_hint')
    boards_with_prebuilt = []
    versions = opts.version_hint.split(',')
    for builder in opts.builder_hint.split(','):
      if not all(cros_util.has_test_image(builder, v) for v in versions):
        logger.warning(
            'builder=%s does not have prebuilt test image for %s, ignore',
            builder, opts.version_hint)
        continue
      boards_with_prebuilt.append(cros_lab_util.normalize_board_name(builder))
    logger.info('boards with prebuilt: %s', boards_with_prebuilt)
    if not boards_with_prebuilt:
      raise errors.ArgumentError(
          '--version_hint',
          'given builders have no prebuilt for %s' % opts.version_hint)
    variants = filter_dimensions_by_board(boards_with_prebuilt, variants)
    if not variants:
      raise errors.NoDutAvailable(
          'Devices with specified constraints have no prebuilt. '
          'Wrong version number?')

  while True:
    # Query every time because each iteration takes a few minutes
    bots = select_available_bots_randomly(
        dimensions, variants, num=opts.parallel, is_busy=False)
    if not bots:
      bots = select_available_bots_randomly(
          dimensions, variants, num=opts.parallel, is_busy=True)
    if not bots:
      raise errors.NoDutAvailable(
          'no bots satisfy constraints; all are in maintenance state?')

    remaining_time = opts.time_limit - (time.time() - t0)
    if remaining_time <= 0:
      break
    timeout = min(120, remaining_time)
    host = asyncio.run(lease_dut_parallelly(opts.duration, bots, timeout))
    if host:
      # Resolve what board we should build during bisection.
      board_to_build = None
      bots = cros_lab_util.swarming_bots_list(['dut_name:' + host])
      host_board = bots[0]['dimensions']['label-board'][0]
      if opts.builder_hint:
        for builder in opts.builder_hint.split(','):
          if cros_lab_util.normalize_board_name(builder) == host_board:
            board_to_build = builder
            break
      else:
        board_to_build = host_board

      return 'ready', host, board_to_build
    time.sleep(1)

  logger.warning('unable to lease DUT in time limit')
  return 'wait', None, None


def cmd_allocate_dut(opts):
  leased_dut = None
  try:
    todo, host, board = do_allocate_dut(opts)
    leased_dut = cros_lab_util.dut_name_to_address(host) if host else None
    result = {'result': todo, 'leased_dut': leased_dut, 'board': board}
    print(json.dumps(result))
  except Exception as e:
    logger.exception('cmd_allocate_dut failed')
    exception_name = e.__class__.__name__
    result = {
        'result': 'failed',
        'exception': exception_name,
        'text': str(e),
    }
    print(json.dumps(result))


def cmd_repair_dut(opts):
  cros_lab_util.repair(opts.dut)


@cli.fatal_error_handler
def main():
  common.init()
  parser = argparse.ArgumentParser()
  cli.patching_argparser_exit(parser)
  common.add_common_arguments(parser)
  subparsers = parser.add_subparsers(
      dest='command', title='commands', metavar='<command>')

  parser_version_info = subparsers.add_parser(
      'version_info',
      help='Query version info of given chromeos build',
      description='Given chromeos `board` and `version`, '
      'print version information of components.')
  parser_version_info.add_argument(
      'board', help='ChromeOS board name, like "samus".')
  parser_version_info.add_argument(
      'version',
      type=cros_util.argtype_cros_version,
      help='ChromeOS version, like "9876.0.0" or "R62-9876.0.0"')
  parser_version_info.add_argument(
      'name',
      nargs='?',
      help='Component name. If specified, output its version string. '
      'Otherwise output all version info as dict in json format.')
  parser_version_info.set_defaults(func=cmd_version_info)

  parser_query_dut_board = subparsers.add_parser(
      'query_dut_board', help='Query board name of given DUT')
  parser_query_dut_board.add_argument('dut')
  parser_query_dut_board.set_defaults(func=cmd_query_dut_board)

  parser_reboot = subparsers.add_parser(
      'reboot',
      help='Reboot a DUT',
      description='Reboot a DUT and verify the reboot is successful.')
  parser_reboot.add_argument('dut')
  parser_reboot.set_defaults(func=cmd_reboot)

  parser_lease_dut = subparsers.add_parser(
      'lease_dut',
      help='Lease a DUT in the lab',
      description='Lease a DUT in the lab. '
      'This is implemented by `skylab lease-dut` with additional checking.')
  # "skylab lease-dut" doesn't take reason, so this is not required=True.
  parser_lease_dut.add_argument('--session', help='session name')
  parser_lease_dut.add_argument('dut')
  parser_lease_dut.add_argument(
      '--duration',
      type=float,
      help='duration in seconds; will be round to minutes')
  parser_lease_dut.set_defaults(func=cmd_lease_dut)

  parser_release_dut = subparsers.add_parser(
      'release_dut',
      help='Release a DUT in the lab',
      description='Release a DUT in the lab. '
      'This is implemented by `skylab release-dut` with additional checking.')
  parser_release_dut.add_argument('--session', help='session name')
  parser_release_dut.add_argument('dut')
  parser_release_dut.set_defaults(func=cmd_release_dut)

  parser_allocate_dut = subparsers.add_parser(
      'allocate_dut',
      help='Allocate a DUT in the lab',
      description='Allocate a DUT in the lab. It will lease a DUT in the lab '
      'for bisecting. The caller (bisect-kit runner) of this command should '
      'retry this command again later if no DUT available now.')
  parser_allocate_dut.add_argument(
      '--session', required=True, help='session name')
  parser_allocate_dut.add_argument(
      '--pool',
      help='Pool to search DUT (default: %(default)s)',
      default=DEFAULT_DUT_POOL)
  group = parser_allocate_dut.add_mutually_exclusive_group(required=True)
  group.add_argument('--board', help='allocation criteria; comma separated')
  group.add_argument('--model', help='allocation criteria; comma separated')
  group.add_argument('--sku', help='allocation criteria; comma separated')
  group.add_argument('--dut_name', help='allocation criteria; comma separated')
  parser_allocate_dut.add_argument(
      '--version_hint', help='chromeos version; comma separated')
  parser_allocate_dut.add_argument(
      '--builder_hint', help='chromeos builder; comma separated')
  # Pubsub ack deadline is 10 minutes (b/143663659). Default 9 minutes with 1
  # minute buffer.
  parser_allocate_dut.add_argument(
      '--time_limit',
      type=int,
      default=9 * 60,
      help='Time limit to attempt lease in seconds (default: %(default)s)')
  parser_allocate_dut.add_argument(
      '--duration',
      type=float,
      help='lease duration in seconds; will be round to minutes')
  parser_allocate_dut.add_argument(
      '--parallel',
      type=int,
      default=1,
      help='Submit multiple lease tasks to speed up (default: %(default)d)')
  parser_allocate_dut.set_defaults(func=cmd_allocate_dut)

  parser_repair_dut = subparsers.add_parser(
      'repair_dut',
      help='Repair a DUT in the lab',
      description='Repair a DUT in the lab. '
      'This is simply wrapper of "deploy repair" with additional checking.')
  parser_repair_dut.add_argument('dut')
  parser_repair_dut.set_defaults(func=cmd_repair_dut)

  opts = parser.parse_args()
  common.config_logging(opts)

  # It's optional by default since python3.
  if not opts.command:
    parser.error('command is missing')
  opts.func(opts)


if __name__ == '__main__':
  main()