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