site_utils/add_detected_host_labels.py

#!/usr/bin/python2

# 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.

Limitations:
 - 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.
   See crosbug.com/38569
"""


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.
SINGLETON_LABEL_PREFIX = ['power:']

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
    try:
        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.',
                         hostname)
        return False
    finally:
        if host:
            host.close()

    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 label.name.startswith(p)]
        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:
                    afe_label.remove_hosts(hosts=[hostname])
        label.add_hosts(hosts=[hostname])

    missing_labels = set(labels) - set([l.name 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.'
        parser.print_help()
        return 0


    if options.silent:
        logging.disable(logging.CRITICAL)

    if options.info_only:
        logging.disable(logging.DEBUG)

    threadpool = pool.ThreadPool()
    afe = frontend.AFE()

    if options.machines:
        hostnames = [m.strip() for m in options.machines.split(',')]
    else:
        hostnames = afe.get_hostnames()
    successes = sum(threadpool.imap_unordered(
                        lambda x: add_missing_labels(afe, x),
                        hostnames))
    attempts = len(hostnames)

    logging.info('Label updating finished. Failed update on %d out of %d '
                 'hosts.', attempts-successes, attempts)

    return 0


if __name__ == '__main__':
    sys.exit(main())