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