blob: 9650ed9abff3455b079960c5ba71ae50044973b3 [file] [log] [blame]
#!/usr/bin/python3
# Copyright 2013 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""A small wrapper script, iterates through
the known hosts and tries to call get_labels()
to discover host functionality, and adds these
detected labels to host.
Limitations:
- Does not keep a count of how many labels were
actually added.
- If a label is added by this script because it
is detected as supported by get_labels, but later becomes
unsupported, this script has no way to know that it
should be removed, so it will remain attached to the host.
See crosbug.com/38569
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from multiprocessing import pool
import logging
import socket
import argparse
import sys
import common
from autotest_lib.server import hosts
from autotest_lib.server import frontend
from autotest_lib.client.common_lib import error
# A list of label prefix that each dut should only have one of such label with
# the given prefix, e.g., a dut can't have both labels of power:battery and
# power:AC_only.
SINGLETON_LABEL_PREFIX = ['power:']
def add_missing_labels(afe, hostname):
"""
Queries the detectable labels supported by the given host,
and adds those labels to the host.
@param afe: A frontend.AFE() instance.
@param hostname: The host to query and update.
@return: True on success.
False on failure to fetch labels or to add any individual label.
"""
host = None
try:
host = hosts.create_host(hostname)
labels = host.get_labels()
except socket.gaierror:
logging.warning('Unable to establish ssh connection to hostname '
'%s. Skipping.', hostname)
return False
except error.AutoservError:
logging.warning('Unable to query labels on hostname %s. Skipping.',
hostname)
return False
finally:
if host:
host.close()
afe_host = afe.get_hosts(hostname=hostname)[0]
label_matches = afe.get_labels(name__in=labels)
for label in label_matches:
singleton_prefixes = [p for p in SINGLETON_LABEL_PREFIX
if label.name.startswith(p)]
if len(singleton_prefixes) == 1:
singleton_prefix = singleton_prefixes[0]
# Delete existing label with `singleton_prefix`
labels_to_delete = [l for l in afe_host.labels
if l.startswith(singleton_prefix) and
not l in labels]
if labels_to_delete:
logging.warning('Removing label %s', labels_to_delete)
afe_labels_to_delete = afe.get_labels(name__in=labels_to_delete)
for afe_label in afe_labels_to_delete:
afe_label.remove_hosts(hosts=[hostname])
label.add_hosts(hosts=[hostname])
missing_labels = set(labels) - set([l.name for l in label_matches])
if missing_labels:
for label in missing_labels:
logging.warning('Unable to add label %s to host %s. '
'Skipping unknown label.', label, hostname)
return False
return True
def main():
""""
Entry point for add_detected_host_labels script.
"""
parser = argparse.ArgumentParser()
parser.add_argument('-s', '--silent', dest='silent', action='store_true',
help='Suppress all but critical logging messages.')
parser.add_argument('-i', '--info', dest='info_only', action='store_true',
help='Suppress logging messages below INFO priority.')
parser.add_argument('-m', '--machines', dest='machines',
help='Comma separated list of machines to check.')
options = parser.parse_args()
if options.silent and options.info_only:
print('The -i and -s flags cannot be used together.')
parser.print_help()
return 0
if options.silent:
logging.disable(logging.CRITICAL)
if options.info_only:
logging.disable(logging.DEBUG)
threadpool = pool.ThreadPool()
afe = frontend.AFE()
if options.machines:
hostnames = [m.strip() for m in options.machines.split(',')]
else:
hostnames = afe.get_hostnames()
successes = sum(threadpool.imap_unordered(
lambda x: add_missing_labels(afe, x),
hostnames))
attempts = len(hostnames)
logging.info('Label updating finished. Failed update on %d out of %d '
'hosts.', attempts-successes, attempts)
return 0
if __name__ == '__main__':
sys.exit(main())