blob: f4e0b22b95165fe50f9f8bd3d472b7d8fcaaa365 [file] [log] [blame] [edit]
#!/usr/bin/python
#
# Copyright (c) 2011 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.
"""Updates all unlocked hosts in Autotest lab in parallel at a given rate.
Used to update all hosts, or only those of a given platform, in the Autotest
lab to a given version. Allows a configurable number of updates to be started in
parallel. Updates can also be staggered to reduce load."""
import logging
import os
import subprocess
import sys
import threading
import time
import traceback
from collections import deque
from optparse import OptionParser
# Default number of hosts to update in parallel.
DEFAULT_CONCURRENCY = 10
# By default do not stagger any of the updates.
DEFAULT_STAGGER = 0
# Default location of ChromeOS checkout.
DEFAULT_GCLIENT_ROOT = '/usr/local/google/home/${USER}/chromeos/chromeos'
# Default path for individual host logs. Each host will have it's own file. E.g.
# <default_log_path>/<host>.log
DEFAULT_LOG_PATH = '/tmp/mass_update_logs/%s/' % time.strftime('%Y-%m-%d-%H-%M',
time.gmtime())
# Location of Autotest cli executable.
AUTOTEST_LOCATION = '/home/chromeos-test/autotest/cli'
# Default time in seconds to sleep while waiting for threads to complete.
DEFAULT_SLEEP = 10
# Amount of time in seconds to wait before declaring an update as failed.
DEFAULT_TIMEOUT = 2400
class MassUpdateStatus():
"""Used to track status for all updates."""
ssh_failures = []
update_failures = []
successful_updates = 0
class UpdateThread(threading.Thread):
"""Responsible for ssh-test, locking, imaging, and unlocking a host.
Uses the atest CLI for host control and the image_to_live script to actually
update the host. Each thread will continue to process hosts until the queue
is empty."""
_SUCCESS = 0 # Update was successful.
_SSH_FAILURE = 1 # Could not SSH to host or related SSH failure.
_UPDATE_FAILURE = 2 # Update failed for any reason other than SSH.
def __init__(self, options, hosts, status):
self._options = options
self._hosts = hosts
self._status = status
self._logger = logging.getLogger()
threading.Thread.__init__(self)
def run(self):
while self._hosts:
host = self._hosts.popleft()
status = UpdateThread._UPDATE_FAILURE
self._logger.info('Updating host %s' % host)
try:
try:
if not CheckSSH(host=host, options=self._options):
status = UpdateThread._SSH_FAILURE
elif LockHost(host) and ImageHost(host=host, options=self._options):
status = UpdateThread._SUCCESS
finally:
if status == UpdateThread._SUCCESS:
self._logger.info(
'Completed update for host %s successfully.' % host)
self._status.successful_updates += 1
elif status == UpdateThread._SSH_FAILURE:
self._logger.info('Failed to SSH to host %s.' % host)
self._status.ssh_failures.append(host)
else:
self._logger.info('Failed to update host %s.' % host)
self._status.update_failures.append(host)
UnlockHost(host)
except:
traceback.print_exc()
self._logger.warning(
'Exception encountered during update. Skipping host %s.' % host)
def CheckSSH(host, options):
"""Uses the ssh_test script to ensure SSH access to a host is available.
Returns true if an SSH connection to the host was successful."""
return subprocess.Popen(
'%s/src/scripts/ssh_test.sh --remote=%s' % (options.gclient, host),
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE).wait() == 0
def ImageHost(host, options):
"""Uses the image_to_live script to update a host.
Returns true if the imaging process was successful."""
log_file = open(os.path.join(options.log, host + '.log'), 'w')
log_file_err = open(os.path.join(options.log, host + '.log.err'), 'w')
exit_code = subprocess.Popen(
('/usr/local/scripts/alarm %d %s/src/scripts/image_to_live.sh '
'--update_url %s --remote %s' % (DEFAULT_TIMEOUT, options.gclient,
options.url, host)),
shell=True,
stdout=log_file,
stderr=log_file_err).wait()
log_file.close()
log_file_err.close()
return exit_code == 0
def LockHost(host):
"""Locks a host using the atest CLI.
Locking a host tells Autotest that the host shouldn't be scheduled for
any other tasks. Returns true if the locking process was successful."""
success = subprocess.Popen(
'%s/atest host mod -l %s' % (AUTOTEST_LOCATION, host),
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE).wait() == 0
if not success:
logging.getLogger().info('Failed to lock host %s.' % host)
return success
def UnlockHost(host):
"""Unlocks a host using the atest CLI.
Unlocking a host tells Autotest that the host is okay to be scheduled
for other tasks. Returns true if the unlocking process was successful."""
success = subprocess.Popen(
'%s/atest host mod -u %s' % (AUTOTEST_LOCATION, host),
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE).wait() == 0
if not success:
logging.getLogger().info('Failed to unlock host %s.' % host)
return success
def GetHostQueue(options):
"""Returns a queue containing unlocked hosts retrieved from the atest CLI.
If options.label has been specified only unlocked hosts with the specified
label will be returned."""
cmd = ('%s/atest host list --unlocked -s Ready -a acl_cros_test'
% AUTOTEST_LOCATION)
if options.label:
cmd += ' -b ' + options.label
# atest host list will return a tabular data set. Use sed to remove the first
# line which contains column labels we don't need. Then since the first column
# contains the host name, use awk to extract it
cmd += " | sed '1d' | awk '{print $1}'"
stdout = subprocess.Popen(cmd,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE).communicate()[0]
return deque(item.strip() for item in stdout.split('\n') if item.strip())
def ParseOptions():
usage = 'usage: %prog --url=<update url> [options]'
parser = OptionParser(usage)
parser.add_option('-b', '--label', dest='label',
help='Only update hosts with the specified label.')
parser.add_option('-c', '--concurrent', dest='concurrent',
default=DEFAULT_CONCURRENCY,
help=('Number of hosts to be updated concurrently. '
'Defaults to %d hosts.') % DEFAULT_CONCURRENCY)
parser.add_option('-g', '--gclient', dest='gclient',
default=DEFAULT_GCLIENT_ROOT,
help=('Location of ChromeOS checkout. defaults to %s'
% DEFAULT_GCLIENT_ROOT))
parser.add_option('-l', '--log', dest='log',
default=DEFAULT_LOG_PATH,
help=('Where to put individual host log files. '
'Defaults to %s' % DEFAULT_LOG_PATH))
parser.add_option('-s', '--stagger', dest='stagger',
default=DEFAULT_STAGGER,
help=('Attempt to stagger updates. Waits the given amount '
'of time in minutes before starting each updater. '
'Updates will still overlap if the value is set as a '
'multiple of the update period.'))
parser.add_option('-u', '--url', dest='url',
help='Update URL. Points to build for updating hosts.')
options = parser.parse_args()[0]
if options.url is None:
parser.error('An update URL must be provided.')
return options
def InitializeLogging():
"""Configure the global logger for time/date stamping console output.
Returns a logger object for convenience."""
logger = logging.getLogger()
logger.setLevel(logging.INFO)
stream_handler = logging.StreamHandler()
stream_handler.setLevel(logging.INFO)
stream_handler.setFormatter(logging.Formatter('%(asctime)s - %(message)s'))
logger.addHandler(stream_handler)
return logger
def main():
options = ParseOptions()
hosts = GetHostQueue(options)
logger = InitializeLogging()
status = MassUpdateStatus()
# Create log folder if it doesn't exist.
if not os.path.exists(options.log):
os.makedirs(options.log)
logger.info('Starting update using URL %s' % options.url)
logger.info('Individual host logs can be found under %s' % options.log)
try:
# Spawn processing threads which will handle lock, update, and unlock.
for i in range(int(options.concurrent)):
UpdateThread(hosts=hosts, options=options, status=status).start()
# Stagger threads if the option has been enabled.
if options.stagger > 0:
time.sleep(int(options.stagger) * 60)
# Wait for all hosts to be processed and threads to complete. NOTE: Not
# using hosts.join() here because it does not behave properly with CTRL-C
# and KeyboardInterrupt.
while len(threading.enumerate()) > 1:
time.sleep(DEFAULT_SLEEP)
except:
traceback.print_exc()
logger.warning(
'Update process aborted. Some machines may be left locked or updating.')
sys.exit(1)
finally:
logger.info(
('Mass updating complete. %d hosts updated successfully, %d failed.' %
(status.successful_updates, len(status.ssh_failures) +
len(status.update_failures))))
logger.info(('-' * 25) + '[ SUMMARY ]' + ('-' * 25))
for host in status.ssh_failures:
logger.info('Failed to SSH to host %s.' % host)
for host in status.update_failures:
logger.info('Failed to update host %s.' % host)
if len(status.ssh_failures) == 0 and len(status.update_failures) == 0:
logger.info('All hosts updated successfully.')
logger.info('-' * 61)
if __name__ == '__main__':
main()