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
  21 ISOLATE_PY = os.path.join(common_lib.SWARMING_DIR, 'isolate.py')
  22 SWARMING_PY = os.path.join(common_lib.SWARMING_DIR, 'swarming.py')
  23
  24
  25 class Error(Exception):
  26   pass
  27
  28
  29 class ConnectionTimeoutError(Error):
  30   pass
  31
  32
  33 class TaskController(object):
  34   """Provisions, configures, and controls a task machine.
  35
  36   This class is an abstraction of a physical task machine. It provides an
  37   end to end API for controlling a task machine. Operations on the task machine
  38   are performed using the instance's "rpc" property. A simple end to end
  39   scenario is as follows:
  40
  41   task = TaskController(...)
  42   task.Create()
  43   task.WaitForConnection()
  44   proc = task.rpc.subprocess.Popen(['ls'])
  45   print task.rpc.subprocess.GetStdout(proc)
  46   task.Release()
  47   """
  48
  49   _task_count = 0
  50   _tasks = []
  51
  52   def __init__(self, isolated_hash, dimensions, priority=100,
  53                idle_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  54                connection_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  55                verbosity='ERROR', name=None, run_id=None):
  56     assert isinstance(dimensions, dict)
  57     type(self)._tasks.append(self)
  58     type(self)._task_count += 1
  59     self.verbosity = verbosity
  60     self._name = name or 'Task%d' % type(self)._task_count
  61     self._priority = priority
  62     self._isolated_hash = isolated_hash
  63     self._idle_timeout_secs = idle_timeout_secs
  64     self._dimensions = dimensions
  65     self._connect_event = threading.Event()
  66     self._connected = False
  67     self._ip_address = None
  68     self._otp = self._CreateOTP()
  69     self._rpc = None
  70
  71     run_id = run_id or datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
  72     self._task_name = '%s/%s/%s' % (
  73         os.path.splitext(sys.argv[0])[0], self._name, run_id)
  74
  75     parser = argparse.ArgumentParser()
  76     parser.add_argument('--isolate-server')
  77     parser.add_argument('--swarming-server')
  78     parser.add_argument('--task-connection-timeout-secs',
  79                         default=common_lib.DEFAULT_TIMEOUT_SECS)
  80     args, _ = parser.parse_known_args()
  81
  82     self._isolate_server = args.isolate_server
  83     self._swarming_server = args.swarming_server
  84     self._connection_timeout_secs = (connection_timeout_secs or
  85                                     args.task_connection_timeout_secs)
  86
  87   @property
  88   def name(self):
  89     return self._name
  90
  91   @property
  92   def otp(self):
  93     return self._otp
  94
  95   @property
  96   def connected(self):
  97     return self._connected
  98
  99   @property
 100   def connect_event(self):
 101     return self._connect_event
 102
 103   @property
 104   def rpc(self):
 105     return self._rpc
 106
 107   @property
 108   def verbosity(self):
 109     return self._verbosity
 110
 111   @verbosity.setter
 112   def verbosity(self, level):
 113     """Sets the verbosity level as a string.
 114
 115     Either a string ('INFO', 'DEBUG', etc) or a logging level (logging.INFO,
 116     logging.DEBUG, etc) is allowed.
 117     """
 118     assert isinstance(level, (str, int))
 119     if isinstance(level, int):
 120       level = logging.getLevelName(level)
 121     self._verbosity = level  #pylint: disable=attribute-defined-outside-init
 122
 123   @classmethod
 124   def ReleaseAllTasks(cls):
 125     for task in cls._tasks:
 126       task.Release()
 127
 128   def _CreateOTP(self):
 129     """Creates the OTP."""
 130     controller_name = socket.gethostname()
 131     test_name = os.path.basename(sys.argv[0])
 132     creation_time = datetime.datetime.utcnow()
 133     otp = 'task:%s controller:%s test:%s creation:%s' % (
 134         self._name, controller_name, test_name, creation_time)
 135     return otp
 136
 137   def Create(self):
 138     """Creates the task machine."""
 139     logging.info('Creating %s', self.name)
 140     self._connect_event.clear()
 141     self._ExecuteSwarming()
 142
 143   def WaitForConnection(self):
 144     """Waits for the task machine to connect.
 145
 146     Raises:
 147       ConnectionTimeoutError if the task doesn't connect in time.
 148     """
 149     logging.info('Waiting for %s to connect with a timeout of %d seconds',
 150                  self._name, self._connection_timeout_secs)
 151     self._connect_event.wait(self._connection_timeout_secs)
 152     if not self._connect_event.is_set():
 153       raise ConnectionTimeoutError('%s failed to connect' % self.name)
 154
 155   def Release(self):
 156     """Quits the task's RPC server so it can release the machine."""
 157     if self._rpc is not None and self._connected:
 158       logging.info('Releasing %s', self._name)
 159       try:
 160         self._rpc.Quit()
 161       except (socket.error, xmlrpclib.Fault):
 162         logging.error('Unable to connect to %s to call Quit', self.name)
 163       self._rpc = None
 164       self._connected = False
 165
 166   def _ExecuteSwarming(self):
 167     """Executes swarming.py."""
 168     cmd = [
 169         'python',
 170         SWARMING_PY,
 171         'trigger',
 172         self._isolated_hash,
 173         '--priority', str(self._priority),
 174         '--task-name', self._task_name,
 175         ]
 176
 177     if self._isolate_server:
 178       cmd.extend(['--isolate-server', self._isolate_server])
 179     if self._swarming_server:
 180       cmd.extend(['--swarming', self._swarming_server])
 181     for key, value in self._dimensions.iteritems():
 182       cmd.extend(['--dimension', key, value])
 183
 184     cmd.extend([
 185         '--',
 186         '--controller', common_lib.MY_IP,
 187         '--otp', self._otp,
 188         '--verbosity', self._verbosity,
 189         '--idle-timeout', str(self._idle_timeout_secs),
 190         ])
 191
 192     self._ExecuteProcess(cmd)
 193
 194   def _ExecuteProcess(self, cmd):
 195     """Executes a process, waits for it to complete, and checks for success."""
 196     logging.debug('Running %s', ' '.join(cmd))
 197     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 198     _, stderr = p.communicate()
 199     if p.returncode != 0:
 200       raise Error(stderr)
 201
 202   def OnConnect(self, ip_address):
 203     """Receives task ip address on connection."""
 204     self._ip_address = ip_address
 205     self._connected = True
 206     self._rpc = common_lib.ConnectToServer(self._ip_address)
 207     logging.info('%s connected from %s', self._name, ip_address)
 208     self._connect_event.set()