Support for different flavors of ZeroMQ.

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()
 

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
-        )
+        )

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)
 

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 = []

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)