reorganized packages, test passing again

simulator/data/feed.py

@@ -0,0 +1,117 @@
+
+from simulator.data.sources.equity import *
+from simulator.backtest.util import *
+import time
+import logging
+
+class DataFeed(object):
+    
+    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.client_register = {}
+        
+        self.data_workers = {}
+        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())
+             
+    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 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) 
+        
+        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('CONFIRMED')
+        
+        self.syncservice.close()
+        self.logger.info("sync'd all datasources and clients")
+       
+    def run(self):   
+        # Prepare our context and sockets
+        self.context = zmq.Context()
+        
+        ds_finished_counter = 0
+             
+        #create the data sink. Based on http://zguide.zeromq.org/py:tasksink2 
+        #see: http://zguide.zeromq.org/py:taskwork2
+        self.data_socket = self.context.socket(zmq.PULL)
+        self.data_socket.bind(self.data_address)
+        
+        #create the feed
+        self.feed_socket = self.context.socket(zmq.PUB)
+        self.feed_socket.bind(self.feed_address)
+        
+        self.data_buffer.out_socket = self.feed_socket
+        
+        #start the data source workers
+        self.start_data_workers()
+        
+        #wait for all feed subscribers
+        self.sync_clients()
+        
+        self.logger.info("entering feed loop on {addr}".format(addr=self.data_address))
+        
+        while True:
+            message = self.data_socket.recv()
+            event = json.loads(message)
+            if(event["type"] == "DONE"):
+                ds_finished_counter += 1
+                if(len(self.data_workers) == ds_finished_counter):
+                    break
+            else:
+                self.data_buffer.append(event[u's'], event)
+                self.data_buffer.send_next()
+            
+                
+        #drain any remaining messages in the buffer
+        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.data_buffer.received_count, m=self.data_buffer.sent_count))
+        self.data_socket.close()
+        self.feed_socket.close()
+        self.context.term()
+        
+            
+       
+        
+        
+
+        
+            

simulator/data/sources/equity.py

@@ -0,0 +1,110 @@
+import datetime
+import zmq
+import pymongo
+import pymongo.json_util
+import json
+import pytz
+import copy
+import multiprocessing
+import logging
+import random
+from pymongo import ASCENDING, DESCENDING
+
+from simulator.backtest.util import *
+
+from simulator.qbt_server import * #connect_db
+
+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 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)
+        
+        #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))
+            
+    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
+        self.data_socket.send(json.dumps(done_msg))   
+        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)
+          
+   
+       
+        
+