Home | History | Annotate | Download | only in site_utils 
      1 #!/usr/bin/python
      2 
      3 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
      4 # Use of this source code is governed by a BSD-style license that can be
      5 # found in the LICENSE file.
      6 """A small wrapper script, iterates through
      7 the known hosts and tries to call get_labels()
      8 to discover host functionality, and adds these
      9 detected labels to host.
     10 
     11 Limitations:
     12  - Does not keep a count of how many labels were
     13    actually added.
     14  - If a label is added by this script because it
     15    is detected as supported by get_labels, but later becomes
     16    unsupported, this script has no way to know that it
     17    should be removed, so it will remain attached to the host.
     18    See crosbug.com/38569
     19 """
     20 
     21 
     22 from multiprocessing import pool
     23 import logging
     24 import socket
     25 import argparse
     26 import sys
     27 
     28 import common
     29 
     30 from autotest_lib.server import hosts
     31 from autotest_lib.server import frontend
     32 from autotest_lib.client.common_lib import error
     33 
     34 
     35 # A list of label prefix that each dut should only have one of such label with
     36 # the given prefix, e.g., a dut can't have both labels of power:battery and
     37 # power:AC_only.
     38 SINGLETON_LABEL_PREFIX = ['power:']
     39 
     40 def add_missing_labels(afe, hostname):
     41     """
     42     Queries the detectable labels supported by the given host,
     43     and adds those labels to the host.
     44 
     45     @param afe: A frontend.AFE() instance.
     46     @param hostname: The host to query and update.
     47 
     48     @return: True on success.
     49              False on failure to fetch labels or to add any individual label.
     50     """
     51     host = None
     52     try:
     53         host = hosts.create_host(hostname)
     54         labels = host.get_labels()
     55     except socket.gaierror:
     56         logging.warning('Unable to establish ssh connection to hostname '
     57                         '%s. Skipping.', hostname)
     58         return False
     59     except error.AutoservError:
     60         logging.warning('Unable to query labels on hostname %s. Skipping.',
     61                          hostname)
     62         return False
     63     finally:
     64         if host:
     65             host.close()
     66 
     67     afe_host = afe.get_hosts(hostname=hostname)[0]
     68 
     69     label_matches = afe.get_labels(name__in=labels)
     70 
     71     for label in label_matches:
     72         singleton_prefixes = [p for p in SINGLETON_LABEL_PREFIX
     73                               if label.name.startswith(p)]
     74         if len(singleton_prefixes) == 1:
     75             singleton_prefix = singleton_prefixes[0]
     76             # Delete existing label with `singleton_prefix`
     77             labels_to_delete = [l for l in afe_host.labels
     78                                 if l.startswith(singleton_prefix) and
     79                                 not l in labels]
     80             if labels_to_delete:
     81                 logging.warning('Removing label %s', labels_to_delete)
     82                 afe_labels_to_delete = afe.get_labels(name__in=labels_to_delete)
     83                 for afe_label in afe_labels_to_delete:
     84                     afe_label.remove_hosts(hosts=[hostname])
     85         label.add_hosts(hosts=[hostname])
     86 
     87     missing_labels = set(labels) - set([l.name for l in label_matches])
     88 
     89     if missing_labels:
     90         for label in missing_labels:
     91             logging.warning('Unable to add label %s to host %s. '
     92                             'Skipping unknown label.', label, hostname)
     93         return False
     94 
     95     return True
     96 
     97 
     98 def main():
     99     """"
    100     Entry point for add_detected_host_labels script.
    101     """
    102 
    103     parser = argparse.ArgumentParser()
    104     parser.add_argument('-s', '--silent', dest='silent', action='store_true',
    105                         help='Suppress all but critical logging messages.')
    106     parser.add_argument('-i', '--info', dest='info_only', action='store_true',
    107                         help='Suppress logging messages below INFO priority.')
    108     parser.add_argument('-m', '--machines', dest='machines',
    109                         help='Comma separated list of machines to check.')
    110     options = parser.parse_args()
    111 
    112     if options.silent and options.info_only:
    113         print 'The -i and -s flags cannot be used together.'
    114         parser.print_help()
    115         return 0
    116 
    117 
    118     if options.silent:
    119         logging.disable(logging.CRITICAL)
    120 
    121     if options.info_only:
    122         logging.disable(logging.DEBUG)
    123 
    124     threadpool = pool.ThreadPool()
    125     afe = frontend.AFE()
    126 
    127     if options.machines:
    128         hostnames = [m.strip() for m in options.machines.split(',')]
    129     else:
    130         hostnames = afe.get_hostnames()
    131     successes = sum(threadpool.imap_unordered(
    132                         lambda x: add_missing_labels(afe, x),
    133                         hostnames))
    134     attempts = len(hostnames)
    135 
    136     logging.info('Label updating finished. Failed update on %d out of %d '
    137                  'hosts.', attempts-successes, attempts)
    138 
    139     return 0
    140 
    141 
    142 if __name__ == '__main__':
    143     sys.exit(main())
    144