diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py index d1561837..ac6339bc 100644 --- a/tests/test_exception_handling.py +++ b/tests/test_exception_handling.py @@ -3,7 +3,8 @@ import zmq from unittest2 import TestCase from collections import defaultdict -from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm +from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm, \ + InitializeTimeoutAlgorithm, TooMuchProcessingAlgorithm from zipline.finance.trading import SIMULATION_STYLE from zipline.core.devsimulator import AddressAllocator from zipline.lines import SimulatedTrading @@ -143,3 +144,35 @@ class ExceptionTestCase(TestCase): # make sure our path shortening is working self.assertEqual(payload['stack'][0]['filename'], '/zipline/lines.py') self.assertEqual(payload['stack'][-1]['filename'], '/zipline/test_algorithms.py') + + def test_initialize_timeout(self): + + self.zipline_test_config['algorithm'] = \ + InitializeTimeoutAlgorithm( + self.zipline_test_config['sid'] + ) + + zipline = SimulatedTrading.create_test_zipline( + **self.zipline_test_config + ) + output, _ = drain_zipline(self, zipline) + self.assertEqual(output[-1]['prefix'], 'EXCEPTION') + payload = output[-1]['payload'] + self.assertEqual(payload['name'],'Timeout') + self.assertEqual(payload['message'], 'Call to initialize timed out') + +# def test_heartbeat(self): + +# self.zipline_test_config['algorithm'] = \ +# TooMuchProcessingAlgorithm( +# self.zipline_test_config['sid'] +# ) +# zipline = SimulatedTrading.create_test_zipline( +# **self.zipline_test_config +# ) +# output, _ = drain_zipline(self, zipline) +# self.assertEqual(output[-1]['prefix'], 'EXCEPTION') +# payload = output[-1]['payload'] +# self.assertEqual(payload['name'],'Timeout') +# self.assertEqual(payload['message'], 'Too much time spent in handle_data call') + diff --git a/tests/test_finance.py b/tests/test_finance.py index e5f26240..94a98da9 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -121,36 +121,6 @@ class FinanceTestCase(TestCase): zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) assert_single_position(self, zipline) - #@timed(DEFAULT_TIMEOUT) - def test_sid_filter(self): - # 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 - order_count = 100 - no_match_sid = 222 - test_algo = TestAlgorithm( - no_match_sid, - order_amount, - order_count - ) - - self.zipline_test_config['trade_count'] = 200 - self.zipline_test_config['algorithm'] = test_algo - - zipline = SimulatedTrading.create_test_zipline( - **self.zipline_test_config - ) - output, transaction_count = drain_zipline(self, zipline) - - #check that the algorithm received no events - self.assertEqual( - 0, - transaction_count, - "The algorithm should not receive any events due to filtering." - ) - - # TODO: write tests for short sales # TODO: write a test to do massive buying or shorting. diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 22afa34f..811f99ac 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -1,9 +1,12 @@ +import signal from logbook import Logger, Processor from datetime import datetime, timedelta from numbers import Integral +from itertools import groupby from zipline import ndict +from zipline.utils.timeout import timeout, heartbeat, Timeout from zipline.gens.transform import StatefulTransform from zipline.finance.trading import TransactionSimulator @@ -13,10 +16,15 @@ from zipline.gens.utils import hash_args log = Logger('Trade Simulation') +# TODO: make these arguments rather than global constants +INIT_TIMEOUT = 5 +HEARTBEAT_INTERVAL = 1 # seconds +MAX_HEARTBEAT_INTERVALS = 15 + class TradeSimulationClient(object): """ - Generator that takes the expected output of a merge, a user - algorithm, a trading environment, and a simulator style as + Generator-style class that takes the expected output of a merge, a + user algorithm, a trading environment, and a simulator style as arguments. Pipes the merge stream through a TransactionSimulator and a PerformanceTracker, which keep track of the current state of our algorithm's simulated universe. Results are fed to the user's @@ -24,7 +32,7 @@ class TradeSimulationClient(object): TransactionSimulator's order book. TransactionSimulator maintains a dictionary from sids to the - unfulfilled orders placed by the user's algorithm. As trade + as-yet unfilled orders placed by the user's algorithm. As trade events arrive, if the algorithm has open orders against the trade's sid, the simulator will fill orders up to 25% of market cap. Applied transactions are added to a txn field on the event @@ -40,9 +48,9 @@ class TradeSimulationClient(object): performance report, which is appended to event's perf_report field. - Fully processed events are run through a batcher generator, which - batches together events with the same dt field into a single event - to be fed to the algo. The portfolio object is repeatedly + Fully processed events are fed to AlgorithmSimulator, which + batches together events with the same dt field into a single + snapshot to be fed to the algo. The portfolio object is repeatedly overwritten so that only the most recent snapshot of the universe is sent to the algo. """ @@ -54,13 +62,14 @@ class TradeSimulationClient(object): self.environment = environment self.style = sim_style self.algo_sim = None - + self.warmup_start = self.environment.prior_day_open self.algo_start = self.environment.first_open def get_hash(self): """ - There should only ever be one TSC in the system. + There should only ever be one TSC in the system, so + we don't bother passing args into the hash. """ return self.__class__.__name__ + hash_args() @@ -92,9 +101,9 @@ class TradeSimulationClient(object): with_portfolio = perf_tracker.transform(with_filled_orders) # Pass the messages from perf along with the trading client's - # state into the algorithm for simulation. We provide the - # trading client so that the algorithm can place new orders - # into the client's order book. + # state into the algorithm for simulation. We provide a + # pointer to the ordering client's internal state so that the + # algorithm can place new orders into the client's order book. self.algo_sim = AlgorithmSimulator( with_portfolio, ordering_client.state, @@ -106,12 +115,17 @@ class TradeSimulationClient(object): # calculated by the performance tracker) at the end of each # day. It will also yield a risk report at the end of the # simulation. + for message in self.algo_sim: yield message class AlgorithmSimulator(object): - def __init__(self, stream_in, order_book, algo, algo_start): + def __init__(self, + stream_in, + order_book, + algo, + algo_start): self.stream_in = stream_in @@ -128,15 +142,30 @@ class AlgorithmSimulator(object): self.algo_start = algo_start # Monkey patch the user algorithm to place orders in the - # TransactionSimulator's order book. + # TransactionSimulator's order book and use our logger. self.algo.set_order(self.order) - self.algo.set_logger(Logger("AlgoLog")) + self.algolog = Logger("AlgoLog") + self.algo.set_logger(self.algolog) + # Handler for heartbeats during calls to handle_data. + def log_heartbeats(beat_count, stackframe): + t = beat_count * HEARTBEAT_INTERVAL + warning = "handle_data has been processing for %i seconds" %t + self.algolog.warn(warning) + + # Context manager that calls log_heartbeats every HEARTBEAT_INTERVAL + # seconds, raising an exception after MAX_HEARTBEATS + self.heartbeat_monitor = heartbeat( + HEARTBEAT_INTERVAL, + MAX_HEARTBEAT_INTERVALS, + frame_handler=log_heartbeats, + timeout_message="Too much time spent in handle_data call" + ) # ============== # Snapshot Setup # ============== - + # The algorithm's universe as of our most recent event. self.universe = ndict() @@ -147,7 +176,7 @@ class AlgorithmSimulator(object): # We don't have a datetime for the current snapshot until we # receive a message. self.simulation_dt = None - self.this_snapshot_dt = None + self.snapshot_dt = None # ============= # Logging Setup @@ -156,7 +185,7 @@ class AlgorithmSimulator(object): # Processor function for injecting the algo_dt into # user prints/logs. def inject_algo_dt(record): - record.extra['algo_dt'] = self.this_snapshot_dt + record.extra['algo_dt'] = self.snapshot_dt self.processor = Processor(inject_algo_dt) # This is a class, which is instantiated later @@ -211,110 +240,61 @@ class AlgorithmSimulator(object): # Capture any output of this generator to stdout and pipe it # to a logbook interface. Also inject the current algo # snapshot time to any log record generated. - with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''): - # Call the user's initialize method. - self.algo.initialize() - for event in self.stream_in: + # Call user's initialize method with a timeout. + with timeout(INIT_TIMEOUT, message="Call to initialize timed out"): + self.algo.initialize() + + # Group together events with the same dt field. This depends on the + # events already being sorted. + for date, snapshot in groupby(self.stream_in, lambda e: e.dt): + + # Set the simulation date to be the first event we see. + # This should only occur once, at the start of the test. + if self.simulation_dt == None: + self.simulation_dt = date + + # Done message has the risk report, so we yield before exiting. + if date == 'DONE': + for event in snapshot: + yield event.perf_message # We're still in the warmup period. Use the event to - # update our universe, but don't start a snapshot or - # pass anything to handle_data. Discard any - # perf messages. - if event.dt != 'DONE' and event.dt < self.algo_start: - self.update_universe(event) - if event.perf_message: - log.info("Discarding perf message because we're in warmup.") - continue + # update our universe, but don't yield any perf messages, + # and don't send a snapshot to handle_data. + elif date < self.algo_start: + for event in snapshot: + del event['perf_message'] + self.update_universe(event) - # Yield any perf messages received to be relayed back to - # the browser. - - if event.perf_message: - yield event.perf_message - del event['perf_message'] - - if event.dt == "DONE": - if self.this_snapshot_dt: - # StopIteration happened mid-snapshot, so we - # have a universe snapshot that is not yet - # processed by the algorithm. - self.simulate_current_snapshot() - - # Break out of the loop, causing us to raise - # StopIteration This needs to be outside the check - # on self.this_snapshot_dt or else getting a DONE - # immediately after a snapshot finishes will cause - # type errors. - break - - # This should only happen for the first event we run. - if self.simulation_dt == None: - self.simulation_dt = event.dt - - # ====================== - # Time Compression Logic - # ====================== - - if self.this_snapshot_dt != None: - self.update_current_snapshot(event) - - # The algorithm has been missing events because it took - # too long processing. Update the universe with data from - # this event, then check if enough time has passed that we - # can start a new snapshot. + # The algo has taken so long to process events that + # its simulated time is later than the event time. + # Update the universe and yield any perf messages + # encountered, but don't call handle_data. + elif date < self.simulation_dt: + for event in snapshot: + # Only yield if we have something interesting to say. + if event.perf_message != None: + yield event.perf_message + # Delete the message before updating so we don't send it + # to the user. + del event['perf_message'] + self.update_universe(event) + + # Regular snapshot. Update the universe and send a snapshot + # to handle data. else: - self.update_universe(event) - if event.dt >= self.simulation_dt: - self.this_snapshot_dt = event.dt - - - - def update_current_snapshot(self, event): - """ - Update our current snapshot of the universe. Call handle_data if - """ - # The new event matches our snapshot dt. Just update the - # universe and move on. - if event.dt == self.this_snapshot_dt: - self.update_universe(event) - - # The new event does not match our snapshot. - else: - self.simulate_current_snapshot() - - # Once we've finished simulating the old snapshot, - # we can update the universe with the new event. - self.update_universe(event) - - # The current event is later than the simulation time, - # which means the algorithm finished quickly enough to - # receive the new event. Start a new snapshot with this - # event's dt. - if event.dt >= self.simulation_dt: - self.this_snapshot_dt = event.dt - - # The algorithm spent enough time processing that it - # missed the new event. Wait to start a new snapshot until - # the events catch up to the algo's simulated dt. - else: - self.this_snapshot_dt = None - - def simulate_current_snapshot(self): - """ - Run the user's algo against our current snapshot and update the algo's - simulated time. - """ - start_tic = datetime.now() - self.algo.handle_data(self.universe) - stop_tic = datetime.now() - - # How long did you take? - delta = stop_tic - start_tic - - # Update the simulation time. - self.simulation_dt = self.this_snapshot_dt + delta + for event in snapshot: + # Only yield if we have something interesting to say. + if event.perf_message != None: + yield event.perf_message + del event['perf_message'] + + self.update_universe(event) + + # Send the current state of the universe to the user's algo. + self.simulate_snapshot(date) def update_universe(self, event): """ @@ -326,3 +306,23 @@ class AlgorithmSimulator(object): # Update our knowledge of this event's sid for field in event.keys(): self.universe[event.sid][field] = event[field] + + def simulate_snapshot(self, date): + """ + Run the user's algo against our current snapshot and update + the algo's simulated time. + """ + # Needs to be set so that we inject the proper date into algo + # log/print lines. + self.snapshot_dt = date + + start_tic = datetime.now() + with self.heartbeat_monitor: + self.algo.handle_data(self.universe) + stop_tic = datetime.now() + + # How long did you take? + delta = stop_tic - start_tic + + # Update the simulation time. + self.simulation_dt = date + delta diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index a7881fa8..f6bdfd4d 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -221,6 +221,56 @@ class DivByZeroAlgorithm(): def get_sid_filter(self): return [self.sid] +class InitializeTimeoutAlgorithm(): + def __init__(self, sid): + self.sid = sid + self.incr = 0 + + def initialize(self): + import time + from zipline.gens.tradesimulation import INIT_TIMEOUT + time.sleep(INIT_TIMEOUT + 1) + + def set_order(self, order_callable): + pass + + def set_logger(self, logger): + pass + + def set_portfolio(self, portfolio): + pass + + def handle_data(self, data): + pass + + def get_sid_filter(self): + return [self.sid] + +class TooMuchProcessingAlgorithm(): + def __init__(self, sid): + self.sid = sid + + def initialize(self): + pass + + def set_order(self, order_callable): + pass + + def set_logger(self, logger): + pass + + def set_portfolio(self, portfolio): + pass + + def handle_data(self, data): + # Unless we're running on some sort of + # supercomputer this will hit timeout. + for i in xrange(100000000): + self.foo = i + + def get_sid_filter(self): + return [self.sid] + class TimeoutAlgorithm(): def __init__(self, sid): diff --git a/zipline/utils/timeout.py b/zipline/utils/timeout.py new file mode 100644 index 00000000..0310a208 --- /dev/null +++ b/zipline/utils/timeout.py @@ -0,0 +1,125 @@ +import signal + +from functools import wraps + +from pprint import pprint as pp +from numbers import Number +from logbook import Logger + +class Timeout(Exception): + + def __init__(self, frame, message=''): + self.frame = frame + self.message = message + +class timeout(object): + """ + Utility to make a function raise TimeoutException if it spends + more than a specified number of seconds executing. Can be used + as a decorator to apply a static timeout to a function, or as + a context manager to dynamically add a timeout to a code block. + """ + + def __init__(self, seconds, message=''): + self.seconds = seconds + self.message = message + assert isinstance(seconds, Number), "Failed to specify a timeout." + assert seconds > 0, "Timeout must be greater than 0" + + def handler(self, signum, frame): + raise Timeout(frame, self.message) + + def __call__(self, fn): + + @wraps(fn) + def call_fn_with_timeout(*args, **kwargs): + # Set the alarm. + signal.signal(signal.SIGALRM, self.handler) + signal.setitimer(signal.ITIMER_REAL, self.seconds, 0) + try: + outval = fn(*args, **kwargs) + + # Deactivate the alarm once we're done so that the + # decorator doesn't have unexpected side-effects later. + # Note that this will still raise Timeout if the + # call to fn takes too long. + finally: + signal.setitimer(signal.ITIMER_REAL, 0, 0) + signal.signal(signal.SIGALRM, signal.SIG_DFL) + + # Return the value of fn if it finished before the alarm. This + # won't execute if the Timeout was raised. + return outval + return call_fn_with_timeout + + def __enter__(self): + # Set the alarm on entrance. + signal.signal(signal.SIGALRM, self.handler) + signal.setitimer(signal.ITIMER_REAL, self.seconds, 0) + + def __exit__(self, type, value, traceback): + # Deactivate the alarm on exit. This will re-raise + # any exceptions raised inside the with block. + signal.signal(signal.SIGALRM, self.handler) + signal.setitimer(signal.ITIMER_REAL, 0, 0) + +class heartbeat(object): + """ + Utility to perform pseudo-heartbeat checks on a single-threaded + function. Calls frame_handler on the current stack frame of the + wrapped function every ``interval`` seconds. After ``max_interval`` + intervals, raises Timeout. Can be used either as a decorator or + a context manager. + """ + def __init__(self, + interval, + max_intervals, + frame_handler=None, + timeout_message=''): + + self.interval = interval + self.max_intervals = max_intervals + self.frame_handler = frame_handler + self.timeout_message = timeout_message + self.count = 0 + + def handler(self, signum, frame): + self.count += 1 + if self.frame_handler: + self.frame_handler(self.count, frame) + + if self.count >= self.max_intervals: + raise Timeout(frame, self.timeout_message) + + def __call__(self, fn): + + @wraps(fn) + def call_fn_with_heartbeat(*args, **kwargs): + # Set a timer to call our handler every ``interval`` seconds. + signal.signal(signal.SIGALRM, self.handler) + signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval) + try: + outval = fn(*args, **kwargs) + + finally: + # Deactivate the timer once we're done so that the + # decorator doesn't have unexpected side-effects later. + signal.setitimer(signal.ITIMER_REAL, 0, 0) + signal.signal(signal.SIGALRM, signal.SIG_DFL) + + # Return the value of fn if it finished without tripping + # an exception. This won't execute if the Timeout or any + # other exception was raised by self.handle. + return outval + return call_fn_with_heartbeat + + def __enter__(self): + # Set a timer to call our handler every N seconds. + signal.signal(signal.SIGALRM, self.handler) + signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval) + + def __exit__(self, type, value, traceback): + # Turn off the timer on exit. This will re-raise any exception raised + # during execution of the with-block + signal.setitimer(signal.ITIMER_REAL, 0, 0) + signal.signal(signal.SIGALRM, signal.SIG_DFL)