testing/legion/task_controller.py

   1 # Copyright 2015 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Defines the task controller library."""
   6
   7 import argparse
   8 import datetime
   9 import logging
  10 import os
  11 import socket
  12 import subprocess
  13 import sys
  14 import tempfile
  15 import threading
  16
  17 #pylint: disable=relative-import
  18 import common_lib
  19 import process
  20 import rpc_server
  21 import jsonrpclib
  22
  23 ISOLATE_PY = os.path.join(common_lib.SWARMING_DIR, 'isolate.py')
  24 SWARMING_PY = os.path.join(common_lib.SWARMING_DIR, 'swarming.py')
  25
  26
  27 class Error(Exception):
  28   pass
  29
  30
  31 class ConnectionTimeoutError(Error):
  32   pass
  33
  34
  35 class TaskController(object):
  36   """Provisions, configures, and controls a task machine.
  37
  38   This class is an abstraction of a physical task machine. It provides an
  39   end to end API for controlling a task machine. Operations on the task machine
  40   are performed using the instance's "rpc" property. A simple end to end
  41   scenario is as follows:
  42
  43   task = TaskController(...)
  44   task.Create()
  45   task.WaitForConnection()
  46   proc = task.rpc.subprocess.Popen(['ls'])
  47   print task.rpc.subprocess.GetStdout(proc)
  48   task.Release()
  49   """
  50
  51   _task_count = 0
  52   _tasks = []
  53
  54   def __init__(self, isolated_hash, dimensions, priority=100,
  55                idle_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  56                connection_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  57                verbosity='ERROR', name=None, run_id=None):
  58     assert isinstance(dimensions, dict)
  59     type(self)._tasks.append(self)
  60     type(self)._task_count += 1
  61     self.verbosity = verbosity
  62     self._name = name or 'Task%d' % type(self)._task_count
  63     self._priority = priority
  64     self._isolated_hash = isolated_hash
  65     self._idle_timeout_secs = idle_timeout_secs
  66     self._dimensions = dimensions
  67     self._connect_event = threading.Event()
  68     self._connected = False
  69     self._ip_address = None
  70     self._otp = self._CreateOTP()
  71     self._rpc = None
  72     self._output_dir = None
  73
  74     run_id = run_id or datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
  75     self._task_name = '%s/%s/%s' % (
  76         os.path.splitext(sys.argv[0])[0], self._name, run_id)
  77
  78     parser = argparse.ArgumentParser()
  79     parser.add_argument('--isolate-server')
  80     parser.add_argument('--swarming-server')
  81     parser.add_argument('--task-connection-timeout-secs',
  82                         default=common_lib.DEFAULT_TIMEOUT_SECS)
  83     args, _ = parser.parse_known_args()
  84
  85     self._isolate_server = args.isolate_server
  86     self._swarming_server = args.swarming_server
  87     self._connection_timeout_secs = (connection_timeout_secs or
  88                                     args.task_connection_timeout_secs)
  89
  90   @property
  91   def name(self):
  92     return self._name
  93
  94   @property
  95   def otp(self):
  96     return self._otp
  97
  98   @property
  99   def connected(self):
 100     return self._connected
 101
 102   @property
 103   def connect_event(self):
 104     return self._connect_event
 105
 106   @property
 107   def rpc(self):
 108     return self._rpc
 109
 110   @property
 111   def verbosity(self):
 112     return self._verbosity
 113
 114   @verbosity.setter
 115   def verbosity(self, level):
 116     """Sets the verbosity level as a string.
 117
 118     Either a string ('INFO', 'DEBUG', etc) or a logging level (logging.INFO,
 119     logging.DEBUG, etc) is allowed.
 120     """
 121     assert isinstance(level, (str, int))
 122     if isinstance(level, int):
 123       level = logging.getLevelName(level)
 124     self._verbosity = level  #pylint: disable=attribute-defined-outside-init
 125
 126   @property
 127   def output_dir(self):
 128     if not self._output_dir:
 129       self._output_dir = self.rpc.GetOutputDir()
 130     return self._output_dir
 131
 132   @classmethod
 133   def ReleaseAllTasks(cls):
 134     for task in cls._tasks:
 135       task.Release()
 136
 137   def Process(self, cmd, *args, **kwargs):
 138     return process.ControllerProcessWrapper(self.rpc, cmd, *args, **kwargs)
 139
 140   def _CreateOTP(self):
 141     """Creates the OTP."""
 142     controller_name = socket.gethostname()
 143     test_name = os.path.basename(sys.argv[0])
 144     creation_time = datetime.datetime.utcnow()
 145     otp = 'task:%s controller:%s test:%s creation:%s' % (
 146         self._name, controller_name, test_name, creation_time)
 147     return otp
 148
 149   def Create(self):
 150     """Creates the task machine."""
 151     logging.info('Creating %s', self.name)
 152     self._connect_event.clear()
 153     self._ExecuteSwarming()
 154
 155   def WaitForConnection(self):
 156     """Waits for the task machine to connect.
 157
 158     Raises:
 159       ConnectionTimeoutError if the task doesn't connect in time.
 160     """
 161     logging.info('Waiting for %s to connect with a timeout of %d seconds',
 162                  self._name, self._connection_timeout_secs)
 163     self._connect_event.wait(self._connection_timeout_secs)
 164     if not self._connect_event.is_set():
 165       raise ConnectionTimeoutError('%s failed to connect' % self.name)
 166
 167   def Release(self):
 168     """Quits the task's RPC server so it can release the machine."""
 169     if self._rpc is not None and self._connected:
 170       logging.info('Copying output-dir files to controller')
 171       self.RetrieveOutputFiles()
 172       logging.info('Releasing %s', self._name)
 173       try:
 174         self._rpc.Quit()
 175       except (socket.error, jsonrpclib.Fault):
 176         logging.error('Unable to connect to %s to call Quit', self.name)
 177       self._rpc = None
 178       self._connected = False
 179
 180   def _ExecuteSwarming(self):
 181     """Executes swarming.py."""
 182     cmd = [
 183         'python',
 184         SWARMING_PY,
 185         'trigger',
 186         self._isolated_hash,
 187         '--priority', str(self._priority),
 188         '--task-name', self._task_name,
 189         ]
 190
 191     if self._isolate_server:
 192       cmd.extend(['--isolate-server', self._isolate_server])
 193     if self._swarming_server:
 194       cmd.extend(['--swarming', self._swarming_server])
 195     for key, value in self._dimensions.iteritems():
 196       cmd.extend(['--dimension', key, value])
 197
 198     cmd.extend([
 199         '--',
 200         '--controller', common_lib.MY_IP,
 201         '--otp', self._otp,
 202         '--verbosity', self._verbosity,
 203         '--idle-timeout', str(self._idle_timeout_secs),
 204         '--output-dir', '${ISOLATED_OUTDIR}'
 205         ])
 206
 207     self._ExecuteProcess(cmd)
 208
 209   def _ExecuteProcess(self, cmd):
 210     """Executes a process, waits for it to complete, and checks for success."""
 211     logging.debug('Running %s', ' '.join(cmd))
 212     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 213     _, stderr = p.communicate()
 214     if p.returncode != 0:
 215       raise Error(stderr)
 216
 217   def OnConnect(self, ip_address):
 218     """Receives task ip address on connection."""
 219     self._ip_address = ip_address
 220     self._connected = True
 221     self._rpc = rpc_server.RpcServer.Connect(self._ip_address)
 222     logging.info('%s connected from %s', self._name, ip_address)
 223     self._connect_event.set()
 224
 225   def RetrieveOutputFiles(self):
 226     """Retrieves all files in the output-dir."""
 227     files = self.rpc.ListDir(self.output_dir)
 228     for fname in files:
 229       remote_path = self.rpc.PathJoin(self.output_dir, fname)
 230       local_name = os.path.join(common_lib.GetOutputDir(),
 231                                 '%s.%s' % (self.name, fname))
 232       contents = self.rpc.ReadFile(remote_path)
 233       with open(local_name, 'wb+') as fh:
 234         fh.write(contents)