mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 1 | __author__ = """Copyright Andy Whitcroft, Martin J. Bligh - 2006, 2007""" |
| 2 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 3 | import sys, os, signal, time, cPickle, logging |
jadmanski | d93d7d2 | 2008-05-29 21:37:29 +0000 | [diff] [blame] | 4 | |
| 5 | from autotest_lib.client.common_lib import error, utils |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 6 | from autotest_lib.client.common_lib.cros import retry |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 7 | |
| 8 | |
showard | 75cdfee | 2009-06-10 17:40:41 +0000 | [diff] [blame] | 9 | # entry points that use subcommand must set this to their logging manager |
| 10 | # to get log redirection for subcommands |
| 11 | logging_manager_object = None |
| 12 | |
| 13 | |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 14 | def parallel(tasklist, timeout=None, return_results=False): |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 15 | """ |
| 16 | Run a set of predefined subcommands in parallel. |
| 17 | |
| 18 | @param tasklist: A list of subcommand instances to execute. |
| 19 | @param timeout: Number of seconds after which the commands should timeout. |
| 20 | @param return_results: If True instead of an AutoServError being raised |
| 21 | on any error a list of the results|exceptions from the tasks is |
| 22 | returned. [default: False] |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 23 | """ |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 24 | run_error = False |
| 25 | for task in tasklist: |
| 26 | task.fork_start() |
mbligh | c3aee0f | 2008-01-17 16:26:39 +0000 | [diff] [blame] | 27 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 28 | remaining_timeout = None |
| 29 | if timeout: |
| 30 | endtime = time.time() + timeout |
mbligh | c3aee0f | 2008-01-17 16:26:39 +0000 | [diff] [blame] | 31 | |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 32 | results = [] |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 33 | for task in tasklist: |
| 34 | if timeout: |
| 35 | remaining_timeout = max(endtime - time.time(), 1) |
| 36 | try: |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 37 | status = task.fork_waitfor(timeout=remaining_timeout) |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 38 | except error.AutoservSubcommandError: |
| 39 | run_error = True |
| 40 | else: |
| 41 | if status != 0: |
| 42 | run_error = True |
mbligh | 158ba7b | 2008-03-07 18:29:12 +0000 | [diff] [blame] | 43 | |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 44 | results.append(cPickle.load(task.result_pickle)) |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 45 | task.result_pickle.close() |
| 46 | |
| 47 | if return_results: |
| 48 | return results |
| 49 | elif run_error: |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 50 | message = 'One or more subcommands failed:\n' |
| 51 | for task, result in zip(tasklist, results): |
| 52 | message += 'task: %s returned/raised: %r\n' % (task, result) |
| 53 | raise error.AutoservError(message) |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 54 | |
| 55 | |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 56 | def parallel_simple(function, arglist, subdir_name_constructor=lambda x: str(x), |
| 57 | log=True, timeout=None, return_results=False): |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 58 | """ |
| 59 | Each element in the arglist used to create a subcommand object, |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 60 | where that arg is used both as a subdir name, and a single argument |
| 61 | to pass to "function". |
mbligh | dd3235b | 2008-01-14 16:44:19 +0000 | [diff] [blame] | 62 | |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 63 | We create a subcommand object for each element in the list, |
| 64 | then execute those subcommand objects in parallel. |
| 65 | |
| 66 | NOTE: As an optimization, if len(arglist) == 1 a subcommand is not used. |
| 67 | |
| 68 | @param function: A callable to run in parallel once per arg in arglist. |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 69 | @param arglist: A list of arguments to be used one per subcommand |
| 70 | @param subdir_name_constructor: A function that returns a name for the |
| 71 | result sub-directory created per subcommand. |
| 72 | Signature is: |
| 73 | subdir_name_constructor(arg) |
| 74 | where arg is the argument passed to function. |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 75 | @param log: If True, output will be written to output in a subdirectory |
| 76 | named after each subcommand's arg. |
| 77 | @param timeout: Number of seconds after which the commands should timeout. |
| 78 | @param return_results: If True instead of an AutoServError being raised |
| 79 | on any error a list of the results|exceptions from the function |
| 80 | called on each arg is returned. [default: False] |
| 81 | |
| 82 | @returns None or a list of results/exceptions. |
| 83 | """ |
mbligh | 26f0d88 | 2009-06-22 18:30:01 +0000 | [diff] [blame] | 84 | if not arglist: |
Ilja H. Friedel | 04be2bd | 2014-05-07 21:29:59 -0700 | [diff] [blame] | 85 | logging.warning('parallel_simple was called with an empty arglist, ' |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 86 | 'did you forget to pass in a list of machines?') |
| 87 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 88 | # Bypass the multithreading if only one machine. |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 89 | if len(arglist) == 1: |
| 90 | arg = arglist[0] |
| 91 | if return_results: |
| 92 | try: |
| 93 | result = function(arg) |
| 94 | except Exception, e: |
| 95 | return [e] |
| 96 | return [result] |
| 97 | else: |
| 98 | function(arg) |
| 99 | return |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 100 | |
| 101 | subcommands = [] |
| 102 | for arg in arglist: |
| 103 | args = [arg] |
Prathmesh Prabhu | d08c86b | 2017-07-21 16:14:33 -0700 | [diff] [blame] | 104 | subdir = subdir_name_constructor(arg) if log else None |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 105 | subcommands.append(subcommand(function, args, subdir)) |
mbligh | 415dc21 | 2009-06-15 21:53:34 +0000 | [diff] [blame] | 106 | return parallel(subcommands, timeout, return_results=return_results) |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 107 | |
| 108 | |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 109 | class subcommand(object): |
| 110 | fork_hooks, join_hooks = [], [] |
| 111 | |
showard | 75cdfee | 2009-06-10 17:40:41 +0000 | [diff] [blame] | 112 | def __init__(self, func, args, subdir = None): |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 113 | # func(args) - the subcommand to run |
| 114 | # subdir - the subdirectory to log results in |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 115 | if subdir: |
| 116 | self.subdir = os.path.abspath(subdir) |
| 117 | if not os.path.exists(self.subdir): |
| 118 | os.mkdir(self.subdir) |
| 119 | self.debug = os.path.join(self.subdir, 'debug') |
| 120 | if not os.path.exists(self.debug): |
| 121 | os.mkdir(self.debug) |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 122 | else: |
| 123 | self.subdir = None |
showard | 75cdfee | 2009-06-10 17:40:41 +0000 | [diff] [blame] | 124 | self.debug = None |
mbligh | d7685d3 | 2007-08-10 22:08:42 +0000 | [diff] [blame] | 125 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 126 | self.func = func |
| 127 | self.args = args |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 128 | self.pid = None |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 129 | self.returncode = None |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 130 | |
| 131 | |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 132 | def __str__(self): |
| 133 | return str('subcommand(func=%s, args=%s, subdir=%s)' % |
| 134 | (self.func, self.args, self.subdir)) |
| 135 | |
| 136 | |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 137 | @classmethod |
| 138 | def register_fork_hook(cls, hook): |
| 139 | """ Register a function to be called from the child process after |
| 140 | forking. """ |
| 141 | cls.fork_hooks.append(hook) |
| 142 | |
| 143 | |
| 144 | @classmethod |
| 145 | def register_join_hook(cls, hook): |
| 146 | """ Register a function to be called when from the child process |
| 147 | just before the child process terminates (joins to the parent). """ |
| 148 | cls.join_hooks.append(hook) |
| 149 | |
| 150 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 151 | def redirect_output(self): |
showard | 75cdfee | 2009-06-10 17:40:41 +0000 | [diff] [blame] | 152 | if self.subdir and logging_manager_object: |
| 153 | tag = os.path.basename(self.subdir) |
| 154 | logging_manager_object.tee_redirect_debug_dir(self.debug, tag=tag) |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 155 | |
| 156 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 157 | def fork_start(self): |
| 158 | sys.stdout.flush() |
| 159 | sys.stderr.flush() |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 160 | r, w = os.pipe() |
jadmanski | 77e8da8 | 2009-04-21 14:26:40 +0000 | [diff] [blame] | 161 | self.returncode = None |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 162 | self.pid = os.fork() |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 163 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 164 | if self.pid: # I am the parent |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 165 | os.close(w) |
| 166 | self.result_pickle = os.fdopen(r, 'r') |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 167 | return |
mbligh | e183681 | 2009-04-17 18:25:14 +0000 | [diff] [blame] | 168 | else: |
| 169 | os.close(r) |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 170 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 171 | # We are the child from this point on. Never return. |
| 172 | signal.signal(signal.SIGTERM, signal.SIG_DFL) # clear handler |
| 173 | if self.subdir: |
| 174 | os.chdir(self.subdir) |
| 175 | self.redirect_output() |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 176 | |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 177 | try: |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 178 | for hook in self.fork_hooks: |
| 179 | hook(self) |
Paul Hobbs | 4ea9564 | 2017-08-20 11:22:57 -0700 | [diff] [blame] | 180 | result = self.func(*self.args) |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 181 | os.write(w, cPickle.dumps(result, cPickle.HIGHEST_PROTOCOL)) |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 182 | exit_code = 0 |
mbligh | a57cc92 | 2009-08-24 22:04:19 +0000 | [diff] [blame] | 183 | except Exception, e: |
| 184 | logging.exception('function failed') |
| 185 | exit_code = 1 |
| 186 | os.write(w, cPickle.dumps(e, cPickle.HIGHEST_PROTOCOL)) |
| 187 | |
| 188 | os.close(w) |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 189 | |
| 190 | try: |
| 191 | for hook in self.join_hooks: |
| 192 | hook(self) |
| 193 | finally: |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 194 | sys.stdout.flush() |
| 195 | sys.stderr.flush() |
jadmanski | 550fdc2 | 2008-11-20 16:32:08 +0000 | [diff] [blame] | 196 | os._exit(exit_code) |
mbligh | b0fab82 | 2007-07-25 16:40:19 +0000 | [diff] [blame] | 197 | |
| 198 | |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 199 | def _handle_exitstatus(self, sts): |
| 200 | """ |
| 201 | This is partially borrowed from subprocess.Popen. |
| 202 | """ |
| 203 | if os.WIFSIGNALED(sts): |
| 204 | self.returncode = -os.WTERMSIG(sts) |
| 205 | elif os.WIFEXITED(sts): |
| 206 | self.returncode = os.WEXITSTATUS(sts) |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 207 | else: |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 208 | # Should never happen |
| 209 | raise RuntimeError("Unknown child exit status!") |
mbligh | c3aee0f | 2008-01-17 16:26:39 +0000 | [diff] [blame] | 210 | |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 211 | if self.returncode != 0: |
| 212 | print "subcommand failed pid %d" % self.pid |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 213 | print "%s" % (self.func,) |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 214 | print "rc=%d" % self.returncode |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 215 | print |
showard | 75cdfee | 2009-06-10 17:40:41 +0000 | [diff] [blame] | 216 | if self.debug: |
| 217 | stderr_file = os.path.join(self.debug, 'autoserv.stderr') |
| 218 | if os.path.exists(stderr_file): |
| 219 | for line in open(stderr_file).readlines(): |
| 220 | print line, |
jadmanski | 0afbb63 | 2008-06-06 21:10:57 +0000 | [diff] [blame] | 221 | print "\n--------------------------------------------\n" |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 222 | raise error.AutoservSubcommandError(self.func, self.returncode) |
| 223 | |
| 224 | |
| 225 | def poll(self): |
| 226 | """ |
| 227 | This is borrowed from subprocess.Popen. |
| 228 | """ |
| 229 | if self.returncode is None: |
| 230 | try: |
| 231 | pid, sts = os.waitpid(self.pid, os.WNOHANG) |
| 232 | if pid == self.pid: |
| 233 | self._handle_exitstatus(sts) |
| 234 | except os.error: |
| 235 | pass |
| 236 | return self.returncode |
| 237 | |
| 238 | |
| 239 | def wait(self): |
| 240 | """ |
| 241 | This is borrowed from subprocess.Popen. |
| 242 | """ |
| 243 | if self.returncode is None: |
| 244 | pid, sts = os.waitpid(self.pid, 0) |
| 245 | self._handle_exitstatus(sts) |
| 246 | return self.returncode |
| 247 | |
| 248 | |
| 249 | def fork_waitfor(self, timeout=None): |
| 250 | if not timeout: |
| 251 | return self.wait() |
| 252 | else: |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 253 | _, result = retry.timeout(self.wait, timeout_sec=timeout) |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 254 | |
David James | 3f638af | 2014-10-25 18:51:26 -0700 | [diff] [blame] | 255 | if result is None: |
| 256 | utils.nuke_pid(self.pid) |
| 257 | print "subcommand failed pid %d" % self.pid |
| 258 | print "%s" % (self.func,) |
| 259 | print "timeout after %ds" % timeout |
| 260 | print |
| 261 | result = self.wait() |
| 262 | |
| 263 | return result |