diff --git a/tests/client.py b/tests/client.py index c2c07560..3ceede6f 100644 --- a/tests/client.py +++ b/tests/client.py @@ -1,5 +1,4 @@ import logging -from gevent_zeromq import zmq import zipline.protocol as zp from zipline.core.component import Component diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py new file mode 100644 index 00000000..6b67c9db --- /dev/null +++ b/tests/test_exception_handling.py @@ -0,0 +1,159 @@ +import zmq + +from unittest2 import TestCase +from collections import defaultdict + +from zipline.test_algorithms import ExceptionAlgorithm +from zipline.finance.trading import SIMULATION_STYLE +from zipline.core.devsimulator import AddressAllocator +from zipline.lines import SimulatedTrading + +from zipline.utils.test_utils import \ + drain_zipline, \ + check, \ + setup_logger, \ + teardown_logger + +DEFAULT_TIMEOUT = 15 # seconds +EXTENDED_TIMEOUT = 90 + +allocator = AddressAllocator(1000) + + +class ExceptionTestCase(TestCase): + + leased_sockets = defaultdict(list) + + def setUp(self): + self.zipline_test_config = { + 'allocator' : allocator, + 'sid' : 133, + 'devel' : False, + 'results_socket' : allocator.lease(1)[0], + 'simulation_style' : SIMULATION_STYLE.FIXED_SLIPPAGE + } + self.ctx = zmq.Context() + setup_logger(self) + + def tearDown(self): + self.ctx.term() + teardown_logger(self) + + def test_exception_in_init(self): + # Simulation + # ---------- + self.zipline_test_config['algorithm'] = \ + ExceptionAlgorithm( + 'initialize', + self.zipline_test_config['sid'] + ) + + zipline = SimulatedTrading.create_test_zipline( + **self.zipline_test_config + ) + output, _ = drain_zipline(self, zipline) + self.assertEqual(len(output), 1) + self.assertEqual(output[-1]['prefix'], 'EXCEPTION') + payload = output[-1]['payload']['stack'] + check(self, payload, INITIALIZE_TB) + + self.assertTrue(zipline.sim.ready()) + self.assertFalse(zipline.sim.exception) + + + def test_exception_in_handle_data(self): + + # Simulation + # ---------- + self.zipline_test_config['algorithm'] = \ + ExceptionAlgorithm( + 'handle_data', + self.zipline_test_config['sid'] + ) + + zipline = SimulatedTrading.create_test_zipline( + **self.zipline_test_config + ) + + output, _ = drain_zipline(self, zipline) + + self.assertEqual(len(output), 1) + self.assertEqual(output[-1]['prefix'], 'EXCEPTION') + payload = output[-1]['payload']['stack'] + check(self, payload, HANDLE_DATA_TB) + + self.assertTrue(zipline.sim.ready()) + self.assertFalse(zipline.sim.exception) + + + # TODO: + # - define more zipline failure modes: exception in other + # components, exception in Monitor, etc. write tests + # for those scenarios. + + + +INITIALIZE_TB =\ +[{'filename': '/zipline/core/component.py', + 'line': 'self._run()', + 'lineno': 204, + 'method': 'run'}, + {'filename': '/zipline/core/component.py', + 'line': 'self.loop()', + 'lineno': 195, + 'method': '_run'}, + {'filename': '/zipline/core/component.py', + 'line': 'self.do_work()', + 'lineno': 235, + 'method': 'loop'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'self.initialize_algo()', + 'lineno': 97, + 'method': 'do_work'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'self.do_op(self.algorithm.initialize)', + 'lineno': 80, + 'method': 'initialize_algo'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'callable_op(*args, **kwargs)', + 'lineno': 206, + 'method': 'do_op'}, + {'filename': '/zipline/test_algorithms.py', + 'line': 'raise Exception("Algo exception in initialize")', + 'lineno': 166, + 'method': 'initialize'}] + + +HANDLE_DATA_TB =\ +[{'filename': '/zipline/core/component.py', + 'line': 'self._run()', + 'lineno': 204, + 'method': 'run'}, + {'filename': '/zipline/core/component.py', + 'line': 'self.loop()', + 'lineno': 195, + 'method': '_run'}, + {'filename': '/zipline/core/component.py', + 'line': 'self.do_work()', + 'lineno': 235, + 'method': 'loop'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'self.process_event(event)', + 'lineno': 116, + 'method': 'do_work'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'self.run_algorithm()', + 'lineno': 164, + 'method': 'process_event'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'self.do_op(self.algorithm.handle_data, data)', + 'lineno': 186, + 'method': 'run_algorithm'}, + {'filename': '/zipline/components/tradesimulation.py', + 'line': 'callable_op(*args, **kwargs)', + 'lineno': 206, + 'method': 'do_op'}, + {'filename': '/zipline/test_algorithms.py', + 'line': 'raise Exception("Algo exception in handle_data")', + 'lineno': 187, + 'method': 'handle_data'}] diff --git a/tests/test_feed.py b/tests/test_feed.py new file mode 100644 index 00000000..bfff6742 --- /dev/null +++ b/tests/test_feed.py @@ -0,0 +1,237 @@ +import os + +import uuid +import msgpack +import pytz + +from unittest2 import TestCase +from pymongo import Connection, ASCENDING +from itertools import izip, izip_longest, permutations, cycle, chain +from datetime import datetime, timedelta +from collections import deque + +from zipline import ndict +from zipline.gens.sort import date_sort, ready, done, queue_is_ready,queue_is_done,\ + pop_oldest +from zipline.gens.utils import hash_args, assert_datasource_protocol,\ + assert_trade_protocol, alternate +from zipline.gens.tradegens import date_gen, SpecificEquityTrades +from zipline.gens.composites import date_sorted_sources + +import zipline.protocol as zp + +class HelperTestCase(TestCase): + + def setUp(self): + pass + + def tearDown(self): + pass + + def test_individual_queue_logic(self): + queue = deque() + # Empty queues are neither done nor ready. + assert not queue_is_ready(queue) + assert not queue_is_done(queue) + + queue.append(to_dt('foo')) + assert queue_is_ready(queue) + assert not queue_is_done(queue) + + + queue.appendleft(to_dt('DONE')) + assert queue_is_ready(queue) + + # Checking done when we have a message after done will trip an assert. + self.assertRaises(AssertionError, queue_is_done, queue) + + queue.pop() + assert queue_is_ready(queue) + assert queue_is_done(queue) + + def test_pop_logic(self): + sources = {} + ids = ['a', 'b', 'c'] + for id in ids: + sources[id] = deque() + + assert not ready(sources) + assert not done(sources) + + # All sources must have a message to be ready/done + sources['a'].append(to_dt("datetime")) + assert not ready(sources) + assert not done(sources) + sources['a'].pop() + + for id in ids: + sources[id].append(to_dt("datetime")) + + assert ready(sources) + assert not done(sources) + + for id in ids: + sources[id].appendleft(to_dt("DONE")) + + # ["DONE", message] will trip an assert in queue_is_done. + assert ready(sources) + self.assertRaises(AssertionError, done, sources) + + for id in ids: + sources[id].pop() + + assert ready(sources) + assert done(sources) + +class DateSortTestCase(TestCase): + + def setUp(self): + pass + + def tearDown(self): + pass + + def run_date_sort(self, events, expected, source_ids): + """ + Take a list of events, their source_ids, and an expected sorting. + Assert that date_sort's output agrees with expected. + """ + sort_gen = date_sort(events, source_ids) + l = list(sort_gen) + assert l == expected + + def test_single_source(self): + source_ids = ['a'] + # 100 events, increasing by a minute at a time. + type = zp.DATASOURCE_TYPE.TRADE + dates = list(date_gen(count = 100)) + dates.append("DONE") + + # [('a', date1, type), ('a', date2, type), ... ('a', "DONE", type)] + event_args = zip(cycle(source_ids), iter(dates), cycle([type])) + + # Turn event_args into proper events. + events = [mock_data_unframe(*args) for args in event_args] + + # We don't expected Feed to yield the last event. + expected = events[:-1] + + event_gen = (e for e in events) + + self.run_date_sort(event_gen, expected, source_ids) + + def test_multi_source(self): + source_ids = ['a', 'b'] + type = zp.DATASOURCE_TYPE.TRADE + + # Set up source 'a'. Outputs 20 events with 2 minute deltas. + delta_a = timedelta(minutes = 2) + dates_a = list(date_gen(delta = delta_a, count = 20)) + dates_a.append("DONE") + + events_a_args = zip(cycle(['a']), iter(dates_a), cycle([type])) + events_a = [mock_data_unframe(*args) for args in events_a_args] + + # Set up source 'b'. Outputs 10 events with 1 minute deltas. + delta_b = timedelta(minutes = 1) + dates_b = list(date_gen(delta = delta_b, count = 10)) + dates_b.append("DONE") + + events_b_args = zip(cycle(['b']), iter(dates_b), cycle([type])) + events_b = [mock_data_unframe(*args) for args in events_b_args] + + # The expected output is all non-DONE events in both a and b, + # sorted first by dt and then by source_id. + non_dones = events_a[:-1] + events_b[:-1] + expected = sorted(non_dones, compare_by_dt_source_id) + + # Alternating between a and b. + interleaved = alternate(iter(events_a), iter(events_b)) + self.run_date_sort(interleaved, expected, source_ids) + + # All of a, then all of b. + + sequential = chain(iter(events_a), iter(events_b)) + self.run_date_sort(sequential, expected, source_ids) + + def test_sorted_sources(self): + + filter = [1,2] + #Set up source a. One hour between events. + args_a = tuple() + kwargs_a = {'sids' : [1,2,3,4], + 'start' : datetime(2012,6,6,0), + 'delta' : timedelta(hours = 1), + 'filter' : filter + } + #Set up source b. One day between events. + args_b = tuple() + kwargs_b = {'sids' : [1,2,3,4], + 'start' : datetime(2012,6,6,0), + 'delta' : timedelta(days = 1), + 'filter' : filter + } + #Set up source c. One minute between events. + args_c = tuple() + kwargs_c = {'sids' : [1,2,3,4], + 'start' : datetime(2012,6,6,0), + 'delta' : timedelta(minutes = 1), + 'filter' : filter + } + # Set up source d. This should produce no events because the + # internal sids don't match the filter. + args_d = tuple() + kwargs_d = {'sids' : [3,4], + 'start' : datetime(2012,6,6,0), + 'delta' : timedelta(minutes = 1), + 'filter' : filter + } + + sources = (SpecificEquityTrades,) * 4 + source_args = (args_a, args_b, args_c, args_d) + source_kwargs = (kwargs_a, kwargs_b, kwargs_c, kwargs_d) + + # Generate our expected source_ids. + zip_args = zip(source_args, source_kwargs) + expected_ids = ["SpecificEquityTrades" + hash_args(*args, **kwargs) + for args, kwargs in zip_args] + + # Pipe our sources into sort. + sort_out = date_sorted_sources(sources, source_args, source_kwargs) + + # Read all the values from sort and assert that they arrive in + # the correct sorting with the expected hash values. + to_list = list(sort_out) + copy = to_list[:] + for e in to_list: + # All events should match one of our expected source_ids. + assert e.source_id in expected_ids + # But none of them should match source_d. + assert e.source_id != hash_args(*args_d, **kwargs_d) + + expected = sorted(copy, compare_by_dt_source_id) + assert to_list == expected + +def mock_data_unframe(source_id, dt, type): + event = ndict() + event.source_id = source_id + event.dt = dt + event.type = type + return event + +def to_dt(val): + return ndict({'dt': val}) + +def compare_by_dt_source_id(x,y): + if x.dt < y.dt: + return -1 + elif x.dt > y.dt: + return 1 + + elif x.source_id < y.source_id: + return -1 + elif x.source_id > y.source_id: + return 1 + + else: + return 0 diff --git a/tests/test_finance.py b/tests/test_finance.py index f841e37f..d108e019 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -2,11 +2,11 @@ Tests for the zipline.finance package """ import pytz +import zmq from unittest2 import TestCase from datetime import datetime, timedelta from collections import defaultdict -from logbook.compat import LoggingHandler from nose.tools import timed @@ -20,28 +20,36 @@ from zipline.lines import SimulatedTrading from zipline.finance.performance import PerformanceTracker from zipline.utils.protocol_utils import ndict from zipline.finance.trading import TransactionSimulator, SIMULATION_STYLE +from zipline.utils.test_utils import \ + drain_zipline, \ + check, \ + setup_logger, \ + teardown_logger,\ + assert_single_position + DEFAULT_TIMEOUT = 15 # seconds EXTENDED_TIMEOUT = 90 allocator = AddressAllocator(1000) - class FinanceTestCase(TestCase): leased_sockets = defaultdict(list) def setUp(self): self.zipline_test_config = { - 'allocator' : allocator, - 'sid' : 133, - 'devel' : True + 'allocator' : allocator, + 'sid' : 133, + #'devel' : True, + 'results_socket' : allocator.lease(1)[0] } - self.log_handler = LoggingHandler() - self.log_handler.push_application() + self.ctx = zmq.Context() + + setup_logger(self) def tearDown(self): - self.log_handler.pop_application() + teardown_logger(self) @timed(DEFAULT_TIMEOUT) def test_factory_daily(self): @@ -109,150 +117,18 @@ class FinanceTestCase(TestCase): self.assertTrue(env.last_close.month == 12) self.assertTrue(env.last_close.day == 31) - # The following two tests appear broken no that the order source is - # non blocking. HUNCH: The trades are streaming through before the orders - # are placed. - - #@timed(EXTENDED_TIMEOUT) - def test_orders(self): - - # Simulation - # ---------- - - self.zipline_test_config['simulation_style'] = \ - SIMULATION_STYLE.FIXED_SLIPPAGE - zipline = SimulatedTrading.create_test_zipline( - **self.zipline_test_config - ) - zipline.simulate(blocking=True) - - self.assertTrue(zipline.sim.ready()) - self.assertFalse(zipline.sim.exception) - - # TODO: Make more assertions about the final state of the components. - self.assertEqual(zipline.sim.feed.pending_messages(), 0, \ - "The feed should be drained of all messages, found {n} remaining." \ - .format(n=zipline.sim.feed.pending_messages())) - - # the trading client should receive one transaction for every - # order placed. - self.assertEqual( - zipline.trading_client.txn_count, - zipline.trading_client.order_count - ) - - - #@timed(DEFAULT_TIMEOUT) - def test_aggressive_buying(self): - - # Simulation - # ---------- - - # TODO: for some reason the orders aren't filled without an extra - # trade. - trade_count = 5000 - self.zipline_test_config['order_count'] = trade_count - 1 - self.zipline_test_config['trade_count'] = trade_count - self.zipline_test_config['order_amount'] = 1 - - # tell the simulator to fill the orders in individual transactions - # matching the order volume exactly. - self.zipline_test_config['simulation_style'] = \ - SIMULATION_STYLE.FIXED_SLIPPAGE - self.zipline_test_config['environment'] = factory.create_trading_environment() - - sid_list = [self.zipline_test_config['sid']] - - self.zipline_test_config['trade_source'] = factory.create_minutely_trade_source( - sid_list, - trade_count, - self.zipline_test_config['environment'] - ) - - zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) - zipline.simulate(blocking=True) - - self.assertTrue(zipline.sim.ready()) - self.assertFalse(zipline.sim.exception) - - self.assertEqual(zipline.sim.feed.pending_messages(), 0, \ - "The feed should be drained of all messages, found {n} remaining." \ - .format(n=zipline.sim.feed.pending_messages())) - - # - # the trading client should receive one transaction for every - # order placed. - self.assertEqual( - zipline.trading_client.txn_count, - zipline.trading_client.order_count - ) - - - @timed(EXTENDED_TIMEOUT) - def test_performance(self): + def test_full_zipline(self): #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( - zipline.sim.feed.pending_messages(), - 0, - "The feed should be drained of all messages, found {n} remaining." \ - .format(n=zipline.sim.feed.pending_messages()) - ) - - self.assertEqual( - zipline.sim.merge.pending_messages(), - 0, - "The merge should be drained of all messages, found {n} remaining." \ - .format(n=zipline.sim.merge.pending_messages()) - ) - - self.assertEqual( - zipline.algorithm.count, - zipline.algorithm.incr, - "The test algorithm should send as many orders as specified.") - - - transaction_sim = zipline.trading_client.txn_sim - self.assertEqual( - transaction_sim.txn_count, - zipline.trading_client.perf.txn_count, - "The perf tracker should handle the same number of transactions \ - as the simulator emits." - ) - - self.assertEqual( - len(zipline.get_positions()), - 1, - "Portfolio should have one position." - ) - - sid = self.zipline_test_config['sid'] - self.assertEqual( - zipline.get_positions()[sid]['sid'], - sid, - "Portfolio should have one position in " + str(sid) - ) - - self.assertEqual( - zipline.sources['SpecificEquityTrades'].count, - self.zipline_test_config['trade_count'], - "The simulated trade source should send all trades." - ) - - self.assertEqual( - zipline.algorithm.frame_count, - self.zipline_test_config['trade_count'], - "The algorithm should receive all trades." - ) + assert_single_position(self, zipline) #@timed(DEFAULT_TIMEOUT) def test_sid_filter(self): - """Ensure the algorithm's filter prevents events from arriving.""" + # Ensure the algorithm's filter prevents events from arriving. # create a test algorithm whose filter will not match any of the # trade events sourced inside the zipline. order_amount = 100 @@ -270,12 +146,15 @@ class FinanceTestCase(TestCase): zipline = SimulatedTrading.create_test_zipline( **self.zipline_test_config ) + output, transaction_count = drain_zipline(self, zipline) + + self.assertTrue(zipline.sim.ready()) + self.assertFalse(zipline.sim.exception) - zipline.simulate(blocking=True) #check that the algorithm received no events self.assertEqual( 0, - test_algo.frame_count, + transaction_count, "The algorithm should not receive any events due to filtering." ) @@ -402,6 +281,9 @@ class FinanceTestCase(TestCase): self.transaction_sim(**params1) def transaction_sim(self, **params): + """ This is a utility method that asserts expected + results for conversion of orders to transactions given a + trade history""" trade_count = params['trade_count'] trade_amount = params['trade_amount'] diff --git a/tests/test_mongods.py b/tests/test_mongods.py new file mode 100644 index 00000000..d9b8dbe5 --- /dev/null +++ b/tests/test_mongods.py @@ -0,0 +1,114 @@ +import os + +import uuid +import msgpack +import pytz + +from unittest2 import TestCase +from pymongo import Connection, ASCENDING +from itertools import izip, izip_longest +from datetime import datetime, timedelta + +from zipline.gens.mongods import create_pymongo_iterator, MongoTradeHistoryGen +from zipline.gens.utils import hash_args, assert_datasource_protocol,\ + assert_trade_protocol, mock_raw_event + +import zipline.protocol as zp + +mongo_conn_args = { + 'mongodb_host': 'localhost', + 'mongodb_port': 27017, +} + +class TempMongo(object): + + def __enter__(self): + self.conn = Connection(mongo_conn_args['mongodb_host'], + mongo_conn_args['mongodb_port']) + + temp_id = 'qexec_test_id' + + self.db = self.conn[temp_id] + + return self + + def __exit__(self, type, value, traceback): + self.conn.drop_database(self.db.name) + +class TestMongoDataGenerator(TestCase): + + def setUp(self): + pass + def tearDown(self): + pass + + def test_create_pymongo_iterator(self): + + with TempMongo() as temp_mongo: + db = temp_mongo.db + coll = db.test + coll.ensure_index([('dt', ASCENDING), ('sid', ASCENDING)]) + + for i in xrange(100): + # sid = 1, dt ranging from 0 to 99 + coll.insert(mock_raw_event(1, i)) + + start_date = 20 + end_date = 50 + filter = {'sid' : [1]} + args = (coll, filter, start_date, end_date) + + cursor = create_pymongo_iterator(*args) + # We filter to only get dt's between 20 and 50 + expected = (mock_raw_event(1, i) for i in xrange(20, 51)) + + # Assert that our iterator returns the expected values. + for cursor_event, expected_event in izip_longest(cursor, expected): + del cursor_event['_id'] + # Easiest way to convert unicode to strings. + cursor_event = msgpack.loads(msgpack.dumps(cursor_event)) + assert expected_event.keys() == cursor_event.keys() + assert expected_event.values() == cursor_event.values() + + def test_MongoTradeHistoryGen(self): + + with TempMongo() as temp_mongo: + db = temp_mongo.db + coll = db.test + coll.ensure_index([('dt', ASCENDING), ('sid', ASCENDING)]) + + start_date = datetime(year = 2012,month=6,day=5,hour=0) + delta = timedelta(hours = 1) + + for i in xrange(100): + # sid = 1, dt's increasing an hour at a time from start + time = start_date + i * delta + coll.insert(mock_raw_event(1, time)) + + # Halfway through the events we added to db. + end_date = start_date + delta * 50 + + filter = {'sid' : [1]} + args = (coll, filter, start_date, end_date) + db_gen = MongoTradeHistoryGen(*args) + + expected_times = (start_date + i * delta for i in xrange(51)) + expected_events = (mock_raw_event(1, t) for t in expected_times) + + # DB events should match the expected events for price, dt, volume, + # and sid. They should also conform to the trade frame protocol. + + for db, expected in izip_longest(db_gen, expected_events): + expected['dt'] = expected['dt'].replace(tzinfo = pytz.utc) + # Check that our output meets the trade protocol. + assert_trade_protocol(db) + + # Check that our output matches expectations + for field in iter(['sid', 'dt', 'price', 'volume']): + assert db[field] == expected[field] + + # Expected output of hash_args: + assert db['source_id'] == \ + 'MongoTradeHistoryGen983a27fd0710414239a5cde71ef5a8fc' + + diff --git a/tests/test_monitor.py b/tests/test_monitor.py index 8b670356..5f55aaee 100644 --- a/tests/test_monitor.py +++ b/tests/test_monitor.py @@ -1,16 +1,15 @@ -import gevent -from logbook.compat import LoggingHandler +from zipline.utils.test_utils import setup_logger, teardown_logger from unittest2 import TestCase, skip from zipline.core.monitor import Controller class TestMonitor(TestCase): def setUp(self): - self.log_handler = LoggingHandler() - self.log_handler.push_application() + setup_logger(self, '/var/log/qexec/qexec.log') + def tearDown(self): - self.log_handler.pop_application() + teardown_logger(self) def test_init(self): pub_socket = 'tcp://127.0.0.1:5000' @@ -25,18 +24,3 @@ class TestMonitor(TestCase): con = Controller(pub_socket, route_socket, ) con.manage([ 'a', 'b', 'c', 'd' ]) - - @skip - def test_poll(self): - from mock_zmq import zmq_synthetic - pub_socket = 'tcp://127.0.0.1:5000' - route_socket = 'tcp://127.0.0.1:5001' - cancel_socket = 'tcp://127.0.0.1:5002' - - con = Controller(pub_socket, route_socket, cancel_socket) - con.manage([ 'a', 'b', 'c', 'd' ]) - con.zmq = zmq_synthetic - con.zmq_flavor = 'green' - - con.period = 0.00001 - gevent.spawn(con.run).join(timeout=con.period) diff --git a/tests/test_optimize.py b/tests/test_optimize.py index 256ee20c..d571041c 100644 --- a/tests/test_optimize.py +++ b/tests/test_optimize.py @@ -13,7 +13,7 @@ EXTENDED_TIMEOUT = 90 allocator = AddressAllocator(1000) -from logbook.compat import LoggingHandler +from zipline.utils.test_utils import setup_logger, teardown_logger class TestUpDown(TestCase): """This unittest verifies that the BuySellAlgorithm in @@ -31,11 +31,11 @@ class TestUpDown(TestCase): 'amplitude' : 30, 'base_price' : 50 } - self.log_handler = LoggingHandler() - self.log_handler.push_application() + setup_logger(self, '/var/log/qexec/qexec.log') + def tearDown(self): - self.log_handler.pop_application() + teardown_logger(self) @skip @timed(DEFAULT_TIMEOUT) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 1fec3e7a..1fe1ce3c 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -2,7 +2,7 @@ from datetime import timedelta from collections import defaultdict from unittest2 import TestCase -from logbook.compat import LoggingHandler +from zipline.utils.test_utils import setup_logger, teardown_logger import zipline.utils.factory as factory from zipline.finance.vwap import DailyVWAP, VWAPTransform @@ -25,11 +25,10 @@ class ZiplineWithTransformsTestCase(TestCase): 'sid' : 133, 'devel' : True } - self.log_handler = LoggingHandler() - self.log_handler.push_application() + setup_logger(self, '/var/log/qexec/qexed.log') def tearDown(self): - self.log_handler.pop_application() + teardown_logger(self) def test_vwap_tnfm(self): zipline = SimulatedTrading.create_test_zipline( @@ -47,8 +46,7 @@ class FinanceTransformsTestCase(TestCase): def setUp(self): self.trading_environment = factory.create_trading_environment() - self.log_handler = LoggingHandler() - self.log_handler.push_application() + setup_logger(self, '/var/log/qexec/qexec.log') def tearDown(self): self.log_handler.pop_application() diff --git a/zipline/components/aggregator.py b/zipline/components/aggregator.py index 6eb99fe9..80d6a087 100644 --- a/zipline/components/aggregator.py +++ b/zipline/components/aggregator.py @@ -83,26 +83,16 @@ class Aggregate(Component): self.drain() self.signal_done() else: - try: - event = self.unframe(message) - except zp.INVALID_DATASOURCE_FRAME as exc: - # Error deserializing - return self.signal_exception(exc) + event = self.unframe(message) + self.append(event) - try: - self.append(event) + if self.is_full(): + event = self.next() - if 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) + if event: + self.send(event) + else: + pass # ------------- # Flow Control diff --git a/zipline/components/tradesimulation.py b/zipline/components/tradesimulation.py index 12b4c8d7..60202854 100644 --- a/zipline/components/tradesimulation.py +++ b/zipline/components/tradesimulation.py @@ -38,6 +38,7 @@ class TradeSimulationClient(Component): self.perf = perf.PerformanceTracker(self.trading_environment) self.zmq_out = None self.results_socket = results_socket + self.algo_initialized = False @property def get_id(self): @@ -56,9 +57,6 @@ class TradeSimulationClient(Component): # initialize with all possible sids. self.perf.set_sids(self.algorithm.get_sid_filter()) - # N.B. Initialize is now called from open, because we - # need to have a socket open for logging. - def open(self): self.result_feed = self.connect_result() if self.results_socket: @@ -70,7 +68,6 @@ class TradeSimulationClient(Component): self.setup_logging(sock) self.perf.publish_to(sock) - self.initialize_algo() def initialize_algo(self): """ Setup loggers for algorithm and run algorithm's own @@ -80,7 +77,8 @@ class TradeSimulationClient(Component): self.algo_log = Logger("AlgoLog") self.algorithm.set_logger(self.algo_log) - self.run_logged_op(self.algorithm.initialize) + self.do_op(self.algorithm.initialize) + self.algo_initialized = True def setup_logging(self, socket = None): sock = socket or self.results_socket @@ -95,6 +93,8 @@ class TradeSimulationClient(Component): self.stdout_capture = stdout_only_pipe def do_work(self): + if not self.algo_initialized: + self.initialize_algo() # see if the poller has results for the result_feed if self.socks.get(self.result_feed) == self.zmq.POLLIN: @@ -183,9 +183,15 @@ class TradeSimulationClient(Component): # data injection pipeline for log rerouting # any fields injected here should be added to # LOG_EXTRA_FIELDS in zipline/protocol.py - self.run_logged_op(self.algorithm.handle_data, data) + self.do_op(self.algorithm.handle_data, data) - def run_logged_op(self, callable_op, *args, **kwargs): + def exception_callback(self, exc_type, exc_value, exc_traceback): + if self.results_socket: + log.info("Sending exception frame") + msg = zp.EXCEPTION_FRAME(exc_traceback) + self.out_socket.send(msg) + + def do_op(self, callable_op, *args, **kwargs): """ Wrap a callable operation with the zmq logbook handler if it exits.""" if self.zmq_out: @@ -221,8 +227,8 @@ class TradeSimulationClient(Component): with log_pipeline.threadbound(), self.stdout_capture(self.logger, ''): self.algorithm.handle_data('data') - def connect_order(self): - return self.connect_push_socket(self.addresses['order_address']) + #def connect_order(self): + # return self.connect_push_socket(self.addresses['order_address']) def order(self, sid, amount): order = zp.ndict({ diff --git a/zipline/core/component.py b/zipline/core/component.py index 841800e1..0e1d04f8 100644 --- a/zipline/core/component.py +++ b/zipline/core/component.py @@ -14,12 +14,9 @@ from setproctitle import setproctitle # pyzmq import zmq -# gevent_zeromq -import gevent_zeromq -# zmq_ctypes -#import zmq_ctypes -from zipline.utils.gpoll import _Poller as GeventPoller +from zipline.core.monitor import PARAMETERS + from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_STATE, \ COMPONENT_FAILURE, CONTROL_FRAME, CONTROL_UNFRAME @@ -27,6 +24,10 @@ log = logbook.Logger('Component') from zipline.exceptions import ComponentNoInit +class KillSignal(Exception): + def __init__(self): + pass + class Component(object): """ @@ -152,41 +153,13 @@ class Component(object): def do_work(self): raise NotImplementedError - def init_zmq(self, flavor): - """ - ZMQ in all flavors. Have it your way. - - mp - Distinct contexts | pyzmq - thread - Same context | pyzmq - green - Same context | gevent_zeromq - pypy - Same context | zmq_ctypes - - """ - - if flavor == 'mp': - self.zmq = zmq - self.context = self.zmq.Context() - self.zmq_poller = self.zmq.Poller - # The the process title so you can watch it in top - setproctitle(self.__class__.__name__) - return - if flavor == 'thread': - self.zmq = zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = self.zmq.Poller - return - if flavor == 'green': - self.zmq = gevent_zeromq.zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = GeventPoller - return - if flavor == 'pypy': - self.zmq = zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = self.zmq.Poller - return - - raise Exception("Unknown ZeroMQ Flavor") + def init_zmq(self): + self.zmq = zmq + self.context = self.zmq.Context() + self.zmq_poller = self.zmq.Poller + # The the process title so you can watch it in top + setproctitle(self.__class__.__name__) + return def _run(self): """ @@ -204,16 +177,15 @@ class Component(object): self.done = False # TODO: use state flag self.sockets = [] - self.init_zmq(self.zmq_flavor) + self.init_zmq() self.setup_poller() - self.open() self.setup_control() + self.open() self.signal_ready() self.lock_ready() - self.wait_ready() # ----------------------- # YOU SHALL NOT PASS!!!!! @@ -231,14 +203,17 @@ class Component(object): try: self._run() except Exception as exc: - exc_info = sys.exc_info() - self.signal_exception(exc) + if not isinstance(exc, KillSignal): + self.signal_exception(exc) + else: + # if we get a kill signal, forcibly close all the + # sockets. + # exc_info = sys.exc_info() + # self.relay_exception(exc_info[0], exc_info[1], exc_info[2]) + self.teardown_sockets() - # Reraise the exception - raise exc_info[0], exc_info[1], exc_info[2] finally: self.shutdown() - self.teardown_sockets() log.info("Exiting %r" % self) def working(self): @@ -314,7 +289,6 @@ class Component(object): # controller that we're done. elif event == CONTROL_PROTOCOL.SHUTDOWN: self.signal_done() - self.shutdown() # ========= # Hard Kill @@ -326,7 +300,9 @@ class Component(object): # In case we didn't receive a ping, send a pre-emptive # pong to the monitor. - elif self.last_ping and time.time() - self.last_ping > 1: + elif hasattr(self, 'control_out') and \ + self.last_ping and \ + time.time() - self.last_ping > 1: # send a ping ahead of schedule pre_pong = time.time() heartbeat_frame = CONTROL_FRAME( @@ -337,8 +313,13 @@ class Component(object): # Echo back the heartbeat identifier to tell the # controller that this component is still alive and # doing work - self.control_out.send(heartbeat_frame) + self.control_out.send(heartbeat_frame, self.zmq.NOBLOCK) self.last_ping = pre_pong + elif self.last_ping and \ + time.time() - self.last_ping > PARAMETERS.MAX_COMPONENT_WAIT: + # monitor is gone without sending the shutdown + # signal, do a hard exit. + self.kill() # ---------------------------- # Cleanup & Modes of Failure @@ -349,6 +330,7 @@ class Component(object): Close all zmq sockets safely. This is universal, no matter where this is running it will need the sockets closed. """ + log.warn("{id} closing all sockets".format(id=self.get_id)) #close all the sockets for sock in self.sockets: sock.close() @@ -360,6 +342,7 @@ class Component(object): Tear down after normal operation. """ if self.on_done: + log.warn("{id} calling done.".format(id=self.get_id)) self.on_done() def kill(self): @@ -369,7 +352,8 @@ class Component(object): Tear down ( fast ) as a mode of failure in the simulation or on service halt. """ - raise NotImplementedError + # sys.exit(1) + raise KillSignal() # ---------------------- # Internal Maintenance @@ -397,33 +381,67 @@ class Component(object): # in the locked quasimode. Respond to HEARTBEAT and GO # messages. + start_wait = time.time() + while self.waiting: - #socks = dict(self.poll.poll(self.heartbeat_timeout)) + socks = dict(self.poll.poll(0)) - msg = self.control_in.recv() - event, payload = CONTROL_UNFRAME(msg) + assert self.control_in, \ + 'Component does not have a control_in socket' - # ==== - # Go - # ==== + if socks.get(self.control_in) == zmq.POLLIN: - # A distributed lock from the controller to ensure - # synchronized start. + msg = self.control_in.recv() + event, payload = CONTROL_UNFRAME(msg) - if event == CONTROL_PROTOCOL.HEARTBEAT: - heartbeat_frame = CONTROL_FRAME( - CONTROL_PROTOCOL.OK, - payload - ) - self.control_out.send(heartbeat_frame) - log.info('Prestart Heartbeat ' + self.get_id) + # ==== + # Go + # ==== + + # A distributed lock from the controller to ensure + # synchronized start. + + if event == CONTROL_PROTOCOL.HEARTBEAT: + heartbeat_frame = CONTROL_FRAME( + CONTROL_PROTOCOL.OK, + payload + ) + self.control_out.send(heartbeat_frame) + log.info('Prestart Heartbeat ' + self.get_id) + + elif event == CONTROL_PROTOCOL.GO: + # Side effectful call from the controller to unlock + # and begin doing work only when the entire topology + # of the system beings to come online + log.info('Unlocking ' + self.__class__.__name__) + self.unlock_ready() + + # ========= + # 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: + self.signal_done() + break + + # ========= + # Hard Kill + # ========= + + # Just exit. + elif event == CONTROL_PROTOCOL.KILL: + self.kill() + break + + elif time.time() - start_wait > PARAMETERS.MAX_COMPONENT_WAIT: + log.info('No go signal from monitor, %s exiting' \ + % self.__class__.__name__) + self.kill() + break - elif event == CONTROL_PROTOCOL.GO: - # Side effectful call from the controller to unlock - # and begin doing work only when the entire topology - # of the system beings to come online - log.info('Unlocking ' + self.__class__.__name__) - self.unlock_ready() def signal_ready(self): log.info(self.__class__.__name__ + ' is ready') @@ -452,7 +470,8 @@ class Component(object): def signal_exception(self, exc=None, scope=None): """ - This is *very* important error tracking handler. + All exceptions inside any component should boil back to + this handler. Will inform the system that the component has failed and how it has failed. @@ -472,16 +491,48 @@ class Component(object): self._exception = exc exc_type, exc_value, exc_traceback = sys.exc_info() trace = ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback)) - sys.stdout.write(trace) - if hasattr(self, 'control_out'): - exception_frame = CONTROL_FRAME( - CONTROL_PROTOCOL.EXCEPTION, - trace - ) - self.control_out.send(exception_frame) + # if a downstream component fails, this component may try + # sending when there are zero connections to the socket, + # which will raise ZMQError(EAGAIN). So, it doesn't make + # sense to relay this exception to Monitor and the rest + # of the zipline. + if isinstance(exc, zmq.ZMQError) and exc.errno == zmq.EAGAIN: + log.warn("{id} raised a ZMQError(EAGAIN) not relaying"\ + .format(id=self.get_id)) + return + + # sys.stdout.write(trace) + log.exception("Unexpected error in run for {id}.".format(id=self.get_id)) + + self.relay_exception(exc_type, exc_value, exc_traceback) + + if hasattr(self, 'control_out') and self.control_out: + try: + log.info('{id} sending exception to controller'.format(id=self.get_id)) + exception_frame = CONTROL_FRAME( + CONTROL_PROTOCOL.EXCEPTION, + trace + ) + self.control_out.send(exception_frame, self.zmq.NOBLOCK) + # The controller should relay the exception back + # to all zipline components. Wait here until the + # notice arrives, and we can assume other zipline + # components have broken out of their message + # loops. + for i in xrange(PARAMETERS.MAX_COMPONENT_WAIT): + self.heartbeat(timeout=1000) + log.warn("{id} never heard back from monitor."\ + .format(id=self.get_id)) + except: + log.exception("Exception waiting for controller reply") + + def relay_exception(self, exc_type, exc_value, exc_traceback): + if hasattr(self, 'exception_callback') and self.exception_callback: + log.info('{id} making exception callback'.format(id=self.get_id)) + self.exception_callback(exc_type, exc_value, exc_traceback) + - #LOGGER.exception("Unexpected error in run for {id}.".format(id=self.get_id)) def signal_done(self): """ @@ -489,21 +540,23 @@ class Component(object): """ self.state_flag = COMPONENT_STATE.DONE + # notify internal work loop that we're done + self.done = True # TODO: use state flag - if self.out_socket: + if hasattr(self, 'out_socket') and self.out_socket: msg = zmq.Message(str(CONTROL_PROTOCOL.DONE)) self.out_socket.send(msg) + if hasattr(self, 'control_out'): + # notify controller we're done + done_frame = CONTROL_FRAME( + CONTROL_PROTOCOL.DONE, + '' + ) - # notify controller we're done - done_frame = CONTROL_FRAME( - CONTROL_PROTOCOL.DONE, - '' - ) - - self.control_out.send(done_frame) - log.info("[%s] sent control done" % self.get_id) + self.control_out.send(done_frame) + log.info("[%s] sent control done" % self.get_id) # there is a narrow race condition where we finish just # after the Monitor accepts our prior heartbeat, but just @@ -511,8 +564,7 @@ class Component(object): # last heartbeat, and wait an unusually long time. self.heartbeat(timeout=5000) - # notify internal work look that we're done - self.done = True # TODO: use state flag + # ----------- @@ -524,9 +576,6 @@ class Component(object): Setup the poller used for multiplexing the incoming data handling sockets. """ - - # Initializes the poller class specified by the flavor of - # ZeroMQ. Either zmq.Poller or gpoll.Poller . self.poll = self.zmq_poller() def bind_data(self): diff --git a/zipline/core/devsimulator.py b/zipline/core/devsimulator.py index 8835cb07..cbce4b66 100644 --- a/zipline/core/devsimulator.py +++ b/zipline/core/devsimulator.py @@ -5,7 +5,6 @@ See :py:method"" import logbook import threading -from zipline.core.simulatorref import SimulatorBase log = logbook.Logger('Dev Simulator') @@ -34,71 +33,4 @@ class AddressAllocator(object): return sockets def reaquire(self, *conn): - pass - - -class Simulator(SimulatorBase): - - zmq_flavor = 'thread' - - def __init__(self, addresses): - # TODO: rethink this - SimulatorBase.__init__(self, addresses) - self.subthreads = [] - self.running = False - - log.warn(DEPRECATION_WARNING) - - @property - def get_id(self): - return 'Simple Simulator' - - def launch_controller(self): - thread = threading.Thread( - target=self.controller.run, - ) - thread.start() - - self.subthreads.append(thread) - return thread - - def simulate(self): - thread = threading.Thread(target=self.run) - thread.start() - - self.subthreads.append(thread) - self.running = True - - return thread - - def did_clean_shutdown(self): - return not any([t.isAlive() for t in self.subthreads]) - - def shutdown(self): - """ - Destroy all tracked components. - """ - - if not self.running: - return - - - for component in self.components.itervalues(): - component.shutdown() - - for thread in self.subthreads: - if thread.is_alive(): - thread._Thread__stop() - - #self.controller.shutdown() - - self.running = False - - assert self.did_clean_shutdown() - - def launch_component(self, component): - thread = threading.Thread(target=component.run) - thread.start() - - self.subthreads.append(thread) - return thread + pass \ No newline at end of file diff --git a/zipline/core/monitor.py b/zipline/core/monitor.py index efcc6314..d3a831f1 100644 --- a/zipline/core/monitor.py +++ b/zipline/core/monitor.py @@ -1,17 +1,15 @@ +import inspect import os import zmq import sys import time -import gevent import itertools import logbook -import gevent_zeromq from setproctitle import setproctitle -from signal import SIGHUP, SIGINT +from signal import SIGHUP, SIGINT, SIGKILL, signal from collections import OrderedDict, Counter -from zipline.utils.gpoll import _Poller as GeventPoller from zipline.protocol import CONTROL_PROTOCOL, CONTROL_FRAME, \ CONTROL_UNFRAME, CONTROL_STATES, INVALID_CONTROL_FRAME \ @@ -42,7 +40,11 @@ log = logbook.Logger('Controller') # the system. PARAMETERS = ndict(dict( - GENERATIONAL_PERIOD = 10, #seconds + # time Monitor will wait for a heartbeat, in seconds + GENERATIONAL_PERIOD = 10, + # time Component will wait for GO and for a heartbeat before + # timing out. + MAX_COMPONENT_WAIT = 20, ALLOWED_SKIPPED_HEARTBEATS = 10, ALLOWED_INVALID_HEARTBEATS = 3, PRESTART_HEARBEATS = 3, @@ -67,9 +69,8 @@ class Controller(object): debug = True period = PARAMETERS.GENERATIONAL_PERIOD - def __init__(self, pub_socket, route_socket, devel=True): + def __init__(self, pub_socket, route_socket): - self.devel = devel self.nosignals = False self.context = None self.zmq = None @@ -96,34 +97,17 @@ class Controller(object): self.missed_beats = Counter() - log.warn("Running Controller in development mode, will ONLY synchronize start.") + # if we are inside a test, we want to skip signalling + # back to the parent process. + self.inside_test = 'nose' in inspect.stack()[-1][1] - def init_zmq(self, flavor): - assert self.zmq_flavor in ['thread', 'mp', 'green'] - if flavor == 'mp': - self.zmq = zmq - self.context = self.zmq.Context() - self.zmq_poller = self.zmq.Poller - if self.devel: - log.warning("USING DEVELOPMENT MODE IN MP CONTEXT NOT RECOMMENDED") - return - if flavor == 'thread': - self.zmq = zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = self.zmq.Poller - return - if flavor == 'green': - self.zmq = gevent_zeromq.zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = GeventPoller - return - if flavor == 'pypy': - self.zmq = zmq - self.context = self.zmq.Context.instance() - self.zmq_poller = self.zmq.Poller - return + def init_zmq(self): + self.zmq = zmq + self.context = self.zmq.Context() + self.zmq_poller = self.zmq.Poller + return def manage(self, topology): """ @@ -157,7 +141,7 @@ class Controller(object): def run(self): self.running = True - self.init_zmq(self.zmq_flavor) + self.init_zmq() setproctitle('Monitor') self.state = CONTROL_STATES.INIT @@ -166,8 +150,8 @@ class Controller(object): # ----------------------- # The last breathe of the interpreter will assume that we've # failed unless we specify otherwise. - if not self.devel: - sys.exitfunc = self.signal_interrupt + log.info('registering exit function') + sys.exitfunc = self.signal_interrupt # We overload this if ( and only if ) the topology exits # cleanly. This prevents failure modes where the monitor # dies. @@ -345,31 +329,17 @@ class Controller(object): if complete: self.send_go() - # If we're running in development stop here - # because our responsibilites are over. The - # zipline will either run to completion or die, - # monitor doesn't care anymore because its all - # threads. - - if self.devel: - log.warn("Shutting down Controller because in devel mode") - #sys.exitfunc = lambda: None - self.shutdown(soft=True) - log.info('Heartbeat (%s, %s)' % (done, complete)) # ================ # Exit Strategies # ================ - if self.zmq_flavor == 'green': - gevent.sleep(0) - # Will also fall out of loop when done, if using # non-freeform topology if done: log.info('Entire topology exited cleanly') - self.shutdown(soft=True) + self.shutdown() # Noop exit func #sys.exitfunc = lambda: None @@ -387,12 +357,12 @@ class Controller(object): we're good. The topology exited cleanly and we can prove it. """ - if not self.nosignals: - ppid = os.getppid() - log.warning("Sending SIGHUP") - os.kill(ppid, SIGHUP) - else: - log.warning("Would SIGHUP here, but disabled") + if self.inside_test: + log.warning("Skipping SIGHUP because we're in a nosetest") + return + ppid = os.getppid() + log.warning("Sending SIGHUP") + os.kill(ppid, SIGHUP) def signal_interrupt(self): """ @@ -400,11 +370,8 @@ class Controller(object): interpreter exits. If the monitor dies the system is considered a failure. """ - if not self.nosignals: - ppid = os.getpid() - os.kill(ppid, SIGINT) - else: - log.warning("Would SIGINT here, but disabled") + ppid = os.getpid() + os.kill(ppid, SIGINT) def beat(self): """ @@ -491,7 +458,7 @@ class Controller(object): def fail_universal(self): # TODO: this requires higher order functionality log.error('System in exception state, shutting down') - self.shutdown(soft=True) + self.shutdown() def fail(self, component): if self.state is CONTROL_STATES.TERMINATE: @@ -523,7 +490,7 @@ class Controller(object): Shutdown the system on failure. """ log.error('System in exception state, shutting down') - self.shutdown(soft=True) + self.kill() def exception(self, component, failure): universal = self.exception_universal @@ -532,7 +499,6 @@ class Controller(object): if component in self.topology or self.freeform: self.error_replay[(component, time.time())] = failure log.error('Component in exception state: %s' % component) - log.error(str(failure)) exception_handlers.get(component, universal)() else: @@ -638,20 +604,22 @@ class Controller(object): for (component, time), error in self.error_replay.iteritems(): log.info('Component Log for -- %s --:\n%s' % (component, error)) - def shutdown(self, hard=False, soft=True): + def kill(self): + if self.state is CONTROL_STATES.TERMINATE: + return - assert hard or soft, """ Must specify kill hard or soft """ + log.info('Hard Shutdown') + self.send_hardkill() + self.state = CONTROL_STATES.TERMINATE + self.alive = False + + + def shutdown(self): if self.state is CONTROL_STATES.TERMINATE: return + log.info('Soft Shutdown') + self.send_softkill() + self.state = CONTROL_STATES.TERMINATE self.alive = False - - if hard and not self.devel: - self.state = CONTROL_STATES.TERMINATE - log.info('Hard Shutdown') - - if soft and not self.devel: - self.state = CONTROL_STATES.TERMINATE - log.info('Soft Shutdown') - self.send_softkill() diff --git a/zipline/core/process.py b/zipline/core/process.py new file mode 100644 index 00000000..b2a01429 --- /dev/null +++ b/zipline/core/process.py @@ -0,0 +1,90 @@ +""" +The process simulator. Each component in a separate +multiprocessing.process. +""" + +import logbook +import multiprocessing +from zipline.core.host import ComponentHost + +log = logbook.Logger('Process Simulator') + +class ProcessSimulator(ComponentHost): + """ + The process simulator. + """ + + zmq_flavor = 'mp' + + def __init__(self, addresses): + ComponentHost.__init__(self, addresses) + self.subprocesses = [] + self.running = False + self.mapping = {} + + def define(self, key, val): + """ + Returns the mapping between a component and its + pid. + """ + self.mapping[key] = val + + @property + def get_id(self): + return 'Multiprocess Simulator' + + # ========= + # Launchers + # ========= + # + # invoked by the host's open() + + def launch_controller(self): + proc = multiprocessing.Process(target=self.controller.run) + proc.start() + self.con = proc + + # Process specific + self.controller_process = proc + self.mapping[proc.pid] = 'Controller' + + def launch_component(self, component): + proc = multiprocessing.Process(target=component.run) + proc.start() + self.subprocesses.append(proc) + + self.mapping[proc.pid] = component.get_id + return proc + + def simulate(self): + """ + Kick off the simulation + """ + self.run() + + def did_clean_shutdown(self): + cleanly = not any([s.is_alive() for s in self.subprocesses]) + if not cleanly: + for process in self.subprocesses: + if process.is_alive(): + log.error('Failed to Yield', self.mapping[process.pid]) + return cleanly + + def shutdown(self, ensure_clean=True): + """ + Shutdown the simulation. + """ + for component in self.components.itervalues(): + component.shutdown() + + for process in self.subprocesses: + process.join(timeout=1) + process.terminate() + + self.controller.shutdown(soft=True) + self.running = False + + self.con.terminate() + + if ensure_clean: + assert self.did_clean_shutdown() diff --git a/zipline/core/simulatorref.py b/zipline/core/simulatorref.py deleted file mode 100644 index abe87379..00000000 --- a/zipline/core/simulatorref.py +++ /dev/null @@ -1,94 +0,0 @@ -""" - -The reference simulator for all of Quantopian infastructure. - -If a subclass does not conform to the API it will fail at -compiletime. - -Subclasses: - - - (partial) zipline.devsimulator.Simulator - - ( full ) qexec.executor.simulator.ProcessSimulator - - ( full ) qexec.executor.simulator.ThreadSimulator - - ( full ) qexec.executor.simulator.GreenletSimulator - -""" - -import abc -from zipline.core.host import ComponentHost - -class SimulatorBase(ComponentHost): - - __metaclass__ = abc.ABCMeta - - def __init__(self, addresses): - """ - Initailizes the simulator. - """ - ComponentHost.__init__(self, addresses) - - @abc.abstractproperty - def get_id(self): - """Human readable name of the simulator.""" - return "Reference Simulator" - - @abc.abstractmethod - def launch_component(self, component): - """ Launch an indvidiaul component in the simulation. """ - raise NotImplementedError - - @abc.abstractmethod - def launch_controller(self): - """ Launch the controller for the simulation. """ - raise NotImplementedError - - @abc.abstractmethod - def simulate(self): - """ Run a simulation. """ - raise NotImplementedError - - @abc.abstractmethod - def shutdown(self): - """ Normal shutdown procedure. """ - raise NotImplementedError - - def cancel(self): - """ Soft shutdown """ - self.controller.shutdown(soft=True) - - def kill(self): - """ Hard shutdown """ - self.controller.shutdown(hard=True) - - # Extension Methods - # ----------------- - # Provided by some simulators, those that do not will degrade - # gracefully. - - # - ``did_clean_shutdown`` - # - ``point_of_failure`` - # - ``launch_debugger`` - - def did_clean_shutdown(self): - """ - Returns True if all the subcomponents in the simulation yielded - cleanly. - """ - return False - - def point_of_failure(self): - """ Returns the point of failure of the code. """ - failures = [ - c for c in self._components.values() - if c.exception - ] - - # Sort by failure time so we can follow the failure - # through the system. - return sorted(failures, key=lambda c: c.fail_time) - - def launch_debugger(self): - """ - Launches a remote debug shell in the context of the failed component. - """ - pass diff --git a/zipline/finance/sources.py b/zipline/finance/sources.py index 0aba2186..b9f818ea 100644 --- a/zipline/finance/sources.py +++ b/zipline/finance/sources.py @@ -40,14 +40,7 @@ class TradeDataSource(DataSource): if event.sid in self.filter['sid']: message = zp.DATASOURCE_FRAME(event) - else: - blank = ndict({ - "type" : zp.DATASOURCE_TYPE.TRADE, - "source_id" : self.get_id - }) - message = zp.DATASOURCE_FRAME(blank) - - self.data_socket.send(message) + self.data_socket.send(message) class RandomEquityTrades(TradeDataSource): @@ -90,7 +83,7 @@ class RandomEquityTrades(TradeDataSource): class SpecificEquityTrades(TradeDataSource): """ - Generates a random stream of trades for testing. + Generates a non-random stream of trades for testing. """ def init(self, event_list): diff --git a/zipline/gens/__init__.py b/zipline/gens/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py new file mode 100644 index 00000000..832909ad --- /dev/null +++ b/zipline/gens/composites.py @@ -0,0 +1,91 @@ +import datetime +from itertools import tee, starmap +from collections import namedtuple + +from zipline.gens.tradegens import SpecificEquityTrades +from zipline.gens.utils import roundrobin, hash_args +from zipline.gens.sort import date_sort +from zipline.gens.merge import merge +from zipline.gens.transform import stateful_transform + +SortBundle = namedtuple("SortBundle", ['source', 'args', 'kwargs']) +MergeBundle = namedtuple("MergeBundle", ['stream', 'tnfm', 'args', 'kwargs']) + +def date_sorted_sources(sources, source_args, source_kwargs): + """ + Takes a list of generator functions, a list of tuples of positional arguments, + and a list of dictionaries of keyword arguments. Packages up all arguments + and passes them into a date_sort. + """ + assert len(sources) == len(source_args) == len(source_kwargs) + # Package up sources and arguments. + + # Create a generator of SortBundle objects to be turned into + # namestrings and generator objects. + bundle_gen = starmap(SortBundle, zip(sources, source_args, source_kwargs)) + + # Load the results of the generator into a tuple so that the + # results can be used twice (once in namestring comprehension, + # once in the generator comprehension for intialized sources. + bundles = tuple(bundle_gen) + + # Calculate namestring hashes to pass to date_sort. + names = [bundle.source.__name__ + hash_args(*bundle.args, **bundle.kwargs) + for bundle in bundles] + # Pass each source its arguments. + initialized = [bundle.source(*bundle.args, **bundle.kwargs) + for bundle in bundles] + + # Convert the list of generators into a flat stream by pulling + # one element at a time from each. + stream_in = roundrobin(*initialized) + + # Guarantee the flat stream will be sorted by date, using source_id as + # tie-breaker, which is fully deterministic (given deterministic string + # representation for all args/kwargs) + return date_sort(stream_in, names) + + +def merged_transforms(sorted_stream, tnfms, tnfm_args, tnfm_kwargs): + """ + A generator that takes the expected output of a date_sort, pipes it + through a given set of transforms, and runs the results throught a + merge to output a unified stream. tnfms should be a list of + pointers to generator functions. tnfm_args should be a list of + tuples, representing the arguments to be passed to each transform. + tnfm_kwargs should be a list of dictionaries representing keyword + arguments to each transform. + """ + + # We should have as many sets of args as we have transforms. + assert len(tnfms) == len(tnfm_args) == len(tnfm_kwargs) + + # Create a copy of the stream for each transform. + split = tee(sorted_stream, len(tnfms)) + + # Package each transform with a stream copy and set of args. Use a list + # so that we can re-use this for calculating hashes. + bundle_gen = starmap(MergeBundle, zip(split, tnfms, tnfm_args, tnfm_kwargs)) + + bundles = tuple(bundle_gen) + # list comprehension to create transform generators from + # bundles + tnfm_gens = [ + stateful_transform( + bundle.stream, + bundle.tnfm, + *bundle.args, + **bundle.kwargs + ) + for bundle in bundles] + + # Generate expected hashes for each transform + hashes = [bundle.tnfm.__name__ + hash_args(*bundle.args, **bundle.kwargs) + for bundle in bundles] + + # Roundrobin the outputs of our transforms to create a single flat stream. + to_merge = roundrobin(*tnfm_gens) + + # Pipe the stream into merge. + merged = merge(to_merge, hashes) + return merged_transforms diff --git a/zipline/gens/merge.py b/zipline/gens/merge.py new file mode 100644 index 00000000..4778ed5b --- /dev/null +++ b/zipline/gens/merge.py @@ -0,0 +1,89 @@ +""" +Generator version of Merge. +""" + +from collections import deque + +from zipline import ndict +from zipline.gens.utils import hash_args, \ + assert_merge_protocol + + +def merge(stream_in, tnfm_ids): + """ + A generator that takes a generator and a list of source_ids. We + maintain an internal queue for each id in source_ids. Once we + have a message from every queue, we pop an event from each queue + and merge them together into an event. We raise an error if we + do not receive the same number of events from all sources. + """ + + assert isinstance(tnfm_ids, list) + + # Set up an internal queue for each expected source. + tnfms = {} + for id in tnfm_ids: + assert isinstance(id, basestring), "Bad source_id %s" % id + tnfms[id] = deque() + + # Process incoming streams. + for message in stream_in: + assert isinstance(message, tuple), \ + "Bad message in merge: %s" %message + assert len(message) == 2 + id, value = message + assert id in tnfm_ids, \ + "Message from unexpected tnfm: %s, %s" % (id, tnfm_ids) + assert isinstance(value, ndict), "Bad message in merge: %s" %message + + tnfms[id].append(value) + + # Only pop messages when we have a pending message from + # all datasources. Stop if all sources have signalled done. + + while ready(tnfms) and not done(tnfms): + message = merge_one(tnfms) + assert_merge_protocol(tnfm_ids, message) + yield message + + # We should have only a done message left in each queue. + for queue in tnfms.itervalues(): + assert len(queue) == 1, "Bad queue in merge on exit: %s" % queue + assert queue[0].dt == "DONE", \ + "Bad last message in merge on exit: %s" % queue + +def merge_one(sources): + output = ndict() + for key, queue in sources.iteritems(): + new_xform = ndict({key: queue.popleft()}) + output.merge(new_xform) + return output + + +#TODO: This is replicated in sort. Probably should be one source file. +def ready(sources): + """ + Feed is ready when every internal queue has at least one message. Note that + this include DONE messages, so done(sources) is True only if ready(sources). + """ + assert isinstance(sources, dict) + return all( (queue_is_ready(source) for source in sources.itervalues()) ) + +def queue_is_ready(queue): + assert isinstance(queue, deque) + return len(queue) > 0 + +def done(sources): + """Feed is done when all internal queues have only a "DONE" message.""" + assert isinstance(sources, dict) + return all( (queue_is_done(source) for source in sources.itervalues()) ) + +def queue_is_done(queue): + assert isinstance(queue, deque) + if len(queue) == 0: + return False + if queue[0].dt == "DONE": + assert len(queue) == 1, "Message after DONE in date_sort: %s" % queue + return True + else: + return False diff --git a/zipline/gens/mongods.py b/zipline/gens/mongods.py new file mode 100644 index 00000000..aa22d4c1 --- /dev/null +++ b/zipline/gens/mongods.py @@ -0,0 +1,96 @@ +""" +Generator-style DataSource that loads from MongoDB. +""" + +import pytz +import logbook +import pymongo + +from pymongo import ASCENDING +from datetime import datetime + +from zipline import ndict +from zipline.gens.utils import hash_args, \ + assert_trade_protocol + +import zipline.protocol as zp + +def MongoTradeHistoryGen(collection, filter, start_date, end_date): + """A generator that takes a pymongo Collection object, a list of + filters, a start date and an end_date and yields ndicts containing + the results of a query to its collection with the given filter, + start, and end. The output is also packaged with a unique + source_id string for downstream sorting + """ + + assert isinstance(collection, pymongo.collection.Collection) + assert isinstance(filter, dict) + assert isinstance(start_date, (datetime)) + assert isinstance(end_date, (datetime)) + + # Set up internal iterator. This outputs raw dictionaries. + iterator = create_pymongo_iterator(collection, filter, start_date, end_date) + + # Create unique identifier string that can be used to break + # sorting ties deterministically + argstring = hash_args(collection, filter, start_date, end_date) + source_id = "MongoTradeHistoryGen" + argstring + + # All datasources + for event in iterator: + # Construct a new event that fulfills the datasource protocol. + event['type'] = zp.DATASOURCE_TYPE.TRADE + event['dt'] = event['dt'].replace(tzinfo=pytz.utc) + event['source_id'] = source_id + + payload = ndict(event) + assert_trade_protocol(payload) + yield payload + +def create_pymongo_iterator(collection, filter, start_date, end_date): + """ + Returns an iterator that spits out raw objects loaded from a + MongoDB collection. + + See the comments on :py:class:`zipline.messaging.DataSource` + for expected content of filter. + """ + log = logbook.Logger("MongoDBQuery") + + # Object that will hold our database query. + spec = {} + + # add the filters from the algorithm. + for name, value in filter.iteritems(): + + # Add the list of sids that we care about. + if name == 'sid': + assert isinstance(value, list) + sid_range = {'sid':{'$in':value}} + spec.update(sid_range) + + # limit the data to the date range [start, end], inclusive + date_range = {'dt':{'$gte': start_date, '$lte': end_date}} + spec.update(date_range) + + fields = ['sid','price','volume','dt'] + + # In our collection, load all objects matching spec. Of those + # objects, get only the fields matching fields, and return the + # loaded objects sorted by dt from least to greatest. + + cursor = collection.find( + fields = fields, + spec = spec, + sort = [("dt",ASCENDING)], + slave_ok = True + ) + + # Optimize the cursor sort to query in 'dt' and 'sid' order. + cursor = cursor.hint([('dt', ASCENDING),('sid', ASCENDING)]) + + # Set up the iterator + iterator = iter(cursor) + log.info("MongoDataSource iterator ready") + + return iterator diff --git a/zipline/gens/sort.py b/zipline/gens/sort.py new file mode 100644 index 00000000..f6ff7a5e --- /dev/null +++ b/zipline/gens/sort.py @@ -0,0 +1,118 @@ +""" +Generator version of Feed. +""" +from collections import deque +from zipline import ndict +from zipline.gens.utils import \ + assert_datasource_unframe_protocol, \ + assert_sort_protocol + +def date_sort(stream_in, source_ids): + """ + A generator that takes a generator and a list of source_ids. We + maintain an internal queue for each id in source_ids. While we + have messages pending from all sources, we pull the earliest + message and yield it. + """ + + assert isinstance(source_ids, (list, tuple)) + + # Set up an internal queue for each expected source. + sources = {} + for id in source_ids: + assert isinstance(id, basestring), "Bad source_id %s" % id + sources[id] = deque() + + # Process incoming streams. + for message in stream_in: + # Incoming messages should be the output of DATASOURCE_UNFRAME. + assert_datasource_unframe_protocol(message), \ + "Bad message in date_sort: %s" % message + + # Only allow messages from sources we expect. + assert message.source_id in sources, "Unexpected source: %s" % message + + sources[message.source_id].append(message) + + # Only pop messages when we have a pending message from + # all datasources. Stop if all sources have signalled done. + + while ready(sources) and not done(sources): + message = pop_oldest(sources) + assert_sort_protocol(message) + yield message + + # We should have only a done message left in each queue. + for queue in sources.itervalues(): + assert len(queue) == 1, "Bad queue in date_sort on exit: %s" % queue + assert queue[0].dt == "DONE", \ + "Bad last message in date_sort on exit: %s" % queue + +def ready(sources): + """ + Feed is ready when every internal queue has at least one + message. Note that this include DONE messages, so done(sources) is + True only if ready(sources). + """ + assert isinstance(sources, dict) + return all( (queue_is_ready(source) for source in sources.itervalues()) ) + +def queue_is_ready(queue): + assert isinstance(queue, deque) + return len(queue) > 0 + +def done(sources): + """Feed is done when all internal queues have only a "DONE" message.""" + assert isinstance(sources, dict) + return all( (queue_is_done(source) for source in sources.itervalues()) ) + +def queue_is_done(queue): + assert isinstance(queue, deque) + if len(queue) == 0: + return False + if queue[0].dt == "DONE": + assert len(queue) == 1, "Message after DONE in date_sort: %s" % queue + return True + else: + return False + +def pop_oldest(sources): + + oldest_event = None + + # Iterate over the dict, checking internal queues for the oldest + # pending event. + + for queue in sources.itervalues(): + current_event = queue[0] + # Skip queues that are done. + if current_event.dt == "DONE": + continue + # Any event is older than nothing. + elif oldest_event == None: + oldest_event = current_event + # Keep the older event. Break ties by source_id. This will + # trip an assert if we have duplicate sources. + else: + oldest_event = older(oldest_event, current_event) + + # Pop the oldest event we found from its queue and return it. + return sources[oldest_event.source_id].popleft() + +# Return the event with the older timestamp. Break ties by source_id. +def older(oldest, current): + assert isinstance(oldest, ndict) + assert isinstance(oldest, ndict) + + # Try to compare by dt. + if oldest.dt < current.dt: + return oldest + elif oldest.dt > current.dt: + return current + # Break ties by source_id. + elif oldest.source_id < current.source_id: + return oldest + elif oldest.source_id > current.source_id: + return current + else: + assert False, "Duplicate event" diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py new file mode 100644 index 00000000..c7ee74f8 --- /dev/null +++ b/zipline/gens/tradegens.py @@ -0,0 +1,130 @@ +""" +Tools to generate trade events without a backing store. Useful for testing +and zipline development +""" +import random +from itertools import chain, cycle, ifilter, izip +from datetime import datetime, timedelta + +from zipline.utils.factory import create_trade +from zipline.gens.utils import hash_args, mock_done + +def date_gen(start = datetime(2012, 6, 6, 0), + delta = timedelta(minutes = 1), + count = 100): + """ + Utility to generate a stream of dates. + """ + return (start + (i * delta) for i in xrange(count)) + +def mock_prices(count, rand = False): + """ + Utility to generate a stream of mock prices. By default + cycles through values from 0.0 to 10.0, n times. Optional + flag to give random values between 0.0 and 10.0 + """ + + if rand: + return (random.uniform(0.0, 10.0) for i in xrange(count)) + else: + return (float(i % 11) for i in xrange(1,count+1)) + +def mock_volumes(count, rand = False): + """ + Utility to generate a set of volumes. By default cycles + through values from 100 to 1000, incrementing by 50. Optional + flag to give random values between 100 and 1000. + """ + if rand: + return (random.randrange(100, 1000) for i in xrange(count)) + else: + return ((i * 50)%900 + 100 for i in xrange(count)) + +def fuzzy_dates(count = 500): + """ + Add +-10 seconds to each event from a date_gen. Note that this + still guarantees sorting, since the default on date_gen is minute + separation of events. + """ + for date in date_gen(count = count): + yield date + timedelta(seconds = random.randint(-10, 10)) + +def SpecificEquityTrades(*args, **config): + """ + Yields all events in event_list that match the given sid_filter. + If no event_list is specified, generates an internal stream of events + to filter. Returns all events if filter is None. + """ + # We shouldn't get any positional arguments. + assert args == () + + # Unpack config dictionary with default values. + count = config.get('count', 500) + sids = config.get('sids', [1, 2]) + start = config.get('start', datetime(2012, 6, 6, 0)) + delta = config.get('delta', timedelta(minutes = 1)) + + # Default to None for event_list and filter. + event_list = config.get('event_list') + filter = config.get('filter') + + arg_string = hash_args(*args, **config) + namestring = "SpecificEquityTrades" + arg_string + # If we have an event_list, ignore the other arguments and use the list. + # TODO: still append our namestring? + if event_list: + unfiltered = (event for event in event_list) + + # Set up iterators for each expected field. + else: + dates = date_gen(count = count, start = start, delta = delta) + prices = mock_prices(count) + volumes = mock_volumes(count) + sids = cycle(sids) + + # Combine the iterators into a single iterator of arguments + arg_gen = izip(sids, prices, volumes, dates) + + # Convert argument packages into events. + unfiltered = (create_trade(*args, source_id = namestring) + for args in arg_gen) + + # If we specified a sid filter, filter out elements that don't match the filter. + if filter: + filtered = ifilter(lambda event: event.sid in filter, unfiltered) + + # Otherwise just use all events. + else: + filtered = unfiltered + + # Add a done message to the end of the stream. For a live + # datasource this would be handled by the containing Component. + out = chain(filtered, [mock_done(namestring)]) + return out + +def RandomEquityTrades(*args, **config): + # We shouldn't get any positional args. + assert args == () + + count = config.get('count', 500) + sids = config.get('sids', [1,2]) + filter = config.get('filter') + + dates = fuzzy_dates(count) + prices = mock_prices(count, rand = True) + volumes = mock_volumes(count, rand = True) + sids = cycle(sids) + + arg_gen = izip(sids, prices, volumes, dates) + + unfiltered = (create_trade(*args) for args in arg_gen) + + if filter: + filtered = ifilter(lambda event: event.sid in filter, unfiltered) + else: + filtered = unfiltered + return filtered + +# if __name__ == "__main__": +# import nose.tools; nose.tools.set_trace() +# trades = SpecificEquityTrades(filter = [1]) diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py new file mode 100644 index 00000000..a03a841a --- /dev/null +++ b/zipline/gens/transform.py @@ -0,0 +1,180 @@ +""" +Generator versions of transforms. +""" +import types + +from datetime import datetime +from collections import deque, defaultdict +from numbers import Number + +from zipline import ndict +from zipline.gens.utils import assert_sort_unframe_protocol, \ + assert_transform_protocol, hash_args + +class Passthrough(object): + """ + Trivial class for forwarding events. + """ + def __init__(self): + pass + + def update(self, event): + assert isinstance(event, ndict),"Bad event in Passthrough: %s" % event + assert event.has_key('sid'), "No sid in Passthrough: %s" % event + assert event.has_key('dt'), "No dt in Passthorughz: %s" % event + return event + +def functional_transform(stream_in, func, *args, **kwargs): + """ + Generic transform generator that takes each message from an in-stream + and yields the output of a function on that message. Not sure how + useful this will be in reality, but good for testing. + """ + assert isinstance(func, types.FunctionType), \ + "Functional" + namestring = func.__name__ + hash_args(*args, **kwargs) + + for message in stream_in: + assert_sort_unframe_protocol(message) + out_value = func(message, *args, **kwargs) + assert_transform_protocol(out_value) + yield(namestring, out_value) + +def stateful_transform(stream_in, tnfm_class, *args, **kwargs): + """ + Generic transform generator that takes each message from an in-stream + and sorts it to a state class. For each call to update, the state + class must produce a message to be fed downstream. + """ + + assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \ + "Stateful transform requires a class." + assert tnfm_class.__dict__.has_key('update'), \ + "Stateful transform requires the class to have an update method" + + # Create an instance of our transform class. + state = tnfm_class(*args, **kwargs) + + # Generate the string associated with this generator's output. + namestring = tnfm_class.__name__ + hash_args(*args, **kwargs) + + for message in stream_in: + assert_sort_unframe_protocol(message) + out_value = state.update(message) + assert_transform_protocol(out_value) + yield (namestring, out_value) + +class MovingAverage(object): + """ + Class that maintains a dictionary from sids to EventWindows + Upon receipt of each message we update the + corresponding window and return the calculated average. + """ + + def __init__(self, delta, fields): + self.delta = delta + self.fields = fields + + # No way to pass arguments to the defaultdict factory, so we + # need to define a method to generate the correct EventWindows. + self.sid_windows = defaultdict(self.create_window) + + def create_window(self): + """Factory method for self.sid_windows.""" + return EventWindow(self.delta, self.fields) + + def update(self, event): + """ + Update the event window for this event's sid. Return an ndict from + tracked fields to averages. + """ + + assert isinstance(event, ndict),"Bad event in MovingAverage: %s" % event + assert event.has_key('sid'), "No sid in MovingAverage: %s" % event + assert event.has_key('dt'), "No dt in MovingAverage: %s" % event + + output = ndict({'sid': event.sid, 'dt': event.dt}) + # This will create a new EventWindow if this is the first + # message for this sid. + window = self.sid_windows[event.sid] + window.update(event) + averages = window.get_averages() + + # Return the calculated averages along with + output.merge(averages) + return output + +class EventWindow(object): + """ + Maintains a list of events that are within a certain timedelta + of the most recent tick. The expected use of this class is to + track events associated with a single sid. We provide simple + functionality for averages, but anything more complicated + should be handled by a containing class. + """ + + def __init__(self, delta, fields): + self.ticks = deque() + self.delta = delta + self.fields = fields + self.totals = defaultdict(float) + + def __len__(self): + return len(self.ticks) + + def update(self, event): + self.assert_well_formed(event) + # Add new event and increment totals. + self.ticks.append(event) + for field in self.fields: + self.totals[field] += event[field] + + # We return a list of all out-of-range events we removed. + out_of_range = [] + + # Clear out expired events, decrementing totals. + # newest oldest + # | | + # V V + + while (self.ticks[-1].dt - self.ticks[0].dt) >= self.delta: + # popleft removes and returns ticks[0] + popped = self.ticks.popleft() + # Decrement totals + for field in self.fields: + self.totals[field] -= popped[field] + # Add the popped element to the list of dropped events. + out_of_range.append(popped) + + return out_of_range + + def average(self, field): + assert field in self.fields + if len(self.ticks) == 0: + return 0.0 + else: + return self.totals[field] / len(self.ticks) + + def get_averages(self): + """ + Return an ndict of all our tracked averages. + """ + out = ndict() + # out.ticks = len(self.ticks) + for field in self.fields: + out[field] = self.average(field) + return out + + def assert_well_formed(self, event): + assert isinstance(event, ndict), "Bad event in EventWindow:%s" % event + assert event.has_key('dt'), "Missing dt in EventWindow:%s" % event + assert isinstance(event.dt, datetime),"Bad dt in EventWindow:%s" % event + if len(self.ticks) > 0: + # Something is wrong if new event is older than previous. + assert event.dt >= self.ticks[-1].dt, \ + "Events arrived out of order in EventWindow: %s -> %s" % (event, self.ticks[0]) + for field in self.fields: + assert event.has_key(field), \ + "Event missing [%s] in EventWindow" % field + assert isinstance(event[field], Number), \ + "Got %s for %s in EventWindow" % (event[field], field) diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py new file mode 100644 index 00000000..e2f859cb --- /dev/null +++ b/zipline/gens/utils.py @@ -0,0 +1,102 @@ +import pytz +import numbers + +from hashlib import md5 +from datetime import datetime +from itertools import izip_longest +from zipline import ndict +from zipline.protocol import DATASOURCE_TYPE + +def mock_raw_event(sid, dt): + event = { + 'sid' : sid, + 'dt' : dt, + 'price' : 1.0, + 'volume' : 1 + } + return event + +def mock_done(source_id): + return ndict({'dt': "DONE", "source_id" : source_id, 'type' : 0}) + +def alternate(g1, g2): + """Specialized version of roundrobin for just 2 generators.""" + for e1, e2 in izip_longest(g1, g2): + if e1 != None: + yield e1 + if e2 != None: + yield e2 + +def roundrobin(*args): + """ + Takes N generators, pulling one element off each until all inputs + are empty. + """ + for elem_tuple in izip_longest(*args): + for value in elem_tuple: + if value != None: + yield value + + +def hash_args(*args, **kwargs): + """Define a unique string for any set of representable args.""" + arg_string = '_'.join([str(arg) for arg in args]) + kwarg_string = '_'.join([str(key) + '=' + str(value) for key, value in kwargs.iteritems()]) + combined = ':'.join([arg_string, kwarg_string]) + + hasher = md5() + hasher.update(combined) + return hasher.hexdigest() + +def assert_datasource_protocol(event): + """Assert that an event meets the protocol for datasource outputs.""" + + assert isinstance(event, ndict) + assert isinstance(event.source_id, basestring) + assert event.type in DATASOURCE_TYPE + + # Done packets have no dt. + if not event.type == DATASOURCE_TYPE.DONE: + assert isinstance(event.dt, datetime) + assert event.dt.tzinfo == pytz.utc + +def assert_trade_protocol(event): + """Assert that an event meets the protocol for datasource TRADE outputs.""" + assert_datasource_protocol(event) + + assert isinstance(event, ndict) + assert event.type == DATASOURCE_TYPE.TRADE + assert isinstance(event.sid, int) + assert isinstance(event.price, numbers.Real) + assert isinstance(event.volume, numbers.Integral) + assert isinstance(event.dt, datetime) + +def assert_datasource_unframe_protocol(event): + """Assert that an event is valid output of zp.DATASOURCE_UNFRAME.""" + assert isinstance(event, ndict) + assert isinstance(event.source_id, basestring) + assert event.type in DATASOURCE_TYPE + assert event.has_key('dt') + +def assert_sort_protocol(event): + """Assert that an event is valid input to zp.FEED_FRAME.""" + assert isinstance(event, ndict) + assert isinstance(event.source_id, basestring) + assert event.type in DATASOURCE_TYPE + assert event.has_key('dt') + +def assert_sort_unframe_protocol(event): + """Same as above.""" + assert isinstance(event, ndict) + assert isinstance(event.source_id, basestring) + assert event.type in DATASOURCE_TYPE + assert event.has_key('dt') + +def assert_transform_protocol(event): + """Transforms should return an ndict to be merged by merge.""" + assert isinstance(event, ndict) + +def assert_merge_protocol(tnfm_ids, message): + """Merge should output an ndict with a field for each id in its transform set.""" + assert isinstance(message, ndict) + assert set(tnfm_ids) == set(message.keys()) diff --git a/zipline/gens/zmq_gens.py b/zipline/gens/zmq_gens.py new file mode 100644 index 00000000..524852a7 --- /dev/null +++ b/zipline/gens/zmq_gens.py @@ -0,0 +1,16 @@ +import zmq + +import zipline.protocol as zp + +def gen_from_zmq(poller, unframe): + """ + A generator that takes an initialized zmq poller and yields + messages from the poller until it gets a zp.CONTROL_PROTOCOL.DONE. + """ + while True: + message = poller.recv() + if message = zp.CONTROL_PROTOCOL.DONE: + yield "DONE" + break + else: + yield unframe(message) diff --git a/zipline/lines.py b/zipline/lines.py index f56c1eb3..c20c1e46 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -69,7 +69,7 @@ from zipline.transforms import BaseTransform from zipline.test_algorithms import TestAlgorithm from zipline.components import TradeSimulationClient -from zipline.core.devsimulator import Simulator +from zipline.core.process import ProcessSimulator from zipline.core.monitor import Controller from zipline.finance.trading import SIMULATION_STYLE @@ -105,8 +105,6 @@ class SimulatedTrading(object): :py:class:`zipline.trading.TradingEnvironment` - allocator: an instance of :py:class:`zipline.simulator.AddressAllocator` - - simulator_class: a :py:class:`zipline.core.host.ComponentHost` - subclass (not an instance) - simulation_style: optional parameter that configures the :py:class:`zipline.finance.trading.TransactionSimulator`. Expects a SIMULATION_STYLE as defined in :py:mod:`zipline.finance.trading` @@ -117,7 +115,6 @@ class SimulatedTrading(object): self.trading_environment = config['trading_environment'] self.sim_style = config.get('simulation_style') - self.devel = config.get('devel', False) self.leased_sockets = [] self.sim_context = None @@ -136,15 +133,11 @@ class SimulatedTrading(object): self.con = Controller( sockets[5], sockets[6], - devel = self.devel ) - # TODO: Not freeform - self.con.manage('freeform') - self.started = False - self.sim = config['simulator_class'](addresses) + self.sim = ProcessSimulator(addresses) self.clients = {} @@ -179,9 +172,6 @@ class SimulatedTrading(object): - order_amount - the number of shares per order, defaults to 100 - trade_count - the number of trades to simulate, defaults to 101 to ensure all orders are processed. - - simulator_class - optional parameter that provides an alternative - subclass of ComponentHost to hold the whole zipline. Defaults to - :py:class:`zipline.simulator.Simulator` - algorithm - optional parameter providing an algorithm. defaults to :py:class:`zipline.test.algorithms.TestAlgorithm` - trade_source - optional parameter to specify trades, if present. @@ -221,11 +211,6 @@ class SimulatedTrading(object): # trade than order trade_count = 101 - if config.has_key('simulator_class'): - simulator_class = config['simulator_class'] - else: - simulator_class = Simulator - simulation_style = config.get('simulation_style') if not simulation_style: simulation_style = SIMULATION_STYLE.FIXED_SLIPPAGE @@ -266,22 +251,13 @@ class SimulatedTrading(object): 'algorithm' : test_algo, 'trading_environment' : trading_environment, 'allocator' : allocator, - 'simulator_class' : simulator_class, 'simulation_style' : simulation_style, 'results_socket' : results_socket, - 'devel' : config.get('devel', False) }) #------------------- zipline.add_source(trade_source) - # Save us from needless debugging - inside_test = 'nose' in inspect.stack()[-1][1] - if inside_test and not config.get('devel', False): - assert False, """ - You need to run the SimulatedTrading inside a test with devel=True - """ - return zipline def add_source(self, source): @@ -366,7 +342,7 @@ class SimulatedTrading(object): def setup_controller(self): """ - Prepare the controller tro manage the topology specified + Prepare the controller to manage the topology specified by this line. """ self.con.manage(self.topology) @@ -377,14 +353,6 @@ class SimulatedTrading(object): self.started = True self.sim_context = self.sim.simulate() - # If we're in development mode then flag all the - # components in the topology as devel so as to indicate - # that they won't poll on the control channels for - # anything other than the synchronized start. - if self.devel: - for component in self.components: - component.devel = True - # If we're using a threaded simulator block on the pool # of thread since we're only ever in a test and we don't # generally monitor the state of the system as a hold at @@ -392,16 +360,8 @@ class SimulatedTrading(object): # TODO: better way of identifying concurrency substrate if blocking: - if self.sim.zmq_flavor == 'thread': - log.debug('Blocking') - for thread in self.sim.subthreads: - #log.debug('Waiting on %r' % thread) - log.debug('Waiting on %r' % thread) - thread.join() - log.debug('Yielded on %r' % thread) - else: - for process in self.sim.subprocesses: - process.join() + for process in self.sim.subprocesses: + process.join() @property def is_success(self): diff --git a/zipline/protocol.py b/zipline/protocol.py index 866e5410..cfe6f611 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -118,16 +118,21 @@ import msgpack import numbers import datetime import pytz +import traceback +import re +import os from collections import namedtuple from utils.protocol_utils import Enum, FrameExceptionFactory, ndict, namelookup -from utils.date_utils import EPOCH, UN_EPOCH +from utils.date_utils import EPOCH, UN_EPOCH, epoch_now # ----------------------- # Control Protocol # ----------------------- +PRODUCTION_PREFIXES = ['PERF', 'RISK', 'EXCEPTION', 'CANCEL'] + INVALID_CONTROL_FRAME = FrameExceptionFactory('CONTROL') CONTROL_STATES = Enum( @@ -226,28 +231,35 @@ def DATASOURCE_FRAME(event): - *ds_type* a string denoting the datasource type. Must be on of: - TRADE + - DONE - (others to follow soon) - *payload* a msgpack string carrying the payload for the frame """ - assert isinstance(event.source_id, basestring) assert isinstance(event.type, int), 'Unexpected type %s' % (event.type) #datasources will send sometimes send empty msgs to feel gaps - if len(event.keys()) == 2: + if (event.type == DATASOURCE_TYPE.EMPTY): return msgpack.dumps(tuple([ event.type, event.source_id, - DATASOURCE_TYPE.EMPTY + "EMPTY" ])) - if(event.type == DATASOURCE_TYPE.TRADE): + elif(event.type == DATASOURCE_TYPE.TRADE): return msgpack.dumps(tuple([ event.type, event.source_id, TRADE_FRAME(event) ])) + + elif(event.type == DATASOURCE_TYPE.DONE): + return msgpack.dumps(tuple([ + event.type, + event.source_id, + "DONE" + ])) else: raise INVALID_DATASOURCE_FRAME(str(event)) @@ -259,8 +271,9 @@ def DATASOURCE_UNFRAME(msg): Returns a dict containing at least: - - source_id - - type + - source_id: instance-unique string + - type: datasource type + - dt: None, 'DONE' or a datetime object other properties are added based on the datasource type: @@ -282,6 +295,8 @@ def DATASOURCE_UNFRAME(msg): child_value = ndict({'dt':None}) elif(ds_type == DATASOURCE_TYPE.TRADE): child_value = TRADE_UNFRAME(payload) + elif(ds_type == DATASOURCE_TYPE.DONE): + child_value = ndict({'dt' : 'DONE'}) else: raise INVALID_DATASOURCE_FRAME(msg) @@ -305,6 +320,7 @@ def FEED_FRAME(event): - source_id - type + - dt """ assert isinstance(event, ndict), 'unknown type %s' % str(event) source_id = event.source_id @@ -319,6 +335,9 @@ def FEED_UNFRAME(msg): #TODO: anything we can do to assert more about the content of the dict? assert isinstance(payload, dict) rval = ndict(payload) + assert rval.source_id + assert rval.type in DATASOURCE_TYPE + assert rval.dt UNPACK_DATE(rval) return rval except TypeError: @@ -503,6 +522,50 @@ def convert_transactions(transactions): def RISK_FRAME(risk): return BT_UPDATE_FRAME('RISK', risk) +def EXCEPTION_FRAME(exception_tb): + stack_list = traceback.extract_tb(exception_tb) + rlist = [] + for stack in stack_list: + filename = shorten_filename(stack[0]) + rstack = { + 'filename' : filename, + 'lineno' : stack[1], + 'method' : stack[2], + 'line' : stack[3] + } + rlist.append(rstack) + result = { + 'date' : epoch_now(), + 'stack' : rlist + } + + return BT_UPDATE_FRAME('EXCEPTION', result) + +def shorten_filename(filename): + if filename == None: + return None + + # check if the path contains zipeline_repo + path_re = r'(?<=zipline_repo).*' + match = re.search(path_re, filename) + + if match: + return match.group(0) + parts = filename.split('zipline_repo') + return parts[1] + else: + # return just the filename. + head, tail = os.path.split(filename) + return tail + +def CANCEL_FRAME(date): + result = { + 'date' : EPOCH(date) + } + + return BT_UPDATE_FRAME('CANCEL', result) + + def BT_UPDATE_FRAME(prefix, payload): """ Frames prepared by RISK_FRAME and PERF_FRAME methods are sent via the same @@ -579,6 +642,7 @@ def tuple_to_date(date_tuple): DATASOURCE_TYPE = Enum( 'TRADE', 'EMPTY', + 'DONE' ) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index aa07fa54..d66cdef0 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -54,7 +54,7 @@ class TestAlgorithm(): at the close of a simulation. """ - def __init__(self, sid, amount, order_count): + def __init__(self, sid, amount, order_count, sid_filter=None): self.count = order_count self.sid = sid self.amount = amount @@ -63,6 +63,10 @@ class TestAlgorithm(): self.order = None self.frame_count = 0 self.portfolio = None + if sid_filter: + self.sid_filter = sid_filter + else: + self.sid_filter = [self.sid] def initialize(self): pass @@ -84,7 +88,7 @@ class TestAlgorithm(): self.incr += 1 def get_sid_filter(self): - return [self.sid] + return self.sid_filter # class HeavyBuyAlgorithm(): @@ -145,7 +149,7 @@ class NoopAlgorithm(object): pass def get_sid_filter(self): - return None + return [] class ExceptionAlgorithm(object): """ @@ -153,8 +157,9 @@ class ExceptionAlgorithm(object): constructor. """ - def __init__(self, throw_from): - self.throw_from == throw_from + def __init__(self, throw_from, sid): + self.throw_from = throw_from + self.sid = sid def initialize(self): if self.throw_from == "initialize": @@ -187,12 +192,12 @@ class ExceptionAlgorithm(object): if self.throw_from == "get_sid_filter": raise Exception("Algo exception in get_sid_filter") else: - return [1] + return [self.sid] class TestPrintAlgorithm(): - def __init__(self): - pass + def __init__(self, sid): + self.sid = sid def initialize(self): print "Initializing..." @@ -211,12 +216,13 @@ class TestPrintAlgorithm(): pass def get_sid_filter(self): - return [1] + return [self.sid] class TestLoggingAlgorithm(): - def __init__(self): + def __init__(self, sid): self.log = None + self.sid = sid def initialize(self): self.log.info("Initializing...") @@ -234,4 +240,4 @@ class TestLoggingAlgorithm(): self.log.info("Handling Data...") def get_sid_filter(self): - return [1] + return [self.sid] diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 6101c3b2..db440891 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -69,9 +69,9 @@ def create_trading_environment(year=2006): return trading_environment -def create_trade(sid, price, amount, datetime): +def create_trade(sid, price, amount, datetime, source_id = "test_factory"): row = zp.ndict({ - 'source_id' : "test_factory", + 'source_id' : source_id, 'type' : zp.DATASOURCE_TYPE.TRADE, 'sid' : sid, 'dt' : datetime, @@ -89,6 +89,8 @@ def get_next_trading_dt(current, interval, trading_calendar): return next + + def create_trade_history(sid, prices, amounts, interval, trading_calendar): trades = [] current = trading_calendar.first_open diff --git a/zipline/utils/gpoll.py b/zipline/utils/gpoll.py deleted file mode 100644 index 8c18d4cf..00000000 --- a/zipline/utils/gpoll.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -This is somewhat legally ambigious, since it technically -hasn't been merged in gevent_zeromq but given that the -author issued it as a Pull Request on a MIT project, -indicates that its probably fine to use. ~Steve -""" - -import zmq -from zmq import * - -from zmq.core.poll import Poller as _original_Poller - -import gevent -from gevent import select -from gevent_zeromq.core import _Socket - -def patch_poller(self): - zmq.Poller = _Poller - -class _Poller(_original_Poller): - """ - Replacement for :class:`zmq.core.Poller` - - Ensures that the greened Poller below is used in calls - to :meth:`zmq.core.Poller.poll`. - """ - - def _get_descriptors(self): - """ - Returns three elements tuple with socket descriptors ready for - gevent.select - """ - rlist = [] - wlist = [] - xlist = [] - - for socket, flags in self.sockets.items(): - if isinstance(socket, _Socket): - fd = socket.getsockopt(FD) - elif isinstance(socket, int): - fd = socket - elif hasattr(socket, 'fileno'): - try: - fd = int(socket.fileno()) - except: - raise ValueError('fileno() must return an valid integer fd') - else: - raise TypeError("Socket must be a 0MQ socket, an integer fd or \ - have a fileno() method: %r" % socket) - - if flags & POLLIN: rlist.append(fd) - if flags & POLLOUT: wlist.append(fd) - if flags & POLLERR: xlist.append(fd) - - return (rlist, wlist, xlist) - - def poll(self, timeout=-1): - """Overridden method to ensure that the green version of Poller is used - - Behaves the same as :meth:`zmq.core.Poller.poll` - """ - - if timeout is None: - timeout = -1 - - timeout = int(timeout) - if timeout < 0: - timeout = -1 - - rlist = None - wlist = None - xlist = None - - if timeout > 0: - tout = gevent.Timeout.start_new(timeout/1000.0) - - try: - # Loop until timeout or events available - while True: - events = super(_Poller, self).poll(0) - if events or timeout == 0: - return events - - # wait for activity on sockets in a green way - if not rlist and not wlist and not xlist: - rlist, wlist, xlist = self._get_descriptors() - - try: - select.select(rlist, wlist, xlist) - except gevent.select.error, ex: - raise ZMQError(*ex.args) - - except gevent.Timeout, t: - if t is not tout: - raise - return [] - finally: - if timeout > 0: - tout.cancel() diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py new file mode 100644 index 00000000..6655c265 --- /dev/null +++ b/zipline/utils/test_utils.py @@ -0,0 +1,142 @@ +import zmq +import time +import zipline.protocol as zp +from datetime import datetime +import blist +from zipline.utils.date_utils import EPOCH +from itertools import izip +from logbook import FileHandler + +def setup_logger(test, path='/var/log/zipline/zipline.log'): + test.log_handler = FileHandler(path) + test.log_handler.push_application() + +def teardown_logger(test): + test.log_handler.pop_application() + +def check_list(test, a, b, label): + test.assertTrue(isinstance(a, (list, blist.blist))) + test.assertTrue(isinstance(b, (list, blist.blist))) + i = 0 + for a_val, b_val in izip(a, b): + check(test, a_val, b_val, label + "[" + str(i) + "]") + + +def check_dict(test, a, b, label): + test.assertTrue(isinstance(a, dict)) + test.assertTrue(isinstance(b, dict)) + for key in a.keys(): + # ignore the extra fields used by dictshield + if key in ['progress']: + continue + test.assertTrue(a.has_key(key), "missing key at: " + label + "." + key) + test.assertTrue(b.has_key(key), "missing key at: " + label + "." + key) + a_val = a[key] + b_val = b[key] + check(test, a_val, b_val, label + "." + key) + + +def check_datetime(test, a, b, label): + test.assertTrue(isinstance(a, datetime)) + test.assertTrue(isinstance(b, datetime)) + test.assertEqual(EPOCH(a), EPOCH(b), "mismatched dates " + label) + + +def check(test, a, b, label=None): + """ + Check equality for arbitrarily nested dicts and lists that terminate + in types that allow direct comparisons (string, ints, floats, datetimes) + """ + if not label: + label = '' + if isinstance(a, dict): + check_dict(test, a, b, label) + elif isinstance(a, (list, blist.blist)): + check_list(test, a, b, label) + elif isinstance(a, datetime): + check_datetime(test, a, b, label) + else: + test.assertEqual(a, b, "mismatch on path: " + label) + + +def drain_zipline(test, zipline): + assert test.ctx, "method expects a valid zmq context" + assert test.zipline_test_config, "method expects a valid test config" + assert isinstance(test.zipline_test_config, dict) + assert test.zipline_test_config['results_socket'], \ + "need to specify a socket address for logs/perf/risk" + test.receiver = create_receiver( + test.zipline_test_config['results_socket'], + test.ctx + ) + # Bind and connect are asynch, so allow time for bind before + # starting the zipline (TSC connects internally). + time.sleep(1) + + # start the simulation + zipline.simulate(blocking=False) + output, transaction_count = drain_receiver(test.receiver) + # some processes will exit after the message stream is + # finished. We block here to avoid collisions with subsequent + # ziplines. + for process in zipline.sim.subprocesses: + process.join() + + return output, transaction_count + +def create_receiver(socket_addr, ctx): + receiver = ctx.socket(zmq.PULL) + receiver.bind(socket_addr) + + return receiver + +def drain_receiver(receiver): + output = [] + transaction_count = 0 + while True: + msg = receiver.recv() + if msg == str(zp.CONTROL_PROTOCOL.DONE): + break + else: + update = zp.BT_UPDATE_UNFRAME(msg) + output.append(update) + if update['prefix'] == 'PERF': + transaction_count += \ + len(update['payload']['daily_perf']['transactions']) + elif update['prefix'] == 'EXCEPTION': + break + + receiver.close() + del receiver + + return output, transaction_count + + +def assert_single_position(test, zipline): + output, transaction_count = drain_zipline(test, zipline) + + test.assertTrue(zipline.sim.ready()) + test.assertFalse(zipline.sim.exception) + + test.assertEqual( + test.zipline_test_config['order_count'], + transaction_count + ) + + # the final message is the risk report, the second to + # last is the final day's results. Positions is a list of + # dicts. + closing_positions = output[-2]['payload']['daily_perf']['positions'] + + test.assertEqual( + len(closing_positions), + 1, + "Portfolio should have one position." + ) + + sid = test.zipline_test_config['sid'] + test.assertEqual( + closing_positions[0]['sid'], + sid, + "Portfolio should have one position in " + str(sid) + ) diff --git a/zipline/utils/zmq_utils.py b/zipline/utils/zmq_utils.py index 5e0a529f..49177ee4 100644 --- a/zipline/utils/zmq_utils.py +++ b/zipline/utils/zmq_utils.py @@ -1,5 +1,5 @@ """ -Misc ZeroMQ utilities. +Misc ZeroMQ experimental tools """ import gevent import msgpack