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