def _iter_zmq(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
backurl = 'tcp://%s:%s' % (config.dbserver.host, socket.port)
task_in_url = 'tcp://%s:%s' % (config.dbserver.host,
config.zworkers.task_in_port)
with Socket(task_in_url, zmq.PUSH, 'connect') as sender:
num_results = 0
for args in self._genargs(backurl):
sender.send((self.task_func, args))
num_results += 1
yield num_results
yield from self.loop(range(num_results), iter(socket))
def _iter_zmq(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
task_in_url = ('tcp://%(master_host)s:%(task_in_port)s' %
config.zworkers)
with Socket(task_in_url, zmq.PUSH, 'connect') as sender:
n = 0
for args in self._genargs(socket.backurl):
sender.send((self.task_func, args))
n += 1
yield n
for _ in range(n):
obj = socket.zsocket.recv_pyobj()
# receive n responses for the n requests sent
yield obj
def _iter_celery(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
logging.info('Using receiver %s', socket.backurl)
results = []
for piks in self._genargs(socket.backurl):
res = safetask.delay(self.task_func, piks)
# populating Starmap.task_ids, used in celery_cleanup
self.task_ids.append(res.task_id)
results.append(res)
num_results = len(results)
yield num_results
it = self.iter_native(results)
isocket = iter(socket)
while num_results:
res = next(isocket)
if self.calc_id and self.calc_id != res.mon.calc_id:
logging.warn(
'Discarding a result from job %d, since this '
'is job %d', res.mon.calc_id, self.calc_id)
continue
err = next(it)
if isinstance(err, Exception): # TaskRevokedError
raise err
num_results -= 1
yield res
def _iter_celery_zmq(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
logging.info('Using receiver %s', socket.backurl)
it = self._iter_celery(socket.backurl)
yield next(it) # number of results
isocket = iter(socket)
for _ in it:
yield next(isocket)
def safely_call(func, args, monitor=dummy_mon):
"""
Call the given function with the given arguments safely, i.e.
by trapping the exceptions. Return a pair (result, exc_type)
where exc_type is None if no exceptions occur, otherwise it
is the exception class and the result is a string containing
error message and traceback.
:param func: the function to call
:param args: the arguments
"""
isgenfunc = inspect.isgeneratorfunction(func)
monitor.operation = 'total ' + func.__name__
if hasattr(args[0], 'unpickle'):
# args is a list of Pickled objects
args = [a.unpickle() for a in args]
if monitor is dummy_mon: # in the DbServer
assert not isgenfunc, func
return Result.new(func, args, monitor)
mon = args[-1]
mon.operation = 'total ' + func.__name__
mon.measuremem = True
if mon is not monitor:
mon.children.append(monitor) # monitor is a child of mon
mon.weight = getattr(args[0], 'weight', 1.) # used in task_info
with Socket(monitor.backurl, zmq.PUSH, 'connect') as zsocket:
msg = check_mem_usage() # warn if too much memory is used
if msg:
zsocket.send(Result(None, mon, msg=msg))
if inspect.isgeneratorfunction(func):
gfunc = func
else:
def gfunc(*args):
yield func(*args)
gobj = gfunc(*args)
for count in itertools.count():
res = Result.new(next, (gobj, ), mon, count=count)
# StopIteration -> TASK_ENDED
try:
zsocket.send(res)
except Exception: # like OverflowError
_etype, exc, tb = sys.exc_info()
err = Result(exc,
mon,
''.join(traceback.format_tb(tb)),
count=count)
zsocket.send(err)
mon.duration = 0
mon.counts = 0
mon.children.clear()
if res.msg == 'TASK_ENDED':
break
def _iter_zmq(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
task_in_url = ('tcp://%(master_host)s:%(task_in_port)s' %
config.zworkers)
with Socket(task_in_url, zmq.PUSH, 'connect') as sender:
num_results = 0
for args in self._genargs(socket.backurl):
sender.send((self.task_func, args))
num_results += 1
yield num_results
isocket = iter(socket)
while num_results:
res = next(isocket)
if self.calc_id and self.calc_id != res.mon.calc_id:
logging.warn(
'Discarding a result from job %d, since this '
'is job %d', res.mon.calc_id, self.calc_id)
continue
num_results -= 1
yield res
def _iter_celery(self):
with Socket(self.receiver, zmq.PULL, 'bind') as socket:
backurl = 'tcp://%s:%s' % (config.dbserver.host, socket.port)
logging.debug('Using receiver %s', backurl)
results = []
for piks in self._genargs(backurl):
res = safetask.delay(self.task_func, piks)
# populating Starmap.task_ids, used in celery_cleanup
self.task_ids.append(res.task_id)
results.append(res)
num_results = len(results)
yield num_results
yield from self.loop(self.iter_native(results), iter(socket))
def submit(self, *args, func=None, monitor=None):
"""
Submit the given arguments to the underlying task
"""
monitor = monitor or self.monitor
func = func or self.task_func
if not hasattr(self, 'socket'): # first time
self.__class__.running_tasks = self.tasks
self.socket = Socket(self.receiver, zmq.PULL, 'bind').__enter__()
monitor.backurl = 'tcp://%s:%s' % (
config.dbserver.host, self.socket.port)
assert not isinstance(args[-1], Monitor) # sanity check
dist = 'no' if self.num_tasks == 1 else self.distribute
if dist != 'no':
args = pickle_sequence(args)
self.sent += numpy.array([len(p) for p in args])
res = getattr(self, dist + '_submit')(func, args, monitor)
self.tasks.append(res)
def submit(self, *args):
"""
Submit the given arguments to the underlying task
"""
global running_tasks
if not hasattr(self, 'socket'): # first time
running_tasks = self.tasks
self.socket = Socket(self.receiver, zmq.PULL, 'bind').__enter__()
self.monitor.backurl = 'tcp://%s:%s' % (config.dbserver.host,
self.socket.port)
assert not isinstance(args[-1], Monitor) # sanity check
# add incremental task number and task weight
self.monitor.task_no = len(self.tasks) + 1
dist = 'no' if self.num_tasks == 1 else self.distribute
if dist != 'no':
args = pickle_sequence(args)
self.sent += numpy.array([len(p) for p in args])
res = getattr(self, dist + '_submit')(args)
self.tasks.append(res)
def safely_call(func, args):
"""
Call the given function with the given arguments safely, i.e.
by trapping the exceptions. Return a pair (result, exc_type)
where exc_type is None if no exceptions occur, otherwise it
is the exception class and the result is a string containing
error message and traceback.
:param func: the function to call
:param args: the arguments
"""
with Monitor('total ' + func.__name__, measuremem=True) as child:
if args and hasattr(args[0], 'unpickle'):
# args is a list of Pickled objects
args = [a.unpickle() for a in args]
if args and isinstance(args[-1], Monitor):
mon = args[-1]
mon.operation = func.__name__
mon.children.append(child) # child is a child of mon
else: # in the DbServer
mon = child
try:
res = Result(func(*args), mon)
except Exception:
_etype, exc, tb = sys.exc_info()
res = Result(exc, mon, ''.join(traceback.format_tb(tb)))
# FIXME: check_mem_usage is disabled here because it's causing
# dead locks in threads when log messages are raised.
# Check is done anyway in other parts of the code
# further investigation is needed
# check_mem_usage(mon) # check if too much memory is used
backurl = getattr(mon, 'backurl', None)
if backurl is None:
return res
with Socket(backurl, zmq.PUSH, 'connect') as zsocket:
try:
zsocket.send(res)
except Exception: # like OverflowError
_etype, exc, tb = sys.exc_info()
err = Result(exc, mon, ''.join(traceback.format_tb(tb)))
zsocket.send(err)
return zsocket.num_sent
def zmq_submit(self, func, args, monitor):
if not hasattr(self, 'sender'):
task_in_url = 'tcp://%s:%s' % (config.dbserver.host,
config.zworkers.task_in_port)
self.sender = Socket(task_in_url, zmq.PUSH, 'connect').__enter__()
return self.sender.send((func, args, self.task_no, monitor))