Alter init methods, poll for control messages.

etc/requirements.txt

@@ -3,3 +3,4 @@ pyzmq==2.1.11
 gevent-zeromq==0.2.2
 msgpack-python==0.1.12
 humanhash==0.0.1
+ujson=1.18

zipline/component.py

@@ -38,6 +38,7 @@ class Component(object):
         self.out_socket        = None
         self.gevent_needed     = False
         self.killed            = False
+        self.controller        = None
         self.heartbeat_timeout = 2000
         self.state_flag        = COMPONENT_STATE.OK # OK | DONE | EXCEPTION
 
@@ -46,6 +47,15 @@ class Component(object):
         self.guid = uuid.uuid4()
         self.huid = humanhash.humanize(self.guid.hex)
 
+        self.init()
+
+    def init(self):
+        """
+        Subclasses should override this to extend the setup for
+        the class. Shouldn't have side effects.
+        """
+        pass
+
     # ------------
     # Core Methods
     # ------------
@@ -57,7 +67,9 @@ class Component(object):
         """
         Tear down after normal operation.
         """
-        raise NotImplementedError
+        #close all the sockets
+        for sock in self.sockets:
+            sock.close()
 
     def kill(self):
         """
@@ -89,9 +101,7 @@ class Component(object):
         self.setup_control()
         self.loop()
 
-        #close all the sockets
-        for sock in self.sockets:
-            sock.close()
+        self.destroy()
 
     def run(self, catch_exceptions=False):
         """
@@ -258,8 +268,8 @@ class Component(object):
         """
         assert self.controller
 
-        self.control_out = self.controller.message_sender()
-        self.control_in = self.controller.message_listener()
+        self.control_out = self.controller.message_sender(context=self.context)
+        self.control_in = self.controller.message_listener(context=self.context)
 
         self.poll.register(self.control_in, self.zmq.POLLIN)
         self.sockets.extend([self.control_in, self.control_out])
@@ -275,7 +285,7 @@ class Component(object):
         self.sync_socket.connect(self.addresses['sync_address'])
         #self.sync_socket.setsockopt(self.zmq.LINGER,0)
 
-        # Explictly, a different poller for obvious reasons.
+        # Explictly a different poller for obvious reasons.
         # I'm not fond of having this poller init'd as a side
         # effect of a method call. Still thinking about where to
         # put it at the moment though...
@@ -305,6 +315,12 @@ class Component(object):
             self.sockets         ,
         )
 
+    def __len__(self):
+        """
+        Some components overload this for debug purposes
+        """
+        raise NotImplementedError
+
     def __repr__(self):
         """
         Return a usefull string representation of the component

zipline/messaging.py

@@ -1,12 +1,11 @@
 """
 Commonly used messaging components.
 """
-import json
-import uuid
 import datetime
+import ujson as json
+
 import zipline.util as qutil
 from zipline.component import Component
-
 from zipline.protocol import CONTROL_PROTOCOL
 
 class ComponentHost(Component):
@@ -18,6 +17,11 @@ class ComponentHost(Component):
     def __init__(self, addresses, gevent_needed=False):
         Component.__init__(self)
         self.addresses = addresses
+        self.gevent_needed  = gevent_needed
+
+        self.init()
+
+    def init(self):
 
         # workaround for defect in threaded use of strptime:
         # http://bugs.python.org/issue11108
@@ -26,13 +30,10 @@ class ComponentHost(Component):
         self.components     = {}
         self.sync_register  = {}
         self.timeout        = datetime.timedelta(seconds=5)
-        self.gevent_needed  = gevent_needed
-        self.heartbeat_timeout = 2000
 
         self.feed           = ParallelBuffer()
         self.merge          = MergedParallelBuffer()
         self.passthrough    = PassthroughTransform()
-        self.controller     = None
 
         #register the feed and the merge
         self.register_components([self.feed, self.merge, self.passthrough])
@@ -112,13 +113,13 @@ class ComponentHost(Component):
             if self.sync_socket in socks and socks[self.sync_socket] == self.zmq.POLLIN:
                 msg = self.sync_socket.recv()
                 parts = msg.split(':')
