From 2d6c4688f980dc345eed9f14ae12bcf280513a6d Mon Sep 17 00:00:00 2001 From: fawce Date: Mon, 19 Mar 2012 23:19:57 -0400 Subject: [PATCH] added a new component property is_blocking. Non-blocking datasources and transforms will not prevent merge or feed from exiting, even if they have not reported DONE. refactored lines.py to provide a test zipline method and to take keyword arguments in the constructor and factory method. --- zipline/finance/trading.py | 12 ++++- zipline/lines.py | 100 +++++++++++++++++++++++++++++++---- zipline/messaging.py | 14 +++-- zipline/test/test_finance.py | 80 ++++++---------------------- 4 files changed, 126 insertions(+), 80 deletions(-) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 1dcd69bc..83d1e763 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -67,6 +67,7 @@ class TradeSimulationClient(qmsg.Component): if msg == str(zp.CONTROL_PROTOCOL.DONE): qutil.LOGGER.info("Client is DONE!") self.run_callbacks() + self.signal_order_done() self.signal_done() return @@ -153,6 +154,14 @@ class OrderDataSource(qmsg.DataSource): def get_type(self): return zp.DATASOURCE_TYPE.ORDER + # + @property + def is_blocking(self): + """ + This datasource is in a loop with the TradingSimulationClient + """ + return False + def open(self): qmsg.DataSource.open(self) self.order_socket = self.bind_order() @@ -177,7 +186,7 @@ class OrderDataSource(qmsg.DataSource): [self.order_socket], #allow half the time of a heartbeat for the order #timeout, so we have time to signal we are done. - timeout=self.heartbeat_timeout/2000 + #timeout=self.heartbeat_timeout/2000 ) @@ -186,6 +195,7 @@ class OrderDataSource(qmsg.DataSource): #no order message means there was a timeout above, #and the client is done sending orders (but isn't #telling us himself!). + qutil.LOGGER.warn("signaling orders done on timeout.") self.signal_done() return diff --git a/zipline/lines.py b/zipline/lines.py index 58e7eed7..619af0d7 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -111,20 +111,28 @@ class SimulatedTrading(object): simulation. """ - def __init__(self, algorithm, trading_environment, allocator): + def __init__(self, **config): """ - :param algorithm: a class that follows the algorithm protocol. Must - have a handle_frame method that accepts a pandas.Dataframe of the - current state of the simulation universe. Must have an order property - which can be set equal to the order method of trading_client. - :param trading_environment: TradingEnvironment object. + :param config: a dict with the following required properties:: + - algorithm: a class that follows the algorithm protocol. Must + have a handle_frame method that accepts a pandas.Dataframe of the + current state of the simulation universe. Must have an order + property which can be set equal to the order method of + trading_client. (TODO: where should this protocol be documented?) + - trading_environment: an instance of + :py:class:`zipline.trading.TradingEnvironment` + - allocator: an instance of + :py:class:`zipline.simulator.AddressAllocator` + - simulator_class: a :py:class:`zipline.messaging.ComponentHost` + subclass (not an instance) """ - self.algorithm = algorithm - self.allocator = allocator + assert isinstance(config, dict) + self.algorithm = config['algorithm'] + self.allocator = config['allocator'] + self.trading_environment = config['trading_environment'] + self.leased_sockets = [] - self.trading_environment = trading_environment self.sim_context = None - self.algorithm = algorithm sockets = self.allocate_sockets(8) addresses = { @@ -142,7 +150,8 @@ class SimulatedTrading(object): logging = qutil.LOGGER ) - self.sim = Simulator(addresses) + + self.sim = config['simulator_class'](addresses) self.clients = {} self.trading_client = TradeSimulationClient(self.trading_environment) @@ -177,6 +186,75 @@ class SimulatedTrading(object): self.trading_client.add_event_callback(self.algorithm.handle_frame) #register the trading_client's order method with the algorithm self.algorithm.set_order(self.trading_client.order) + + @staticmethod + def create_test_zipline(**config): + """ + :param config: A configuration object that is a dict with:: + - environment - a \ + :py:class:`zipline.finance.trading.TradeSimulationClient` + - allocator - a :py:class:`zipline.simulator.AddressAllocator` + - sid - an integer, which will be used as the security ID. + - order_count - the number of orders the test algo will place, + defaults to 100 + - trade_count - the number of trades to simulate, defaults to 100 + - simulator_class - optional parameter that provides an alternative + subclass of ComponentHost to hold the whole zipline. Defaults to + :py:class:`zipline.simulator.Simulator` + """ + assert isinstance(config, dict) + trading_environment = config['environment'] + allocator = config['allocator'] + sid = config['sid'] + if config.has_key('order_count'): + order_count = config['order_count'] + else: + order_count = 100 + + if config.has_key('trade_count'): + trade_count = config['trade_count'] + else: + trade_count = 100 + + if config.has_key('simulator_class'): + simulator_class = config['simulator_class'] + else: + simulator_class = Simulator + + #------------------- + # Trade Source + #------------------- + sids = [sid] + #------------------- + trade_source = factory.create_daily_trade_source( + sids, + trade_count, + trading_environment + ) + #------------------- + # Create the Algo + #------------------- + order_amount = 100 + #------------------- + test_algo = TestAlgorithm( + sid, + order_amount, + order_count + ) + #------------------- + # Simulation + #------------------- + zipline = SimulatedTrading(**{ + 'algorithm':test_algo, + 'trading_environment':trading_environment, + 'allocator':allocator, + 'simulator_class':simulator_class + }) + #------------------- + + zipline.add_source(trade_source) + + return zipline def add_source(self, source): assert isinstance(source, zmsg.DataSource) diff --git a/zipline/messaging.py b/zipline/messaging.py index 921daafb..b87ca2ec 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -65,7 +65,6 @@ class ComponentHost(Component): communication with them. """ assert isinstance(components, list) - for component in components: component.gevent_needed = self.gevent_needed @@ -78,9 +77,13 @@ class ComponentHost(Component): if isinstance(component, DataSource): self.feed.add_source(component.get_id) + if not component.is_blocking: + self.feed.ds_finished_counter +=1 if isinstance(component, BaseTransform): self.merge.add_source(component.get_id) - + if not component.is_blocking: + self.feed.ds_finished_counter +=1 + def unregister_component(self, component_id): del self.components[component_id] del self.sync_register[component_id] @@ -220,6 +223,7 @@ class ParallelBuffer(Component): if len(self.data_buffer) == self.ds_finished_counter: #drain any remaining messages in the buffer + qutil.LOGGER.debug("draining feed") self.drain() self.signal_done() else: @@ -261,7 +265,7 @@ class ParallelBuffer(Component): """ Send the (chronologically) next message in the buffer. """ - if(not(self.is_full() or self.draining)): + if not (self.is_full() or self.draining): return event = self.next() @@ -433,6 +437,10 @@ class BaseTransform(Component): def get_type(self): return COMPONENT_TYPE.CONDUIT + @property + def is_blocking(self): + return True + def open(self): """ Establishes zmq connections. diff --git a/zipline/test/test_finance.py b/zipline/test/test_finance.py index a7273eae..0e2bce14 100644 --- a/zipline/test/test_finance.py +++ b/zipline/test/test_finance.py @@ -45,6 +45,12 @@ class FinanceTestCase(TestCase): ) self.allocator = allocator + + self.zipline_test_config = { + 'allocator':self.allocator, + 'sid':133, + 'environment':self.trading_environment + } @timed(DEFAULT_TIMEOUT) def test_trade_feed_protocol(self): @@ -151,38 +157,9 @@ class FinanceTestCase(TestCase): @timed(DEFAULT_TIMEOUT) def test_orders(self): - - # Just verify sending and receiving orders. - # -------------- - # - SID=133 - sids = [133] - trade_count = 100 - trade_source = factory.create_daily_trade_source( - sids, - trade_count, - self.trading_environment - ) - - # Create the Algo - #------------------- - order_amount = 100 - order_count = 10 - test_algo = TestAlgorithm( - SID, - order_amount, - order_count - ) - # Simulation # ---------- - zipline = SimulatedTrading( - test_algo, - self.trading_environment, - self.allocator - ) - - zipline.add_source(trade_source) + zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) zipline.simulate(blocking=True) self.assertTrue(zipline.sim.ready()) @@ -196,38 +173,10 @@ class FinanceTestCase(TestCase): @timed(DEFAULT_TIMEOUT) def test_performance(self): - - # verify order -> transaction -> portfolio position. - # -------------- - SID=133 - sids = [133] - trade_count = 100 - trade_source = factory.create_daily_trade_source( - sids, - trade_count, - self.trading_environment - ) - - # Create the Algo - #------------------- - order_amount = 100 - order_count = 25 - test_algo = TestAlgorithm( - SID, - order_amount, - order_count - ) - - # Simulation - # ---------- - zipline = SimulatedTrading( - test_algo, - self.trading_environment, - self.allocator - ) - - zipline.add_source(trade_source) - + #provide enough trades to ensure all orders are filled. + self.zipline_test_config['order_count'] = 100 + self.zipline_test_config['trade_count'] = 200 + zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) zipline.simulate(blocking=True) self.assertEqual( @@ -245,14 +194,14 @@ class FinanceTestCase(TestCase): ) self.assertEqual( - test_algo.count, - test_algo.incr, + zipline.algorithm.count, + zipline.algorithm.incr, "The test algorithm should send as many orders as specified.") order_source = zipline.sources[zp.FINANCE_COMPONENT.ORDER_SOURCE] self.assertEqual( order_source.sent_count, - test_algo.count, + zipline.algorithm.count, "The order source should have sent as many orders as the algo." ) @@ -270,6 +219,7 @@ class FinanceTestCase(TestCase): "Portfolio should have one position." ) + SID = self.zipline_test_config['sid'] self.assertEqual( zipline.get_positions()[SID]['sid'], SID,