| #!/usr/bin/python2 |
| |
| from __future__ import print_function |
| |
| import argparse |
| import logging |
| import multiprocessing |
| import subprocess |
| import sys |
| from multiprocessing.pool import ThreadPool |
| |
| import common |
| from autotest_lib.server import frontend |
| from autotest_lib.site_utils.lib import infra |
| |
| DEPLOY_SERVER_LOCAL = ('/usr/local/autotest/site_utils/deploy_server_local.py') |
| POOL_SIZE = 124 |
| |
| |
| def _filter_servers(servers): |
| """Filter a set of servers to those that should be deployed to.""" |
| non_push_roles = {'devserver', 'crash_server', 'reserve'} |
| for s in servers: |
| if s['status'] == 'repair_required': |
| continue |
| if s['status'] == 'backup': |
| continue |
| if set(s['roles']) & non_push_roles: |
| continue |
| yield s |
| |
| |
| def discover_servers(afe): |
| """Discover the in-production servers to update. |
| |
| Returns the set of servers from serverdb that are in production and should |
| be updated. This filters out servers in need of repair, or servers of roles |
| that are not yet supported by deploy_server / deploy_server_local. |
| |
| @param afe: Server to contact with RPC requests. |
| |
| @returns: A set of server hostnames. |
| """ |
| # Example server details.... |
| # { |
| # 'hostname': 'server1', |
| # 'status': 'backup', |
| # 'roles': ['drone', 'scheduler'], |
| # 'attributes': {'max_processes': 300} |
| # } |
| rpc = frontend.AFE(server=afe) |
| servers = rpc.run('get_servers') |
| |
| return {s['hostname'] for s in _filter_servers(servers)} |
| |
| |
| def _parse_arguments(args): |
| """Parse command line arguments. |
| |
| @param args: The command line arguments to parse. (usually sys.argv[1:]) |
| |
| @returns A tuple of (argparse.Namespace populated with argument values, |
| list of extra args to pass to deploy_server_local). |
| """ |
| parser = argparse.ArgumentParser( |
| formatter_class=argparse.RawDescriptionHelpFormatter, |
| description='Run deploy_server_local on a bunch of servers. Extra ' |
| 'arguments will be passed through.', |
| epilog=('Update all servers:\n' |
| ' deploy_server.py -x --afe cautotest\n' |
| '\n' |
| 'Update one server:\n' |
| ' deploy_server.py <server> -x\n' |
| )) |
| |
| parser.add_argument('-x', action='store_true', |
| help='Actually perform actions. If not supplied, ' |
| 'script does nothing.') |
| parser.add_argument('--afe', |
| help='The AFE server used to get servers from server_db,' |
| 'e.g, cautotest. Used only if no SERVER specified.') |
| parser.add_argument('servers', action='store', nargs='*', metavar='SERVER') |
| |
| return parser.parse_known_args() |
| |
| |
| def _update_server(server, extra_args=[]): |
| """Run deploy_server_local for given server. |
| |
| @param server: hostname to update. |
| @param extra_args: args to be passed in to deploy_server_local. |
| |
| @return: A tuple of (server, success, output), where: |
| server: Name of the server. |
| sucess: True if update succeeds, False otherwise. |
| output: A string of the deploy_server_local script output |
| including any errors. |
| """ |
| cmd = ('%s %s' % |
| (DEPLOY_SERVER_LOCAL, ' '.join(extra_args))) |
| success = False |
| try: |
| output = infra.execute_command(server, cmd) |
| success = True |
| except subprocess.CalledProcessError as e: |
| output = e.output |
| |
| return server, success, output |
| |
| def _update_in_parallel(servers, extra_args=[]): |
| """Update a group of servers in parallel. |
| |
| @param servers: A list of servers to update. |
| @param options: Options for the push. |
| |
| @returns A dictionary from server names that failed to the output |
| of the update script. |
| """ |
| # Create a list to record all the finished servers. |
| manager = multiprocessing.Manager() |
| finished_servers = manager.list() |
| |
| do_server = lambda s: _update_server(s, extra_args) |
| |
| # The update actions run in parallel. If any update failed, we should wait |
| # for other running updates being finished. Abort in the middle of an update |
| # may leave the server in a bad state. |
| pool = ThreadPool(POOL_SIZE) |
| try: |
| results = pool.map_async(do_server, servers) |
| pool.close() |
| |
| # Track the updating progress for current group of servers. |
| incomplete_servers = set() |
| server_names = set([s[0] for s in servers]) |
| while not results.ready(): |
| incomplete_servers = sorted(set(servers) - set(finished_servers)) |
| print('Not finished yet. %d servers in this group. ' |
| '%d servers are still running:\n%s\n' % |
| (len(servers), len(incomplete_servers), incomplete_servers)) |
| # Check the progress every 20s |
| results.wait(20) |
| |
| # After update finished, parse the result. |
| failures = {} |
| for server, success, output in results.get(): |
| if not success: |
| failures[server] = output |
| |
| return failures |
| |
| finally: |
| pool.terminate() |
| pool.join() |
| |
| |
| def main(args): |
| """Entry point to deploy_server.py |
| |
| @param args: The command line arguments to parse. (usually sys.argv) |
| |
| @returns The system exit code. |
| """ |
| options, extra_args = _parse_arguments(args[1:]) |
| # Remove all the handlers from the root logger to get rid of the handlers |
| # introduced by the import packages. |
| logging.getLogger().handlers = [] |
| logging.basicConfig(level=logging.DEBUG) |
| |
| servers = options.servers |
| if not servers: |
| if not options.afe: |
| print('No servers or afe specified. Aborting') |
| return 1 |
| print('Retrieving servers from %s..' % options.afe) |
| servers = discover_servers(options.afe) |
| print('Retrieved servers were: %s' % servers) |
| |
| if not options.x: |
| print('Doing nothing because -x was not supplied.') |
| print('servers: %s' % options.servers) |
| print('extra args for deploy_server_local: %s' % extra_args) |
| return 0 |
| |
| failures = _update_in_parallel(servers, extra_args) |
| |
| if not failures: |
| print('Completed all updates successfully.') |
| return 0 |
| |
| print('The following servers failed, with the following output:') |
| for s, o in failures.iteritems(): |
| print('======== %s ========' % s) |
| print(o) |
| |
| print('The servers that failed were:') |
| print('\n'.join(failures.keys())) |
| print('\n\nTo retry on failed servers, run the following command:') |
| retry_cmd = [args[0], '-x'] + failures.keys() + extra_args |
| print(' '.join(retry_cmd)) |
| return 1 |
| |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main(sys.argv)) |