+                sync_id, status = parts
 
+                # TODO: move into frame protocol
                 if len(parts) != 2:
                     qutil.LOGGER.info("got bad confirm: {msg}".format(msg=msg))
                     continue
 
-                sync_id, status = parts
-
                 if status == str(CONTROL_PROTOCOL.DONE): # TODO: other way around
                     qutil.LOGGER.info("{id} is DONE".format(id=sync_id))
                     self.unregister_component(sync_id)
@@ -152,6 +153,8 @@ class ParallelBuffer(Component):
         self.data_buffer            = {}
         self.ds_finished_counter    = 0
 
+    def init(self):
+        pass
 
     @property
     def get_id(self):
@@ -168,12 +171,17 @@ class ParallelBuffer(Component):
         # wait for synchronization reply from the host
         socks = dict(self.poll.poll(self.heartbeat_timeout)) #timeout after 2 seconds.
 
+        if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
+            msg = self.control_in.recv()
+
         if self.pull_socket in socks and socks[self.pull_socket] == self.zmq.POLLIN:
             message = self.pull_socket.recv()
+
             if message == str(CONTROL_PROTOCOL.DONE):
                 self.ds_finished_counter += 1
+
                 if len(self.data_buffer) == self.ds_finished_counter:
-                     #drain any remaining messages in the buffer
+                    #drain any remaining messages in the buffer
                     self.drain()
                     self.signal_done()
             else:
@@ -181,13 +189,6 @@ class ParallelBuffer(Component):
                 self.append(event[u'id'], event)
                 self.send_next()
 
-    def __len__(self):
-        """
-        Buffer's length is same as internal map holding separate
-        sorted arrays of events keyed by source id.
-        """
-        return len(self.data_buffer)
-
     def append(self, source_id, value):
         """
         Add an event to the buffer for the source specified by
@@ -255,6 +256,13 @@ class ParallelBuffer(Component):
             self.feed_socket.send(json.dumps(event), self.zmq.NOBLOCK)
             self.sent_count += 1
 
+    def __len__(self):
+        """
+        Buffer's length is same as internal map holding separate
+        sorted arrays of events keyed by source id.
+        """
+        return len(self.data_buffer)
+
 
 class MergedParallelBuffer(ParallelBuffer):
     """
@@ -264,6 +272,15 @@ class MergedParallelBuffer(ParallelBuffer):
     def __init__(self):
         ParallelBuffer.__init__(self)
 
+        self.init()
+
+    def init(self):
+        pass
+
+    @property
+    def get_id(self):
+        return "MERGE"
+
     def open(self):
         self.pull_socket = self.bind_merge()
         self.feed_socket = self.bind_result()
@@ -283,28 +300,31 @@ class MergedParallelBuffer(ParallelBuffer):
                 result[source] = cur['value']
         return result
 
-    @property
-    def get_id(self):
-        return "MERGE"
-
 
 class BaseTransform(Component):
-    """Top level execution entry point for the transform::
+    """
+    Top level execution entry point for the transform::
 
-            - connects to the feed socket to subscribe to events
-            - connets to the result socket (most oftened bound by a TransformsMerge) to PUSH transforms
-            - processes all messages received from feed, until DONE message received
-            - pushes all transforms
-            - sends DONE to result socket, closes all sockets and context
+        - connects to the feed socket to subscribe to events
+        - connets to the result socket (most oftened bound by a TransformsMerge) to PUSH transforms
+        - processes all messages received from feed, until DONE message received
+        - pushes all transforms
+        - sends DONE to result socket, closes all sockets and context
 
     Parent class for feed transforms. Subclass and override transform
-    method to create a new derived value from the combined feed."""
+    method to create a new derived value from the combined feed.
+    """
 
     def __init__(self, name):
         Component.__init__(self)
-        self.state         = {}
+        self.state = {}
         self.state['name'] = name
 
+        self.init()
+
+    def init(self):
+        pass
+
     @property
     def get_id(self):
         return self.state['name']
