From 6dbbcfaf6634de820c0c11a3903ec7f27bc67712 Mon Sep 17 00:00:00 2001 From: Stephen Diehl Date: Tue, 3 Jul 2012 07:20:58 -0400 Subject: [PATCH] Replay selective fsm2 changes to new branch. --- zipline/components/aggregator.py | 95 +++++++++++++++----------------- zipline/components/feed.py | 37 +++++++------ zipline/components/merge.py | 14 ++--- zipline/core/controlled.py | 69 +++++++++++++++++++++++ 4 files changed, 139 insertions(+), 76 deletions(-) create mode 100644 zipline/core/controlled.py diff --git a/zipline/components/aggregator.py b/zipline/components/aggregator.py index a7bd6c82..451fc28a 100644 --- a/zipline/components/aggregator.py +++ b/zipline/components/aggregator.py @@ -12,12 +12,12 @@ Abstract base class for Feed and Merge. import logbook import zipline.protocol as zp - from zipline.core.component import Component -from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE, \ - CONTROL_FRAME, CONTROL_UNFRAME -from zipline.utils.protocol_utils import Enum +from zipline.core.controlled import do_handle_control_events +from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE from zipline.transitions import WorkflowMeta +from zipline.utils.protocol_utils import Enum + log = logbook.Logger('Aggregate') @@ -34,6 +34,10 @@ AGGREGATE_TRANSITIONS = dict( do_drain = (READY , DRAINING) , ) +# ========= +# Component +# ========= + class Aggregate(Component): """ Abstract superclass to Merge & Feed. Acts on two sockets @@ -41,7 +45,7 @@ class Aggregate(Component): - pull_socket - feed_socket - Both use ``data_buffer`` for buffering. + Both use ``sources`` for buffering. Feed and Merge define these differently. """ @@ -53,6 +57,9 @@ class Aggregate(Component): def get_type(self): return COMPONENT_TYPE.CONDUIT + def add_source(self, source_id): + self.sources[source_id] = [] + # ------------- # Core Methods # ------------- @@ -61,39 +68,26 @@ class Aggregate(Component): # wait for synchronization reply from the host socks = dict(self.poll.poll(self.heartbeat_timeout)) - # TODO: Abstract this out, maybe on base component - if socks.get(self.control_in) == self.zmq.POLLIN: - msg = self.control_in.recv() - event, payload = CONTROL_UNFRAME(msg) - - # -- Heartbeat -- - if event == CONTROL_PROTOCOL.HEARTBEAT: - # Heart outgoing - heartbeat_frame = CONTROL_FRAME( - CONTROL_PROTOCOL.OK, - payload - ) - self.control_out.send(heartbeat_frame) - - # -- Soft Kill -- - elif event == CONTROL_PROTOCOL.SHUTDOWN: - self.signal_done() - self.shutdown() - - # -- Hard Kill -- - elif event == CONTROL_PROTOCOL.KILL: - self.kill() - + # ---------------- + # Control Dispatch + # ---------------- + #do_handle_control_events(self, socks) + # ------------- + # Work Dispatch + # ------------- if socks.get(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 + if len(self.sources) == self.ds_finished_counter: + # Drain any remaining messages in the buffer log.debug("Draining Feed") + + self.state = DRAINING + self.drain() self.signal_done() else: @@ -105,7 +99,15 @@ class Aggregate(Component): try: self.append(event) - self.send_next() + + if not (self.is_full() or self.draining): + event = self.next() + + if event: + self.send(event) + else: + pass + except zp.INVALID_DATASOURCE_FRAME as exc: # Invalid message return self.signal_exception(exc) @@ -118,30 +120,25 @@ class Aggregate(Component): """ Send all messages in the buffer. """ - self.state = DRAINING while self.pending_messages() > 0: - self.send_next() + event = self.next() + if event: + self.send(event) - def send_next(self): + def send(self, event): """ Send the (chronologically) next message in the buffer. """ - if not (self.is_full() or self.draining): - return - - event = self.next() - - if event: - self.feed_socket.send(self.frame(event), self.zmq.NOBLOCK) - self.sent_counters[event.source_id] += 1 - self.sent_count += 1 + self.feed_socket.send(self.frame(event), self.zmq.NOBLOCK) + self.sent_counters[event.source_id] += 1 + self.sent_count += 1 def is_full(self): """ Indicates whether the buffer has messages in buffer for all un-DONE, blocking sources. """ - for source_id, events in self.data_buffer.iteritems(): + for source_id, events in self.sources.iteritems(): if len(events) == 0: return False return True @@ -152,19 +149,13 @@ class Aggregate(Component): buffer. """ total = 0 - for events in self.data_buffer.itervalues(): + for events in self.sources.itervalues(): total += len(events) return total - def add_source(self, source_id): - """ - Add a data source to the buffer. - """ - self.data_buffer[source_id] = [] - 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) + return len(self.sources) diff --git a/zipline/components/feed.py b/zipline/components/feed.py index f51e515d..719a99a6 100644 --- a/zipline/components/feed.py +++ b/zipline/components/feed.py @@ -1,5 +1,5 @@ import logbook -from collections import Counter +from collections import defaultdict, Counter from zipline.components.aggregator import Aggregate, \ AGGREGATE_STATES, AGGREGATE_TRANSITIONS @@ -28,7 +28,7 @@ class Feed(Aggregate): self.received_count = 0 self.ds_finished_counter = 0 - self.data_buffer = {} + self.sources = defaultdict(list) # source_id -> integer count self.sent_counters = Counter() @@ -71,7 +71,7 @@ class Feed(Aggregate): Add an event to the buffer for the source specified by source_id. """ - self.data_buffer[event.source_id].append(event) + self.sources[event.source_id].append(event) self.recv_counters[event.source_id] += 1 self.received_count += 1 @@ -84,24 +84,27 @@ class Feed(Aggregate): if not(self.is_full() or self.draining): return - cur_source = None earliest_source = None earliest_event = None - #iterate over the queues of events from all sources - #(1 queue per datasource) - for events in self.data_buffer.itervalues(): - if len(events) == 0: - continue - cur_source = events - first_in_list = events[0] - if first_in_list.dt == None: - #this is a filler event, discard - events.pop(0) + # iterate over the queues of source from all sources + # (1 queue per datasource) + + for source in self.sources.itervalues(): + if len(source) == 0: continue - if (earliest_event == None) or (first_in_list.dt <= earliest_event.dt): - earliest_event = first_in_list - earliest_source = cur_source + head = source[0] + + if head.dt == None: + #this is a filler event, discard + source.pop(0) + continue + + if (earliest_event == None) or (head.dt <= earliest_event.dt): + earliest_event = head + earliest_source = source if earliest_event != None: return earliest_source.pop(0) + else: + return False diff --git a/zipline/components/merge.py b/zipline/components/merge.py index bbb9dc6a..6dc7661e 100644 --- a/zipline/components/merge.py +++ b/zipline/components/merge.py @@ -2,7 +2,7 @@ import zipline.protocol as zp from zipline.components.aggregator import Aggregate, \ AGGREGATE_STATES, AGGREGATE_TRANSITIONS -from collections import Counter +from collections import defaultdict, Counter class Merge(Aggregate): """ @@ -19,9 +19,7 @@ class Merge(Aggregate): self.draining = False self.ds_finished_counter = 0 - # Depending on the size of this, might want to use a data - # structure with better asymptotics. - self.data_buffer = {} + self.sources = defaultdict(list) # source_id -> integer count self.sent_counters = Counter() @@ -61,7 +59,7 @@ class Merge(Aggregate): source_id. """ - self.data_buffer[event.keys()[0]].append(event) + self.sources[event.keys()[0]].append(event) self.received_count += 1 def next(self): @@ -73,8 +71,10 @@ class Merge(Aggregate): return #get the raw event from the passthrough transform. - result = self.data_buffer[zp.TRANSFORM_TYPE.PASSTHROUGH].pop(0).PASSTHROUGH - for source, events in self.data_buffer.iteritems(): + passthrough = self.sources[zp.TRANSFORM_TYPE.PASSTHROUGH] + result = passthrough.pop(0).PASSTHROUGH + + for source, events in self.sources.iteritems(): if source == zp.TRANSFORM_TYPE.PASSTHROUGH: continue if len(events) > 0: diff --git a/zipline/core/controlled.py b/zipline/core/controlled.py new file mode 100644 index 00000000..08c7345c --- /dev/null +++ b/zipline/core/controlled.py @@ -0,0 +1,69 @@ +""" +Poller logic for a component which is controlled by the monitor, this is +largely universal and thus we break it out into a seperate module and +splice it into the dispatch loops for each component instance. + +Example usage:: + + def do_work(): + socks = self.poll.poll() + + # Handle control events + do_handle_control_events() + + # Handle other events + if socks.get(socket) == zmq.POLLIN: + ... +""" + +import zmq +from zipline.core.component import Component +from zipline.protocol import CONTROL_PROTOCOL, CONTROL_FRAME, CONTROL_UNFRAME + +def do_handle_control_events(cls, poller): + assert isinstance(cls, Component) + assert cls.control_in, 'Component does not have a control_in socket' + + if poller.get(cls.control_in) == zmq.POLLIN: + msg = cls.control_in.recv() + event, payload = CONTROL_UNFRAME(msg) + + # =========== + # Heartbeat + # =========== + + # The controller will send out a single number packed in + # a CONTROL_FRAME with ``heartbeat`` event every + # (n)-seconds. The component then has n seconds to + # respond to it. If not then it will be considered as + # malfunctioning or maybe CPU bound. + + if event == CONTROL_PROTOCOL.HEARTBEAT: + # Heart outgoing + heartbeat_frame = CONTROL_FRAME( + CONTROL_PROTOCOL.OK, + payload + ) + # Echo back the heartbeat identifier to tell the + # controller that this component is still alive and + # doing work + cls.control_out.send(heartbeat_frame) + + # ========= + # Soft Kill + # ========= + + # Try and clean up properly and send out any reports or + # data that are done during a clean shutdown. Inform the + # controller that we're done. + elif event == CONTROL_PROTOCOL.SHUTDOWN: + cls.signal_done() + cls.shutdown() + + # ========= + # Hard Kill + # ========= + + # Just exit. + elif event == CONTROL_PROTOCOL.KILL: + cls.kill()