Merge pull request #41 from quantopian/cancellation2

Cancellation2

zipline/component.py

@@ -10,6 +10,7 @@ import uuid
 import time
 import socket
 import gevent
+import traceback
 import humanhash
 
 # pyzmq
@@ -305,11 +306,11 @@ class Component(object):
 
         self._exception = exc
         exc_type, exc_value, exc_traceback = sys.exc_info()
-        self.stack_trace = exc_traceback
+        trace = '\n>>>'.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
 
         exception_frame = CONTROL_FRAME(
             CONTROL_PROTOCOL.EXCEPTION,
-            str(exc)
+            trace
         )
         self.control_out.send(exception_frame)
 
@@ -492,7 +493,8 @@ class Component(object):
         """
         The descriptive name of the component.
         """
-        return 'UNKNOWN COMPONENT'
+        # Prevents the bug that Thomas ran into
+        raise NotImplementedError
 
     @property
     def get_type(self):

zipline/messaging.py

@@ -227,7 +227,7 @@ class Feed(Component):
 
             # -- Soft Kill --
             elif event == CONTROL_PROTOCOL.SHUTDOWN:
-                self.done()
+                self.signal_done()
                 self.shutdown()
 
             # -- Hard Kill --
@@ -496,7 +496,7 @@ class BaseTransform(Component):
 
             # -- Soft Kill --
             elif event == CONTROL_PROTOCOL.SHUTDOWN:
-                self.done()
+                self.signal_done()
                 self.shutdown()
 
             # -- Hard Kill --

zipline/monitor.py

@@ -3,13 +3,14 @@ import gevent
 import itertools
 # pyzmq
 import zmq
-# gevent_zeromq
 import gevent_zeromq
-# zmq_ctypes
-#import zmq_ctypes
+
+from collections import OrderedDict
 
 from protocol import CONTROL_PROTOCOL, CONTROL_FRAME, \
-    CONTROL_UNFRAME, CONTROL_STATES, INVALID_CONTROL_FRAME
+    CONTROL_UNFRAME, CONTROL_STATES, INVALID_CONTROL_FRAME \
+
+states = CONTROL_STATES
 
 from gpoll import _Poller as GeventPoller
 
@@ -103,6 +104,17 @@ from gpoll import _Poller as GeventPoller
 #  | 0 | | 0 | | 0 | | 0 |
 #  +---+ +---+ +---+ +---+
 
+INIT, SOURCES_READY, RUNNING, TERMINATE = CONTROL_STATES
+
+state_transitions = frozenset([
+    (-1            , INIT),
+    (INIT          , SOURCES_READY),
+    (SOURCES_READY , RUNNING),
+    (INIT          , TERMINATE),
+    (SOURCES_READY , TERMINATE),
+    (RUNNING       , TERMINATE),
+])
+
 class UnknownChatter(Exception):
     def __init__(self, name):
         self.named = name
@@ -149,27 +161,26 @@ class Controller(object):
 
     def __init__(self, pub_socket, route_socket, logging = None):
 
-        self.context = None
-        self.zmq = None
+        self.context    = None
+        self.zmq        = None
         self.zmq_poller = None
 
-        polling = False
-
-        self.polling = polling
-
+        self.running = False
+        self.polling = False
         self.tracked = set()
         self.responses = set()
 
-        self.ctime = 0
-        self.tic = time.time()
+        self.ctime    = 0
+        self.tic      = time.time()
         self.freeform = False
+        self._state   = -1
 
         self.associated = []
 
-        self.pub_socket = pub_socket
+        self.pub_socket   = pub_socket
         self.route_socket = route_socket
 
-        self.error_replay = {}
+        self.error_replay = OrderedDict()
 
         if logging:
             self.logging = logging
@@ -182,23 +193,23 @@ class Controller(object):
         assert self.zmq_flavor in ['thread', 'mp', 'green']
 
         if flavor == 'mp':
-            self.zmq = zmq
-            self.context = self.zmq.Context()
+            self.zmq        = zmq
+            self.context    = self.zmq.Context()
             self.zmq_poller = self.zmq.Poller
             return
         if flavor == 'thread':
-            self.zmq = zmq
-            self.context = self.zmq.Context.instance()
+            self.zmq        = zmq
+            self.context    = self.zmq.Context.instance()
             self.zmq_poller = self.zmq.Poller
             return
         if flavor == 'green':
