def __init__(self, context, options): self._instances.append(self) self.context = context self.options = options self.mongodb = AsyncConnection(options.mongo_uri) # External sockets src_sock = context.socket(zmq.PUSH) sink_sock = context.socket(zmq.PULL) self._cli_sock = context.socket(zmq.XREP) self._cli_sock.bind(options.zmr['req_bind']) # Load some options, bind some sockets self._uri = options.zmr['worker_uri'] self._src_port = options.zmr.get('src_port', None) self._sink_port = options.zmr.get('sink_port', None) if self._src_port: src_sock.bind('%s:%s' % (self._uri, self._src_port)) else: self._src_port = src_sock.bind_to_random_port(self._uri) if self._sink_port: sink_sock.bind('%s:%s' % (self._uri, self._sink_port)) else: self._sink_port = sink_sock.bind_to_random_port(self._uri) #Job Manager self.job_manager = JobManager( src_sock, sink_sock, options.zmr['outstanding_chunks'], options.zmr['max_chunk_timeout']) # Internal request handling sockets self._req_worker_sock = context.socket(zmq.XREQ) self._req_worker_sock.bind('inproc://req-workers') self._terminate = False
class ZMRRouter(object): _instances = [] def __init__(self, context, options): self._instances.append(self) self.context = context self.options = options self.mongodb = AsyncConnection(options.mongo_uri) # External sockets src_sock = context.socket(zmq.PUSH) sink_sock = context.socket(zmq.PULL) self._cli_sock = context.socket(zmq.XREP) self._cli_sock.bind(options.zmr['req_bind']) # Load some options, bind some sockets self._uri = options.zmr['worker_uri'] self._src_port = options.zmr.get('src_port', None) self._sink_port = options.zmr.get('sink_port', None) if self._src_port: src_sock.bind('%s:%s' % (self._uri, self._src_port)) else: self._src_port = src_sock.bind_to_random_port(self._uri) if self._sink_port: sink_sock.bind('%s:%s' % (self._uri, self._sink_port)) else: self._sink_port = sink_sock.bind_to_random_port(self._uri) #Job Manager self.job_manager = JobManager( src_sock, sink_sock, options.zmr['outstanding_chunks'], options.zmr['max_chunk_timeout']) # Internal request handling sockets self._req_worker_sock = context.socket(zmq.XREQ) self._req_worker_sock.bind('inproc://req-workers') self._terminate = False def start(self): # gevent.spawn(self.sink_handler) for i in xrange(self.options.zmr['request_greenlets']): gevent.spawn(self.request_handler, i, self.context, 'inproc://req-workers') util.gevent_queue(self._cli_sock, self._req_worker_sock) def serve_forever(self): while True: if self._terminate: break gevent.sleep(1) 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()