|
|
# 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.
|
|
|
|
|
|
from i18n import _
|
|
|
import errno, os, signal, sys, threading, util
|
|
|
|
|
|
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
|
|
|
|
|
|
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
|
|
|
|
|
|
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)
|
|
|
oldhandler = signal.getsignal(signal.SIGINT)
|
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
|
pids, problem = [], [0]
|
|
|
for pargs in partition(args, workers):
|
|
|
pid = os.fork()
|
|
|
if pid == 0:
|
|
|
signal.signal(signal.SIGINT, oldhandler)
|
|
|
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)
|
|
|
except: # re-raises (close enough for debugging anyway)
|
|
|
try:
|
|
|
ui.traceback()
|
|
|
finally:
|
|
|
os._exit(255)
|
|
|
pids.append(pid)
|
|
|
pids.reverse()
|
|
|
os.close(wfd)
|
|
|
fp = os.fdopen(rfd, 'rb', 0)
|
|
|
def killworkers():
|
|
|
# if one worker bails, there's no good reason to wait for the rest
|
|
|
for p in pids:
|
|
|
try:
|
|
|
os.kill(p, signal.SIGTERM)
|
|
|
except OSError, err:
|
|
|
if err.errno != errno.ESRCH:
|
|
|
raise
|
|
|
def waitforworkers():
|
|
|
for _ in pids:
|
|
|
st = _exitstatus(os.wait()[1])
|
|
|
if st and not problem[0]:
|
|
|
problem[0] = st
|
|
|
killworkers()
|
|
|
t = threading.Thread(target=waitforworkers)
|
|
|
t.start()
|
|
|
def cleanup():
|
|
|
signal.signal(signal.SIGINT, oldhandler)
|
|
|
t.join()
|
|
|
status = problem[0]
|
|
|
if status:
|
|
|
if status < 0:
|
|
|
os.kill(os.getpid(), -status)
|
|
|
sys.exit(status)
|
|
|
try:
|
|
|
for line in fp:
|
|
|
l = line.split(' ', 1)
|
|
|
yield int(l[0]), l[1][:-1]
|
|
|
except: # re-raises
|
|
|
killworkers()
|
|
|
cleanup()
|
|
|
raise
|
|
|
cleanup()
|
|
|
|
|
|
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)
|
|
|
|
|
|
if os.name != 'nt':
|
|
|
_platformworker = _posixworker
|
|
|
_exitstatus = _posixexitstatus
|
|
|
|
|
|
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]
|
|
|
|