@@ -325,7 +345,11 @@ class BaseTransform(Component):
             - call transform (subclass' method) on event
             - send the transformed event
         """
-        socks = dict(self.poll.poll(self.heartbeat_timeout)) #timeout after 2 seconds.
+        socks = dict(self.poll.poll(self.heartbeat_timeout))
+
+        if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
+            msg = self.control_in.recv()
+
         if self.feed_socket in socks and socks[self.feed_socket] == self.zmq.POLLIN:
             message = self.feed_socket.recv()
             if message == str(CONTROL_PROTOCOL.DONE):
@@ -360,6 +384,9 @@ class PassthroughTransform(BaseTransform):
     def __init__(self):
         BaseTransform.__init__(self, "PASSTHROUGH")
 
+    def init(self):
+        pass
+
     def transform(self, event):
         return {'value':event}
 
@@ -373,6 +400,9 @@ class DataSource(Component):
     def __init__(self, source_id):
         Component.__init__(self)
         self.id        = source_id
+        self.init()
+
+    def init(self):
         self.cur_event = None
 
     @property

zipline/test/client.py

@@ -1,18 +1,22 @@
-import json
+import ujson as json
+
 import zipline.util as qutil
 import zipline.messaging as qmsg
-
 from zipline.protocol import CONTROL_PROTOCOL
 
 class TestClient(qmsg.Component):
 
     def __init__(self, utest, expected_msg_count=0):
         qmsg.Component.__init__(self)
-        self.received_count     = 0
-        self.expected_msg_count = expected_msg_count
+
         self.utest              = utest
+        self.expected_msg_count = expected_msg_count
+
+        self.init()
+
+    def init(self):
+        self.received_count     = 0
         self.prev_dt            = None
-        self.heartbeat_timeout = 2000
 
     @property
     def get_id(self):
@@ -24,6 +28,9 @@ class TestClient(qmsg.Component):
     def do_work(self):
         socks = dict(self.poll.poll(self.heartbeat_timeout))
 
+        if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
+            msg = self.control_in.recv()
+
         if self.data_feed in socks and socks[self.data_feed] == self.zmq.POLLIN:   
             msg = self.data_feed.recv()
 
@@ -37,10 +44,12 @@ class TestClient(qmsg.Component):
 
             self.received_count += 1
             event = json.loads(msg)
+
             if(self.prev_dt != None):
-                if(not event['dt'] >= self.prev_dt):
+                if not event['dt'] >= self.prev_dt:
                     raise Exception("Message out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt))
 
             self.prev_dt = event['dt']
+
             if(self.received_count % 100 == 0):
                 qutil.LOGGER.info("received {n} messages".format(n=self.received_count))

zipline/transforms/technical.py

@@ -5,6 +5,7 @@ TODO: add trailing stop
 """
 
 import datetime
+
 from zipline.messaging import BaseTransform
 import zipline.util as qutil
 
@@ -16,18 +17,25 @@ class MovingAverage(BaseTransform):
 
     def __init__(self, name, days):
         BaseTransform.__init__(self, name)
+        self.window = datetime.timedelta(days = days)
+
+        self.init()
+
+    def init(self):
         self.events         = []
         self.current_total  = 0
-        self.window         = datetime.timedelta(days = days)
 
     def transform(self, event):
-        """Update the moving average with the latest data point."""
+        """
+        Update the moving average with the latest data point.
+        """
 
         self.events.append(event)
         self.current_total += event['price']
         event_date = qutil.parse_date(event['dt'])
 
         index = 0
+
         for cur_event in self.events:
             cur_date = qutil.parse_date(cur_event['dt'])
             if(cur_date - event_date):
@@ -37,11 +45,10 @@ class MovingAverage(BaseTransform):
             else:
                 break
 
-        if(len(self.events) == 0):
+        if len(self.events) == 0:
             return 0.0
 
         self.average = self.current_total/len(self.events)
-
         self.state['value'] = self.average
-        return self.state
 
+        return self.state