added a test package, converted the hacky test code into cases, fixed bunch off issues with transforms and merge.

This commit is contained in:
fawce
2012-01-18 21:18:04 -05:00
parent e250a9ecf1
commit 9a91b6f95d
8 changed files with 295 additions and 231 deletions
+28 -2
View File
@@ -5,6 +5,9 @@ Small classes to assist with db access, timezone calculations, and so on.
import datetime
import pytz
import json
import logging
import uuid
import zmq
class DocWrap():
"""
@@ -114,7 +117,7 @@ class ParallelBuffer(object):
class MergedParallelBuffer(ParallelBuffer):
def merge_next(self):
def next(self):
if(not(self.is_full() or self.draining)):
return
@@ -123,4 +126,27 @@ class MergedParallelBuffer(ParallelBuffer):
if(len(events) > 0):
result[source] = events.pop(0)
return result
return result
class FeedSync(object):
def __init__(self, feed, name):
self.feed = feed
self.id = "{name}-{id}".format(name=name, id=uuid.uuid1())
self.feed.register_sync(self.id)
self.logger = logging.getLogger()
#self.logger.info("registered {id} with feed".format(id=self.id))
def confirm(self):
context = zmq.Context()
#synchronize with feed
sync_socket = context.socket(zmq.REQ)
sync_socket.connect(self.feed.sync_address)
# send a synchronization request to the feed
sync_socket.send(self.id)
# wait for synchronization reply from the feed
sync_socket.recv()
sync_socket.close()
context.term()
self.logger.info("sync'd feed from {id}".format(id = self.id))
+28 -18
View File
@@ -6,24 +6,27 @@ import logging
class DataFeed(object):
def __init__(self, db, subscriber_count):
def __init__(self, config):
self.logger = logging.getLogger()
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.db = db
self.client_register = {}
self.data_workers = {}
emt1 = EquityMinuteTrades(133, self.db, self.data_address, self.sync_address, 1)
self.data_workers[1] = emt1
emt2 = EquityMinuteTrades(134, self.db, self.data_address, self.sync_address, 2)
self.data_workers[2] = emt2
self.config = config
for name, info in config.iteritems():
if(info['class'] == "EquityMinuteTrades"):
emt = EquityMinuteTrades(info['sid'], self, name)
self.data_workers[name] = emt
elif(info['class'] == "RandomEquityTrades"):
ret = RandomEquityTrades(info['sid'], self, name, info['count'])
self.data_workers[name] = ret
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():
@@ -31,21 +34,29 @@ class DataFeed(object):
source.start()
self.logger.info("ds processes launched")
def register_sync(self, sync_id):
self.client_register[sync_id] = "UNCONFIRMED"
def registration_complete(self):
for sync_id, status in self.client_register.iteritems():
if status == "UNCONFIRMED":
return False
return True
def sync_clients(self):
# Socket to receive signals
self.logger.info("waiting for all datasources and clients to be ready")
self.syncservice = self.context.socket(zmq.REP)
self.syncservice.bind(self.sync_address)
subscribers = 1
while subscribers <= self.sync_count:
self.logger.info("sync'ing {count} of {total}".format(count=subscribers, total=self.sync_count))
while not self.registration_complete():
# wait for synchronization request
msg = self.syncservice.recv()
self.client_register[msg] = "CONFIRMED"
#self.logger.info("confirmed {id}".format(id=msg))
# send synchronization reply
self.syncservice.send('')
subscribers += 1
self.syncservice.send('CONFIRMED')
self.syncservice.close()
self.logger.info("sync'd all datasources and clients")
@@ -62,7 +73,7 @@ class DataFeed(object):
self.data_socket.bind(self.data_address)
#create the feed
self.feed_socket = self.context.socket(zmq.PUSH)
self.feed_socket = self.context.socket(zmq.PUB)
self.feed_socket.bind(self.feed_address)
self.data_buffer.out_socket = self.feed_socket
@@ -92,7 +103,6 @@ class DataFeed(object):
#send the DONE message
self.feed_socket.send("DONE")
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()
+72 -44
View File
@@ -7,65 +7,53 @@ import pytz
import copy
import multiprocessing
import logging
import random
from pymongo import ASCENDING, DESCENDING
from backtest.util import *
from qbt_server import connect_db
class EquityMinuteTrades(object):
def __init__(self, sid, db, data_address, sync_address, source_id):
self.sid = sid
self.db = db
self.source_id = source_id
self.logger = logging.getLogger()
self.sync_address = sync_address
self.data_address = data_address
self.logger.info("data address is {ds}".format(ds=data_address))
class DataSource(object):
def __init__(self, feed, source_id):
self.source_id = source_id
self.logger = logging.getLogger()
self.feed = feed
self.sync = FeedSync(self.feed, str(source_id))
self.data_address = self.feed.data_address
self.logger.info("data address is {ds}".format(ds=self.feed.data_address))
self.cur_event = None
def start(self):
self.proc = multiprocessing.Process(target=self.run)
self.proc.start()
def run(self):
self.logger.info("starting data source:{sid} on {addr}".format(sid=self.sid, addr=self.data_address))
self.context = zmq.Context()
#synchronize with feed
sync_socket = self.context.socket(zmq.REQ)
sync_socket.connect(self.sync_address)
# send a synchronization request to the feed
sync_socket.send('')
# wait for synchronization reply from the feed
sync_socket.recv()
sync_socket.close()
def open(self):
self.logger.info("starting data source:{source_id} on {addr}".format(source_id=self.source_id, addr=self.feed.data_address))
self.context = zmq.Context()
#create the data sink. Based on http://zguide.zeromq.org/py:tasksink2
self.data_socket = self.context.socket(zmq.PUSH)
self.data_socket.connect(self.data_address)
eventQS = self.db.equity.trades.minute.find(fields=["sid","price","volume","dt"],
spec={"sid":self.sid},
sort=[("dt",ASCENDING)],
slave_ok=True)
self.logger.info("found {count} events".format(count=eventQS.count()))
for doc in eventQS:
doc_dt = doc['dt'].replace(tzinfo = pytz.utc)
doc_dt_str = format_date(doc_dt)
event = copy.copy(doc)
event['dt'] = doc_dt_str
event['s'] = self.source_id
event['type'] = 'event'
del(event['_id'])
#signal we are ready
self.sync.confirm()
def run(self):
try:
self.open()
self.send_all()
self.close()
except Exception as err:
self.logger.error(err, "Unexpected failure running datasource - {name}.".format(name=name))
#send this event to feed address
#self.logger.info("sending {event}".format(event=event))
self.data_socket.send(json.dumps(event))
def send(self, event):
event['s'] = self.source_id
event['type'] = 'event'
self.data_socket.send(json.dumps(event))
def close(self):
done_msg = {}
done_msg['type'] = 'DONE'
done_msg['s'] = self.source_id
@@ -73,9 +61,49 @@ class EquityMinuteTrades(object):
self.data_socket.close()
self.context.term()
self.logger.info("finished processing data source")
class EquityMinuteTrades(DataSource):
def __init__(self, sid, feed, source_id):
self.sid = sid
self.connection, self.db = connect_db()
DataSource.__init__(self, feed, source_id)
def send_all(self):
eventQS = self.db.equity.trades.minute.find(fields=["sid","price","volume","dt"],
spec={"sid":self.sid},
sort=[("dt",ASCENDING)],
slave_ok=True)
self.logger.info("found {count} events".format(count=eventQS.count()))
for doc in eventQS:
doc_dt = doc['dt'].replace(tzinfo = pytz.utc)
doc_dt_str = format_date(doc_dt)
event = copy.copy(doc)
event['dt'] = doc_dt_str
del(event['_id'])
self.send(event)
class RandomEquityTrades(DataSource):
def __init__(self, sid, feed, source_id, count):
DataSource.__init__(self, feed, source_id)
self.count = count
self.sid = sid
def send_all(self):
trade_start = datetime.datetime.now()
minute = datetime.timedelta(minutes=1)
price = random.uniform(5.0,50.0)
for i in range(self.count):
price = price + random.uniform(-0.05,0.05)
event = {'sid':self.sid, 'dt':format_date(trade_start + (minute * i)),'price':price, 'volume':random.randrange(100,10000,100)}
self.send(event)
+79 -67
View File
@@ -3,13 +3,14 @@ import logging
import datetime
import json
import config
import copy
import multiprocessing
from backtest.util import *
class Transform(object):
"""Parent class for feed transforms. Subclass to create a new derived value from the combined feed."""
def __init__(self, feed_address, result_address, sync_address, config_dict):
def __init__(self, feed, config_dict, result_address):
"""
feed_address - zmq socket address, Transform will CONNECT a PULL socket and receive messages until "DONE" is received.
result_address - zmq socket address, Transform will CONNECT a PUSH socket and send messaes until feed_socket receives "DONE"
@@ -19,42 +20,38 @@ class Transform(object):
the last transform in a series should be server=True so that clients can connect.
"""
self.logger = logging.getLogger()
self.feed_address = feed_address
self.feed = feed
self.feed_address = feed.feed_address
self.result_address = result_address
self.sync_address = sync_address
self.config = config.Config(config_dict)
self.name = self.config.get_string('name')
self.sync = FeedSync(feed, self.name)
self.state = {}
self.state['name'] = self.name
self.received_count = 0
self.sent_count = 0
def run(self):
self.open()
self.process_all()
self.close()
def open(self):
self.context = zmq.Context()
self.logger.info("starting {name} transform".format(name = self.name))
#create the feed PULL.
self.feed_socket = self.context.socket(zmq.PULL)
#create the feed SUB.
self.feed_socket = self.context.socket(zmq.SUB)
self.feed_socket.connect(self.feed_address)
self.feed_socket.setsockopt(zmq.SUBSCRIBE,'')
#create the result PUSH
self.result_socket = self.context.socket(zmq.PUSH)
self.result_socket.connect(self.result_address)
self.logger.info("sync'ing feed from {name}".format(name = self.name))
#synchronize with feed
sync_socket = self.context.socket(zmq.REQ)
sync_socket.connect(self.sync_address)
# send a synchronization request to the feed
sync_socket.send('')
# wait for synchronization reply from the feed
sync_socket.recv()
sync_socket.close()
def process_all(self):
self.logger.info("starting {name} event loop".format(name = self.name))
self.run_loop()
def run_loop(self):
self.sync.confirm()
while True:
message = self.feed_socket.recv()
@@ -65,9 +62,11 @@ class Transform(object):
event = json.loads(message)
cur_state = self.update(event)
cur_state['dt'] = event['dt']
cur_state['name'] = self.name
self.result_socket.send(json.dumps(cur_state))
self.sent_count += 1
def close(self):
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()
@@ -80,8 +79,8 @@ class Transform(object):
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)
def __init__(self, feed, props, result_address="tcp://127.0.0.1:20202"):
Transform.__init__(self, feed, props, result_address)
self.events = []
self.window = datetime.timedelta(days = self.config.get_integer('days'),
@@ -119,15 +118,15 @@ class MergedTransformsFeed(Transform):
""" Merge data feed and array of transform feeds into a single result vector.
PULL from feed
PULL from child transforms
PUSH to client
PUSH merged message to client
"""
def __init__(self, feed_address, result_address, sync_address, props):
def __init__(self, feed, props):
"""
config - must have an entry for 'transforms':array of dicts, which are convertedto configs.
"""
Transform.__init__(self, feed_address, result_address, sync_address, props)
Transform.__init__(self, feed, props, "tcp://127.0.0.1:20202")
self.transform_address = "tcp://127.0.0.1:{port}".format(port=10104)
self.transform_socket = None
self.create_transforms(self.config.transforms)
@@ -138,10 +137,12 @@ class MergedTransformsFeed(Transform):
for props in configs:
class_name = props['class']
if(class_name == 'MovingAverage'):
mavg = MovingAverage(self.feed_address, self.transform_address, self.sync_address, props)
mavg = MovingAverage(self.feed, props, self.transform_address)
self.transforms[mavg.name] = mavg
self.data_buffer = ParallelBuffer(self.transforms.keys())
keys = copy.copy(self.transforms.keys())
keys.append("feed") #for the raw feed
self.data_buffer = MergedParallelBuffer(keys)
for name, transform in self.transforms.iteritems():
self.logger.info("starting {name}".format(name=name))
@@ -149,51 +150,62 @@ class MergedTransformsFeed(Transform):
proc.start()
self.buffers = {}
for name, transform in self.transforms:
for name, transform in self.transforms.iteritems():
self.buffers[name] = []
def get_socket(self):
def open(self):
self.context = zmq.Context()
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):
self.logger.info("starting {name} transform".format(name = self.name))
#create the feed SUB.
self.feed_socket = self.context.socket(zmq.SUB)
self.feed_socket.connect(self.feed_address)
self.feed_socket.setsockopt(zmq.SUBSCRIBE,'')
#create the result PUSH
self.result_socket = self.context.socket(zmq.PUSH)
self.result_socket.bind(self.result_address)
#create the transform PULL.
self.transform_socket = self.context.socket(zmq.PULL)
self.transform_socket.bind(self.transform_address)
self.data_buffer.out_socket = self.result_socket
# Initialize poll set
self.poller = zmq.Poller()
self.poller.register(self.feed_socket, zmq.POLLIN)
self.poller.register(self.transform_socket, zmq.POLLIN)
def close(self):
self.transform_socket.close()
Transform.close(self)
def process_all(self):
self.sync.confirm()
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))
socks = dict(self.poller.poll())
self.feed_socket.close()
self.result_socket.close()
self.context.term()
def update(self, event):
state = {}
state['feed'] = event
count = 0
while count < len(self.transforms):
message = self.get_socket().recv()
if(message == "DONE"):
return "DONE"
data = json.loads(message)
state[data['name']] = data
if self.feed_socket in socks and socks[self.feed_socket] == zmq.POLLIN:
message = self.feed_socket.recv()
self.received_count += 1
if(message != "DONE"):
event = json.loads(message)
self.data_buffer.append("feed",event)
return state
if self.transform_socket in socks and socks[self.transform_socket] == zmq.POLLIN:
t_message = self.transform_socket.recv()
if(message != "DONE"):
t_event = json.loads(t_message)
self.data_buffer.append(t_event['name'], t_event)
self.data_buffer.send_next()
#drain any remaining messages in the buffer
self.data_buffer.drain()
#signal to client that we're done
self.result_socket.send("DONE")
self.logger.info("Transform {name} recieved {r} and sent {s}".format(name=self.name, r=self.data_buffer.received_count, s=self.data_buffer.sent_count))
+17 -23
View File
@@ -8,11 +8,13 @@ from backtest.util import *
class TestClient(object):
def __init__(self,feed_address, sync_address, bind=False):
self.logger = logging.getLogger()
self.feed_address = feed_address
self.sync_address = sync_address
self.bind = bind
def __init__(self,feed, address, bind=False):
self.logger = logging.getLogger()
self.feed = feed
self.address = address
self.sync = FeedSync(feed, "testclient")
self.bind = bind
self.received_count = 0
def run(self):
@@ -22,41 +24,33 @@ class TestClient(object):
self.data_feed = self.context.socket(zmq.PULL)
if(self.bind):
self.logger.info("binding to {feed_address}".format(feed_address=self.feed_address))
self.data_feed.bind(self.feed_address)
self.logger.info("binding to {address}".format(address=self.address))
self.data_feed.bind(self.address)
else:
self.data_feed.connect(self.feed_address)
self.logger.info("connecting to {address}".format(address=self.address))
self.data_feed.connect(self.address)
self.logger.info("synchronizing with data feed")
#synchronize with feed
sync_socket = self.context.socket(zmq.REQ)
sync_socket.connect(self.sync_address)
# send a synchronization request to the feed
sync_socket.send('')
# wait for synchronization reply from the feed
sync_socket.recv()
sync_socket.close()
self.sync.confirm()
self.logger.info("Starting the client loop")
counter = 0
prev_dt = None
while True:
msg = self.data_feed.recv()
counter += 1
if(msg == "DONE"):
self.logger.info("DONE!")
break
self.received_count += 1
self.logger.info(msg)
event = json.loads(msg)
if(prev_dt != None):
if(not event['dt'] >= prev_dt):
raise Exception("message arrived out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt))
prev_dt = event['dt']
if(counter % 100 == 0):
self.logger.info("received {n} messages".format(n=counter))
if(self.received_count % 100 == 0):
self.logger.info("received {n} messages".format(n=self.received_count))
self.logger.info("received {n} messages".format(n=counter))
self.logger.info("received {n} messages".format(n=self.received_count))
self.data_feed.close()
self.context.term()
+5 -77
View File
@@ -1,78 +1,6 @@
import zmq
import logging
import tornado
from data.sources.equity import *
from data.feed import *
from data.transforms import MergedTransformsFeed, MovingAverage
import unittest
from test.testmessaging import *
from qbt_server import connect_db
from qbt_client import TestClient
def datafeed():
connection, db = connect_db()
logger = logging.getLogger()
feed = DataFeed(db, 2)
#feed.run()
feed_proc = multiprocessing.Process(target=feed.run)
feed_proc.start()
config = {}
config['name'] = '**merged feed**'
config['transforms'] = [{'name':'mavg1', 'class':'MovingAverage', 'hours':1},{'name':'mavg2', 'class':'MovingAverage', 'hours':2}]
result_address = "tcp://127.0.0.1:20202"
mavg = MovingAverage(feed.feed_address, result_address, feed.sync_address, config['transforms'][0])
mavg_proc = multiprocessing.Process(target=mavg.run)
logger.info("about to launch moving average")
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()
#subscribe a client directly to the consolidated feed
#client = TestClient(feed.feed_address, feed.sync_address)
logger.info("starting the client")
#subscribe a client to the transformed feed
client = TestClient(result_address, feed.sync_address, bind=True)
client.run()
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)
proc2 = multiprocessing.Process(target=pub)
proc1.start()
proc2.start()
def sub():
context = zmq.Context()
controller = context.socket(zmq.PULL)
controller.connect("tcp://127.0.0.1:10099")
#controller.setsockopt(zmq.SUBSCRIBE, '')
while True:
try:
message = controller.recv()
print message
except zmq.ZMQError as err:
if err.errno != zmq.EAGAIN:
raise err
def pub():
context = zmq.Context()
controller = context.socket(zmq.PUSH)
controller.bind("tcp://127.0.0.1:10099")
while True:
controller.send("HELLO3")
if __name__ == "__main__":
tornado.options.parse_command_line()
datafeed()
if __name__ == '__main__':
suite = unittest.TestLoader().loadTestsFromTestCase(MessagingTestCase)
unittest.TextTestRunner(verbosity=2).run(suite)
View File
+66
View File
@@ -0,0 +1,66 @@
import unittest
import zmq
import logging
import tornado
from data.sources.equity import *
from data.feed import *
from data.transforms import MergedTransformsFeed, MovingAverage
from qbt_client import TestClient
class MessagingTestCase(unittest.TestCase):
def setUp(self):
tornado.options.parse_command_line()
self.total_data_count = 800
self.feed_config = {'emt1':{'sid':133, 'class':'RandomEquityTrades', 'count':400},
'emt2':{'sid':134, 'class':'RandomEquityTrades', 'count':400}}
self.feed = DataFeed(self.feed_config)
self.feed_proc = multiprocessing.Process(target=self.feed.run)
self.config = {}
self.config['name'] = '**merged feed**'
self.config['transforms'] = [{'name':'mavg1', 'class':'MovingAverage', 'hours':1},{'name':'mavg2', 'class':'MovingAverage', 'hours':2}]
def dtest_client(self):
#subscribe a client to the transformed feed
client = TestClient(self.feed, self.feed.feed_address)
feed_proc = multiprocessing.Process(target=self.feed.run)
feed_proc.start()
client.run()
self.assertEqual(self.feed.data_buffer.pending_messages(), 0, "The feed should be drained of all messages, found {n} remaining.".format(n=self.feed.data_buffer.pending_messages()))
self.assertEqual(self.total_data_count, client.received_count, "The client should have received ({n}) the same number of messages as the feed sent ({m}).".format(n=client.received_count, m=self.total_data_count))
def dtest_moving_average_to_client(self):
mavg = MovingAverage(self.feed, self.config['transforms'][0])
mavg_proc = multiprocessing.Process(target=mavg.run)
mavg_proc.start()
client = TestClient(self.feed, mavg.result_address, bind=True)
feed_proc = multiprocessing.Process(target=self.feed.run)
feed_proc.start()
client.run()
self.assertEqual(self.feed.data_buffer.pending_messages(), 0, "The feed should be drained of all messages.")
self.assertEqual(self.total_data_count, client.received_count, "The client should have received the same number of messages as the feed sent.")
def test_merged_to_client(self):
merger = MergedTransformsFeed(self.feed, self.config)
merger_proc = multiprocessing.Process(target=merger.run)
merger_proc.start()
client = TestClient(self.feed, merger.result_address)
feed_proc = multiprocessing.Process(target=self.feed.run)
feed_proc.start()
client.run()
self.assertEqual(self.feed.data_buffer.pending_messages(), 0, "The feed should be drained of all messages.")
self.assertEqual(self.total_data_count, client.received_count, "The client should have received the same number of messages as the feed sent.")