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 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 35 | def add_missing_labels(afe, hostname): |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 36 | """ |
| 37 | Queries the detectable labels supported by the given host, |
| 38 | and adds those labels to the host. |
| 39 | |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 40 | @param afe: A frontend.AFE() instance. |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 41 | @param hostname: The host to query and update. |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 42 | |
| 43 | @return: True on success. |
| 44 | False on failure to fetch labels or to add any individual label. |
| 45 | """ |
| 46 | |
| 47 | try: |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 48 | host = hosts.create_host(hostname) |
| 49 | labels = host.get_labels() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 50 | except socket.gaierror: |
| 51 | logging.warning('Unable to establish ssh connection to hostname ' |
| 52 | '%s. Skipping.', hostname) |
| 53 | return False |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 54 | except error.AutoservError: |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 55 | logging.warning('Unable to query labels on hostname %s. Skipping.', |
| 56 | hostname) |
| 57 | return False |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 58 | finally: |
| 59 | host.close() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 60 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 61 | label_matches = afe.get_labels(name__in=labels) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 62 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 63 | for label in label_matches: |
| 64 | label.add_hosts(hosts=[hostname]) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 65 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 66 | missing_labels = set(labels) - set([l.name for l in label_matches]) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 67 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 68 | if missing_labels: |
| 69 | for label in missing_labels: |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 70 | logging.warning('Unable to add label %s to host %s. ' |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 71 | 'Skipping unknown label.', label, hostname) |
| 72 | return False |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 73 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 74 | return True |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 75 | |
| 76 | |
| 77 | def main(): |
| 78 | """" |
| 79 | Entry point for add_detected_host_labels script. |
| 80 | """ |
| 81 | |
| 82 | parser = argparse.ArgumentParser() |
| 83 | parser.add_argument('-s', '--silent', dest='silent', action='store_true', |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 84 | help='Suppress all but critical logging messages.') |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 85 | parser.add_argument('-i', '--info', dest='info_only', action='store_true', |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 86 | help='Suppress logging messages below INFO priority.') |
| 87 | parser.add_argument('-m', '--machines', dest='machines', |
| 88 | help='Comma separated list of machines to check.') |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 89 | options = parser.parse_args() |
| 90 | |
| 91 | if options.silent and options.info_only: |
| 92 | print 'The -i and -s flags cannot be used together.' |
| 93 | parser.print_help() |
| 94 | return 0 |
| 95 | |
| 96 | |
| 97 | if options.silent: |
| 98 | logging.disable(logging.CRITICAL) |
| 99 | |
| 100 | if options.info_only: |
| 101 | logging.disable(logging.DEBUG) |
| 102 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 103 | threadpool = pool.ThreadPool() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 104 | afe = frontend.AFE() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 105 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 106 | if options.machines: |
| 107 | hostnames = [m.strip() for m in options.machines.split(',')] |
| 108 | else: |
| 109 | hostnames = afe.get_hostnames() |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 110 | |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 111 | successes = sum(threadpool.imap_unordered( |
| 112 | lambda x: add_missing_labels(afe, x), |
| 113 | hostnames)) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 114 | attempts = len(hostnames) |
| 115 | |
| 116 | logging.info('Label updating finished. Failed update on %d out of %d ' |
Alex Miller | 6f84a79 | 2014-02-06 16:07:56 -0800 | [diff] [blame] | 117 | 'hosts.', attempts-successes, attempts) |
Aviv Keshet | 53bd44e | 2013-01-31 13:08:41 -0800 | [diff] [blame] | 118 | |
| 119 | return 0 |
| 120 | |
| 121 | |
| 122 | if __name__ == '__main__': |
J. Richard Barnette | ecdbd19 | 2013-10-02 17:44:08 -0700 | [diff] [blame] | 123 | sys.exit(main()) |