mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-08 10:21:19 +08:00
Merge branch 'master' of github.com:quantopian/zipline
This commit is contained in:
fawce
+50
-19
@@ -8,17 +8,16 @@ class Controller(object):
|
||||
polling = False
|
||||
debug = False
|
||||
|
||||
def __init__(self, pull_socket, pub_socket, context=None, logging = None):
|
||||
def __init__(self, pull_socket, pub_socket, logging = None):
|
||||
|
||||
self._ctx = None
|
||||
|
||||
self.associated = []
|
||||
|
||||
if not context:
|
||||
self._ctx = zmq.Context()
|
||||
else:
|
||||
self._ctx = context
|
||||
|
||||
self.pull_socket = pull_socket
|
||||
self.push_socket = pull_socket # same port
|
||||
self.pub_socket = pub_socket
|
||||
self.sub_socket = pub_socket # same port
|
||||
|
||||
if logging:
|
||||
self.logging = logging
|
||||
@@ -30,18 +29,26 @@ class Controller(object):
|
||||
self.success = 0
|
||||
self.failed = 0
|
||||
|
||||
def run(self, debug=False):
|
||||
def run(self, debug=False, context=None):
|
||||
self.polling = True
|
||||
|
||||
#if debug:
|
||||
return self._poll()
|
||||
#else:
|
||||
#return self._poll_fast()
|
||||
if not context:
|
||||
self._ctx = zmq.Context()
|
||||
else:
|
||||
self._ctx = context
|
||||
|
||||
if debug:
|
||||
return self._poll_fast() # the c loop
|
||||
else:
|
||||
return self._poll() # use a python loop
|
||||
|
||||
def _poll_fast(self):
|
||||
"""
|
||||
C version of the polling forwarder.
|
||||
"""
|
||||
self.pull = self._ctx.socket(zmq.PULL)
|
||||
self.pub = self._ctx.socket(zmq.PUB)
|
||||
|
||||
zmq.device(zmq.FORWARDER, self.pull, self.pub)
|
||||
|
||||
def _poll(self):
|
||||
@@ -60,7 +67,9 @@ class Controller(object):
|
||||
|
||||
while self.polling:
|
||||
try:
|
||||
self.pub.send(self.pull.recv())
|
||||
msg = self.pull.recv()
|
||||
print msg
|
||||
self.pub.send(msg)
|
||||
except KeyboardInterrupt:
|
||||
self.polling = False
|
||||
break
|
||||
@@ -75,24 +84,36 @@ class Controller(object):
|
||||
self.failed += 1
|
||||
continue
|
||||
|
||||
def message_sender(self):
|
||||
# -------------------
|
||||
# Hooks for Endpoints
|
||||
# -------------------
|
||||
|
||||
def message_sender(self, context=None):
|
||||
"""
|
||||
Spin off a socket used for sending messages to this
|
||||
controller.
|
||||
"""
|
||||
s = self._ctx.socket(zmq.PUSH)
|
||||
s.connect(self.pull_socket)
|
||||
|
||||
if not context:
|
||||
context = zmq.Context()
|
||||
|
||||
s = context.socket(zmq.PUSH)
|
||||
s.connect(self.push_socket)
|
||||
self.associated.append(s)
|
||||
return s
|
||||
|
||||
def message_listener(self):
|
||||
def message_listener(self, context = None, filters=None):
|
||||
"""
|
||||
Spin off a socket used for receiving messages from this
|
||||
controller.
|
||||
"""
|
||||
s = self._ctx.socket(zmq.SUB)
|
||||
s.connect(self.pub_socket)
|
||||
s.setsockopt(zmq.SUBSCRIBE, '')
|
||||
|
||||
if not context:
|
||||
context = zmq.Context()
|
||||
|
||||
s = context.socket(zmq.SUB)
|
||||
s.connect(self.sub_socket)
|
||||
s.setsockopt(zmq.SUBSCRIBE, filters or '')
|
||||
self.associated.append(s)
|
||||
return s
|
||||
|
||||
@@ -116,3 +137,13 @@ class Controller(object):
|
||||
return
|
||||
return float(self.success) / (self.success + self.failed)
|
||||
|
||||
if __name__ == '__main__':
|
||||
print 'Running on ',\
|
||||
'tcp://127.0.0.1:5000', \
|
||||
'tcp://127.0.0.1:5001',
|
||||
|
||||
controller = Controller(
|
||||
'tcp://127.0.0.1:5000',
|
||||
'tcp://127.0.0.1:5001',
|
||||
)
|
||||
controller.run()
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
Dummy simulator backported from Qexec for development on Zipline.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from multiprocessing import Process
|
||||
from unittest2 import TestCase
|
||||
|
||||
from zipline.monitor import Controller
|
||||
|
||||
import gevent
|
||||
from gevent_zeromq import zmq
|
||||
|
||||
ctx = zmq.Context()
|
||||
|
||||
class TestControlProtocol(TestCase):
|
||||
|
||||
def setUpController(self):
|
||||
self.controller.run()
|
||||
|
||||
def setUp(self):
|
||||
self.controller = Controller(
|
||||
'tcp://127.0.0.1:5000',
|
||||
'tcp://127.0.0.1:5001',
|
||||
)
|
||||
self.control_proc = Process(target=self.setUpController)
|
||||
self.control_proc.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.control_proc.terminate()
|
||||
ctx.destroy()
|
||||
|
||||
def asyncMessage(self, socket):
|
||||
return socket.recv()
|
||||
|
||||
def send_and_receive(self, push, sub, message='\x01'):
|
||||
msg = gevent.spawn(sub.recv)
|
||||
push.send(message)
|
||||
gevent.sleep(0) # explicit gevent yield
|
||||
msg.join()
|
||||
self.assertEqual(msg.value, message)
|
||||
|
||||
def test_control_message(self):
|
||||
|
||||
sub = self.controller.message_listener(context=ctx)
|
||||
message = gevent.spawn(self.asyncMessage, sub)
|
||||
|
||||
push = self.controller.message_sender(context=ctx)
|
||||
|
||||
# Don't like introducing time constants but because of
|
||||
# the way gevent scheduler works zmq will often send all
|
||||
# the messages off before the other thread even gets to
|
||||
# listening.
|
||||
self.send_and_receive(push, sub)
|
||||
sub.close()
|
||||
push.close()
|
||||
|
||||
def test_control_delivery(self):
|
||||
# Assert that the number of messages sent on the wire is
|
||||
# the number of messages received, ie we don't drop any.
|
||||
# This is of course depenendent on the topology of the
|
||||
# listeners being fixed. Which normally it isn't.
|
||||
|
||||
sub = self.controller.message_listener(context=ctx)
|
||||
message = gevent.spawn(self.asyncMessage, sub)
|
||||
|
||||
push = self.controller.message_sender(context=ctx)
|
||||
|
||||
# Don't like introducing time constants but because of
|
||||
# the way gevent scheduler works zmq will often send all
|
||||
# the messages off before the other thread even gets to
|
||||
# listening.
|
||||
for i in xrange(25):
|
||||
self.send_and_receive(push, sub)
|
||||
|
||||
sub.close()
|
||||
push.close()
|
||||
@@ -0,0 +1,189 @@
|
||||
import uuid
|
||||
import copy
|
||||
import atexit
|
||||
import pickle
|
||||
|
||||
from datetime import datetime
|
||||
from collections import defaultdict
|
||||
|
||||
from UserDict import DictMixin
|
||||
|
||||
class Snapshot(object, DictMixin):
|
||||
"""
|
||||
A snapshot in time of a history container.
|
||||
"""
|
||||
|
||||
def __init__(self, state, version, ts):
|
||||
self.version = version
|
||||
self.timestamp = ts
|
||||
self._state = state
|
||||
|
||||
def keys(self):
|
||||
return self._state.keys()
|
||||
|
||||
def values(self):
|
||||
return self._state.values()
|
||||
|
||||
def items(self):
|
||||
return self._state.items()
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._state.__getitem__(key)
|
||||
|
||||
def has_key(self, key):
|
||||
return self._state.has_key(key)
|
||||
|
||||
def copy(self):
|
||||
return copy.copy(self._state)
|
||||
|
||||
class History(object, DictMixin):
|
||||
"""
|
||||
A duck-typed dictionary that tracks its time evolution.
|
||||
|
||||
Worth noting this not a particuarly high-performance
|
||||
data structure due to the copious amount of copying going on.
|
||||
"""
|
||||
|
||||
def __init__(self, default=None):
|
||||
if default:
|
||||
initial = defaultdict(default)
|
||||
else:
|
||||
initial = {}
|
||||
|
||||
self.version = 0
|
||||
self.changeset = [('CREATE', None)]
|
||||
self.current = Snapshot(initial, version=self.version, ts=datetime.now())
|
||||
self._history = [self.current]
|
||||
|
||||
def items(self, version=-1):
|
||||
return self._history[version].items()
|
||||
|
||||
def keys(self, version=-1):
|
||||
return self._history[version].keys()
|
||||
|
||||
def rollback(self, version):
|
||||
pass
|
||||
|
||||
def event(self, tup):
|
||||
self.changeset.append(tup)
|
||||
|
||||
def __getitem__(self, key, version=-1):
|
||||
return self._history[version].__getitem__(key)
|
||||
|
||||
def __setitem__(self, key, val):
|
||||
if self.current.has_key(key):
|
||||
self.changeset.append(('CHANGE', key))
|
||||
else:
|
||||
self.changeset.append(('ADD', key))
|
||||
|
||||
state = self.current.copy()
|
||||
state[key] = val
|
||||
|
||||
self.version += 1
|
||||
self.current = Snapshot(state, self.version, datetime.now())
|
||||
self._history.append(self.current)
|
||||
|
||||
def __delitem__(self, key):
|
||||
self.changeset.append(('REMOVE', key))
|
||||
|
||||
state = self.current.copy()
|
||||
del state[key]
|
||||
|
||||
self.version += 1
|
||||
self.current = Snapshot(state, self.version, datetime.now())
|
||||
self._history.append(self.current)
|
||||
|
||||
def history(self):
|
||||
for change in self.changeset:
|
||||
print change
|
||||
|
||||
def __repr__(self):
|
||||
return ':'.join(['historical', self.current._state.__repr__()])
|
||||
|
||||
SocketHistory = History()
|
||||
ContextHistory = History()
|
||||
|
||||
def patch_zmq(_zmq=None):
|
||||
"""
|
||||
Monkey patch zeromq to allow for socket tracking.
|
||||
"""
|
||||
if _zmq:
|
||||
zmq = _zmq
|
||||
else:
|
||||
import zmq
|
||||
|
||||
_Context = zmq.Context
|
||||
_Socket = zmq.Socket
|
||||
|
||||
class TrackedSocket(zmq.Socket):
|
||||
|
||||
def __init__(self, context, socket_type):
|
||||
self.context = context
|
||||
self.uuid = str(uuid.uuid4())
|
||||
SocketHistory[self.uuid] = self
|
||||
_Socket.__init__(self, context, socket_type)
|
||||
|
||||
def connect(self, address):
|
||||
SocketHistory.event(('CONNECT', self.uuid, address))
|
||||
_Socket.connect(self, address)
|
||||
|
||||
def bind(self, address):
|
||||
SocketHistory.event(('BIND', self.uuid, address))
|
||||
_Socket.bind(self, address)
|
||||
|
||||
def close(self, *args, **kwargs):
|
||||
del SocketHistory[self.uuid]
|
||||
_Socket.close(self, *args, **kwargs)
|
||||
|
||||
def setsockopt(self, option, optval):
|
||||
if option == zmq.IDENTITY:
|
||||
old = SocketHistory[self.uuid]
|
||||
SocketHistory[optval] = old
|
||||
del SocketHistory[self.uuid]
|
||||
self.uuid = optval
|
||||
|
||||
_Socket.setsockopt(self, option, optval)
|
||||
|
||||
class TrackedContext(zmq.Context):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.sockets = {}
|
||||
_Context.__init__(self, *args, **kwargs)
|
||||
self.uuid = str(uuid.uuid4())
|
||||
ContextHistory[self.uuid] = self
|
||||
|
||||
def socket(self, socket_type):
|
||||
sock = TrackedSocket(self, socket_type)
|
||||
ContextHistory.event(('EMBED', self.uuid, sock.uuid))
|
||||
self.sockets[sock.uuid] = sock
|
||||
return sock
|
||||
|
||||
def name(self, name):
|
||||
"""
|
||||
Name the context. Is a superset of the vanilla pyzmq
|
||||
API.
|
||||
"""
|
||||
old = ContextHistory[self.context.uuid]
|
||||
ContextHistory[name] = old
|
||||
del ContextHistory[self.context.uuid]
|
||||
self.uuid = name
|
||||
|
||||
def term(self, *args, **kwargs):
|
||||
for uid, sock in self.sockets.iteritems():
|
||||
if not sock.closed:
|
||||
del SocketHistory[sock.uuid]
|
||||
del ContextHistory[self.uuid]
|
||||
_Context.term(self, *args, **kwargs)
|
||||
|
||||
def destroy(self, *args, **kwargs):
|
||||
ContextHistory.event(('DESTROY', self.uuid))
|
||||
_Context.destroy(self, *args, **kwargs)
|
||||
|
||||
zmq.Context = TrackedContext
|
||||
zmq.Socket = TrackedSocket
|
||||
return TrackedContext, TrackedSocket
|
||||
|
||||
def track_to_file(f):
|
||||
def write_track():
|
||||
pickle.dump(SocketHistory.changeset, file(f, 'wb+'))
|
||||
atexit.register(write_track)
|
||||
Reference in New Issue
Block a user