Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 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 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 22 | from multiprocessing import pool |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 23 | import logging |
| 24 | import socket |
| 25 | import argparse |
| 26 | import sys |
| 27 | |
| 28 | import common |
| 29 | |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 30 | from autotest_lib.server import hosts |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 31 | from autotest_lib.server import frontend |
| 32 | from autotest_lib.client.common_lib import error |
| 33 | |
| 34 | |
Dan Shi | 1f753ba | 2015-09-30 15:28:38 -0700 | [diff] [blame] | 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 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 40 | def add_missing_labels(afe, hostname): |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 41 | """ |
| 42 | Queries the detectable labels supported by the given host, |
| 43 | and adds those labels to the host. |
| 44 | |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 45 | @param afe: A frontend.AFE() instance. |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 46 | @param hostname: The host to query and update. |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 47 | |
| 48 | @return: True on success. |
| 49 | False on failure to fetch labels or to add any individual label. |
| 50 | """ |
Prashanth Balasubramanian | 1e5b345 | 2014-10-16 09:57:39 -0700 | [diff] [blame] | 51 | host = None |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 52 | try: |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 53 | host = hosts.create_host(hostname) |
| 54 | labels = host.get_labels() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 55 | except socket.gaierror: |
| 56 | logging.warning('Unable to establish ssh connection to hostname ' |
| 57 | '%s. Skipping.', hostname) |
| 58 | return False |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 59 | except error.AutoservError: |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 60 | logging.warning('Unable to query labels on hostname %s. Skipping.', |
| 61 | hostname) |
| 62 | return False |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 63 | finally: |
Prashanth Balasubramanian | 1e5b345 | 2014-10-16 09:57:39 -0700 | [diff] [blame] | 64 | if host: |
| 65 | host.close() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 66 | |
Dan Shi | 1f753ba | 2015-09-30 15:28:38 -0700 | [diff] [blame] | 67 | afe_host = afe.get_hosts(hostname=hostname)[0] |
| 68 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 69 | label_matches = afe.get_labels(name__in=labels) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 70 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 71 | for label in label_matches: |
Dan Shi | 1f753ba | 2015-09-30 15:28:38 -0700 | [diff] [blame] | 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]) |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 85 | label.add_hosts(hosts=[hostname]) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 86 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 87 | missing_labels = set(labels) - set([l.name for l in label_matches]) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 88 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 89 | if missing_labels: |
| 90 | for label in missing_labels: |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 91 | logging.warning('Unable to add label %s to host %s. ' |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 92 | 'Skipping unknown label.', label, hostname) |
| 93 | return False |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 94 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 95 | return True |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 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', |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 105 | help='Suppress all but critical logging messages.') |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 106 | parser.add_argument('-i', '--info', dest='info_only', action='store_true', |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 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.') |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 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 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 124 | threadpool = pool.ThreadPool() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 125 | afe = frontend.AFE() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 126 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 127 | if options.machines: |
| 128 | hostnames = [m.strip() for m in options.machines.split(',')] |
| 129 | else: |
| 130 | hostnames = afe.get_hostnames() |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 131 | successes = sum(threadpool.imap_unordered( |
| 132 | lambda x: add_missing_labels(afe, x), |
| 133 | hostnames)) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 134 | attempts = len(hostnames) |
| 135 | |
| 136 | logging.info('Label updating finished. Failed update on %d out of %d ' |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 137 | 'hosts.', attempts-successes, attempts) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 138 | |
| 139 | return 0 |
| 140 | |
| 141 | |
| 142 | if __name__ == '__main__': |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 143 | sys.exit(main()) |