-            self.zmq = gevent_zeromq.zmq
-            self.context = self.zmq.Context.instance()
+            self.zmq        = gevent_zeromq.zmq
+            self.context    = self.zmq.Context.instance()
             self.zmq_poller = GeventPoller
             return
         if flavor == 'pypy':
-            self.zmq = zmq
-            self.context = self.zmq.Context.instance()
+            self.zmq        = zmq
+            self.context    = self.zmq.Context.instance()
             self.zmq_poller = self.zmq.Poller
             return
 
@@ -217,19 +228,24 @@ class Controller(object):
             self.freeform = False
             self.topology = frozenset(topology)
 
-        default_states = [
-            CONTROL_STATES.RUNNING,
-            CONTROL_STATES.SHUTDOWN,
-            CONTROL_STATES.TERMINATE,
-        ]
-
-        self.states = states or default_states
         self.polling = True
+        self.state = CONTROL_STATES.INIT
 
-        # Start off in RUNNING, state
-        self.state = self.states[0]
+    @property
+    def state(self):
+        return self._state
+
+    @state.setter
+    def state(self, new):
+        old, self._state = self._state, new
+
+        if (old, new) not in state_transitions:
+            raise RuntimeError("[Controller] Invalid State Transition : %s -> %s" %(old, new))
+        else:
+            self.logging.info("[Controller] State Transition : %s -> %s" %(old, new))
 
     def run(self):
+        self.running = True
         self.init_zmq(self.zmq_flavor)
 
         try:
@@ -256,6 +272,9 @@ class Controller(object):
     # -------------
 
     def send_heart(self):
