mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 15:38:11 +08:00
fawce
155 lines
3.8 KiB
Python
155 lines
3.8 KiB
Python
"""
|
|
Misc ZeroMQ experimental tools
|
|
"""
|
|
import gevent
|
|
import msgpack
|
|
import numpy
|
|
from numpy import dtype
|
|
from pandas import DataFrame
|
|
from gevent_zeromq import zmq
|
|
|
|
from contextlib import closing
|
|
|
|
class ZmqDone(Exception):
|
|
|
|
def __init__(self, socket, frame):
|
|
self.ident = socket.identity
|
|
self.frame = str(frame)
|
|
|
|
def __str__(self):
|
|
return 'Socket ( %s ) finished with frame ( %s )' % \
|
|
( self.ident, self.frame )
|
|
|
|
class zs(object):
|
|
"""
|
|
A wrapper for the *very* common pattern of reading from a
|
|
upstream socket until you get a DONE or EXCEPTION frame.
|
|
|
|
# Eliminates all the boilerplate serialization logic
|
|
# and error handling cases into 3 lines.
|
|
|
|
halts = (ERROR_FRAME, CLOSE_FRAME)
|
|
stream = zs(socket, halts)
|
|
|
|
stream.on_error(YouFailAtFailing)
|
|
|
|
for msg in stream:
|
|
print msg
|
|
|
|
"""
|
|
|
|
def __init__(self, socket, halts, srl=msgpack):
|
|
self._socket = socket
|
|
self.exc_case = halts[0]
|
|
self.done_case = halts[1]
|
|
|
|
self.loads = srl.loads
|
|
self.halt_method = 'exception'
|
|
self.exception = ZmqDone
|
|
self.function = None
|
|
|
|
def __iter__(self):
|
|
self.last = msg = self.loads(self._socket.recv())
|
|
|
|
if msg == self.exc_case:
|
|
return self.halt()
|
|
|
|
if msg == self.done_case:
|
|
raise StopIteration
|
|
|
|
yield msg
|
|
|
|
def last(self):
|
|
return self.last
|
|
|
|
def halt(self):
|
|
if self.halt_method == 'exception':
|
|
raise self.exception
|
|
elif self.halt_method == 'function':
|
|
return self.function()
|
|
|
|
def on_error(self, callee):
|
|
|
|
if isinstance(callee, Exception):
|
|
self.halt_method = 'exception'
|
|
self.exception = callee
|
|
else:
|
|
self.halt_method = 'function'
|
|
self.function = callee
|
|
|
|
def ZmqConsole(sock_typ, socket_addr, sock_conn=None, context=None):
|
|
"""
|
|
A utility to drop into a ZeroMQ pdb console and inspect
|
|
messages as they come through. If you just want to pipe to
|
|
stdout, don't use this.
|
|
"""
|
|
|
|
context = context or zmq.Context.instance()
|
|
socket = context.socket(zmq.PULL)
|
|
socket.bind(socket_addr)
|
|
|
|
def console():
|
|
while True:
|
|
msg = socket.recv_pyobj()
|
|
print msg
|
|
import pdb; pdb.set_trace()
|
|
|
|
return gevent.spawn(console)
|
|
|
|
class NumpyChannel(zmq.Socket):
|
|
|
|
def recv_pandas(self, flags=0, copy=True, track=False):
|
|
|
|
# Pandas Metadata
|
|
index, columns, dtype_name, shape = msgpack.loads(self.recv(flags=flags))
|
|
|
|
# Pandas ndarray
|
|
ndbuffer = self.recv(flags=flags, copy=copy, track=track)
|
|
buf = buffer(ndbuffer)
|
|
|
|
ndarray = numpy.frombuffer(buf, dtype=dtype(dtype_name)).reshape(shape)
|
|
return DataFrame(data=ndarray, index=index,
|
|
columns=columns, dtype=dtype_name)
|
|
|
|
def send_pandas(self, df, flags=0, copy=True, track=False):
|
|
|
|
# Pandas Metadata
|
|
index = df.index.tolist()
|
|
columns = df.columns.tolist()
|
|
dtype_name = df.values.dtype.name
|
|
shape = df.values.shape
|
|
|
|
# Pandas ndarray
|
|
ndarray = df.values
|
|
|
|
metadata = msgpack.dumps((index, columns, dtype_name, shape))
|
|
|
|
self.send(metadata, flags|zmq.SNDMORE)
|
|
return self.send(ndarray, flags, copy=copy, track=track)
|
|
|
|
if __name__ == '__main__':
|
|
|
|
from numpy.random import randn
|
|
df = DataFrame(randn(5,5))
|
|
|
|
ctx = zmq.Context.instance()
|
|
|
|
def send():
|
|
pub = NumpyChannel(ctx, zmq.PUSH)
|
|
pub.bind('inproc://a')
|
|
|
|
for i in xrange(100):
|
|
pub.send_pandas(df, copy=False)
|
|
|
|
def recv():
|
|
sub = NumpyChannel(ctx, zmq.PULL)
|
|
sub.connect('inproc://a')
|
|
|
|
for i in xrange(100):
|
|
sub.recv_pandas(copy=False)
|
|
|
|
gevent.joinall([
|
|
gevent.spawn(send),
|
|
gevent.spawn(recv)
|
|
])
|