def start(self, job):
q = None
for p in self.parts:
p.input = q
p.output = q = Queue(1000)
for p in self.parts:
p._greenlet = gevent.spawn_link_exception(util.ending_request(job._mongodb), p.run)
def run(self):
for cname, values in self.iter:
ci = MRCommitter(
self._job, self._output_type, cname,
self._job.reducefun(cname))
ci.input = Queue(1000)
ci._greenlet = gevent.spawn_link_exception(util.ending_request(self._job._mongodb), ci.run_wrapper)
for v in values:
ci.input.put(v[1])
ci.input.put(StopIteration)
ci._greenlet.join()
def request_handler(self, i, context, uri):
sock = context.socket(zmq.REP)
sock.connect(uri)
while True:
try:
obj = util.recv_bson(sock)
log.debug('Request %s in %s', obj, i)
command = obj.get('$command')
log.debug('Req %s in %s', command, i)
if command == 'echo':
util.send_bson(sock, dict(message=obj))
elif command == 'terminate':
self._terminate = True
elif command in ('mapreduce', 'xmapreduce', 'basic'):
job = Job.from_request(self, obj)
self.job_manager.alloc_job(job)
log.debug('Start job %s', job.id)
if obj.get('wait'):
job.run()
self.job_manager.dealloc_job(job.id)
else:
gevent.spawn(util.ending_request(self.mongodb), job.run)
util.send_bson(sock, dict(job_id=job.id, status=job.status))
continue
elif command == 'status':
try:
jid = obj['job_id']
status = self.job_manager.job_status(jid)
util.send_bson(sock, dict(job_id=jid, status=status))
if status['state'] in ('complete', 'error'):
self.job_manager.dealloc_job(jid)
except:
log.exception('Error getting status for job')
util.send_bson(sock, dict(
job_id=obj['job_id'], status=dict(
state='UNKNOWN')))
continue
elif command == 'config':
util.send_bson(sock, dict(
self.options.zmr,
src_port=self._src_port,
sink_port=self._sink_port))
except:
log.exception('Error in request handler')
util.send_bson(sock, dict(error=traceback.format_exc()))
continue
finally:
self.mongodb.end_request()
def run(self, requests):
for req_num in xrange(requests):
connection = ZM.doc_session.db.connection
header = util.recv_bson(self._src)
log.debug('RECV %s', header.keys())
more = self._src.getsockopt(zmq.RCVMORE)
parts = []
while more:
if header['compress']:
parts.append(zlib.decompress(self._src.recv()))
else:
parts.append(self._src.recv())
more = self._src.getsockopt(zmq.RCVMORE)
with util.ending_request(connection):
self.handle_message(header, parts)
r = resource.getrusage(resource.RUSAGE_SELF)
if r.ru_maxrss > self._suicide_level:
log.error('Suicide (painlessly) worker at RSS %s / %s',
r.ru_maxrss, self._suicide_level)
break
log.info('Closing worker after %d requests', req_num)
self._src.close()
self._sink.close()
