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