diff --git a/zipline/component.py b/zipline/component.py index 119138e3..d33ae44b 100644 --- a/zipline/component.py +++ b/zipline/component.py @@ -9,8 +9,16 @@ import sys import uuid import time import socket +import gevent import humanhash +# pyzmq +import zmq +# gevent_zeromq +import gevent_zeromq +# zmq_ctypes +#import zmq_ctypes + from datetime import datetime import zipline.util as qutil @@ -61,8 +69,8 @@ class Component(object): self.zmq = None self.context = None self.addresses = None + self.out_socket = None - self.gevent_needed = False self.killed = False self.controller = None self.heartbeat_timeout = 2000 @@ -126,27 +134,51 @@ class Component(object): def do_work(self): raise NotImplementedError + def init_zmq(self, flavor): + """ + ZMQ in all flavors. Have it your way. + + mp - Distinct contexts | pyzmq + thread - Same context | pyzmq + green - Same context | gevent_zeromq + pypy - Same context | zmq_ctypes + + """ + + if flavor == 'mp': + self.zmq = zmq + self.context = self.zmq.Context() + return + if flavor == 'thread': + self.zmq = zmq + self.context = self.zmq.Context.instance() + return + if flavor == 'green': + self.zmq = gevent_zeromq.zmq + self.context = self.zmq.Context.instance() + return + if flavor == 'pypy': + self.zmq = zmq + self.context = self.zmq.Context.instance() + return + + import pdb; pdb.set_trace() + raise Exception("Unknown ZeroMQ Flavor") + def _run(self): self.start_tic = time.time() self.done = False # TODO: use state flag self.sockets = [] - if self.gevent_needed: - qutil.LOGGER.info("Loading gevent specific zmq for {id}".format(id=self.get_id)) - import gevent_zeromq - self.zmq = gevent_zeromq.zmq - else: - import zmq - self.zmq = zmq - - self.context = self.zmq.Context.instance() + self.init_zmq(self.zmq_flavor) self.setup_poller() self.open() self.setup_sync() self.setup_control() + self.loop() self.shutdown() diff --git a/zipline/lines.py b/zipline/lines.py index 211a4f39..50f3b787 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -309,7 +309,8 @@ class SimulatedTrading(object): self.sim_context.join() def shutdown(self): - self.allocator.reaquire(*self.leased_sockets) + pass + #self.allocator.reaquire(*self.leased_sockets) #-------------------------------- # Component property accessors @@ -333,4 +334,4 @@ class ZiplineException(Exception): return "Unexpected exception {line}: {msg}".format( line=self.name, msg=self.message - ) \ No newline at end of file + ) diff --git a/zipline/messaging.py b/zipline/messaging.py index 12724023..1b7ffbdf 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -16,15 +16,17 @@ class ComponentHost(Component): start, and then wait for all components to be finished. """ - def __init__(self, addresses, gevent_needed=False): + def __init__(self, addresses): Component.__init__(self) self.addresses = addresses - self.gevent_needed = gevent_needed self.running = False self.init() def init(self): + assert hasattr(self, 'zmq_flavor'), \ + """ You must specify a flavor of ZeroMQ for all + ComponentHost subclasses. """ # Component Registry, keyed by get_id # ---------------------- @@ -67,10 +69,12 @@ class ComponentHost(Component): assert isinstance(components, list) for component in components: - component.gevent_needed = self.gevent_needed component.addresses = self.addresses component.controller = self.controller + # Hosts share their zmq flavor with hosted components + component.zmq_flavor = self.zmq_flavor + self._components[component.guid] = component self.components[component.get_id] = component self.sync_register[component.get_id] = datetime.datetime.utcnow() @@ -97,13 +101,8 @@ class ComponentHost(Component): self.sync_socket = self.context.socket(self.zmq.REP) self.sync_socket.bind(self.addresses['sync_address']) - # There is a namespace collision between three classes - # which use the self.poller property to mean different - # things. - # ===================================================== self.sync_poller = self.zmq.Poller() self.sync_poller.register(self.sync_socket, self.zmq.POLLIN) - # ===================================================== self.sockets.append(self.sync_socket) diff --git a/zipline/simulator.py b/zipline/simulator.py index 99185c7c..8d6480b5 100644 --- a/zipline/simulator.py +++ b/zipline/simulator.py @@ -26,9 +26,11 @@ class AddressAllocator(object): def reaquire(self, *conn): pass -# + class Simulator(ComponentHost): + zmq_flavor = 'thread' + def __init__(self, addresses): ComponentHost.__init__(self, addresses) self.subthreads = [] diff --git a/zipline/zmq_utils.py b/zipline/zmq_utils.py index e13a5e34..eeb04411 100644 --- a/zipline/zmq_utils.py +++ b/zipline/zmq_utils.py @@ -1,7 +1,94 @@ +""" +Misc ZeroMQ utilities. +""" import gevent from gevent_zeromq import zmq +from contextlib import closing + +class ZmqTimeout(object): + def __init__(self, socket): + self._socket = socket + + def __enter__(self): + pass + + def __exit__(self): + self._socket.close() + +class ZmqDone(object): + + def __init__(self, socket, frame): + self.ident = socket.identity + self.frame = str(frame) + + def __str__(self): + return 'Socket ( %s ) finished with frame ( %s )' % \ + ( self.ident, self.frame ) + +class zs(object): + """ + A wrapper for the *very* common pattern of reading from a + upstream socket until you get a DONE or EXCEPTION frame. + + # Eliminates all the boilerplate serialization logic + # and error handling cases into 3 lines. + + halts = (ERROR_FRAME, CLOSE_FRAME) + stream = zs(socket, halts) + + stream.on_error(YouFailAtFailing) + + for msg in stream: + print msg + + """ + + def __init__(self, socket, halts, srl=msgpack): + self._socket = socket + self.exc_case = halts[0] + self.done_case = halts[1] + + self.loads = srl.loads + self.halt_method = 'exception' + self.exception = ZmqDone + self.function = None + + def __iter__(self): + self.last = msg = self.loads(self._socket.recv()) + + if msg == self.exc_case: + return self.halt() + + if msg == self.done_case: + raise StopIteration + + yield msg + + def last(self): + return self.last + + def halt(self): + if self.halt_method == 'exception': + raise self.exception + elif self.halt_method == 'function': + return self.function() + + def on_error(self, callee): + + if isinstance(callee, Exception): + self.halt_method = 'exception' + self.exception = callee + else: + self.halt_method = 'function' + self.function = callee + def ZmqConsole(sock_typ, socket_addr, sock_conn=None, context=None): + """ + A utility to drop into a ZeroMQ pdb console and inspect + messages as they come through. If you just want to pipe to + stdout, don't use this. + """ context = context or zmq.Context.instance() socket = context.socket(zmq.PULL)