src/tools/swarming_client/tools/run_on_bots.py

   1 #!/usr/bin/env python
   2 # Copyright 2014 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 """Automated maintenance tool to run a script on bots.
   7
   8 To use this script, write a self-contained python script (use a .par if
   9 necessary), specify it on the command line and it will be packaged and triggered
  10 on all the swarming bots corresponding to the --dimension filters specified, or
  11 all the bots if no filter is specified.
  12 """
  13
  14 __version__ = '0.1'
  15
  16 import os
  17 import tempfile
  18 import shutil
  19 import subprocess
  20 import sys
  21
  22 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  23
  24 # Must be first import.
  25 import parallel_execution
  26
  27 from third_party import colorama
  28 from third_party.depot_tools import fix_encoding
  29 from utils import tools
  30
  31
  32 def get_bot_list(swarming_server, dimensions, dead_only):
  33   """Returns a list of swarming bots."""
  34   cmd = [
  35     sys.executable, 'swarming.py', 'query',
  36     '--swarming', swarming_server,
  37     '--bare',
  38   ]
  39   for k, v in dimensions.iteritems():
  40     cmd.extend(('--dimension', k, v))
  41   if dead_only:
  42     cmd.append('--dead-only')
  43   return subprocess.check_output(cmd, cwd=ROOT_DIR).splitlines()
  44
  45
  46 def archive(isolate_server, script):
  47   """Archives the tool and return the sha-1."""
  48   base_script = os.path.basename(script)
  49   isolate = {
  50     'variables': {
  51       'command': ['python', base_script],
  52       'isolate_dependency_tracked': [base_script],
  53     },
  54   }
  55   tempdir = tempfile.mkdtemp(prefix='run_on_bots')
  56   try:
  57     isolate_file = os.path.join(tempdir, 'tool.isolate')
  58     isolated_file = os.path.join(tempdir, 'tool.isolated')
  59     with open(isolate_file, 'wb') as f:
  60       f.write(str(isolate))
  61     shutil.copyfile(script, os.path.join(tempdir, base_script))
  62     cmd = [
  63       sys.executable, 'isolate.py', 'archive',
  64       '--isolate-server', isolate_server,
  65       '-i', isolate_file,
  66       '-s', isolated_file,
  67     ]
  68     return subprocess.check_output(cmd, cwd=ROOT_DIR).split()[0]
  69   finally:
  70     shutil.rmtree(tempdir)
  71
  72
  73 def run_serial(
  74     swarming_server, isolate_server, priority, deadline, repeat, isolated_hash,
  75     name, bots):
  76   """Runs the task one at a time.
  77
  78   This will be mainly bound by task scheduling latency, especially if the slaves
  79   are busy and the priority is low.
  80   """
  81   result = 0
  82   now = parallel_execution.timestamp()
  83   for i in xrange(repeat):
  84     for bot in bots:
  85       # Use an unique task name to ensure the task is executed.
  86       suffix = '/%d' % i if repeat > 1 else ''
  87       task_name = parallel_execution.unique_task_to_name(
  88           name, {'id': bot}, isolated_hash, now) + suffix
  89       cmd = [
  90         sys.executable, 'swarming.py', 'run',
  91         '--swarming', swarming_server,
  92         '--isolate-server', isolate_server,
  93         '--priority', priority,
  94         '--deadline', deadline,
  95         '--dimension', 'id', bot,
  96         '--task-name', task_name,
  97         isolated_hash,
  98       ]
  99       r = subprocess.call(cmd, cwd=ROOT_DIR)
 100       result = max(r, result)
 101   return result
 102
 103
 104 def run_parallel(
 105     swarming_server, isolate_server, priority, deadline, repeat, isolated_hash,
 106     name, bots):
 107   now = parallel_execution.timestamp()
 108   tasks = []
 109   for i in xrange(repeat):
 110     suffix = '/%d' % i if repeat > 1 else ''
 111     tasks.extend(
 112         (
 113           parallel_execution.unique_task_to_name(
 114               name, {'id': bot}, isolated_hash, now) + suffix,
 115           isolated_hash,
 116           {'id': bot},
 117         ) for bot in bots)
 118   extra_args = ['--priority', priority, '--deadline', deadline]
 119   print('Using priority %s' % priority)
 120   for failed_task in parallel_execution.run_swarming_tasks_parallel(
 121       swarming_server, isolate_server, extra_args, tasks):
 122     _name, dimensions, stdout = failed_task
 123     print('%sFailure: %s%s\n%s' % (
 124       colorama.Fore.RED, dimensions, colorama.Fore.RESET, stdout))
 125
 126
 127 def main():
 128   parser = parallel_execution.OptionParser(
 129       usage='%prog [options] script.py', version=__version__)
 130   parser.add_option(
 131       '--serial', action='store_true',
 132       help='Runs the task serially, to be used when debugging problems since '
 133            'it\'s slow')
 134   parser.add_option(
 135       '--repeat', type='int', default=1,
 136       help='Runs the task multiple time on each bot, meant to be used as a '
 137            'load test')
 138   options, args = parser.parse_args()
 139
 140   if len(args) != 1:
 141     parser.error(
 142         'Must pass one python script to run. Use --help for more details')
 143
 144   if not options.priority:
 145     parser.error(
 146         'Please provide the --priority option. Either use a very low number\n'
 147         'so the task completes as fast as possible, or an high number so the\n'
 148         'task only runs when the bot is idle.')
 149
 150   # 1. Query the bots list.
 151   bots = get_bot_list(options.swarming, options.dimensions, False)
 152   print('Found %d bots to process' % len(bots))
 153   if not bots:
 154     return 1
 155
 156   dead_bots = get_bot_list(options.swarming, options.dimensions, True)
 157   if dead_bots:
 158     print('Warning: found %d dead bots' % len(dead_bots))
 159
 160   # 2. Archive the script to run.
 161   isolated_hash = archive(options.isolate_server, args[0])
 162   print('Running %s' % isolated_hash)
 163
 164   # 3. Trigger the tasks.
 165   name = os.path.basename(args[0])
 166   if options.serial:
 167     return run_serial(
 168         options.swarming,
 169         options.isolate_server,
 170         str(options.priority),
 171         str(options.deadline),
 172         options.repeat,
 173         isolated_hash,
 174         name,
 175         bots)
 176
 177   return run_parallel(
 178       options.swarming,
 179       options.isolate_server,
 180       str(options.priority),
 181       str(options.deadline),
 182       options.repeat,
 183       isolated_hash,
 184       name,
 185       bots)
 186
 187
 188 if __name__ == '__main__':
 189   fix_encoding.fix_encoding()
 190   tools.disable_buffering()
 191   colorama.init()
 192   sys.exit(main())