• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2# Copyright 2015 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Install an initial test image on a set of DUTs.
7
8The methods in this module are meant for two nominally distinct use
9cases that share a great deal of code internally.  The first use
10case is for deployment of DUTs that have just been placed in the lab
11for the first time.  The second use case is for use after repairing
12a servo.
13
14Newly deployed DUTs may be in a somewhat anomalous state:
15  * The DUTs are running a production base image, not a test image.
16    By extension, the DUTs aren't reachable over SSH.
17  * The DUTs are not necessarily in the AFE database.  DUTs that
18    _are_ in the database should be locked.  Either way, the DUTs
19    cannot be scheduled to run tests.
20  * The servos for the DUTs need not be configured with the proper
21    board.
22
23More broadly, it's not expected that the DUT will be working at the
24start of this operation.  If the DUT isn't working at the end of the
25operation, an error will be reported.
26
27The script performs the following functions:
28  * Configure the servo for the target board, and test that the
29    servo is generally in good order.
30  * For the full deployment case, install dev-signed RO firmware
31    from the designated stable test image for the DUTs.
32  * For both cases, use servo to install the stable test image from
33    USB.
34  * If the DUT isn't in the AFE database, add it.
35
36The script imposes these preconditions:
37  * Every DUT has a properly connected servo.
38  * Every DUT and servo have proper DHCP and DNS configurations.
39  * Every servo host is up and running, and accessible via SSH.
40  * There is a known, working test image that can be staged and
41    installed on the target DUTs via servo.
42  * Every DUT has the same board.
43  * For the full deployment case, every DUT must be in dev mode,
44    and configured to allow boot from USB with ctrl+U.
45
46The implementation uses the `multiprocessing` module to run all
47installations in parallel, separate processes.
48
49"""
50
51import atexit
52from collections import namedtuple
53import functools
54import json
55import logging
56import multiprocessing
57import os
58import shutil
59import sys
60import tempfile
61import time
62import traceback
63
64from chromite.lib import gs
65
66import common
67from autotest_lib.client.common_lib import error
68from autotest_lib.client.common_lib import host_states
69from autotest_lib.client.common_lib import time_utils
70from autotest_lib.client.common_lib import utils
71from autotest_lib.server import constants
72from autotest_lib.server import frontend
73from autotest_lib.server import hosts
74from autotest_lib.server.cros.dynamic_suite.constants import VERSION_PREFIX
75from autotest_lib.server.hosts import afe_store
76from autotest_lib.server.hosts import servo_host
77from autotest_lib.site_utils.deployment import commandline
78from autotest_lib.site_utils.stable_images import assign_stable_images
79
80
81_LOG_FORMAT = '%(asctime)s | %(levelname)-10s | %(message)s'
82
83_DEFAULT_POOL = constants.Labels.POOL_PREFIX + 'suites'
84
85_DIVIDER = '\n============\n'
86
87_LOG_BUCKET_NAME = 'chromeos-install-logs'
88
89_OMAHA_STATUS = 'gs://chromeos-build-release-console/omaha_status.json'
90
91# Lock reasons we'll pass when locking DUTs, depending on the
92# host's prior state.
93_LOCK_REASON_EXISTING = 'Repairing or deploying an existing host'
94_LOCK_REASON_NEW_HOST = 'Repairing or deploying a new host'
95
96_ReportResult = namedtuple('_ReportResult', ['hostname', 'message'])
97
98
99class _NoAFEServoPortError(Exception):
100    """Exception when there is no servo port stored in the AFE."""
101
102
103class _MultiFileWriter(object):
104
105    """Group file objects for writing at once."""
106
107    def __init__(self, files):
108        """Initialize _MultiFileWriter.
109
110        @param files  Iterable of file objects for writing.
111        """
112        self._files = files
113
114    def write(self, s):
115        """Write a string to the files.
116
117        @param s  Write this string.
118        """
119        for file in self._files:
120            file.write(s)
121
122
123def _get_upload_log_path(arguments):
124    return 'gs://{bucket}/{name}'.format(
125        bucket=_LOG_BUCKET_NAME,
126        name=commandline.get_default_logdir_name(arguments))
127
128
129def _upload_logs(dirpath, gspath):
130    """Upload report logs to Google Storage.
131
132    @param dirpath  Path to directory containing the logs.
133    @param gspath   Path to GS bucket.
134    """
135    ctx = gs.GSContext()
136    ctx.Copy(dirpath, gspath, recursive=True)
137
138
139def _get_omaha_build(board):
140    """Get the currently preferred Beta channel build for `board`.
141
142    Open and read through the JSON file provided by GoldenEye that
143    describes what version Omaha is currently serving for all boards
144    on all channels.  Find the entry for `board` on the Beta channel,
145    and return that version string.
146
147    @param board  The board to look up from GoldenEye.
148
149    @return Returns a Chrome OS version string in standard form
150            R##-####.#.#.  Will return `None` if no Beta channel
151            entry is found.
152    """
153    ctx = gs.GSContext()
154    omaha_status = json.loads(ctx.Cat(_OMAHA_STATUS))
155    omaha_board = board.replace('_', '-')
156    for e in omaha_status['omaha_data']:
157        if (e['channel'] == 'beta' and
158                e['board']['public_codename'] == omaha_board):
159            milestone = e['chrome_version'].split('.')[0]
160            build = e['chrome_os_version']
161            return 'R%s-%s' % (milestone, build)
162    return None
163
164
165def _update_build(afe, report_log, arguments):
166    """Update the stable_test_versions table.
167
168    This calls the `set_stable_version` RPC call to set the stable
169    repair version selected by this run of the command.  Additionally,
170    this updates the stable firmware for the board.  The repair version
171    is selected from three possible versions:
172      * The stable test version currently in the AFE database.
173      * The version Omaha is currently serving as the Beta channel
174        build.
175      * The version supplied by the user.
176    The actual version selected will be whichever of these three is
177    the most up-to-date version.
178
179    The stable firmware version will be set to whatever firmware is
180    bundled in the selected repair image. If the selected repair image bundles
181    firmware for more than one model, then the firmware for every model in the
182    build will be updated.
183
184    This function will log information about the available versions
185    prior to selection.  After selection the repair and firmware
186    versions slected will be logged.
187
188    @param afe          AFE object for RPC calls.
189    @param report_log   File-like object for logging report output.
190    @param arguments    Command line arguments with options.
191
192    @return Returns the version selected.
193    """
194    # Gather the current AFE and Omaha version settings, and report them
195    # to the user.
196    cros_version_map = afe.get_stable_version_map(afe.CROS_IMAGE_TYPE)
197    fw_version_map = afe.get_stable_version_map(afe.FIRMWARE_IMAGE_TYPE)
198    afe_cros = cros_version_map.get_version(arguments.board)
199    afe_fw = fw_version_map.get_version(arguments.board)
200    omaha_cros = _get_omaha_build(arguments.board)
201    report_log.write('AFE    version is %s.\n' % afe_cros)
202    report_log.write('Omaha  version is %s.\n' % omaha_cros)
203    report_log.write('AFE   firmware is %s.\n' % afe_fw)
204    cros_version = afe_cros
205
206    # Check whether we should upgrade the repair build to either
207    # the Omaha or the user's requested build.  If we do, we must
208    # also update the firmware version.
209    if (omaha_cros is not None
210            and (cros_version is None or
211                 utils.compare_versions(cros_version, omaha_cros) < 0)):
212        cros_version = omaha_cros
213    if arguments.build and arguments.build != cros_version:
214        if (cros_version is None
215                or utils.compare_versions(cros_version, arguments.build) < 0):
216            cros_version = arguments.build
217        else:
218            report_log.write('Selected version %s is too old; '
219                             'using version %s'
220                             % (arguments.build, cros_version))
221
222    afe_fw_versions = {arguments.board: afe_fw}
223    fw_versions = assign_stable_images.get_firmware_versions(
224        cros_version_map, arguments.board, cros_version)
225    # At this point `cros_version` is our new repair build, and
226    # `fw_version` is our new target firmware.  Call the AFE back with
227    # updates as necessary.
228    if not arguments.nostable:
229        if cros_version != afe_cros:
230            cros_version_map.set_version(arguments.board, cros_version)
231
232            if fw_versions != afe_fw_versions:
233                for model, fw_version in fw_versions.iteritems():
234                    if fw_version is not None:
235                        fw_version_map.set_version(model, fw_version)
236                    else:
237                        fw_version_map.delete_version(model)
238
239    # Report the new state of the world.
240    report_log.write(_DIVIDER)
241    report_log.write('Repair CrOS version for board %s is now %s.\n' %
242                     (arguments.board, cros_version))
243    for model, fw_version in fw_versions.iteritems():
244        report_log.write('Firmware version for model %s is now %s.\n' %
245                         (model, fw_version))
246    return cros_version
247
248
249def _create_host(hostname, afe, afe_host):
250    """Create a CrosHost object for a DUT to be installed.
251
252    @param hostname  Hostname of the target DUT.
253    @param afe       A frontend.AFE object.
254    @param afe_host  AFE Host object for the DUT.
255    """
256    machine_dict = {
257            'hostname': hostname,
258            'afe_host': afe_host,
259            'host_info_store': afe_store.AfeStore(hostname, afe),
260    }
261    servo_args = hosts.CrosHost.get_servo_arguments({})
262    return hosts.create_host(machine_dict, servo_args=servo_args)
263
264
265def _try_lock_host(afe_host):
266    """Lock a host in the AFE, and report whether it succeeded.
267
268    The lock action is logged regardless of success; failures are
269    logged if they occur.
270
271    @param afe_host AFE Host instance to be locked.
272
273    @return `True` on success, or `False` on failure.
274    """
275    try:
276        logging.warning('Locking host now.')
277        afe_host.modify(locked=True,
278                        lock_reason=_LOCK_REASON_EXISTING)
279    except Exception as e:
280        logging.exception('Failed to lock: %s', e)
281        return False
282    return True
283
284
285def _try_unlock_host(afe_host):
286    """Unlock a host in the AFE, and report whether it succeeded.
287
288    The unlock action is logged regardless of success; failures are
289    logged if they occur.
290
291    @param afe_host AFE Host instance to be unlocked.
292
293    @return `True` on success, or `False` on failure.
294    """
295    try:
296        logging.warning('Unlocking host.')
297        afe_host.modify(locked=False, lock_reason='')
298    except Exception as e:
299        logging.exception('Failed to unlock: %s', e)
300        return False
301    return True
302
303
304def _update_host_attributes(afe, hostname, host_attrs):
305    """Update the attributes for a given host.
306
307    @param afe          AFE object for RPC calls.
308    @param hostname     Host name of the DUT.
309    @param host_attrs   Dictionary with attributes to be applied to the
310                        host.
311    """
312    # Grab the servo hostname/port/serial from `host_attrs` if supplied.
313    # For new servo V4 deployments, we require the user to supply the
314    # attributes (because there are no appropriate defaults).  So, if
315    # none are supplied, we assume it can't be V4, and apply the
316    # defaults for servo V3.
317    host_attr_servo_host = host_attrs.get(servo_host.SERVO_HOST_ATTR)
318    host_attr_servo_port = host_attrs.get(servo_host.SERVO_PORT_ATTR)
319    host_attr_servo_serial = host_attrs.get(servo_host.SERVO_SERIAL_ATTR)
320    servo_hostname = (host_attr_servo_host or
321                      servo_host.make_servo_hostname(hostname))
322    servo_port = (host_attr_servo_port or
323                  str(servo_host.ServoHost.DEFAULT_PORT))
324    afe.set_host_attribute(servo_host.SERVO_HOST_ATTR,
325                           servo_hostname,
326                           hostname=hostname)
327    afe.set_host_attribute(servo_host.SERVO_PORT_ATTR,
328                           servo_port,
329                           hostname=hostname)
330    if host_attr_servo_serial:
331        afe.set_host_attribute(servo_host.SERVO_SERIAL_ATTR,
332                               host_attr_servo_serial,
333                               hostname=hostname)
334
335
336def _get_afe_host(afe, hostname, host_attrs, arguments):
337    """Get an AFE Host object for the given host.
338
339    If the host is found in the database, return the object
340    from the RPC call with the updated attributes in host_attr_dict.
341
342    If no host is found, create one with appropriate servo
343    attributes and the given board label.
344
345    @param afe          AFE object for RPC calls.
346    @param hostname     Host name of the DUT.
347    @param host_attrs   Dictionary with attributes to be applied to the
348                        host.
349    @param arguments    Command line arguments with options.
350
351    @return A tuple of the afe_host, plus a flag. The flag indicates
352            whether the Host should be unlocked if subsequent operations
353            fail.  (Hosts are always unlocked after success).
354    """
355    hostlist = afe.get_hosts([hostname])
356    unlock_on_failure = False
357    if hostlist:
358        afe_host = hostlist[0]
359        if not afe_host.locked:
360            if _try_lock_host(afe_host):
361                unlock_on_failure = True
362            else:
363                raise Exception('Failed to lock host')
364        if afe_host.status not in host_states.IDLE_STATES:
365            if unlock_on_failure and not _try_unlock_host(afe_host):
366                raise Exception('Host is in use, and failed to unlock it')
367            raise Exception('Host is in use by Autotest')
368        # This host was pre-existing; if the user didn't supply
369        # attributes, don't update them, because the defaults may
370        # not be correct.
371        if host_attrs:
372            _update_host_attributes(afe, hostname, host_attrs)
373    else:
374        afe_host = afe.create_host(hostname,
375                                   locked=True,
376                                   lock_reason=_LOCK_REASON_NEW_HOST)
377        afe_host.add_labels([constants.Labels.BOARD_PREFIX + arguments.board])
378        _update_host_attributes(afe, hostname, host_attrs)
379    afe_host = afe.get_hosts([hostname])[0]
380    return afe_host, unlock_on_failure
381
382
383def _install_firmware(host):
384    """Install dev-signed firmware after removing write-protect.
385
386    At start, it's assumed that hardware write-protect is disabled,
387    the DUT is in dev mode, and the servo's USB stick already has a
388    test image installed.
389
390    The firmware is installed by powering on and typing ctrl+U on
391    the keyboard in order to boot the the test image from USB.  Once
392    the DUT is booted, we run a series of commands to install the
393    read-only firmware from the test image.  Then we clear debug
394    mode, and shut down.
395
396    @param host   Host instance to use for servo and ssh operations.
397    """
398    servo = host.servo
399    # First power on.  We sleep to allow the firmware plenty of time
400    # to display the dev-mode screen; some boards take their time to
401    # be ready for the ctrl+U after power on.
402    servo.get_power_state_controller().power_off()
403    servo.switch_usbkey('dut')
404    servo.get_power_state_controller().power_on()
405    time.sleep(10)
406    # Dev mode screen should be up now:  type ctrl+U and wait for
407    # boot from USB to finish.
408    servo.ctrl_u()
409    if not host.wait_up(timeout=host.USB_BOOT_TIMEOUT):
410        raise Exception('DUT failed to boot in dev mode for '
411                        'firmware update')
412    # Disable software-controlled write-protect for both FPROMs, and
413    # install the RO firmware.
414    for fprom in ['host', 'ec']:
415        host.run('flashrom -p %s --wp-disable' % fprom,
416                 ignore_status=True)
417    host.run('chromeos-firmwareupdate --mode=factory')
418    # Get us out of dev-mode and clear GBB flags.  GBB flags are
419    # non-zero because boot from USB was enabled.
420    host.run('/usr/share/vboot/bin/set_gbb_flags.sh 0',
421             ignore_status=True)
422    host.run('crossystem disable_dev_request=1',
423             ignore_status=True)
424    host.halt()
425
426
427def _install_test_image(host, arguments):
428    """Install a test image to the DUT.
429
430    Install a stable test image on the DUT using the full servo
431    repair flow.
432
433    @param host       Host instance for the DUT being installed.
434    @param arguments  Command line arguments with options.
435    """
436    # Don't timeout probing for the host usb device, there could be a bunch
437    # of servos probing at the same time on the same servo host.  And
438    # since we can't pass None through the xml rpcs, use 0 to indicate None.
439    if not host.servo.probe_host_usb_dev(timeout=0):
440        raise Exception('No USB stick detected on Servo host')
441    try:
442        if not arguments.noinstall:
443            if not arguments.nostage:
444                host.servo.image_to_servo_usb(
445                        host.stage_image_for_servo())
446            if arguments.full_deploy:
447                _install_firmware(host)
448            host.servo_install()
449    except error.AutoservRunError as e:
450        logging.exception('Failed to install: %s', e)
451        raise Exception('chromeos-install failed')
452    finally:
453        host.close()
454
455
456def _install_and_update_afe(afe, hostname, host_attrs, arguments):
457    """Perform all installation and AFE updates.
458
459    First, lock the host if it exists and is unlocked.  Then,
460    install the test image on the DUT.  At the end, unlock the
461    DUT, unless the installation failed and the DUT was locked
462    before we started.
463
464    If installation succeeds, make sure the DUT is in the AFE,
465    and make sure that it has basic labels.
466
467    @param afe          AFE object for RPC calls.
468    @param hostname     Host name of the DUT.
469    @param host_attrs   Dictionary with attributes to be applied to the
470                        host.
471    @param arguments    Command line arguments with options.
472    """
473    afe_host, unlock_on_failure = _get_afe_host(afe, hostname, host_attrs,
474                                                arguments)
475    try:
476        host = _create_host(hostname, afe, afe_host)
477        _install_test_image(host, arguments)
478        host.labels.update_labels(host)
479        platform_labels = afe.get_labels(
480                host__hostname=hostname, platform=True)
481        if not platform_labels:
482            platform = host.get_platform()
483            new_labels = afe.get_labels(name=platform)
484            if not new_labels:
485                afe.create_label(platform, platform=True)
486            afe_host.add_labels([platform])
487        version = [label for label in afe_host.labels
488                       if label.startswith(VERSION_PREFIX)]
489        if version:
490            afe_host.remove_labels(version)
491    except Exception as e:
492        if unlock_on_failure and not _try_unlock_host(afe_host):
493            logging.error('Failed to unlock host!')
494        raise
495
496    if not _try_unlock_host(afe_host):
497        raise Exception('Install succeeded, but failed to unlock the DUT.')
498
499
500def _install_dut(arguments, host_attr_dict, hostname):
501    """Deploy or repair a single DUT.
502
503    @param arguments       Command line arguments with options.
504    @param host_attr_dict  Dict mapping hostnames to attributes to be
505                           stored in the AFE.
506    @param hostname        Host name of the DUT to install on.
507
508    @return On success, return `None`.  On failure, return a string
509            with an error message.
510    """
511    # In some cases, autotest code that we call during install may
512    # put stuff onto stdout with 'print' statements.  Most notably,
513    # the AFE frontend may print 'FAILED RPC CALL' (boo, hiss).  We
514    # want nothing from this subprocess going to the output we
515    # inherited from our parent, so redirect stdout and stderr, before
516    # we make any AFE calls.  Note that this is reasonable because we're
517    # in a subprocess.
518
519    logpath = os.path.join(arguments.logdir, hostname + '.log')
520    logfile = open(logpath, 'w')
521    sys.stderr = sys.stdout = logfile
522    _configure_logging_to_file(logfile)
523
524    afe = frontend.AFE(server=arguments.web)
525    try:
526        _install_and_update_afe(afe, hostname,
527                                host_attr_dict.get(hostname, {}),
528                                arguments)
529    except Exception as e:
530        logging.exception('Original exception: %s', e)
531        return str(e)
532    return None
533
534
535def _report_hosts(report_log, heading, host_results_list):
536    """Report results for a list of hosts.
537
538    To improve visibility, results are preceded by a header line,
539    followed by a divider line.  Then results are printed, one host
540    per line.
541
542    @param report_log         File-like object for logging report
543                              output.
544    @param heading            The header string to be printed before
545                              results.
546    @param host_results_list  A list of _ReportResult tuples
547                              to be printed one per line.
548    """
549    if not host_results_list:
550        return
551    report_log.write(heading)
552    report_log.write(_DIVIDER)
553    for result in host_results_list:
554        report_log.write('{result.hostname:30} {result.message}\n'
555                         .format(result=result))
556    report_log.write('\n')
557
558
559def _report_results(afe, report_log, hostnames, results):
560    """Gather and report a summary of results from installation.
561
562    Segregate results into successes and failures, reporting
563    each separately.  At the end, report the total of successes
564    and failures.
565
566    @param afe          AFE object for RPC calls.
567    @param report_log   File-like object for logging report output.
568    @param hostnames    List of the hostnames that were tested.
569    @param results      List of error messages, in the same order
570                        as the hostnames.  `None` means the
571                        corresponding host succeeded.
572    """
573    successful_hosts = []
574    success_reports = []
575    failure_reports = []
576    for result, hostname in zip(results, hostnames):
577        if result is None:
578            successful_hosts.append(hostname)
579        else:
580            failure_reports.append(_ReportResult(hostname, result))
581    if successful_hosts:
582        afe.reverify_hosts(hostnames=successful_hosts)
583        for h in afe.get_hosts(hostnames=successful_hosts):
584            for label in h.labels:
585                if label.startswith(constants.Labels.POOL_PREFIX):
586                    result = _ReportResult(h.hostname,
587                                           'Host already in %s' % label)
588                    success_reports.append(result)
589                    break
590            else:
591                h.add_labels([_DEFAULT_POOL])
592                result = _ReportResult(h.hostname,
593                                       'Host added to %s' % _DEFAULT_POOL)
594                success_reports.append(result)
595    report_log.write(_DIVIDER)
596    _report_hosts(report_log, 'Successes', success_reports)
597    _report_hosts(report_log, 'Failures', failure_reports)
598    report_log.write(
599        'Installation complete:  %d successes, %d failures.\n' %
600        (len(success_reports), len(failure_reports)))
601
602
603def _clear_root_logger_handlers():
604    """Remove all handlers from root logger."""
605    root_logger = logging.getLogger()
606    for h in root_logger.handlers:
607        root_logger.removeHandler(h)
608
609
610def _configure_logging_to_file(logfile):
611    """Configure the logging module for `install_duts()`.
612
613    @param log_file  Log file object.
614    """
615    _clear_root_logger_handlers()
616    handler = logging.StreamHandler(logfile)
617    formatter = logging.Formatter(_LOG_FORMAT, time_utils.TIME_FMT)
618    handler.setFormatter(formatter)
619    root_logger = logging.getLogger()
620    root_logger.addHandler(handler)
621
622
623def _get_used_servo_ports(servo_hostname, afe):
624    """
625    Return a list of used servo ports for the given servo host.
626
627    @param servo_hostname:  Hostname of the servo host to check for.
628    @param afe:             AFE instance.
629
630    @returns a list of used ports for the given servo host.
631    """
632    used_ports = []
633    host_list = afe.get_hosts_by_attribute(
634            attribute=servo_host.SERVO_HOST_ATTR, value=servo_hostname)
635    for host in host_list:
636        afe_host = afe.get_hosts(hostname=host)
637        if afe_host:
638            servo_port = afe_host[0].attributes.get(servo_host.SERVO_PORT_ATTR)
639            if servo_port:
640                used_ports.append(int(servo_port))
641    return used_ports
642
643
644def _get_free_servo_port(servo_hostname, used_servo_ports, afe):
645    """
646    Get a free servo port for the servo_host.
647
648    @param servo_hostname:    Hostname of the servo host.
649    @param used_servo_ports:  Dict of dicts that contain the list of used ports
650                              for the given servo host.
651    @param afe:               AFE instance.
652
653    @returns a free servo port if servo_hostname is non-empty, otherwise an
654        empty string.
655    """
656    used_ports = []
657    servo_port = servo_host.ServoHost.DEFAULT_PORT
658    # If no servo hostname was specified we can assume we're dealing with a
659    # servo v3 or older deployment since the servo hostname can be
660    # inferred from the dut hostname (by appending '-servo' to it).  We only
661    # need to find a free port if we're using a servo v4 since we can use the
662    # default port for v3 and older.
663    if not servo_hostname:
664        return ''
665    # If we haven't checked this servo host yet, check the AFE if other duts
666    # used this servo host and grab the ports specified for them.
667    elif servo_hostname not in used_servo_ports:
668        used_ports = _get_used_servo_ports(servo_hostname, afe)
669    else:
670        used_ports = used_servo_ports[servo_hostname]
671    used_ports.sort()
672    if used_ports:
673        # Range is taken from servod.py in hdctools.
674        start_port = servo_host.ServoHost.DEFAULT_PORT
675        end_port = start_port - 99
676        # We'll choose first port available in descending order.
677        for port in xrange(start_port, end_port - 1, -1):
678            if port not in used_ports:
679              servo_port = port
680              break
681    used_ports.append(servo_port)
682    used_servo_ports[servo_hostname] = used_ports
683    return servo_port
684
685
686def _get_afe_servo_port(host_info, afe):
687    """
688    Get the servo port from the afe if it matches the same servo host hostname.
689
690    @param host_info   HostInfo tuple (hostname, host_attr_dict).
691
692    @returns Servo port (int) if servo host hostname matches the one specified
693    host_info.host_attr_dict, otherwise None.
694
695    @raises _NoAFEServoPortError: When there is no stored host info or servo
696        port host attribute in the AFE for the given host.
697    """
698    afe_hosts = afe.get_hosts(hostname=host_info.hostname)
699    if not afe_hosts:
700        raise _NoAFEServoPortError
701
702    servo_port = afe_hosts[0].attributes.get(servo_host.SERVO_PORT_ATTR)
703    afe_servo_host = afe_hosts[0].attributes.get(servo_host.SERVO_HOST_ATTR)
704    host_info_servo_host = host_info.host_attr_dict.get(
705        servo_host.SERVO_HOST_ATTR)
706
707    if afe_servo_host == host_info_servo_host and servo_port:
708        return int(servo_port)
709    else:
710        raise _NoAFEServoPortError
711
712
713def _get_host_attributes(host_info_list, afe):
714    """
715    Get host attributes if a hostname_file was supplied.
716
717    @param host_info_list   List of HostInfo tuples (hostname, host_attr_dict).
718
719    @returns Dict of attributes from host_info_list.
720    """
721    host_attributes = {}
722    # We need to choose servo ports for these hosts but we need to make sure
723    # we don't choose ports already used. We'll store all used ports in a
724    # dict of lists where the key is the servo_host and the val is a list of
725    # ports used.
726    used_servo_ports = {}
727    for host_info in host_info_list:
728        host_attr_dict = host_info.host_attr_dict
729        # If the host already has an entry in the AFE that matches the same
730        # servo host hostname and the servo port is set, use that port.
731        try:
732            host_attr_dict[servo_host.SERVO_PORT_ATTR] = _get_afe_servo_port(
733                host_info, afe)
734        except _NoAFEServoPortError:
735            host_attr_dict[servo_host.SERVO_PORT_ATTR] = _get_free_servo_port(
736                host_attr_dict[servo_host.SERVO_HOST_ATTR], used_servo_ports,
737                afe)
738        host_attributes[host_info.hostname] = host_attr_dict
739    return host_attributes
740
741
742def install_duts(argv, full_deploy):
743    """Install a test image on DUTs, and deploy them.
744
745    This handles command line parsing for both the repair and
746    deployment commands.  The two operations are largely identical;
747    the main difference is that full deployment includes flashing
748    dev-signed firmware on the DUT prior to installing the test
749    image.
750
751    @param argv         Command line arguments to be parsed.
752    @param full_deploy  If true, do the full deployment that includes
753                        flashing dev-signed RO firmware onto the DUT.
754    """
755    # Override tempfile.tempdir.  Some of the autotest code we call
756    # will create temporary files that don't get cleaned up.  So, we
757    # put the temp files in our results directory, so that we can
758    # clean up everything in one fell swoop.
759    tempfile.tempdir = tempfile.mkdtemp()
760    # MALCOLM:
761    #   Be comforted.
762    #   Let's make us med'cines of our great revenge,
763    #   To cure this deadly grief.
764    atexit.register(shutil.rmtree, tempfile.tempdir)
765
766    arguments = commandline.parse_command(argv, full_deploy)
767    if not arguments:
768        sys.exit(1)
769    sys.stderr.write('Installation output logs in %s\n' % arguments.logdir)
770
771    # We don't want to distract the user with logging output, so we catch
772    # logging output in a file.
773    logging_file_path = os.path.join(arguments.logdir, 'debug.log')
774    logfile = open(logging_file_path, 'w')
775    _configure_logging_to_file(logfile)
776
777    report_log_path = os.path.join(arguments.logdir, 'report.log')
778    with open(report_log_path, 'w') as report_log_file:
779        report_log = _MultiFileWriter([report_log_file, sys.stdout])
780        afe = frontend.AFE(server=arguments.web)
781        current_build = _update_build(afe, report_log, arguments)
782        host_attr_dict = _get_host_attributes(arguments.host_info_list, afe)
783        install_pool = multiprocessing.Pool(len(arguments.hostnames))
784        install_function = functools.partial(_install_dut, arguments,
785                                             host_attr_dict)
786        results_list = install_pool.map(install_function, arguments.hostnames)
787        _report_results(afe, report_log, arguments.hostnames, results_list)
788
789        gspath = _get_upload_log_path(arguments)
790        report_log.write('Logs will be uploaded to %s\n' % (gspath,))
791
792    try:
793        _upload_logs(arguments.logdir, gspath)
794    except Exception as e:
795        upload_failure_log_path = os.path.join(arguments.logdir,
796                                               'gs_upload_failure.log')
797        with open(upload_failure_log_path, 'w') as file:
798            traceback.print_exc(limit=None, file=file)
799        sys.stderr.write('Failed to upload logs;'
800                         ' failure details are stored in {}.\n'
801                         .format(upload_failure_log_path))
802