| #!/usr/bin/env python2 |
| # Copyright 2019 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 |
| from __future__ import unicode_literals |
| from __future__ import print_function |
| |
| import collections |
| import datetime |
| import io |
| import json |
| import os |
| import subprocess |
| import tempfile |
| import time |
| import shutil |
| import sys |
| import types |
| import itertools |
| |
| import common |
| |
| _THIS_FILE = os.path.abspath(__file__) |
| _THIS_DIR = os.path.dirname(_THIS_FILE) |
| |
| _SKYLAB_EXE = 'skylab' |
| |
| __all__ = ['migrate', 'setup'] |
| |
| _TEMPPATH = object() |
| |
| _FAILED_STEP_SENTINEL = object() |
| |
| _LITERAL_MAP = { |
| 'True': True, |
| 'False': False, |
| 'None': None, |
| } |
| |
| |
| TEXT = (unicode, str) |
| |
| |
| def find_atest_path(): |
| """Get the path to the 'atest' executable. |
| |
| @return : path to 'atest' executable |
| """ |
| atest_exe = os.path.join(_THIS_DIR, 'atest') |
| assert os.path.exists(atest_exe) |
| return atest_exe |
| |
| |
| _ATEST_EXE = find_atest_path() |
| |
| |
| def strip_suffix(str, suffix): |
| if str.endswith(suffix): |
| return str[:-len(suffix)] |
| else: |
| return str |
| |
| |
| def call_with_tempfile(cmd, lines): |
| """Execute command requiring a temporary file and return a CommandOutput struct. |
| |
| @param cmd : the components of the argv to be executed. |
| The magical value _TEMPPATH will be replaced with the path |
| to the temporary file. |
| @param lines : the lines of content to write to the temporary file |
| |
| @returns : CommandOutput struct containing output as list of lines |
| and the exit status |
| """ |
| if isinstance(cmd, (str, unicode)): |
| raise TypeError('cmd cannot be str or unicode') |
| assert not isinstance(lines, (str, unicode)) |
| with tempfile.NamedTemporaryFile() as fh: |
| for line in lines: |
| fh.write(line) |
| if line.endswith('\n'): |
| pass |
| else: |
| fh.write('\n') |
| fh.flush() |
| assert os.path.exists(fh.name) |
| cmd = [(x if x is not _TEMPPATH else fh.name) for x in cmd] |
| try: |
| output = subprocess.check_output(cmd) |
| if isinstance(output, (bytes, unicode)): |
| output = output.splitlines() |
| return CommandOutput( |
| exit_code=0, output=[x.decode('utf-8') for x in output]) |
| except subprocess.CalledProcessError as e: |
| return CommandOutput( |
| exit_code=e.returncode, |
| output=[x.decode('utf-8') for x in e.output.splitlines()]) |
| |
| |
| |
| # accepts: string |
| # returns: string but with exactly one trailing newline |
| def _one_trailing_newline(s): |
| s = s.rstrip("\n") |
| return s + "\n" |
| |
| # accepts: shell command, rest of args |
| # returns: exit_status, stdout, stderr |
| def shell_capture_all(cmd, *rest): |
| shellcmd = ("bash", "-c", cmd, "bash",) + rest |
| pr = subprocess.Popen( |
| shellcmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| ) |
| stdout, stderr = pr.communicate() |
| return pr.returncode, stdout, stderr |
| |
| |
| # accepts: shell command, rest of args |
| # returns: exit_status, stdout, stderr |
| def shell_capture_all_no_stdin(cmd, *rest): |
| with open(os.devnull) as null: |
| shellcmd = ("bash", "-c", cmd, "bash",) + rest |
| pr = subprocess.Popen( |
| shellcmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| stdin=null, |
| ) |
| stdout, stderr = pr.communicate() |
| return pr.returncode, stdout, stderr |
| |
| |
| # accepts: shell command, lines of temporary file |
| # returns: exit_status, stdout, stderr |
| def shell_capture_all_with_tempfile(cmd, lines): |
| exit_status = stdout = stderr = None |
| if lines in (bytes, unicode): |
| raise TypeError("lines must not be text-like (%s)" % type(lines)) |
| filename = None |
| with tempfile.NamedTemporaryFile(delete=False) as fh: |
| filename = fh.name |
| for line in lines: |
| fh.write(_one_trailing_newline(line)) |
| try: |
| exit_status, stdout, stderr = shell_capture_all_no_stdin(cmd, filename) |
| finally: |
| os.unlink(filename) |
| return exit_status, stdout, stderr |
| |
| |
| CommandOutput = collections.namedtuple('CommandOutput', ['output', 'exit_code']) |
| |
| |
| def _nontrivially_pairwise_disjoint(*sets): |
| """If there are any items present in more than one set, then 'sets' is not pairwise disjoint. |
| |
| If there are exactly zero or one sets, then there are no pairs of sets |
| and therefore the pairwise disjoint condition will always hold |
| regardless of the set contents. Therefore, calling |
| _nontrivially_pairwise_disjoint |
| with fewer than 2 sets probably indicates a logic error and will result |
| in an exception being thrown. |
| |
| Example: [{1}, {2}, set(), {3, 4, 5}, set()] |
| CounterExample: [{1, 2}, {2, 3}] |
| |
| @param sets: a sequence of sets |
| @return: whether the sets are pairwise disjoint |
| """ |
| if len(sets) in (0, 1): |
| raise ValueError( |
| 'a collection of 0 or 1 sets is trivially pairwise disjoint.') |
| combined = set() |
| sum_len_set = 0 |
| for set_ in sets: |
| combined.update(set_) |
| sum_len_set += len(set_) |
| assert len(combined) <= sum_len_set |
| return len(combined) == sum_len_set |
| |
| |
| MigrateDutCommandStatus = collections.namedtuple('MigrateDutCommandStatus', [ |
| 'success', 'failure', 'needs_add_to_skylab', 'needs_drone', 'needs_rename' |
| ]) |
| |
| AddToSkylabInventoryAndDroneStatus = collections.namedtuple( |
| 'AddToSkylabInventoryAndDroneStatus', |
| ['complete', 'without_drone', 'not_started']) |
| |
| RenameCommandStatus = collections.namedtuple('RenameCommandStatus', |
| ['renamed', 'not_renamed']) |
| |
| LockCommandStatus = collections.namedtuple('LockCommandStatus', |
| ['locked', 'not_locked', 'tries']) |
| |
| MigrationPlan = collections.namedtuple('MigrationPlan', ['transfer', 'retain']) |
| |
| |
| class MigrationException(Exception): |
| """Raised when migration fails""" |
| pass |
| |
| |
| def stderr_log(*args, **kwargs): |
| return print(*args, file=sys.stderr, **kwargs) |
| |
| |
| def _humantime(): |
| return tuple(datetime.datetime.now().timetuple())[:6] |
| |
| |
| def _migration_json_summary(failed_step=_FAILED_STEP_SENTINEL, |
| plan=None, |
| not_locked=None, |
| migrate_status=None, |
| unconditionally_migrate_status=None): |
| assert isinstance(plan, MigrationPlan) |
| assert not isinstance(not_locked, (str, unicode)) |
| assert isinstance(failed_step, (types.NoneType, unicode)) |
| assert isinstance(migrate_status, (types.NoneType, MigrateDutCommandStatus)) |
| assert isinstance(unconditionally_migrate_status, MigrateDutCommandStatus) |
| |
| def merge_attrs(fieldname, struct1, struct2=None): |
| merged = set() |
| if struct1: |
| merged.update(getattr(struct1, fieldname)) |
| if struct2: |
| merged.update(getattr(struct2, fieldname)) |
| return sorted(merged) |
| |
| |
| out = { |
| 'locked_success': (failed_step is None), |
| 'failed_step': failed_step, |
| 'plan': { |
| 'transfer': merge_attrs('transfer', plan), |
| 'retain': merge_attrs('retain', plan), |
| }, |
| 'duts': { |
| 'migrated': |
| merge_attrs('success', migrate_status, unconditionally_migrate_status), |
| 'not_locked': |
| list(sorted(set(not_locked))), |
| 'needs_add_to_skylab': |
| merge_attrs('needs_add_to_skylab', migrate_status, unconditionally_migrate_status), |
| 'needs_drone': |
| merge_attrs('needs_drone', migrate_status, unconditionally_migrate_status), |
| 'needs_rename': |
| merge_attrs('needs_rename', migrate_status, unconditionally_migrate_status), |
| } |
| } |
| return out |
| |
| |
| class AtestCmd(object): |
| """Helper functions for executing 'atest' commands""" |
| |
| @staticmethod |
| def brief_info_cmd(): |
| """Command line for getting per-host info. |
| |
| @return : list of strings to be executed as external command |
| """ |
| return [_ATEST_EXE, 'host', 'list', '--parse', '-M', _TEMPPATH] |
| |
| @staticmethod |
| def brief_info(hostnames=None): |
| """Run brief info command. |
| |
| @return : iterator of dictionaries describing each hostname |
| """ |
| hostnames = hostnames or set() |
| items = call_with_tempfile(AtestCmd.brief_info_cmd(), hostnames).output |
| for item in AtestCmd.brief_info_filter(items): |
| yield item |
| |
| @staticmethod |
| def brief_info_filter(stream): |
| """Filter lines of output from 'atest host list...'. |
| |
| @return : iterator of fields |
| """ |
| for line in stream: |
| line = line.rstrip() |
| if line: |
| fields = line.split('|') |
| # if the line of output has exactly zero or one |
| # |-delimited sections, then it is not a description |
| # of a DUT. Silently discard such lines. |
| if len(fields) in (0, 1): |
| continue |
| # trim labels entry if it exists |
| if fields[-1].startswith('Labels='): |
| fields.pop() |
| d = {} |
| for f in fields: |
| k, _, v = f.partition('=') |
| # if the value associated with a key is a Python literal |
| # such as True, False, or None, replace it with the |
| # corresponding Python value. |
| # otherwise, use the original string. |
| d[k] = _LITERAL_MAP.get(v, v) |
| yield d |
| |
| @staticmethod |
| def rename_cmd(for_migration=True): |
| """Generate command line arguments for 'rename'. |
| |
| @return : command line arguments |
| """ |
| name_flag = '--for-migration' if for_migration else '--for-rollback' |
| return [ |
| _ATEST_EXE, 'host', 'rename', '--non-interactive', name_flag, |
| '--parse', '-M', _TEMPPATH |
| ] |
| |
| @staticmethod |
| def rename(hostnames=None, for_migration=True): |
| """Rename a list of hosts. |
| |
| @return : iterator of successfully renamed hosts |
| """ |
| hostnames = hostnames or set() |
| |
| to_migrate_hostnames = set() |
| already_migrated_hostnames = set() |
| |
| for hostname in hostnames: |
| if hostname.endswith("-migrated-do-not-use"): |
| already_migrated_hostnames.add(hostname) |
| else: |
| to_migrate_hostnames.add(hostname) |
| |
| stderr_log('begin rename', time.time(), _humantime()) |
| items = call_with_tempfile( |
| AtestCmd.rename_cmd(for_migration=for_migration), |
| lines=to_migrate_hostnames).output |
| |
| out = list(AtestCmd.rename_filter(items)) |
| out_seen = set(out) |
| |
| # out and already_migrated_hostnames should be disjoint |
| # but if they aren't we still don't want to list the same |
| # hostname twice |
| for hostname in already_migrated_hostnames: |
| if hostname not in out_seen: |
| out.append(hostname) |
| |
| stderr_log('end rename', time.time(), _humantime()) |
| return out |
| |
| @staticmethod |
| def rename_filter(stream): |
| """Process each item of output from `atest host rename...`. |
| |
| @return : iterator of successfully renamed hosts |
| """ |
| for item in stream: |
| row = [x.strip() for x in item.strip().split()] |
| if len(row) == 3: |
| src, sep, dest = row |
| # dest has the 'migrated-do-not-use' suffix |
| # use src! |
| if sep != 'to': |
| continue |
| yield src |
| |
| @staticmethod |
| def statjson_cmd(hostname=None): |
| """Command line for generating json for hostname. |
| |
| @return : command line |
| """ |
| return [_ATEST_EXE, 'host', 'statjson', '--', hostname] |
| |
| @staticmethod |
| def statjson(hostname=None): |
| """Run the command for getting the host json. |
| |
| @return : 'atest host statjson' output as parsed json. |
| """ |
| cmd = AtestCmd.statjson_cmd(hostname=hostname) |
| (out, err, exit_status) = capture_all(cmd) |
| if exit_status == 0: |
| try: |
| return json.loads(out.decode('utf-8')) |
| except ValueError: |
| sys.stderr.write(out) |
| sys.stderr.write("\n\n") |
| return None |
| else: |
| if exit_status: |
| if "Failed to stat:" in err: |
| assert "Unknown host" in err |
| return None |
| else: |
| assert "unexpected failure" |
| |
| @staticmethod |
| def atest_lock_cmd(reason=None): |
| """Generate command for 'atest host mod --lock'. |
| |
| @return : command line |
| """ |
| return [ |
| _ATEST_EXE, 'host', 'mod', '--lock', '-r', reason, '-M', _TEMPPATH |
| ] |
| |
| @staticmethod |
| def atest_lock(reason=None, hostnames=None): |
| """Try to lock hostnames via 'atest host mod --lock'. |
| |
| @return : Nothing |
| """ |
| hostnames = hostnames or set() |
| assert isinstance(reason, unicode) |
| cmd = AtestCmd.atest_lock_cmd(reason=reason) |
| # NOTE: attempting to lock a host can fail because the host |
| # is already locked. Therefore, atest_lock always succeeds |
| # regardless of the exit status of the command. |
| call_with_tempfile(cmd, hostnames) |
| |
| @staticmethod |
| def atest_lock_filter(stream): |
| """Take lines from 'atest host mod --lock' and emit a stream of hostnames. |
| |
| The first line "Locked hosts:" is removed. We trim the whitespace of the |
| other lines. |
| |
| Input: |
| Locked Hosts: |
| A |
| B |
| C |
| |
| Output: |
| A |
| B |
| C |
| """ |
| for x in stream: |
| if x.lower().startswith('locked host'): |
| continue |
| else: |
| yield x.strip() |
| |
| @staticmethod |
| def atest_unlock_cmd(): |
| """Generate command for 'atest host mod --unlock'.""" |
| return [_ATEST_EXE, 'host', 'mod', '--unlock', '-M', _TEMPPATH] |
| |
| @staticmethod |
| def atest_unlock(reason=None, hostnames=None): |
| """Unlock hostnames via 'atest host mod --unlock'. |
| |
| @return : iterator of successfully unlocked hosts |
| """ |
| hostnames = hostnames or set() |
| cmd = AtestCmd.atest_unlock_cmd() |
| items = call_with_tempfile(cmd, hostnames).output |
| for item in AtestCmd.atest_unlock_filter(items): |
| yield item |
| |
| @staticmethod |
| def atest_unlock_filter(stream): |
| """Take lines from 'atest host mod --unlock' and emit a stream of hostnames. |
| |
| The first line "Unlocked hosts:" is removed. We trim the whitespace of |
| the other lines. |
| |
| Input: |
| Unlocked Hosts: |
| A |
| B |
| C |
| |
| Output: |
| A |
| B |
| C |
| """ |
| for x in stream: |
| if x.lower().startswith('unlocked host'): |
| continue |
| else: |
| yield x.strip() |
| |
| @staticmethod |
| def atest_get_migration_plan_cmd(ratio): |
| """Generate command for 'atest host get_migration_plan --mlist ...'""" |
| return [ |
| _ATEST_EXE, 'host', 'get_migration_plan', '--ratio', |
| unicode(ratio), '--mlist', _TEMPPATH |
| ] |
| |
| @staticmethod |
| def atest_get_migration_plan(ratio, hostnames=[]): |
| # optimizations in case the ratio is 1 or 0 |
| hostnames = hostnames or set() |
| if ratio == 0: |
| return { |
| 'transfer': [], |
| 'retain': hostnames, |
| } |
| if ratio == 1: |
| return { |
| 'transfer': hostnames, |
| 'retain': [], |
| } |
| cmd = AtestCmd.atest_get_migration_plan_cmd(ratio) |
| output = call_with_tempfile(cmd, hostnames).output |
| out = json.loads(''.join(output)) |
| return out |
| |
| |
| def trywith(exn_type, f, *args, **kwargs): |
| out = None |
| exn = None |
| try: |
| out = f(*args, **kwargs) |
| except exn_type as e: |
| exn = e |
| return (out, exn) |
| |
| |
| |
| def backtick(*args, **kwargs): |
| output = None |
| exit_status = None |
| out, exn = trywith(subprocess.CalledProcessError, subprocess.check_output, *args, **kwargs) |
| if exn is None: |
| output = out |
| exit_status = 0 |
| else: |
| output = exn.output |
| exit_status = exn.returncode |
| return (output, exit_status) |
| |
| |
| def capture_all(*args, **kwargs): |
| proc = subprocess.Popen( |
| *args, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| **kwargs |
| ) |
| out, err = proc.communicate() |
| return (out, err, proc.returncode) |
| |
| |
| |
| # accepts: iterable of hostnames |
| # returns: { |
| # good: hostnames in autotest |
| # bad: hostnames not in autotest |
| # } |
| def autotest_status(hostnames): |
| os.environ["ATEST"] = _ATEST_EXE |
| status, out, err = shell_capture_all_with_tempfile('"${ATEST}" host list --hostnames-only --mlist "$1"', hostnames) |
| good = [] |
| bad = [] |
| # process the bad lines |
| for errline in err: |
| # skip preamble |
| if errline.startswith("Unknown host"): |
| continue |
| bad.append(errline.strip()) |
| # process the good lines |
| for goodline in good: |
| good.append(good.strip()) |
| return { |
| "good": good, |
| "bad": bad |
| } |
| |
| # accepts: iterable of hostnames |
| # returns: { |
| # good: hostnames in skylab |
| # bad: hostnames not in skylab |
| # } |
| def skylab_status(hostnames): |
| os.environ["SKYLAB"] = _SKYLAB_EXE |
| good = [] |
| bad = [] |
| for hostname in hostnames: |
| os.environ["HOSTNAME"] = hostname |
| status, out, err = shell_capture_all_no_stdin('"${SKYLAB}" dut-info "${HOSTNAME}"') |
| # TODO(gregorynisbet): make error checking more robust here to see why exactly |
| # we couldn't get info on the DUT. |
| if status == 0: |
| good.append(hostname) |
| else: |
| bad.append(hostname) |
| return { |
| "good": good, |
| "bad": bad, |
| } |
| |
| |
| |
| # accepts: iterable of hostnames |
| # returns: { |
| # good: hostnames with no issues |
| # not_renamed: hostnames in skylab but not renamed in autotest |
| # not_in_skylab: hostnames that are not in skylab |
| # not_renamed_not_in_skylab: hostnames that aren't renamed or in skylab |
| # } |
| def hostname_migrated_status(hostnames): |
| migrated_map = {} |
| for hostname in hostnames: |
| migrated_map[hostname + "-migrated-do-not-use"] = hostname |
| |
| atest_out = autotest_status(hostnames) |
| atest_out_good = set(atest_out["good"]) |
| atest_out_bad = set(atest_out["bad"]) |
| |
| skylab_out = skylab_status(hostnames) |
| skylab_out_good = set(skylab_out["good"]) |
| skylab_out_bad = set(skylab_out["bad"]) |
| |
| atest_renamed_out = autotest_status(list(migrated_map)) |
| atest_renamed_out_good = set(atest_renamed_out["good"]) |
| atest_renamed_out_bad = set(atest_renamed_out["bad"]) |
| |
| good = [] |
| not_renamed = [] |
| not_in_skylab = [] |
| not_renamed_not_in_skylab = [] |
| |
| for hostname in hostnames: |
| # hostname flags is a string listing the undesirable properties |
| # associated with that particular hostname |
| # A -- host has bad autotest status, either not renamed or old hostname present |
| # S -- not migrated to skylab |
| hostname_flags = set() |
| if hostname in atest_out_good: |
| hostname_flags.add("A") |
| if hostname in atest_renamed_out_bad: |
| hostname_flags.add("A") |
| if hostname in skylab_out_bad: |
| hostname_flags.add("S") |
| |
| if hostname_flags == set(): |
| good.append(hostname) |
| elif hostname_flags == {"A"}: |
| not_renamed.append(hostname) |
| elif hostname_flags == {"S"}: |
| not_in_skylab.append(hostname) |
| elif hostname_flags == {"A", "S"}: |
| not_renamed_not_in_skylab.append(hostname) |
| else: |
| assert False, ("impossible, unexpected set %s" % hostname_flags) |
| |
| return { |
| "good": good, |
| "not_renamed": not_renamed, |
| "not_in_skylab": not_in_skylab, |
| "not_renamed_not_in_skylab": not_renamed_not_in_skylab, |
| } |
| |
| |
| class SkylabCmd(object): |
| """Helper functions for executing Skylab commands""" |
| |
| ADD_MANY_DUTS_CMD = (_SKYLAB_EXE, 'quick-add-duts') |
| |
| @staticmethod |
| def add_one_dut_cmd(): |
| """Create the skylab command line invocation for adding a single DUT.""" |
| return [ |
| _SKYLAB_EXE, |
| 'add-dut', |
| '-skip-image-download', |
| '-skip-install-firmware', |
| '-skip-install-os', |
| '-specs-file', |
| _TEMPPATH, |
| ] |
| |
| @staticmethod |
| def add_one_dut(add_dut_content): |
| """Add one dut to skylab.""" |
| stderr_log('begin add_one_dut', time.time(), _humantime()) |
| cmd = SkylabCmd.add_one_dut_cmd() |
| out = call_with_tempfile(cmd, [json.dumps(add_dut_content)]) |
| stderr_log('end add_one_dut', time.time(), _humantime()) |
| return out |
| |
| @staticmethod |
| def assign_one_dut_cmd(hostname=None): |
| """Command line for assigning a single DUT to a randomly chosen drone.""" |
| # by default, skylab assign-dut will pick a random drone |
| return [_SKYLAB_EXE, 'assign-dut', '--', hostname] |
| |
| @staticmethod |
| def add_many_duts(dut_contents): |
| stderr_log('begin add_many_duts', time.time(), _humantime()) |
| for dut_content in dut_contents: |
| stderr_log("add many DUTs: ", str(dut_content)[:80] + "...") |
| """Add multiple DUTs to skylab at once. |
| |
| @param dut_contents: a sequence of JSON-like objects describing DUTs as |
| used by `skylab add-dut` and `skylab quick-add-dut` |
| |
| @returns : nothing |
| """ |
| # TODO(gregorynisbet) -- how fine-grained does the error reporting need |
| # to be? is it possible for some duts to be |
| # successfully migrated and others not? |
| # The action performed by `skylab quick-add-duts` |
| # is idempotent, so trying multiple times is not |
| # necessarily a problem. |
| td = tempfile.mkdtemp() |
| try: |
| paths = [] |
| for i in range(len(dut_contents)): |
| path_ = os.path.join(td, str(i)) |
| with open(path_, 'w') as fh: |
| json.dump(dut_contents[i], fh) |
| paths.append(path_) |
| cmd = list(SkylabCmd.ADD_MANY_DUTS_CMD) + paths |
| print("log command") |
| stderr_log(cmd) |
| print("capture_all") |
| # ignore cases where the hostname doesn't exist |
| (out, err, exit_status) = capture_all(cmd) |
| if exit_status != 0: |
| if "Failed to stat:" in err: |
| assert "Unknown host" in err |
| # then do nothing |
| |
| # shutil.rmtree(td, ignore_errors=True) |
| finally: |
| stderr_log('end add_many_duts', time.time(), _humantime()) |
| |
| |
| class Migration(object): |
| |
| @staticmethod |
| def migration_plan(ratio, hostnames=None): |
| hostnames = hostnames or set() |
| plan = AtestCmd.atest_get_migration_plan( |
| ratio=ratio, hostnames=hostnames) |
| return MigrationPlan(transfer=plan['transfer'], retain=plan['retain']) |
| |
| @staticmethod |
| def lock(hostnames=None, reason=None, retries=3): |
| """Lock a list of hostnames with retries. |
| """ |
| hostnames = hostnames or set() |
| assert isinstance(reason, unicode) |
| to_lock = set(hostnames) |
| for _ in range(retries): |
| AtestCmd.atest_lock(hostnames=to_lock.copy(), reason=reason) |
| |
| @staticmethod |
| def ensure_lock(hostnames=None): |
| """Without changing the state of a DUT, determine which are locked. |
| |
| @return : LockCommandStatus |
| """ |
| hostnames = hostnames or set() |
| dut_infos = AtestCmd.brief_info(hostnames=hostnames) |
| all_hosts = set(hostnames) |
| confirmed_locked = set() |
| for dut_info in dut_infos: |
| locked = dut_info['Locked'] |
| assert locked in (True, False) |
| if locked: |
| confirmed_locked.add(dut_info['Host']) |
| return LockCommandStatus( |
| locked=confirmed_locked, |
| not_locked=(all_hosts - confirmed_locked), |
| tries=None, |
| ) |
| |
| @staticmethod |
| def rename(hostnames=None, for_migration=True, retries=1): |
| """Rename a list of hosts with retry. |
| |
| @return : {"renamed": renamed hosts, "not-renamed": not renamed |
| hosts} |
| """ |
| hostnames = hostnames or set() |
| all_hosts = set(hostnames) |
| needs_rename = all_hosts.copy() |
| for _ in range(retries): |
| for successfully_renamed in AtestCmd.rename( |
| hostnames=needs_rename.copy(), for_migration=for_migration): |
| needs_rename.discard(successfully_renamed) |
| out = RenameCommandStatus( |
| renamed=(all_hosts - needs_rename), |
| not_renamed=needs_rename, |
| ) |
| return out |
| |
| @staticmethod |
| def add_to_skylab_inventory_and_drone(use_quick_add, hostnames=None, rename_retries=3): |
| """@returns : AddToSkylabInventoryAndDroneStatus""" |
| hostnames = hostnames or set() |
| assert not isinstance(hostnames, (unicode, bytes)) |
| stderr_log('begin add hostnames to inventory', time.time(), |
| _humantime()) |
| all_hosts = set(hostnames) |
| moved = set() |
| with_drone = set() |
| |
| if use_quick_add: |
| stderr_log("quick add path", time.time(), _humantime()) |
| dut_contents = [] |
| good_hostnames = [] |
| for hostname in hostnames: |
| out_json = dut_contents.append(AtestCmd.statjson(hostname=hostname)) |
| if out_json is None: |
| pass |
| else: |
| good_hostnames.append(out_json) |
| # SkylabCmd.add_many_duts does not check for whether the action was successful |
| # we use hostname_migrated_status to check whether the duts we were supposed to migrate |
| # were actually migrated |
| SkylabCmd.add_many_duts(dut_contents=dut_contents) |
| |
| # strip the migrated suffix when checking the status of each of the hostnames |
| truncated_hostnames = [strip_suffix(hostname, "-migrated-do-not-use") for hostname in hostnames] |
| status_out = hostname_migrated_status(truncated_hostnames) |
| |
| # anything in the good state or missing rename is fine |
| complete = status_out["good"] + status_out["not_renamed"] |
| # anything where the status indicates that the entity is not in skylab yet is not fine |
| not_started = status_out["not_in_skylab"] + status_out["not_renamed_not_in_skylab"] |
| |
| return AddToSkylabInventoryAndDroneStatus( |
| complete=complete, |
| without_drone=set(), |
| not_started=not_started, |
| ) |
| |
| else: |
| stderr_log("slow add path", time.time(), _humantime()) |
| for hostname in hostnames: |
| if hostname not in moved: |
| skylab_dut_descr = AtestCmd.statjson(hostname=hostname) |
| stderr_log("processing hostname", hostname) |
| status = SkylabCmd.add_one_dut(add_dut_content=skylab_dut_descr) |
| if status.exit_code != 0: |
| continue |
| moved.add(hostname) |
| with_drone.add(hostname) |
| |
| out = AddToSkylabInventoryAndDroneStatus( |
| complete=with_drone, |
| without_drone=(moved - with_drone), |
| not_started=((all_hosts - moved) - with_drone), |
| ) |
| stderr_log('end add hostnames to inventory', time.time(), _humantime()) |
| return out |
| |
| @staticmethod |
| def migrate_known_good_duts_until_max_duration_sync(use_quick_add, |
| hostnames=None, |
| max_duration=60 * 60, |
| min_ready_intervals=10, |
| interval_len=0): |
| """Take a list of DUTs and attempt to migrate them when they aren't busy. |
| |
| @param hostnames : list of hostnames |
| @param max_duration : when to stop trying to safely migrate duts |
| @param min_ready_intervals : the minimum number of intervals that a DUT |
| must have a good status |
| @param interval_len : the length in seconds of interval |
| @param use_quick_add : whether to use skylab quick-add-duts. |
| |
| @returns : {"success": successfuly migrated DUTS, "failure": |
| non-migrated DUTS} |
| """ |
| hostnames = hostnames or set() |
| assert interval_len is not None |
| stderr_log('begin migrating only ready DUTs', time.time(), _humantime()) |
| start = time.time() |
| stop = start + max_duration |
| good_intervals = collections.Counter() |
| need_to_move = set(hostnames) |
| successfully_moved = set() |
| needs_add_to_skylab = set() |
| needs_drone = set() |
| needs_rename = set() |
| while time.time() < stop: |
| if not need_to_move: |
| break |
| ready_to_move = set() |
| # determine which duts have been in a good state for min_ready_intervals |
| for record in AtestCmd.brief_info(hostnames=need_to_move.copy()): |
| hostname = record['Host'] |
| if record['Status'] not in {'Running', 'Provisioning'}: |
| good_intervals[hostname] += 1 |
| else: |
| del good_intervals[hostname] |
| if good_intervals[hostname] >= min_ready_intervals: |
| ready_to_move.add(hostname) |
| need_to_move.discard(hostname) |
| # move the ready to move duts now |
| # any dut that is declared ready to move at this point will definitely |
| # reach a terminal state |
| skylab_summary = Migration.add_to_skylab_inventory_and_drone( |
| hostnames=ready_to_move, use_quick_add=use_quick_add) |
| needs_add_to_skylab.update(skylab_summary.not_started) |
| needs_drone.update(skylab_summary.without_drone) |
| # rename the autotest entry all at once |
| rename_summary = Migration.rename( |
| hostnames=skylab_summary.complete, for_migration=True) |
| needs_rename.update(rename_summary.not_renamed) |
| successfully_moved.update(rename_summary.renamed) |
| time.sleep(interval_len) |
| out = MigrateDutCommandStatus( |
| success=successfully_moved, |
| failure=(need_to_move | needs_add_to_skylab | needs_drone |
| | needs_rename), |
| needs_add_to_skylab=needs_add_to_skylab, |
| needs_drone=needs_drone, |
| needs_rename=needs_rename, |
| ) |
| stderr_log('end migrating only ready DUTs', time.time(), _humantime()) |
| return out |
| |
| @staticmethod |
| def migrate_duts_unconditionally(hostnames, use_quick_add): |
| """regardless of the DUTs' status, forcibly migrate all the DUTs to skylab. |
| |
| @returns: MigrateDutCommandStatus |
| """ |
| hostnames = hostnames or set() |
| assert not isinstance(hostnames, (unicode, bytes)) |
| stderr_log('begin unconditional migration', time.time(), _humantime()) |
| successfully_moved = set() |
| needs_add_to_skylab = set() |
| needs_drone = set() |
| needs_rename = set() |
| skylab_summary = Migration.add_to_skylab_inventory_and_drone( |
| hostnames=hostnames, use_quick_add=use_quick_add) |
| needs_add_to_skylab.update(skylab_summary.not_started) |
| needs_drone.update(skylab_summary.without_drone) |
| rename_summary = Migration.rename( |
| hostnames=skylab_summary.complete, for_migration=True) |
| successfully_moved.update(rename_summary.renamed) |
| needs_rename.update(rename_summary.not_renamed) |
| needs_rename.discard(rename_summary.not_renamed) |
| out = MigrateDutCommandStatus( |
| success=successfully_moved, |
| failure=(needs_drone | needs_rename | needs_add_to_skylab), |
| needs_add_to_skylab=needs_add_to_skylab, |
| needs_drone=needs_drone, |
| needs_rename=needs_rename, |
| ) |
| stderr_log('end unconditional migration', time.time(), _humantime()) |
| return out |
| |
| @staticmethod |
| def migrate(hostnames=None, |
| ratio=1, |
| reason=None, |
| max_duration=None, |
| interval_len=None, |
| min_ready_intervals=10, |
| immediately=None, |
| use_quick_add=False): |
| """Migrate duts from autotest to skylab. |
| |
| @param ratio : ratio of DUTs in hostnames to migrate. |
| @param hostnames : hostnames to migrate |
| @param reason : the reason to give for providing the migration |
| @param interval_len : length of time between checks for DUT readiness |
| @param max_duration : the grace period to allow DUTs to finish their |
| tasks |
| @param min_ready_intervals : minimum number of intervals before a device |
| is healthy |
| |
| @return : nothing |
| """ |
| hostnames = hostnames or set() |
| assert isinstance(reason, (unicode, bytes)) |
| assert interval_len is not None |
| assert max_duration is not None |
| assert immediately is not None |
| reason = reason if isinstance(reason, |
| unicode) else reason.decode('utf-8') |
| # log the parameters of the migration |
| stderr_log('begin migrate', time.time(), _humantime()) |
| stderr_log('number of hostnames', len(hostnames), time.time(), _humantime()) |
| stderr_log('ratio', ratio, time.time(), _humantime()) |
| stderr_log('max_duration', max_duration, time.time(), _humantime()) |
| stderr_log('atest', _ATEST_EXE, time.time(), _humantime()) |
| stderr_log('skylab', _SKYLAB_EXE, time.time(), _humantime()) |
| stderr_log('minimum number of intervals', min_ready_intervals, time.time(), _humantime()) |
| stderr_log('immediately', immediately, time.time(), _humantime()) |
| stderr_log('use_quick_add', use_quick_add, time.time(), _humantime()) |
| |
| all_hosts = tuple(hostnames) |
| plan = Migration.migration_plan(ratio=ratio, hostnames=all_hosts) |
| Migration.lock(hostnames=plan.transfer, reason=reason) |
| failed_step = _FAILED_STEP_SENTINEL |
| ensure_lock_status = Migration.ensure_lock(hostnames=plan.transfer) |
| if ensure_lock_status.not_locked: |
| failed_step = 'lock' |
| to_migrate = plan.transfer |
| migrate_status = None |
| if not immediately: |
| migrate_status = \ |
| Migration.migrate_known_good_duts_until_max_duration_sync( |
| hostnames=to_migrate, |
| max_duration=max_duration, |
| min_ready_intervals=min_ready_intervals, |
| interval_len=interval_len, |
| use_quick_add=use_quick_add) |
| to_migrate = migrate_status.failure |
| unconditionally_migrate_status = Migration.migrate_duts_unconditionally( |
| use_quick_add=use_quick_add, |
| hostnames=to_migrate, |
| ) |
| failed_step = None |
| out = _migration_json_summary( |
| failed_step=failed_step, |
| plan=plan, |
| not_locked=ensure_lock_status.not_locked, |
| migrate_status=migrate_status, |
| unconditionally_migrate_status=unconditionally_migrate_status, |
| ) |
| stderr_log('end migrate', time.time(), _humantime()) |
| return out |
| |
| |
| |
| # accepts: iterable |
| # returns: item or None, ok (true if item is real, false otherwise) |
| def next_safe(it): |
| it = iter(it) |
| try: |
| return next(it), True |
| except StopIteration: |
| return None, False |
| |
| |
| # accepts: n (stride length), it (iterable) |
| # returns: iterator of arrays of n items each |
| def natatime(n, it): |
| it = iter(it) |
| while True: |
| out = [] |
| for i in range(n): |
| item, ok = next_safe(it) |
| if ok: |
| out.append(item) |
| else: |
| break |
| if len(out): |
| yield out |
| else: |
| return |
| |
| |
| |
| |
| def migrate(batch_size=None, hostnames=None, **kwargs): |
| if batch_size is None: |
| return Migration.migrate(hostnames=hostnames, **kwargs) |
| if batch_size is not None: |
| it = natatime(n=batch_size, it=hostnames) |
| out = [] |
| for batch in it: |
| res = Migration.migrate(hostnames=hostnames, **kwargs) |
| json.dumps(res, sys.stderr, indent=4) |
| out.append(res) |
| return out |
| |
| |
| def setup(atest_exe=None, skylab_exe=None): |
| """Configure the module-scoped path to atest and skylab executables.""" |
| if atest_exe is not None: |
| _ATEST_EXE = atest_exe |
| if skylab_exe is not None: |
| _SKYLAB_EXE = skylab_exe |
