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     self._platform = None
  74
  75     run_id = run_id or datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
  76     self._task_name = '%s/%s/%s' % (
  77         os.path.splitext(sys.argv[0])[0], self._name, run_id)
  78
  79     parser = argparse.ArgumentParser()
  80     parser.add_argument('--isolate-server')
  81     parser.add_argument('--swarming-server')
  82     parser.add_argument('--task-connection-timeout-secs',
  83                         default=common_lib.DEFAULT_TIMEOUT_SECS)
  84     args, _ = parser.parse_known_args()
  85
  86     self._isolate_server = args.isolate_server
  87     self._swarming_server = args.swarming_server
  88     self._connection_timeout_secs = (connection_timeout_secs or
  89                                     args.task_connection_timeout_secs)
  90
  91   @property
  92   def name(self):
  93     return self._name
  94
  95   @property
  96   def otp(self):
  97     return self._otp
  98
  99   @property
 100   def connected(self):
 101     return self._connected
 102
 103   @property
 104   def connect_event(self):
 105     return self._connect_event
 106
 107   @property
 108   def rpc(self):
 109     return self._rpc
 110
 111   @property
 112   def verbosity(self):
 113     return self._verbosity
 114
 115   @verbosity.setter
 116   def verbosity(self, level):
 117     """Sets the verbosity level as a string.
 118
 119     Either a string ('INFO', 'DEBUG', etc) or a logging level (logging.INFO,
 120     logging.DEBUG, etc) is allowed.
 121     """
 122     assert isinstance(level, (str, int))
 123     if isinstance(level, int):
 124       level = logging.getLevelName(level)
 125     self._verbosity = level  #pylint: disable=attribute-defined-outside-init
 126
 127   @property
 128   def output_dir(self):
 129     if not self._output_dir:
 130       self._output_dir = self.rpc.GetOutputDir()
 131     return self._output_dir
 132
 133   @property
 134   def platform(self):
 135     if not self._platform:
 136       self._platform = self._rpc.GetPlatform()
 137     return self._platform
 138
 139   @classmethod
 140   def ReleaseAllTasks(cls):
 141     for task in cls._tasks:
 142       task.Release()
 143
 144   def Process(self, cmd, *args, **kwargs):
 145     return process.ControllerProcessWrapper(self.rpc, cmd, *args, **kwargs)
 146
 147   def _CreateOTP(self):
 148     """Creates the OTP."""
 149     controller_name = socket.gethostname()
 150     test_name = os.path.basename(sys.argv[0])
 151     creation_time = datetime.datetime.utcnow()
 152     otp = 'task:%s controller:%s test:%s creation:%s' % (
 153         self._name, controller_name, test_name, creation_time)
 154     return otp
 155
 156   def Create(self):
 157     """Creates the task machine."""
 158     logging.info('Creating %s', self.name)
 159     self._connect_event.clear()
 160     self._ExecuteSwarming()
 161
 162   def WaitForConnection(self):
 163     """Waits for the task machine to connect.
 164
 165     Raises:
 166       ConnectionTimeoutError if the task doesn't connect in time.
 167     """
 168     logging.info('Waiting for %s to connect with a timeout of %d seconds',
 169                  self._name, self._connection_timeout_secs)
 170     self._connect_event.wait(self._connection_timeout_secs)
 171     if not self._connect_event.is_set():
 172       raise ConnectionTimeoutError('%s failed to connect' % self.name)
 173
 174   def Release(self):
 175     """Quits the task's RPC server so it can release the machine."""
 176     if self._rpc is not None and self._connected:
 177       logging.info('Copying output-dir files to controller')
 178       self.RetrieveOutputFiles()
 179       logging.info('Releasing %s', self._name)
 180       try:
 181         self._rpc.Quit()
 182       except (socket.error, jsonrpclib.Fault):
 183         logging.error('Unable to connect to %s to call Quit', self.name)
 184       self._rpc = None
 185       self._connected = False
 186
 187   def _ExecuteSwarming(self):
 188     """Executes swarming.py."""
 189     cmd = [
 190         'python',
 191         SWARMING_PY,
 192         'trigger',
 193         self._isolated_hash,
 194         '--priority', str(self._priority),
 195         '--task-name', self._task_name,
 196         ]
 197
 198     if self._isolate_server:
 199       cmd.extend(['--isolate-server', self._isolate_server])
 200     if self._swarming_server:
 201       cmd.extend(['--swarming', self._swarming_server])
 202     for key, value in self._dimensions.iteritems():
 203       cmd.extend(['--dimension', key, value])
 204
 205     cmd.extend([
 206         '--',
 207         '--controller', common_lib.MY_IP,
 208         '--otp', self._otp,
 209         '--verbosity', self._verbosity,
 210         '--idle-timeout', str(self._idle_timeout_secs),
 211         '--output-dir', '${ISOLATED_OUTDIR}'
 212         ])
 213
 214     self._ExecuteProcess(cmd)
 215
 216   def _ExecuteProcess(self, cmd):
 217     """Executes a process, waits for it to complete, and checks for success."""
 218     logging.debug('Running %s', ' '.join(cmd))
 219     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 220     _, stderr = p.communicate()
 221     if p.returncode != 0:
 222       raise Error(stderr)
 223
 224   def OnConnect(self, ip_address):
 225     """Receives task ip address on connection."""
 226     self._ip_address = ip_address
 227     self._connected = True
 228     self._rpc = rpc_server.RpcServer.Connect(self._ip_address)
 229     logging.info('%s connected from %s', self._name, ip_address)
 230     self._connect_event.set()
 231
 232   def RetrieveOutputFiles(self):
 233     """Retrieves all files in the output-dir."""
 234     files = self.rpc.ListDir(self.output_dir)
 235     for fname in files:
 236       remote_path = self.rpc.PathJoin(self.output_dir, fname)
 237       local_name = os.path.join(common_lib.GetOutputDir(),
 238                                 '%s.%s' % (self.name, fname))
 239       contents = self.rpc.ReadFile(remote_path)
 240       with open(local_name, 'wb+') as fh:
 241         fh.write(contents)