#!/usr/bin/env python2
#
# Copyright 2015 Google INc. All Rights Reserved.
"""This module controls locking and unlocking of test machines."""
from __future__ import print_function
import argparse
import getpass
import os
import sys
import traceback
from cros_utils import logger
from cros_utils import machines
class AFELockException(Exception):
"""Base class for exceptions in this module."""
class MachineNotPingable(AFELockException):
"""Raised when machine does not respond to ping."""
class MissingHostInfo(AFELockException):
"""Raised when cannot find info about machine on machine servers."""
class UpdateNonLocalMachine(AFELockException):
"""Raised when user requests to add/remove a ChromeOS HW Lab machine.."""
class DuplicateAdd(AFELockException):
"""Raised when user requests to add a machine that's already on the server."""
class UpdateServerError(AFELockException):
"""Raised when attempt to add/remove a machine from local server fails."""
class LockingError(AFELockException):
"""Raised when server fails to lock/unlock machine as requested."""
class DontOwnLock(AFELockException):
"""Raised when user attmepts to unlock machine locked by someone else."""
# This should not be raised if the user specified '--force'
class NoAFEServer(AFELockException):
"""Raised when cannot find/access the autotest server."""
class AFEAccessError(AFELockException):
"""Raised when cannot get information about lab machine from lab server."""
class AFELockManager(object):
"""Class for locking/unlocking machines vie Autotest Front End servers.
This class contains methods for checking the locked status of machines
on both the ChromeOS HW Lab AFE server and a local AFE server. It also
has methods for adding/removing machines from the local server, and for
changing the lock status of machines on either server. For the ChromeOS
HW Lab, it only allows access to the toolchain team lab machines, as
defined in toolchain-utils/crosperf/default_remotes. By default it will
look for a local server on chrotomation2.svl.corp.google.com, but an
alternative local AFE server can be supplied, if desired.
!!!IMPORTANT NOTE!!! The AFE server can only be called from the main
thread/process of a program. If you launch threads and try to call it
from a thread, you will get an error. This has to do with restrictions
in the Python virtual machine (and signal handling) and cannot be changed.
"""
LOCAL_SERVER = 'chrotomation2.svl.corp.google.com'
def __init__(self,
remotes,
force_option,
chromeos_root,
local_server,
use_local=True,
log=None):
"""Initializes an AFELockManager object.
Args:
remotes: A list of machine names or ip addresses to be managed. Names
and ip addresses should be represented as strings. If the list is
empty, the lock manager will get all known machines.
force_option: A Boolean indicating whether or not to force an unlock of
a machine that was locked by someone else.
chromeos_root: The ChromeOS chroot to use for the autotest scripts.
local_server: A string containing the name or ip address of the machine
that is running an AFE server, which is to be used for managing
machines that are not in the ChromeOS HW lab.
local: A Boolean indicating whether or not to use/allow a local AFE
server to be used (see local_server argument).
use_local: Use the local server instead of the official one.
log: If not None, this is the logger object to be used for writing out
informational output messages. It is expected to be an instance of
Logger class from cros_utils/logger.py.
"""
self.chromeos_root = chromeos_root
self.user = getpass.getuser()
self.logger = log or logger.GetLogger()
autotest_path = os.path.join(chromeos_root,
'src/third_party/autotest/files')
sys.path.append(chromeos_root)
sys.path.append(autotest_path)
sys.path.append(os.path.join(autotest_path, 'server', 'cros'))
# We have to wait to do these imports until the paths above have
# been fixed.
# pylint: disable=import-error
from client import setup_modules
setup_modules.setup(
base_path=autotest_path, root_module_name='autotest_lib')
from dynamic_suite import frontend_wrappers
self.afe = frontend_wrappers.RetryingAFE(
timeout_min=30, delay_sec=10, debug=False, server='cautotest')
self.local = use_local
self.machines = list(set(remotes)) or []
self.toolchain_lab_machines = self.GetAllToolchainLabMachines()
if self.machines and self.AllLabMachines():
self.local = False
if not self.local:
self.local_afe = None
else:
dargs = {}
dargs['server'] = local_server or AFELockManager.LOCAL_SERVER
# Make sure local server is pingable.
error_msg = ('Local autotest server machine %s not responding to ping.' %
dargs['server'])
self.CheckMachine(dargs['server'], error_msg)
self.local_afe = frontend_wrappers.RetryingAFE(
timeout_min=30, delay_sec=10, debug=False, **dargs)
if not self.machines:
self.machines = self.toolchain_lab_machines + self.GetAllNonlabMachines()
self.force = force_option
def AllLabMachines(self):
"""Check to see if all machines being used are HW Lab machines."""
all_lab = True
for m in self.machines:
if m not in self.toolchain_lab_machines:
all_lab = False
break
return all_lab
def CheckMachine(self, machine, error_msg):
"""Verifies that machine is responding to ping.
Args:
machine: String containing the name or ip address of machine to check.
error_msg: Message to print if ping fails.
Raises:
MachineNotPingable: If machine is not responding to 'ping'
"""
if not machines.MachineIsPingable(machine, logging_level='none'):
cros_machine = machine + '.cros'
if not machines.MachineIsPingable(cros_machine, logging_level='none'):
raise MachineNotPingable(error_msg)
def MachineIsKnown(self, machine):
"""Checks to see if either AFE server knows the given machine.
Args:
machine: String containing name or ip address of machine to check.
Returns:
Boolean indicating if the machine is in the list of known machines for
either AFE server.
"""
if machine in self.toolchain_lab_machines:
return True
elif self.local_afe and machine in self.GetAllNonlabMachines():
return True
return False
def GetAllToolchainLabMachines(self):
"""Gets a list of all the toolchain machines in the ChromeOS HW lab.
Returns:
A list of names of the toolchain machines in the ChromeOS HW lab.
"""
machines_file = os.path.join(
os.path.dirname(__file__), 'crosperf', 'default_remotes')
machine_list = []
with open(machines_file, 'r') as input_file:
lines = input_file.readlines()
for line in lines:
_, remotes = line.split(':')
remotes = remotes.strip()
for r in remotes.split():
machine_list.append(r.strip())
return machine_list
def GetAllNonlabMachines(self):
"""Gets a list of all known machines on the local AFE server.
Returns:
A list of the names of the machines on the local AFE server.
"""
non_lab_machines = []
if self.local_afe:
non_lab_machines = self.local_afe.get_hostnames()
return non_lab_machines
def PrintStatusHeader(self, is_lab_machine):
"""Prints the status header lines for machines.
Args:
is_lab_machine: Boolean indicating whether to print HW Lab header or
local machine header (different spacing).
"""
if is_lab_machine:
print('\nMachine (Board)\t\t\t\t\tStatus')
print('---------------\t\t\t\t\t------\n')
else:
print('\nMachine (Board)\t\tStatus')
print('---------------\t\t------\n')
def RemoveLocalMachine(self, m):
"""Removes a machine from the local AFE server.
Args:
m: The machine to remove.
Raises:
MissingHostInfo: Can't find machine to be removed.
"""
if self.local_afe:
host_info = self.local_afe.get_hosts(hostname=m)
if host_info:
host_info = host_info[0]
host_info.delete()
else:
raise MissingHostInfo('Cannot find/delete machine %s.' % m)
def AddLocalMachine(self, m):
"""Adds a machine to the local AFE server.
Args:
m: The machine to be added.
"""
if self.local_afe:
error_msg = 'Machine %s is not responding to ping.' % m
self.CheckMachine(m, error_msg)
self.local_afe.create_host(m)
def AddMachinesToLocalServer(self):
"""Adds one or more machines to the local AFE server.
Verify that the requested machines are legal to add to the local server,
i.e. that they are not ChromeOS HW lab machines, and they are not already
on the local server. Call AddLocalMachine for each valid machine.
Raises:
DuplicateAdd: Attempt to add a machine that is already on the server.
UpdateNonLocalMachine: Attempt to add a ChromeOS HW lab machine.
UpdateServerError: Something went wrong while attempting to add a
machine.
"""
for m in self.machines:
for cros_name in [m, m + '.cros']:
if cros_name in self.toolchain_lab_machines:
raise UpdateNonLocalMachine(
'Machine %s is already in the ChromeOS HW'
'Lab. Cannot add it to local server.' % cros_name)
host_info = self.local_afe.get_hosts(hostname=m)
if host_info:
raise DuplicateAdd('Machine %s is already on the local server.' % m)
try:
self.AddLocalMachine(m)
self.logger.LogOutput('Successfully added %s to local server.' % m)
except Exception as e:
traceback.print_exc()
raise UpdateServerError(
'Error occurred while attempting to add %s. %s' % (m, str(e)))
def RemoveMachinesFromLocalServer(self):
"""Removes one or more machines from the local AFE server.
Verify that the requested machines are legal to remove from the local
server, i.e. that they are not ChromeOS HW lab machines. Call
RemoveLocalMachine for each valid machine.
Raises:
UpdateServerError: Something went wrong while attempting to remove a
machine.
"""
for m in self.machines:
for cros_name in [m, m + '.cros']:
if cros_name in self.toolchain_lab_machines:
raise UpdateNonLocalMachine(
'Machine %s is in the ChromeOS HW Lab. '
'This script cannot remove lab machines.' % cros_name)
try:
self.RemoveLocalMachine(m)
self.logger.LogOutput('Successfully removed %s from local server.' % m)
except Exception as e:
traceback.print_exc()
raise UpdateServerError('Error occurred while attempting to remove %s '
'(%s).' % (m, str(e)))
def ListMachineStates(self, machine_states):
"""Gets and prints the current status for a list of machines.
Prints out the current status for all of the machines in the current
AFELockManager's list of machines (set when the object is initialized).
Args:
machine_states: A dictionary of the current state of every machine in
the current AFELockManager's list of machines. Normally obtained by
calling AFELockManager::GetMachineStates.
"""
local_machines = []
printed_hdr = False
for m in machine_states:
cros_name = m + '.cros'
if (m in self.toolchain_lab_machines or
cros_name in self.toolchain_lab_machines):
name = m if m in self.toolchain_lab_machines else cros_name
if not printed_hdr:
self.PrintStatusHeader(True)
printed_hdr = True
state = machine_states[m]
if state['locked']:
print('%s (%s)\tlocked by %s since %s' %
(name, state['board'], state['locked_by'], state['lock_time']))
else:
print('%s (%s)\tunlocked' % (name, state['board']))
else:
local_machines.append(m)
if local_machines:
self.PrintStatusHeader(False)
for m in local_machines:
state = machine_states[m]
if state['locked']:
print('%s (%s)\tlocked by %s since %s' %
(m, state['board'], state['locked_by'], state['lock_time']))
else:
print('%s (%s)\tunlocked' % (m, state['board']))
def UpdateLockInAFE(self, should_lock_machine, machine):
"""Calls an AFE server to lock/unlock a machine.
Args:
should_lock_machine: Boolean indicating whether to lock the machine (True)
or unlock the machine (False).
machine: The machine to update.
Raises:
LockingError: An error occurred while attempting to update the machine
state.
"""
action = 'lock'
if not should_lock_machine:
action = 'unlock'
kwargs = {'locked': should_lock_machine}
kwargs['lock_reason'] = 'toolchain user request (%s)' % self.user
cros_name = machine + '.cros'
if cros_name in self.toolchain_lab_machines:
machine = cros_name
if machine in self.toolchain_lab_machines:
m = machine.split('.')[0]
afe_server = self.afe
else:
m = machine
afe_server = self.local_afe
try:
afe_server.run(
'modify_hosts',
host_filter_data={'hostname__in': [m]},
update_data=kwargs)
except Exception as e:
traceback.print_exc()
raise LockingError('Unable to %s machine %s. %s' % (action, m, str(e)))
def UpdateMachines(self, lock_machines):
"""Sets the locked state of the machines to the requested value.
The machines updated are the ones in self.machines (specified when the
class object was intialized).
Args:
lock_machines: Boolean indicating whether to lock the machines (True) or
unlock the machines (False).
Returns:
A list of the machines whose state was successfully updated.
"""
updated_machines = []
for m in self.machines:
self.UpdateLockInAFE(lock_machines, m)
# Since we returned from self.UpdateLockInAFE we assume the request
# succeeded.
if lock_machines:
self.logger.LogOutput('Locked machine(s) %s.' % m)
else:
self.logger.LogOutput('Unlocked machine(s) %s.' % m)
updated_machines.append(m)
return updated_machines
def _InternalRemoveMachine(self, machine):
"""Remove machine from internal list of machines.
Args:
machine: Name of machine to be removed from internal list.
"""
# Check to see if machine is lab machine and if so, make sure it has
# ".cros" on the end.
cros_machine = machine
if machine.find('rack') > 0 and machine.find('row') > 0:
if machine.find('.cros') == -1:
cros_machine = cros_machine + '.cros'
self.machines = [
m for m in self.machines if m != cros_machine and m != machine
]
def CheckMachineLocks(self, machine_states, cmd):
"""Check that every machine in requested list is in the proper state.
If the cmd is 'unlock' verify that every machine is locked by requestor.
If the cmd is 'lock' verify that every machine is currently unlocked.
Args:
machine_states: A dictionary of the current state of every machine in
the current AFELockManager's list of machines. Normally obtained by
calling AFELockManager::GetMachineStates.
cmd: The user-requested action for the machines: 'lock' or 'unlock'.
Raises:
DontOwnLock: The lock on a requested machine is owned by someone else.
"""
for k, state in machine_states.iteritems():
if cmd == 'unlock':
if not state['locked']:
self.logger.LogWarning('Attempt to unlock already unlocked machine '
'(%s).' % k)
self._InternalRemoveMachine(k)
if state['locked'] and state['locked_by'] != self.user:
raise DontOwnLock('Attempt to unlock machine (%s) locked by someone '
'else (%s).' % (k, state['locked_by']))
elif cmd == 'lock':
if state['locked']:
self.logger.LogWarning(
'Attempt to lock already locked machine (%s)' % k)
self._InternalRemoveMachine(k)
def HasAFEServer(self, local):
"""Verifies that the AFELockManager has appropriate AFE server.
Args:
local: Boolean indicating whether we are checking for the local server
(True) or for the global server (False).
Returns:
A boolean indicating if the AFELockManager has the requested AFE server.
"""
if local:
return self.local_afe is not None
else:
return self.afe is not None
def GetMachineStates(self, cmd=''):
"""Gets the current state of all the requested machines.
Gets the current state of all the requested machines, both from the HW lab
sever and from the local server. Stores the data in a dictionary keyed
by machine name.
Args:
cmd: The command for which we are getting the machine states. This is
important because if one of the requested machines is missing we raise
an exception, unless the requested command is 'add'.
Returns:
A dictionary of machine states for all the machines in the AFELockManager
object.
Raises:
NoAFEServer: Cannot find the HW Lab or local AFE server.
AFEAccessError: An error occurred when querying the server about a
machine.
"""
if not self.HasAFEServer(False):
raise NoAFEServer('Error: Cannot connect to main AFE server.')
if self.local and not self.HasAFEServer(True):
raise NoAFEServer('Error: Cannot connect to local AFE server.')
machine_list = {}
for m in self.machines:
host_info = None
cros_name = m + '.cros'
if (m in self.toolchain_lab_machines or
cros_name in self.toolchain_lab_machines):
mod_host = m.split('.')[0]
host_info = self.afe.get_hosts(hostname=mod_host)
if not host_info:
raise AFEAccessError('Unable to get information about %s from main'
' autotest server.' % m)
else:
host_info = self.local_afe.get_hosts(hostname=m)
if not host_info and cmd != 'add':
raise AFEAccessError('Unable to get information about %s from '
'local autotest server.' % m)
if host_info:
host_info = host_info[0]
name = host_info.hostname
values = {}
values['board'] = host_info.platform if host_info.platform else '??'
values['locked'] = host_info.locked
if host_info.locked:
values['locked_by'] = host_info.locked_by
values['lock_time'] = host_info.lock_time
else:
values['locked_by'] = ''
values['lock_time'] = ''
machine_list[name] = values
else:
machine_list[m] = {}
return machine_list
def Main(argv):
"""Parse the options, initialize lock manager and dispatch proper method.
Args:
argv: The options with which this script was invoked.
Returns:
0 unless an exception is raised.
"""
parser = argparse.ArgumentParser()
parser.add_argument(
'--list',
dest='cmd',
action='store_const',
const='status',
help='List current status of all known machines.')
parser.add_argument(
'--lock',
dest='cmd',
action='store_const',
const='lock',
help='Lock given machine(s).')
parser.add_argument(
'--unlock',
dest='cmd',
action='store_const',
const='unlock',
help='Unlock given machine(s).')
parser.add_argument(
'--status',
dest='cmd',
action='store_const',
const='status',
help='List current status of given machine(s).')
parser.add_argument(
'--add_machine',
dest='cmd',
action='store_const',
const='add',
help='Add machine to local machine server.')
parser.add_argument(
'--remove_machine',
dest='cmd',
action='store_const',
const='remove',
help='Remove machine from the local machine server.')
parser.add_argument(
'--nolocal',
dest='local',
action='store_false',
default=True,
help='Do not try to use local machine server.')
parser.add_argument(
'--remote', dest='remote', help='machines on which to operate')
parser.add_argument(
'--chromeos_root',
dest='chromeos_root',
required=True,
help='ChromeOS root to use for autotest scripts.')
parser.add_argument(
'--local_server',
dest='local_server',
default=None,
help='Alternate local autotest server to use.')
parser.add_argument(
'--force',
dest='force',
action='store_true',
default=False,
help='Force lock/unlock of machines, even if not'
' current lock owner.')
options = parser.parse_args(argv)
if not options.remote and options.cmd != 'status':
parser.error('No machines specified for operation.')
if not os.path.isdir(options.chromeos_root):
parser.error('Cannot find chromeos_root: %s.' % options.chromeos_root)
if not options.cmd:
parser.error('No operation selected (--list, --status, --lock, --unlock,'
' --add_machine, --remove_machine).')
machine_list = []
if options.remote:
machine_list = options.remote.split()
lock_manager = AFELockManager(machine_list, options.force,
options.chromeos_root, options.local_server,
options.local)
machine_states = lock_manager.GetMachineStates(cmd=options.cmd)
cmd = options.cmd
if cmd == 'status':
lock_manager.ListMachineStates(machine_states)
elif cmd == 'lock':
if not lock_manager.force:
lock_manager.CheckMachineLocks(machine_states, cmd)
lock_manager.UpdateMachines(True)
elif cmd == 'unlock':
if not lock_manager.force:
lock_manager.CheckMachineLocks(machine_states, cmd)
lock_manager.UpdateMachines(False)
elif cmd == 'add':
lock_manager.AddMachinesToLocalServer()
elif cmd == 'remove':
lock_manager.RemoveMachinesFromLocalServer()
return 0
if __name__ == '__main__':
sys.exit(Main(sys.argv[1:]))