From e250a9ecf1a7718cef1b32851a26f372f4a95017 Mon Sep 17 00:00:00 2001 From: fawce Date: Wed, 18 Jan 2012 12:08:07 -0500 Subject: [PATCH] split merging code into a new utility class --- backtest/util.py | 75 +++++++++++++++++++++++++++- data/feed.py | 80 ++++++++++-------------------- data/transforms.py | 121 +++++++++++++++++++++++++++++---------------- test.py | 10 ++-- 4 files changed, 184 insertions(+), 102 deletions(-) diff --git a/backtest/util.py b/backtest/util.py index 39dec830..082b1ed3 100644 --- a/backtest/util.py +++ b/backtest/util.py @@ -4,6 +4,7 @@ Small classes to assist with db access, timezone calculations, and so on. import datetime import pytz +import json class DocWrap(): """ @@ -50,4 +51,76 @@ def format_date(dt): if(dt == None): return None dt_str = dt.strftime('%Y/%m/%d-%H:%M:%S') + "." + str(dt.microsecond / 1000) - return dt_str \ No newline at end of file + return dt_str + + +class ParallelBuffer(object): + """ holds several queues of events by key, allows retrieval in date order or by merging""" + def __init__(self, key_list): + self.out_socket = None + self.sent_count = 0 + self.received_count = 0 + self.draining = False + self.data_buffer = {} + for key in key_list: + self.data_buffer[key] = [] + + def append(self, key, value): + self.data_buffer[key].append(value) + self.received_count += 1 + + def next(self): + if(not(self.is_full() or self.draining)): + return + + cur = None + earliest = None + for source, events in self.data_buffer.iteritems(): + if len(events) == 0: + continue + cur = events + if(earliest == None) or (cur[0]['dt'] <= earliest[0]['dt']): + earliest = cur + + if(earliest != None): + return earliest.pop(0) + + def is_full(self): + for source, events in self.data_buffer.iteritems(): + if (len(events) == 0): + return False + return True + + def pending_messages(self): + total = 0 + for source, events in self.data_buffer.iteritems(): + total += len(events) + return total + + def drain(self): + self.draining = True + while(self.pending_messages() > 0): + self.send_next() + + def send_next(self): + if(not(self.is_full() or self.draining)): + return + + event = self.next() + if(event != None): + self.out_socket.send(json.dumps(event)) + self.sent_count += 1 + + +class MergedParallelBuffer(ParallelBuffer): + + def merge_next(self): + if(not(self.is_full() or self.draining)): + return + + result = {} + for source, events in self.data_buffer.iteritems(): + if(len(events) > 0): + result[source] = events.pop(0) + + return result \ No newline at end of file diff --git a/data/feed.py b/data/feed.py index 6c9a63e1..5e3847a1 100644 --- a/data/feed.py +++ b/data/feed.py @@ -1,5 +1,6 @@ from data.sources.equity import * +from backtest.util import * import time import logging @@ -7,29 +8,27 @@ class DataFeed(object): def __init__(self, db, subscriber_count): self.logger = logging.getLogger() - self.db = db - self.data_workers = {} + self.data_address = "tcp://127.0.0.1:{port}".format(port=10101) self.sync_address = "tcp://127.0.0.1:{port}".format(port=10102) self.feed_address = "tcp://127.0.0.1:{port}".format(port=10103) - self.data_buffer = {} #source_id -> [] - self.subscriber_count = subscriber_count - - self.received_count = 0 - self.sent_count = 0 - - def start_data_workers(self): - """Start a sub-process for each datasource.""" + self.db = db + self.data_workers = {} emt1 = EquityMinuteTrades(133, self.db, self.data_address, self.sync_address, 1) self.data_workers[1] = emt1 - self.data_buffer[1] = [] - emt1.start() emt2 = EquityMinuteTrades(134, self.db, self.data_address, self.sync_address, 2) - self.data_workers[2] = emt1 - self.data_buffer[2] = [] - emt2.start() + self.data_workers[2] = emt2 + self.data_buffer = ParallelBuffer(self.data_workers.keys()) + self.sync_count = subscriber_count + len(self.data_workers) + + + def start_data_workers(self): + """Start a sub-process for each datasource.""" + for source_id, source in self.data_workers.iteritems(): + self.logger.info("starting {id}".format(id=source_id)) + source.start() self.logger.info("ds processes launched") def sync_clients(self): @@ -38,10 +37,10 @@ class DataFeed(object): self.syncservice = self.context.socket(zmq.REP) self.syncservice.bind(self.sync_address) + subscribers = 1 - total = self.subscriber_count + len(self.data_workers) - while subscribers <= total: - self.logger.info("sync'ing {count} of {total}".format(count=subscribers, total=total)) + while subscribers <= self.sync_count: + self.logger.info("sync'ing {count} of {total}".format(count=subscribers, total=self.sync_count)) # wait for synchronization request msg = self.syncservice.recv() # send synchronization reply @@ -66,6 +65,8 @@ class DataFeed(object): self.feed_socket = self.context.socket(zmq.PUSH) self.feed_socket.bind(self.feed_address) + self.data_buffer.out_socket = self.feed_socket + #start the data source workers self.start_data_workers() @@ -82,54 +83,25 @@ class DataFeed(object): if(len(self.data_workers) == ds_finished_counter): break else: - self.data_buffer[event[u's']].append(event) - self.received_count = self.received_count + 1 - self.send_next() + self.data_buffer.append(event[u's'], event) + self.data_buffer.send_next() #drain any remaining messages in the buffer - while(self.pending_messages() > 0): - self.send_next(drain=True) + self.data_buffer.drain() #send the DONE message self.feed_socket.send("DONE") - self.logger.info("received {n} messages, sent {m} messages".format(n=self.received_count, m=self.sent_count)) + self.logger.info("received {n} messages, sent {m} messages".format(n=self.data_buffer.received_count, m=self.data_buffer.sent_count)) self.data_socket.close() self.feed_socket.close() self.context.term() - def send_next(self, drain=False): - if(not(self.buffers_full() or drain)): - return - - cur = None - earliest = None - for source, events in self.data_buffer.iteritems(): - if len(events) == 0: - continue - cur = events - if(earliest == None) or (cur[0]['dt'] <= earliest[0]['dt']): - earliest = cur - - if(earliest != None): - event = earliest.pop(0) - self.feed_socket.send(json.dumps(event)) - self.sent_count += 1 + - def buffers_full(self): - for source, events in self.data_buffer.iteritems(): - if (len(events) == 0): - return False - return True - - def pending_messages(self): - total = 0 - for source, events in self.data_buffer.iteritems(): - total += len(events) - return total + - - \ No newline at end of file + \ No newline at end of file diff --git a/data/transforms.py b/data/transforms.py index 9300af69..919b427f 100644 --- a/data/transforms.py +++ b/data/transforms.py @@ -4,7 +4,7 @@ import datetime import json import config import multiprocessing -from backtest import util +from backtest.util import * class Transform(object): """Parent class for feed transforms. Subclass to create a new derived value from the combined feed.""" @@ -52,6 +52,9 @@ class Transform(object): sync_socket.close() self.logger.info("starting {name} event loop".format(name = self.name)) + self.run_loop() + + def run_loop(self): while True: message = self.feed_socket.recv() @@ -73,8 +76,46 @@ class Transform(object): def update(self, event): return {} + -class Merge(Transform): +class MovingAverage(Transform): + + def __init__(self, feed_address, result_address, sync_address, props, server=False): + Transform.__init__(self, feed_address, result_address, sync_address, props) + self.events = [] + + self.window = datetime.timedelta(days = self.config.get_integer('days'), + seconds = self.config.get_integer('seconds'), + microseconds = self.config.get_integer('microseconds'), + milliseconds = self.config.get_integer('milliseconds'), + minutes = self.config.get_integer('minutes'), + hours = self.config.get_integer('hours'), + weeks = self.config.get_integer('weeks')) + + + + + def update(self, event): + self.events.append(event) + + #filter the event list to the window length. + self.events = [x for x in self.events if (parse_date(x['dt']) - parse_date(event['dt'])) <= self.window] + + if(len(self.events) == 0): + return 0.0 + + total = 0.0 + for event in self.events: + total += event['price'] + + self.average = total/len(self.events) + + self.state['avg'] = self.average + + return self.state + + +class MergedTransformsFeed(Transform): """ Merge data feed and array of transform feeds into a single result vector. PULL from feed PULL from child transforms @@ -100,17 +141,46 @@ class Merge(Transform): mavg = MovingAverage(self.feed_address, self.transform_address, self.sync_address, props) self.transforms[mavg.name] = mavg + self.data_buffer = ParallelBuffer(self.transforms.keys()) + for name, transform in self.transforms.iteritems(): self.logger.info("starting {name}".format(name=name)) proc = multiprocessing.Process(target=transform.run) proc.start() + self.buffers = {} + for name, transform in self.transforms: + self.buffers[name] = [] + def get_socket(self): if(self.transform_socket == None): #create the feed PULL. self.transform_socket = self.context.socket(zmq.PULL) self.transform_socket.bind(self.transform_address) + return self.transform_socket + + def run_loop(self): + + while True: + #get original feed message + message = self.feed_socket.recv() + self.received_count += 1 + if(message == "DONE"): + self.result_socket.send("DONE") + break; + event = json.loads(message) + self.data_buffer.append(event['name'], event) + merged_event = self.data_buffer.merge_next() + if(merged_event != None): + self.result_socket.send(json.dumps(merged_event)) + self.sent_count += 1 + + self.logger.info("Transform {name} recieved {r} and sent {s}".format(name=self.name, r=self.received_count, s=self.sent_count)) + + self.feed_socket.close() + self.result_socket.close() + self.context.term() def update(self, event): @@ -118,47 +188,12 @@ class Merge(Transform): state['feed'] = event count = 0 - while count < len(transforms): - message = get_socket().recv + while count < len(self.transforms): + message = self.get_socket().recv() + if(message == "DONE"): + return "DONE" data = json.loads(message) state[data['name']] = data - return state - - - -class MovingAverage(Transform): - - def __init__(self, feed_address, result_address, sync_address, props, server=False): - Transform.__init__(self, feed_address, result_address, sync_address, props) - self.events = [] - - self.window = datetime.timedelta(days = self.config.get_integer('days'), - seconds = self.config.get_integer('seconds'), - microseconds = self.config.get_integer('microseconds'), - milliseconds = self.config.get_integer('milliseconds'), - minutes = self.config.get_integer('minutes'), - hours = self.config.get_integer('hours'), - weeks = self.config.get_integer('weeks')) - - - - - def update(self, event): - self.events.append(event) - - #filter the event list to the window length. - self.events = [x for x in self.events if (util.parse_date(x['dt']) - util.parse_date(event['dt'])) <= self.window] - - if(len(self.events) == 0): - return 0.0 - - total = 0.0 - for event in self.events: - total += event['price'] - - self.average = total/len(self.events) - - self.state['avg'] = self.average - - return self.state \ No newline at end of file + return state + \ No newline at end of file diff --git a/test.py b/test.py index 7d3c2d7b..c1d91cb7 100644 --- a/test.py +++ b/test.py @@ -3,7 +3,7 @@ import logging import tornado from data.sources.equity import * from data.feed import * -from data.transforms import Merge, MovingAverage +from data.transforms import MergedTransformsFeed, MovingAverage from qbt_server import connect_db from qbt_client import TestClient @@ -11,7 +11,8 @@ from qbt_client import TestClient def datafeed(): connection, db = connect_db() logger = logging.getLogger() - feed = DataFeed(db, 2) #one moving average, one client + feed = DataFeed(db, 2) + #feed.run() feed_proc = multiprocessing.Process(target=feed.run) feed_proc.start() @@ -27,6 +28,7 @@ def datafeed(): mavg_proc.start() #merger = Merge(feed.feed_address, result_address, feed.sync_address, config) + #merger.run() #merger_proc = multiprocessing.Process(target=merger.run) #merger_proc.start() @@ -39,8 +41,8 @@ def datafeed(): client = TestClient(result_address, feed.sync_address, bind=True) client.run() - logger.info("feed has {pending} messages".format(pending=feed.pending_messages())) - assert(feed.pending_messages() == 0) + logger.info("feed has {pending} messages".format(pending=feed.data_buffer.pending_messages())) + assert(feed.data_buffer.pending_messages() == 0) def pubsub(): proc1 = multiprocessing.Process(target=sub)