| import fcntl, os, select, time |
| from subprocess import Popen, PIPE |
| |
| # Run a series of subprocesses. Try to keep up to a certain number going in |
| # parallel at any given time. Enforce time limits. |
| # |
| # This is implemented using non-blocking I/O, and so is Unix-specific. |
| # |
| # We assume that, if a task closes its standard error, then it's safe to |
| # wait for it to terminate. So an ill-behaved task that closes its standard |
| # output and then hangs will hang us, as well. However, as it takes special |
| # effort to close one's standard output, this seems unlikely to be a |
| # problem in practice. |
| class TaskPool(object): |
| |
| # A task we should run in a subprocess. Users should subclass this and |
| # fill in the methods as given. |
| class Task(object): |
| def __init__(self): |
| self.pipe = None |
| self.start_time = None |
| |
| # Record that this task is running, with |pipe| as its Popen object, |
| # and should time out at |deadline|. |
| def start(self, pipe, deadline): |
| self.pipe = pipe |
| self.deadline = deadline |
| |
| # Return a shell command (a string or sequence of arguments) to be |
| # passed to Popen to run the task. The command will be given |
| # /dev/null as its standard input, and pipes as its standard output |
| # and error. |
| def cmd(self): |
| raise NotImplementedError |
| |
| # TaskPool calls this method to report that the process wrote |
| # |string| to its standard output. |
| def onStdout(self, string): |
| raise NotImplementedError |
| |
| # TaskPool calls this method to report that the process wrote |
| # |string| to its standard error. |
| def onStderr(self, string): |
| raise NotImplementedError |
| |
| # TaskPool calls this method to report that the process terminated, |
| # yielding |returncode|. |
| def onFinished(self, returncode): |
| raise NotImplementedError |
| |
| # TaskPool calls this method to report that the process timed out and |
| # was killed. |
| def onTimeout(self): |
| raise NotImplementedError |
| |
| # If a task output handler (onStdout, onStderr) throws this, we terminate |
| # the task. |
| class TerminateTask(Exception): |
| pass |
| |
| def __init__(self, tasks, cwd='.', job_limit=4, timeout=150): |
| self.pending = iter(tasks) |
| self.cwd = cwd |
| self.job_limit = job_limit |
| self.timeout = timeout |
| self.next_pending = self.get_next_pending() |
| |
| # Set self.next_pending to the next task that has not yet been executed. |
| def get_next_pending(self): |
| try: |
| return self.pending.next() |
| except StopIteration: |
| return None |
| |
| def run_all(self): |
| # The currently running tasks: a set of Task instances. |
| running = set() |
| with open(os.devnull, 'r') as devnull: |
| while True: |
| while len(running) < self.job_limit and self.next_pending: |
| t = self.next_pending |
| p = Popen(t.cmd(), bufsize=16384, |
| stdin=devnull, stdout=PIPE, stderr=PIPE, |
| cwd=self.cwd) |
| |
| # Put the stdout and stderr pipes in non-blocking mode. See |
| # the post-'select' code below for details. |
| flags = fcntl.fcntl(p.stdout, fcntl.F_GETFL) |
| fcntl.fcntl(p.stdout, fcntl.F_SETFL, flags | os.O_NONBLOCK) |
| flags = fcntl.fcntl(p.stderr, fcntl.F_GETFL) |
| fcntl.fcntl(p.stderr, fcntl.F_SETFL, flags | os.O_NONBLOCK) |
| |
| t.start(p, time.time() + self.timeout) |
| running.add(t) |
| self.next_pending = self.get_next_pending() |
| |
| # If we have no tasks running, and the above wasn't able to |
| # start any new ones, then we must be done! |
| if not running: |
| break |
| |
| # How many seconds do we have until the earliest deadline? |
| now = time.time() |
| secs_to_next_deadline = max(min([t.deadline for t in running]) - now, 0) |
| |
| # Wait for output or a timeout. |
| stdouts_and_stderrs = ([t.pipe.stdout for t in running] |
| + [t.pipe.stderr for t in running]) |
| (readable,w,x) = select.select(stdouts_and_stderrs, [], [], secs_to_next_deadline) |
| finished = set() |
| terminate = set() |
| for t in running: |
| # Since we've placed the pipes in non-blocking mode, these |
| # 'read's will simply return as many bytes as are available, |
| # rather than blocking until they have accumulated the full |
| # amount requested (or reached EOF). The 'read's should |
| # never throw, since 'select' has told us there was |
| # something available. |
| if t.pipe.stdout in readable: |
| output = t.pipe.stdout.read(16384) |
| if output != "": |
| try: |
| t.onStdout(output) |
| except TerminateTask: |
| terminate.add(t) |
| if t.pipe.stderr in readable: |
| output = t.pipe.stderr.read(16384) |
| if output != "": |
| try: |
| t.onStderr(output) |
| except TerminateTask: |
| terminate.add(t) |
| else: |
| # We assume that, once a task has closed its stderr, |
| # it will soon terminate. If a task closes its |
| # stderr and then hangs, we'll hang too, here. |
| t.pipe.wait() |
| t.onFinished(t.pipe.returncode) |
| finished.add(t) |
| # Remove the finished tasks from the running set. (Do this here |
| # to avoid mutating the set while iterating over it.) |
| running -= finished |
| |
| # Terminate any tasks whose handlers have asked us to do so. |
| for t in terminate: |
| t.pipe.terminate() |
| t.pipe.wait() |
| running.remove(t) |
| |
| # Terminate any tasks which have missed their deadline. |
| finished = set() |
| for t in running: |
| if now >= t.deadline: |
| t.pipe.terminate() |
| t.pipe.wait() |
| t.onTimeout() |
| finished.add(t) |
| # Remove the finished tasks from the running set. (Do this here |
| # to avoid mutating the set while iterating over it.) |
| running -= finished |
| return None |
| |
| def get_cpu_count(): |
| """ |
| Guess at a reasonable parallelism count to set as the default for the |
| current machine and run. |
| """ |
| # Python 2.6+ |
| try: |
| import multiprocessing |
| return multiprocessing.cpu_count() |
| except (ImportError,NotImplementedError): |
| pass |
| |
| # POSIX |
| try: |
| res = int(os.sysconf('SC_NPROCESSORS_ONLN')) |
| if res > 0: |
| return res |
| except (AttributeError,ValueError): |
| pass |
| |
| # Windows |
| try: |
| res = int(os.environ['NUMBER_OF_PROCESSORS']) |
| if res > 0: |
| return res |
| except (KeyError, ValueError): |
| pass |
| |
| return 1 |
| |
| if __name__ == '__main__': |
| # Test TaskPool by using it to implement the unique 'sleep sort' algorithm. |
| def sleep_sort(ns, timeout): |
| sorted=[] |
| class SortableTask(TaskPool.Task): |
| def __init__(self, n): |
| super(SortableTask, self).__init__() |
| self.n = n |
| def start(self, pipe, deadline): |
| super(SortableTask, self).start(pipe, deadline) |
| def cmd(self): |
| return ['sh', '-c', 'echo out; sleep %d; echo err>&2' % (self.n,)] |
| def onStdout(self, text): |
| print '%d stdout: %r' % (self.n, text) |
| def onStderr(self, text): |
| print '%d stderr: %r' % (self.n, text) |
| def onFinished(self, returncode): |
| print '%d (rc=%d)' % (self.n, returncode) |
| sorted.append(self.n) |
| def onTimeout(self): |
| print '%d timed out' % (self.n,) |
| |
| p = TaskPool([SortableTask(_) for _ in ns], job_limit=len(ns), timeout=timeout) |
| p.run_all() |
| return sorted |
| |
| print repr(sleep_sort([1,1,2,3,5,8,13,21,34], 15)) |