+        if not self.running:
+            return
+
         heartbeat_frame = CONTROL_FRAME(
             CONTROL_PROTOCOL.HEARTBEAT,
             str(self.ctime)
@@ -263,6 +282,9 @@ class Controller(object):
         self.pub.send(heartbeat_frame)
 
     def send_hardkill(self):
+        if not self.running:
+            return
+
         kill_frame = CONTROL_FRAME(
             CONTROL_PROTOCOL.KILL,
             ''
@@ -270,6 +292,9 @@ class Controller(object):
         self.pub.send(kill_frame)
 
     def send_softkill(self):
+        if not self.running:
+            return
+
         soft_frame = CONTROL_FRAME(
             CONTROL_PROTOCOL.SHUTDOWN,
             ''
@@ -282,16 +307,35 @@ class Controller(object):
 
     def _poll(self):
 
+        assert self.route_socket
+        assert self.pub_socket
+        assert self.cancel_socket
+
+        # -- Publish --
+        # =============
         self.pub = self.context.socket(self.zmq.PUB)
         self.pub.bind(self.pub_socket)
 
+        # -- Cancel --
+        # =============
+        assert isinstance(self.cancel_socket,basestring), self.cancel_socket
+        self.cancel = self.context.socket(self.zmq.REP)
+        self.cancel.connect(self.cancel_socket)
+
+        # -- Router --
+        # =============
         self.router = self.context.socket(self.zmq.ROUTER)
         self.router.bind(self.route_socket)
 
-        self.associated.extend([self.pub, self.router])
 
         poller = self.zmq.Poller()
         poller.register(self.router, self.zmq.POLLIN)
+        poller.register(self.cancel, self.zmq.POLLIN)
+
+        self.associated += [self.pub, self.router, self.cancel]
+
+        # TODO: actually do this
+        self.state = CONTROL_STATES.SOURCES_READY
 
         buffer = []
 
@@ -325,11 +369,20 @@ class Controller(object):
                         self.logging.error('Invalid frame', rawmessage)
                         pass
 
+                if self.cancel in socks and socks[self.cancel] == self.zmq.POLLIN:
+                    self.logging.info('[Controller] Received Cancellation')
+                    rawmessage = self.cancel.recv()
+                    self.shutdown(soft=True)
+                    break
+
             self.beat()
 
             if self.zmq_flavor == 'green':
                 gevent.sleep(0)
 
+            if self.state is CONTROL_STATES.TERMINATE:
+                break
+
             if not self.polling:
                 break
 
@@ -359,35 +412,89 @@ class Controller(object):
         for component in bad:
             self.fail(component)
 
-    # ------------------
-    # Component Handlers
-    # ------------------
+    # --------------
+    # Init Handlers
+    # --------------
+
+    def new_source(self):
+        if self.state is CONTROL_STATES.RUNNING:
+            self.state = SOURCES_READY
+
+    def new_universal(self):
+        pass
 
     # The various "states of being that a component can inform us
     # of
     def new(self, component):
-        self.logging.info('[Controller] Alive "%s" ' % component)
+        if self.state is CONTROL_STATES.TERMINATE:
+            return
+
+        self.logging.info('[Controller] Now Tracking "%s" ' % component)
+
+        universal = self.new_universal
+        init_handlers = {
+            'FEED' : self.new_source,
+        }
 
         if component in self.topology or self.freeform:
+            init_handlers.get(component, universal)()
             self.tracked.add(component)
         else:
             # Some sort of socket collision has occured, this is
             # a very bad failure mode.
             raise UnknownChatter(component)
 
+    # ------------------
+    # Epic Fail Handling
+    # ------------------
+
+    def fail_universal(self):
+        pass
+        # TODO: this requires higher order functionality
+        #self.logging.error('[Controller] System in exception state, shutting down')
+        #self.shutdown(soft=True)
+
     def fail(self, component):
-        self.logging.info('[Controller] Component "%s" timed out' % component)
-        self.tracked.remove(component)
+        if self.state is CONTROL_STATES.TERMINATE:
+            return
+
+        universal = self.fail_universal
+        fail_handlers = { }
+
+        if component in self.topology or self.freeform:
+            self.logging.info('[Controller] Component "%s" timed out' % component)
+            self.tracked.remove(component)
+            fail_handlers.get(component, universal)()
+
+    # -------------------
+    # Completion Handling
+    # -------------------
 
     def done(self, component):
-        # TODO: This will be what we ship off to vbench at some
-        # point...
-        # print component finished at self.ctime
         self.logging.info('[Controller] Component "%s" done.' % component)
 
+    # --------------
+    # Error Handling
+    # --------------
+
+    def exception_universal(self):
+        """
+        Shutdown the system on failure.
+        """
+        self.logging.error('[Controller] System in exception state, shutting down')
+        self.shutdown(soft=True)
+
     def exception(self, component, failure):
-        self.error_replay[time.time()] = failure
-        self.logging.error('Component "%s" in exception state' % component)
+        universal = self.exception_universal
+        exception_handlers = { }
+
+        if component in self.topology or self.freeform:
+            self.error_replay[(component, time.time())] = failure
+            self.logging.error('[Controller] Component "%s" in exception state' % component)
+
+            exception_handlers.get(component, universal)()
+        else:
+            raise UnknownChatter(component)
 
     # -----------------
     # Protocol Handling
@@ -462,6 +569,11 @@ class Controller(object):
         self.associated.append(s)
         return s
 
+    def do_error_replay(self):
+        for (component, time), error in self.error_replay.iteritems():
+            self.logging.info('[Controller] Error Log for -- %s --:\n%s' %
+                (component, error))
+
     def shutdown(self, hard=False, soft=True, context=None):
 
         if not self.polling:
@@ -472,7 +584,7 @@ class Controller(object):
         assert hard or soft, """ Must specify kill hard or soft """
 
         if hard:
-            self.state = CONTROL_STATES.SHUTDOWN
+            self.state = CONTROL_STATES.TERMINATE
 
             self.logging.info('[Controller] Hard Shutdown')
 
@@ -488,18 +600,22 @@ class Controller(object):
             #for asoc in self.associated:
                 #asoc.close()
 
+        self.do_error_replay()
+
 if __name__ == '__main__':
 
-    print 'Running on ',\
-        'tcp://127.0.0.1:5000', \
-        'tcp://127.0.0.1:5001',
+    print 'Running on '\
+        'tcp://127.0.0.1:5000 '\
+        'tcp://127.0.0.1:5001 '
 
     controller = Controller(
         'tcp://127.0.0.1:5000',
         'tcp://127.0.0.1:5001',
     )
+    controller.zmq_flavor = 'green'
+
     controller.manage(
         'freeform',
         []
     )
-    controller.run('green')
+    controller.run()

zipline/protocol.py

@@ -126,9 +126,6 @@ from collections import namedtuple
 from protocol_utils import Enum, FrameExceptionFactory, namedict
 from date_utils import EPOCH, UN_EPOCH
 
-#import ujson
-#import ultrajson_numpy
-
 # -----------------------
 # Control Protocol
 # -----------------------
@@ -136,9 +133,10 @@ from date_utils import EPOCH, UN_EPOCH
 INVALID_CONTROL_FRAME = FrameExceptionFactory('CONTROL')
 
 CONTROL_STATES = Enum(
+    'INIT',
+    'SOURCES_READY',
     'RUNNING',
-    'SHUTDOWN',  # a soft kill
-    'TERMINATE', # a hard kill
+    'TERMINATE',
 )
 
 CONTROL_PROTOCOL = Enum(
@@ -149,6 +147,7 @@ CONTROL_PROTOCOL = Enum(
     'OK'        , # 3 - rep
     'DONE'      , # 4 - rep
     'EXCEPTION' , # 5 - rep
+    'SIGNAL'    , # 6 - rep
 )
 
 def CONTROL_FRAME(event, payload):
@@ -174,18 +173,6 @@ def CONTROL_UNFRAME(msg):
     except ValueError:
         raise INVALID_CONTROL_FRAME(msg)
 
-# -----------------------
-# Heartbeat Protocol
-# -----------------------
-
-# These encode the msgpack equivelant of 1 and 2. The heartbeat
-# frame should only be 1 byte on the wire.
-
-HEARTBEAT_PROTOCOL = namedict({
-    'REQ' : b'\x01',
-    'REP' : b'\x02',
-})
-
 # -----------------------
 # Component State
 # -----------------------

zipline/protocol_utils.py

@@ -12,6 +12,7 @@ def Enum(*options):
     """
     class cstruct(Structure):
         _fields_ = [(o, c_ubyte) for o in options]
+        __iter__ = lambda s: iter(range(len(options)))
     return cstruct(*range(len(options)))
 
 def FrameExceptionFactory(name):

zipline/transitions.py

@@ -0,0 +1,87 @@
+import types
+from collections import Container, Hashable, Callable
+
+class Any(object): pass
+
+class Workflow(Container, Callable):
+
+    def __init__(self, states, transitions, initial_state):
+        self.simple = set()
+        self.complx = []
+
+        if isinstance(states[0], tuple):
+            self.groups = {b for _,b in states}
+        else:
+            self.groups = set()
+
+        matcher = lambda b: lambda f,t : t == b
+
+        for (a, b) in transitions.itervalues():
+            if a is Any:
+                self.complx.append(matcher(b))
+            if isinstance(a, Hashable) and isinstance(b, Hashable):
+                self.simple.add((a,b))
+
+    def __call__(self, **kwargs):
+        if 'group' in kwargs:
+            return self.groups
+
+    def __contains__(self, state):
+        if state in self.simple:
+            return True
+        for match in self.complx:
+            if match(*state):
+                return True
+        else:
+            return False
+
+class Flowable:
+
+    @property
+    def state(self):
+        if not hasattr(self, '_state'):
+            self._state = self.initial_state
+        else:
+            return self._state
+
+    @state.setter
+    def state(self, new):
+        if not hasattr(self, '_state'):
+            self._state = self.initial_state
+
+        old = self._state
+
+        if (old, new) in self.workflow:
+            self._state = new
+        else:
+            raise RuntimeError("Invalid State Transition : %s -> %s" %(old, new))
+
+class WorkflowMeta(type):
+    """
+    Base metaclass component workflows.
+    """
+
+    def __new__(cls, name, mro, attrs):
+
+        state         = attrs.get('states', None)
+        transitions   = attrs.get('transitions', None)
+        initial_state = attrs.get('initial_state', None)
+
+        if attrs.get('workflow'):
+            raise RuntimeError('`workflow` is a reserved attribute.')
+
+        if not state:
+            raise RuntimeError('Must specify states')
+
+        if not transitions:
+            raise RuntimeError('Must specify transitions')
+
+        if not transitions:
+            raise RuntimeError('Must specify initial_state')
+
+        new_class = super(WorkflowMeta, cls).__new__(
+            cls, name, mro+(Flowable,), attrs
+        )
+        new_class.workflow = Workflow(state, transitions, initial_state)
+
+        return new_class