add host_get_migration_plan class to host.py
Bug: chromium:976324
Change-Id: I6716aa2ecb16bc31d3f0f3dea44a87dbe93a34ca
Reviewed-on: https://chromium-review.googlesource.com/1701369
Tested-by: Gregory Nisbet <gregorynisbet@google.com>
Commit-Ready: Gregory Nisbet <gregorynisbet@google.com>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Gregory Nisbet <gregorynisbet@google.com>
diff --git a/cli/host.py b/cli/host.py
index 62c6f10..2771c84 100644
--- a/cli/host.py
+++ b/cli/host.py
@@ -24,13 +24,18 @@
import random
import re
import socket
+import time
from autotest_lib.cli import action_common, rpc, topic_common, skylab_utils
-from autotest_lib.cli.skylab_json_utils import process_labels
+from autotest_lib.cli import fair_partition
from autotest_lib.client.bin import utils as bin_utils
+from autotest_lib.cli.skylab_json_utils import process_labels
from autotest_lib.client.common_lib import error, host_protections
from autotest_lib.server import frontend, hosts
from autotest_lib.server.hosts import host_info
+from autotest_lib.server.lib.status_history import HostJobHistory
+from autotest_lib.server.lib.status_history import UNUSED, WORKING
+from autotest_lib.server.lib.status_history import BROKEN, UNKNOWN
try:
@@ -345,6 +350,60 @@
self.print_dict(attributes, 'Host Attributes', line_before=True)
+class host_get_migration_plan(host_stat):
+ """atest host get_migration_plan --mlist <file>|<hosts>"""
+ usage_action = "get_migration_plan"
+
+ def __init__(self):
+ super(host_get_migration_plan, self).__init__()
+ self.parser.add_option("--ratio", default=0.5, type=float, dest="ratio")
+ self.add_skylab_options()
+
+ def parse(self):
+ (options, leftover) = super(host_get_migration_plan, self).parse()
+ self.ratio = options.ratio
+ return (options, leftover)
+
+ def execute(self):
+ afe = frontend.AFE()
+ results = super(host_get_migration_plan, self).execute()
+ working = []
+ non_working = []
+ for stats, _, _, _ in results:
+ assert len(stats) == 1
+ stats = stats[0]
+ hostname = stats["hostname"]
+ now = time.time()
+ history = HostJobHistory.get_host_history(
+ afe=afe,
+ hostname=hostname,
+ start_time=now,
+ end_time=now - 24 * 60 * 60,
+ )
+ dut_status, _ = history.last_diagnosis()
+ if dut_status in [UNUSED, WORKING]:
+ working.append(hostname)
+ elif dut_status == BROKEN:
+ non_working.append(hostname)
+ elif dut_status == UNKNOWN:
+ # if it's unknown, randomly assign it to working or
+ # nonworking, since we don't know.
+ # The two choices aren't actually equiprobable, but it
+ # should be fine.
+ random.choice([working, non_working]).append(hostname)
+ else:
+ raise ValueError("unknown status %s" % dut_status)
+ working_transfer, working_retain = fair_partition.partition(working, self.ratio)
+ non_working_transfer, non_working_retain = \
+ fair_partition.partition(non_working, self.ratio)
+ return {
+ "transfer": working_transfer + non_working_transfer,
+ "retain": working_retain + non_working_retain,
+ }
+
+ def output(self, results):
+ print json.dumps(results, indent=4, sort_keys=True)
+
class host_statjson(host_stat):
"""atest host statjson --mlist <file>|<hosts>