site_utils/chromeos_proxy/swarming_bots.py

#! /usr/bin/python

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

"""
Manage swarming bots.

* Launch bots, e.g. 200 bots:
    $ swarming_bots.py launch --working_dir WORKING_DIR --bot_count 200

* Kill working bots:
    $ swarming_bots.py kill

* The hierachy of working dir is like
  WORKING_DIR
    |-- bot_0
    |   |-- 092b5bd4562f579711823f61e311de37247c853a-cacert.pem
    |   |-- bot_config.log
    |   |-- swarming_bot.log
    |   |-- swarming_bot.zip
    |-- bot_1
        |-- 092b5bd4562f579711823f61e311de37247c853a-cacert.pem
        |-- bot_config.log
        |-- swarming_bot.log
        |-- swarming_bot.zip
  Note bot_config.py:get_dimensions() will rely on the the bot number
  in the path to generate bot id.

* TODO (fdeng):
  ** Better management of each individual bot given a bot id.
    *** Kill a bot given a bot id.
    *** Launch a bot given a bot id.
    *** Restart a bot given a bot id.
  ** Kill a bot gracefully.
"""
import argparse
import logging
import os
import shutil
import subprocess
import sys
import time
import urllib

import common

from autotest_lib.client.common_lib import global_config


DEFAULT_SWARMING_PROXY = global_config.global_config.get_config_value(
        'CROS', "swarming_proxy", default=None)
BOT_FILENAME = 'swarming_bot.zip'


class BotManagementError(Exception):
    """Raised for any bot management related error."""


def _parse_args(args):
    """Parse args.

    @param args: Argument list passed from main.

    @return: A tuple with the parsed args, as returned by parser.parse_args.
    """
    if not args:
        print 'Too few arguments, try --help'
        sys.exit(1)

    parser = argparse.ArgumentParser(
            description='Launch swarming bots on a autotest server')
    subparsers = parser.add_subparsers()
    # launch parser
    launch_parser = subparsers.add_parser(
            'launch', help='Launch swarming bots.')
    launch_parser.required = False
    launch_parser.set_defaults(which='launch')
    launch_parser.add_argument(
            '--bot_count', type=int, dest='bot_count',required=True,
            help='Number of bots to launch ')
    launch_parser.add_argument(
            '--working_dir', type=str, dest='working_dir', required=True,
            help='A working directory where bots will store files '
                 'generated at runtime')
    launch_parser.add_argument(
            '--swarming_proxy', type=str, dest='swarming_proxy',
            default=DEFAULT_SWARMING_PROXY,
            help='The URL of the swarming instance to talk to, '
                 'Default to the one specified in global config')
    # kill parser
    kill_parser = subparsers.add_parser(
            'kill', help='Kill current running swarming bots.')
    kill_parser.required = False
    kill_parser.set_defaults(which='kill')

    return parser.parse_args(args)


def launch_bot(bot_dir, swarming_proxy):
    """Launch a swarming bot.

    @param bot_dir: The working directory for the bot.
                    The directory is used to store bot code,
                    log file, and any file generated by the bot
                    at run time.
    @param swarming_proxy: URL to the swarming instance.

    """
    logging.debug('Bootstrap bot in %s', bot_dir)
    if os.path.exists(bot_dir):
        shutil.rmtree(bot_dir)
    os.makedirs(bot_dir)
    dest = os.path.join(bot_dir, BOT_FILENAME)
    logging.debug('Getting bot code from: %s/bot_code', swarming_proxy)
    urllib.urlretrieve('%s/bot_code' % swarming_proxy, dest)
    cmd = [sys.executable, BOT_FILENAME]
    logging.debug('Calling command: %s', cmd)
    process = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=bot_dir)
    logging.debug('Created bot (pid: %d)', process.pid)


def launch(bot_count, working_dir, swarming_proxy):
    """Launch |bot_count| number of swarming bots.

    Each bot will have a seperate working dir,
    the name is like 'WORKING_DIR/bot_3'.
    bot_config.py will depend on the bot number generated here
    to generate bot id.

    @param bot_count: the number of bots to launch.
    @param working_dir: The root working directory for all bots.
    @param swarming_proxy: URL to the swarming instance.

    """
    wdir = os.path.abspath(os.path.expanduser(working_dir))
    for i in range(0, bot_count):
        bot_dir = os.path.join(wdir, 'bot_%s' % i)
        launch_bot(bot_dir, swarming_proxy)
    # If we let the process exit immediately, the last process won't
    # be launched sometimes. So sleep for 3 secs.
    logging.debug('Wait 3 seconds for process creation to complete.')
    time.sleep(3)


def killall():
    """Kill all the running bots."""
    # TODO(fdeng): We should record PID and use python's
    # os module to manage the processes.
    cmd =  ('ps ax|grep \'swarming_bot.*zip\'|grep -v grep|'
            'awk {\'print $1\'}|xargs -i kill {}')

    logging.debug('Killing all bots, cmd: %s', cmd)
    process = subprocess.Popen(cmd, shell=True)
    process.wait()
    if process.returncode != 0:
        raise BotManagementError(
                'Kill bots failed, cmd return code %d: %s',
                process.returncode, cmd)


def main(args):
    """Main.

    @args: A list of system arguments.
    """
    logging.basicConfig(level=logging.DEBUG)
    args = _parse_args(args)
    if args.which == 'launch':
        if not args.swarming_proxy:
            logging.error(
                    'No swarming proxy instance specified. '
                    'Specify swarming_proxy in [CROS] in shadow_config, '
                    'or use --swarming_proxy')
            return 1
        if not args.swarming_proxy.startswith('https://'):
            swarming_proxy = 'https://' + args.swarming_proxy
        else:
            swarming_proxy = args.swarming_proxy
        logging.info('Connecting to %s', swarming_proxy)
        launch(args.bot_count, args.working_dir, args.swarming_proxy)
    elif args.which == 'kill':
        killall()


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))