worker.py
140 lines
| 3.8 KiB
| text/x-python
|
PythonLexer
/ mercurial / worker.py
|
r18635 | # worker.py - master-slave parallelism support | ||
# | ||||
# Copyright 2013 Facebook, Inc. | ||||
# | ||||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
|
r18636 | from i18n import _ | ||
|
r18638 | import os, signal, sys, util | ||
|
r18635 | |||
def countcpus(): | ||||
'''try to count the number of CPUs on the system''' | ||||
# posix | ||||
try: | ||||
n = int(os.sysconf('SC_NPROCESSORS_ONLN')) | ||||
if n > 0: | ||||
return n | ||||
except (AttributeError, ValueError): | ||||
pass | ||||
# windows | ||||
try: | ||||
n = int(os.environ['NUMBER_OF_PROCESSORS']) | ||||
if n > 0: | ||||
return n | ||||
except (KeyError, ValueError): | ||||
pass | ||||
return 1 | ||||
|
r18636 | |||
def _numworkers(ui): | ||||
s = ui.config('worker', 'numcpus') | ||||
if s: | ||||
try: | ||||
n = int(s) | ||||
if n >= 1: | ||||
return n | ||||
except ValueError: | ||||
raise util.Abort(_('number of cpus must be an integer')) | ||||
return min(max(countcpus(), 4), 32) | ||||
if os.name == 'posix': | ||||
_startupcost = 0.01 | ||||
else: | ||||
_startupcost = 1e30 | ||||
def worthwhile(ui, costperop, nops): | ||||
'''try to determine whether the benefit of multiple processes can | ||||
outweigh the cost of starting them''' | ||||
linear = costperop * nops | ||||
workers = _numworkers(ui) | ||||
benefit = linear - (_startupcost * workers + linear / workers) | ||||
return benefit >= 0.15 | ||||
|
r18637 | |||
|
r18638 | def worker(ui, costperarg, func, staticargs, args): | ||
'''run a function, possibly in parallel in multiple worker | ||||
processes. | ||||
returns a progress iterator | ||||
costperarg - cost of a single task | ||||
func - function to run | ||||
staticargs - arguments to pass to every invocation of the function | ||||
args - arguments to split into chunks, to pass to individual | ||||
workers | ||||
''' | ||||
if worthwhile(ui, costperarg, len(args)): | ||||
return _platformworker(ui, func, staticargs, args) | ||||
return func(*staticargs + (args,)) | ||||
def _posixworker(ui, func, staticargs, args): | ||||
rfd, wfd = os.pipe() | ||||
workers = _numworkers(ui) | ||||
|
r18708 | oldhandler = signal.getsignal(signal.SIGINT) | ||
signal.signal(signal.SIGINT, signal.SIG_IGN) | ||||
|
r18638 | for pargs in partition(args, workers): | ||
pid = os.fork() | ||||
if pid == 0: | ||||
|
r18708 | signal.signal(signal.SIGINT, oldhandler) | ||
|
r18638 | try: | ||
os.close(rfd) | ||||
for i, item in func(*(staticargs + (pargs,))): | ||||
os.write(wfd, '%d %s\n' % (i, item)) | ||||
os._exit(0) | ||||
except KeyboardInterrupt: | ||||
os._exit(255) | ||||
os.close(wfd) | ||||
fp = os.fdopen(rfd, 'rb', 0) | ||||
def cleanup(): | ||||
# python 2.4 is too dumb for try/yield/finally | ||||
signal.signal(signal.SIGINT, oldhandler) | ||||
|
r18707 | problem = None | ||
|
r18638 | for i in xrange(workers): | ||
|
r18707 | pid, st = os.wait() | ||
st = _exitstatus(st) | ||||
if st and not problem: | ||||
problem = st | ||||
if problem: | ||||
if problem < 0: | ||||
os.kill(os.getpid(), -problem) | ||||
sys.exit(problem) | ||||
|
r18638 | try: | ||
for line in fp: | ||||
l = line.split(' ', 1) | ||||
yield int(l[0]), l[1][:-1] | ||||
except: # re-raises | ||||
cleanup() | ||||
raise | ||||
cleanup() | ||||
|
r18707 | def _posixexitstatus(code): | ||
'''convert a posix exit status into the same form returned by | ||||
os.spawnv | ||||
returns None if the process was stopped instead of exiting''' | ||||
if os.WIFEXITED(code): | ||||
return os.WEXITSTATUS(code) | ||||
elif os.WIFSIGNALED(code): | ||||
return -os.WTERMSIG(code) | ||||
|
r18638 | if os.name != 'nt': | ||
_platformworker = _posixworker | ||||
|
r18707 | _exitstatus = _posixexitstatus | ||
|
r18638 | |||
|
r18637 | def partition(lst, nslices): | ||
'''partition a list into N slices of equal size''' | ||||
n = len(lst) | ||||
chunk, slop = n / nslices, n % nslices | ||||
end = 0 | ||||
for i in xrange(nslices): | ||||
start = end | ||||
end = start + chunk | ||||
if slop: | ||||
end += 1 | ||||
slop -= 1 | ||||
yield lst[start:end] | ||||