blob: 89b57b2f99ea48324364ebf577de30d3aca5b7c2 [file] [log] [blame]
# Copyright (c) 2013 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.
"""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.
- 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.
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.
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
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.',
return False
if host:
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 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:
missing_labels = set(labels) - set([ 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.'
return 0
if options.silent:
if options.info_only:
threadpool = pool.ThreadPool()
afe = frontend.AFE()
if options.machines:
hostnames = [m.strip() for m in options.machines.split(',')]
hostnames = afe.get_hostnames()
successes = sum(threadpool.imap_unordered(
lambda x: add_missing_labels(afe, x),
attempts = len(hostnames)'Label updating finished. Failed update on %d out of %d '
'hosts.', attempts-successes, attempts)
return 0
if __name__ == '__main__':