##// END OF EJS Templates
added zmq controller/engine entry points
added zmq controller/engine entry points

File last commit:

r3548:48fdd06d
r3550:3ab94ec8
Show More
client.py
639 lines | 21.9 KiB | text/x-python | PythonLexer
#!/usr/bin/env python
"""A semi-synchronous Client for the ZMQ controller"""
import time
from pprint import pprint
from IPython.external.decorator import decorator
import streamsession as ss
import zmq
from zmq.eventloop import ioloop, zmqstream
from remotenamespace import RemoteNamespace
from view import DirectView
from dependency import Dependency, depend, require
def _push(ns):
globals().update(ns)
def _pull(keys):
g = globals()
if isinstance(keys, (list,tuple)):
return map(g.get, keys)
else:
return g.get(keys)
def _clear():
globals().clear()
def execute(code):
exec code in globals()
# decorators for methods:
@decorator
def spinfirst(f,self,*args,**kwargs):
self.spin()
return f(self, *args, **kwargs)
@decorator
def defaultblock(f, self, *args, **kwargs):
block = kwargs.get('block',None)
block = self.block if block is None else block
saveblock = self.block
self.block = block
ret = f(self, *args, **kwargs)
self.block = saveblock
return ret
class AbortedTask(object):
def __init__(self, msg_id):
self.msg_id = msg_id
# @decorator
# def checktargets(f):
# @wraps(f)
# def checked_method(self, *args, **kwargs):
# self._build_targets(kwargs['targets'])
# return f(self, *args, **kwargs)
# return checked_method
# class _ZMQEventLoopThread(threading.Thread):
#
# def __init__(self, loop):
# self.loop = loop
# threading.Thread.__init__(self)
#
# def run(self):
# self.loop.start()
#
class Client(object):
"""A semi-synchronous client to the IPython ZMQ controller
Attributes
----------
ids : set
a set of engine IDs
requesting the ids attribute always synchronizes
the registration state. To request ids without synchronization,
use _ids
history : list of msg_ids
a list of msg_ids, keeping track of all the execution
messages you have submitted
outstanding : set of msg_ids
a set of msg_ids that have been submitted, but whose
results have not been received
results : dict
a dict of all our results, keyed by msg_id
block : bool
determines default behavior when block not specified
in execution methods
Methods
-------
spin : flushes incoming results and registration state changes
control methods spin, and requesting `ids` also ensures up to date
barrier : wait on one or more msg_ids
execution methods: apply/apply_bound/apply_to
legacy: execute, run
query methods: queue_status, get_result
control methods: abort, kill
"""
_connected=False
_engines=None
registration_socket=None
query_socket=None
control_socket=None
notification_socket=None
queue_socket=None
task_socket=None
block = False
outstanding=None
results = None
history = None
debug = False
def __init__(self, addr, context=None, username=None, debug=False):
if context is None:
context = zmq.Context()
self.context = context
self.addr = addr
if username is None:
self.session = ss.StreamSession()
else:
self.session = ss.StreamSession(username)
self.registration_socket = self.context.socket(zmq.PAIR)
self.registration_socket.setsockopt(zmq.IDENTITY, self.session.session)
self.registration_socket.connect(addr)
self._engines = {}
self._ids = set()
self.outstanding=set()
self.results = {}
self.history = []
self.debug = debug
self.session.debug = debug
self._notification_handlers = {'registration_notification' : self._register_engine,
'unregistration_notification' : self._unregister_engine,
}
self._queue_handlers = {'execute_reply' : self._handle_execute_reply,
'apply_reply' : self._handle_apply_reply}
self._connect()
@property
def ids(self):
self._flush_notifications()
return self._ids
def _update_engines(self, engines):
for k,v in engines.iteritems():
eid = int(k)
self._engines[eid] = bytes(v) # force not unicode
self._ids.add(eid)
def _build_targets(self, targets):
if targets is None:
targets = self._ids
elif isinstance(targets, str):
if targets.lower() == 'all':
targets = self._ids
else:
raise TypeError("%r not valid str target, must be 'all'"%(targets))
elif isinstance(targets, int):
targets = [targets]
return [self._engines[t] for t in targets], list(targets)
def _connect(self):
"""setup all our socket connections to the controller"""
if self._connected:
return
self._connected=True
self.session.send(self.registration_socket, 'connection_request')
idents,msg = self.session.recv(self.registration_socket,mode=0)
if self.debug:
pprint(msg)
msg = ss.Message(msg)
content = msg.content
if content.status == 'ok':
if content.queue:
self.queue_socket = self.context.socket(zmq.PAIR)
self.queue_socket.setsockopt(zmq.IDENTITY, self.session.session)
self.queue_socket.connect(content.queue)
if content.task:
self.task_socket = self.context.socket(zmq.PAIR)
self.task_socket.setsockopt(zmq.IDENTITY, self.session.session)
self.task_socket.connect(content.task)
if content.notification:
self.notification_socket = self.context.socket(zmq.SUB)
self.notification_socket.connect(content.notification)
self.notification_socket.setsockopt(zmq.SUBSCRIBE, "")
if content.query:
self.query_socket = self.context.socket(zmq.PAIR)
self.query_socket.setsockopt(zmq.IDENTITY, self.session.session)
self.query_socket.connect(content.query)
if content.control:
self.control_socket = self.context.socket(zmq.PAIR)
self.control_socket.setsockopt(zmq.IDENTITY, self.session.session)
self.control_socket.connect(content.control)
self._update_engines(dict(content.engines))
else:
self._connected = False
raise Exception("Failed to connect!")
#### handlers and callbacks for incoming messages #######
def _register_engine(self, msg):
content = msg['content']
eid = content['id']
d = {eid : content['queue']}
self._update_engines(d)
self._ids.add(int(eid))
def _unregister_engine(self, msg):
# print 'unregister',msg
content = msg['content']
eid = int(content['id'])
if eid in self._ids:
self._ids.remove(eid)
self._engines.pop(eid)
def _handle_execute_reply(self, msg):
# msg_id = msg['msg_id']
parent = msg['parent_header']
msg_id = parent['msg_id']
if msg_id not in self.outstanding:
print "got unknown result: %s"%msg_id
else:
self.outstanding.remove(msg_id)
self.results[msg_id] = ss.unwrap_exception(msg['content'])
def _handle_apply_reply(self, msg):
# pprint(msg)
# msg_id = msg['msg_id']
parent = msg['parent_header']
msg_id = parent['msg_id']
if msg_id not in self.outstanding:
print "got unknown result: %s"%msg_id
else:
self.outstanding.remove(msg_id)
content = msg['content']
if content['status'] == 'ok':
self.results[msg_id] = ss.unserialize_object(msg['buffers'])
elif content['status'] == 'aborted':
self.results[msg_id] = AbortedTask(msg_id)
elif content['status'] == 'resubmitted':
pass # handle resubmission
else:
self.results[msg_id] = ss.unwrap_exception(content)
def _flush_notifications(self):
"flush incoming notifications of engine registrations"
msg = self.session.recv(self.notification_socket, mode=zmq.NOBLOCK)
while msg is not None:
if self.debug:
pprint(msg)
msg = msg[-1]
msg_type = msg['msg_type']
handler = self._notification_handlers.get(msg_type, None)
if handler is None:
raise Exception("Unhandled message type: %s"%msg.msg_type)
else:
handler(msg)
msg = self.session.recv(self.notification_socket, mode=zmq.NOBLOCK)
def _flush_results(self, sock):
"flush incoming task or queue results"
msg = self.session.recv(sock, mode=zmq.NOBLOCK)
while msg is not None:
if self.debug:
pprint(msg)
msg = msg[-1]
msg_type = msg['msg_type']
handler = self._queue_handlers.get(msg_type, None)
if handler is None:
raise Exception("Unhandled message type: %s"%msg.msg_type)
else:
handler(msg)
msg = self.session.recv(sock, mode=zmq.NOBLOCK)
def _flush_control(self, sock):
"flush incoming control replies"
msg = self.session.recv(sock, mode=zmq.NOBLOCK)
while msg is not None:
if self.debug:
pprint(msg)
msg = self.session.recv(sock, mode=zmq.NOBLOCK)
###### get/setitem ########
def __getitem__(self, key):
if isinstance(key, int):
if key not in self.ids:
raise IndexError("No such engine: %i"%key)
return DirectView(self, key)
if isinstance(key, slice):
indices = range(len(self.ids))[key]
ids = sorted(self._ids)
key = [ ids[i] for i in indices ]
# newkeys = sorted(self._ids)[thekeys[k]]
if isinstance(key, (tuple, list, xrange)):
_,targets = self._build_targets(list(key))
return DirectView(self, targets)
else:
raise TypeError("key by int/iterable of ints only, not %s"%(type(key)))
############ begin real methods #############
def spin(self):
"""flush incoming notifications and execution results."""
if self.notification_socket:
self._flush_notifications()
if self.queue_socket:
self._flush_results(self.queue_socket)
if self.task_socket:
self._flush_results(self.task_socket)
if self.control_socket:
self._flush_control(self.control_socket)
@spinfirst
def queue_status(self, targets=None, verbose=False):
"""fetch the status of engine queues
Parameters
----------
targets : int/str/list of ints/strs
the engines on which to execute
default : all
verbose : bool
whether to return lengths only, or lists of ids for each element
"""
targets = self._build_targets(targets)[1]
content = dict(targets=targets)
self.session.send(self.query_socket, "queue_request", content=content)
idents,msg = self.session.recv(self.query_socket, 0)
if self.debug:
pprint(msg)
return msg['content']
@spinfirst
@defaultblock
def clear(self, targets=None, block=None):
"""clear the namespace in target(s)"""
targets = self._build_targets(targets)[0]
print targets
for t in targets:
self.session.send(self.control_socket, 'clear_request', content={},ident=t)
error = False
if self.block:
for i in range(len(targets)):
idents,msg = self.session.recv(self.control_socket,0)
if self.debug:
pprint(msg)
if msg['content']['status'] != 'ok':
error = msg['content']
if error:
return error
@spinfirst
@defaultblock
def abort(self, msg_ids = None, targets=None, block=None):
"""abort the Queues of target(s)"""
targets = self._build_targets(targets)[0]
print targets
if isinstance(msg_ids, basestring):
msg_ids = [msg_ids]
content = dict(msg_ids=msg_ids)
for t in targets:
self.session.send(self.control_socket, 'abort_request',
content=content, ident=t)
error = False
if self.block:
for i in range(len(targets)):
idents,msg = self.session.recv(self.control_socket,0)
if self.debug:
pprint(msg)
if msg['content']['status'] != 'ok':
error = msg['content']
if error:
return error
@spinfirst
@defaultblock
def kill(self, targets=None, block=None):
"""Terminates one or more engine processes."""
targets = self._build_targets(targets)[0]
print targets
for t in targets:
self.session.send(self.control_socket, 'kill_request', content={},ident=t)
error = False
if self.block:
for i in range(len(targets)):
idents,msg = self.session.recv(self.control_socket,0)
if self.debug:
pprint(msg)
if msg['content']['status'] != 'ok':
error = msg['content']
if error:
return error
@defaultblock
def execute(self, code, targets='all', block=None):
"""executes `code` on `targets` in blocking or nonblocking manner.
Parameters
----------
code : str
the code string to be executed
targets : int/str/list of ints/strs
the engines on which to execute
default : all
block : bool
whether or not to wait until done
"""
# block = self.block if block is None else block
# saveblock = self.block
# self.block = block
result = self.apply(execute, (code,), targets=targets, block=block, bound=True)
# self.block = saveblock
return result
def run(self, code, block=None):
"""runs `code` on an engine.
Calls to this are load-balanced.
Parameters
----------
code : str
the code string to be executed
block : bool
whether or not to wait until done
"""
result = self.apply(execute, (code,), targets=None, block=block, bound=False)
return result
def _apply_balanced(self, f, args, kwargs, bound=True, block=None,
after=None, follow=None):
"""the underlying method for applying functions in a load balanced
manner."""
block = block if block is not None else self.block
bufs = ss.pack_apply_message(f,args,kwargs)
content = dict(bound=bound)
msg = self.session.send(self.task_socket, "apply_request",
content=content, buffers=bufs)
msg_id = msg['msg_id']
self.outstanding.add(msg_id)
self.history.append(msg_id)
if block:
self.barrier(msg_id)
return self.results[msg_id]
else:
return msg_id
def _apply_direct(self, f, args, kwargs, bound=True, block=None, targets=None,
after=None, follow=None):
"""Then underlying method for applying functions to specific engines."""
block = block if block is not None else self.block
queues,targets = self._build_targets(targets)
print queues
bufs = ss.pack_apply_message(f,args,kwargs)
if isinstance(after, Dependency):
after = after.as_dict()
elif after is None:
after = []
if isinstance(follow, Dependency):
follow = follow.as_dict()
elif follow is None:
follow = []
subheader = dict(after=after, follow=follow)
content = dict(bound=bound)
msg_ids = []
for queue in queues:
msg = self.session.send(self.queue_socket, "apply_request",
content=content, buffers=bufs,ident=queue, subheader=subheader)
msg_id = msg['msg_id']
self.outstanding.add(msg_id)
self.history.append(msg_id)
msg_ids.append(msg_id)
if block:
self.barrier(msg_ids)
else:
if len(msg_ids) == 1:
return msg_ids[0]
else:
return msg_ids
if len(msg_ids) == 1:
return self.results[msg_ids[0]]
else:
result = {}
for target,mid in zip(targets, msg_ids):
result[target] = self.results[mid]
return result
def apply(self, f, args=None, kwargs=None, bound=True, block=None, targets=None,
after=None, follow=None):
"""calls f(*args, **kwargs) on a remote engine(s), returning the result.
if self.block is False:
returns msg_id or list of msg_ids
else:
returns actual result of f(*args, **kwargs)
"""
# enforce types of f,args,kwrags
args = args if args is not None else []
kwargs = kwargs if kwargs is not None else {}
if not callable(f):
raise TypeError("f must be callable, not %s"%type(f))
if not isinstance(args, (tuple, list)):
raise TypeError("args must be tuple or list, not %s"%type(args))
if not isinstance(kwargs, dict):
raise TypeError("kwargs must be dict, not %s"%type(kwargs))
options = dict(bound=bound, block=block, after=after, follow=follow)
if targets is None:
return self._apply_balanced(f, args, kwargs, **options)
else:
return self._apply_direct(f, args, kwargs, targets=targets, **options)
def push(self, ns, targets=None, block=None):
"""push the contents of `ns` into the namespace on `target`"""
if not isinstance(ns, dict):
raise TypeError("Must be a dict, not %s"%type(ns))
result = self.apply(_push, (ns,), targets=targets, block=block,bound=True)
return result
@spinfirst
def pull(self, keys, targets=None, block=True):
"""pull objects from `target`'s namespace by `keys`"""
result = self.apply(_pull, (keys,), targets=targets, block=block, bound=True)
return result
def barrier(self, msg_ids=None, timeout=-1):
"""waits on one or more `msg_ids`, for up to `timeout` seconds.
Parameters
----------
msg_ids : int, str, or list of ints and/or strs
ints are indices to self.history
strs are msg_ids
default: wait on all outstanding messages
timeout : float
a time in seconds, after which to give up.
default is -1, which means no timeout
Returns
-------
True : when all msg_ids are done
False : timeout reached, msg_ids still outstanding
"""
tic = time.time()
if msg_ids is None:
theids = self.outstanding
else:
if isinstance(msg_ids, (int, str)):
msg_ids = [msg_ids]
theids = set()
for msg_id in msg_ids:
if isinstance(msg_id, int):
msg_id = self.history[msg_id]
theids.add(msg_id)
self.spin()
while theids.intersection(self.outstanding):
if timeout >= 0 and ( time.time()-tic ) > timeout:
break
time.sleep(1e-3)
self.spin()
return len(theids.intersection(self.outstanding)) == 0
@spinfirst
def get_results(self, msg_ids,status_only=False):
"""returns the result of the execute or task request with `msg_id`"""
if not isinstance(msg_ids, (list,tuple)):
msg_ids = [msg_ids]
theids = []
for msg_id in msg_ids:
if isinstance(msg_id, int):
msg_id = self.history[msg_id]
theids.append(msg_id)
content = dict(msg_ids=theids, status_only=status_only)
msg = self.session.send(self.query_socket, "result_request", content=content)
zmq.select([self.query_socket], [], [])
idents,msg = self.session.recv(self.query_socket, zmq.NOBLOCK)
if self.debug:
pprint(msg)
# while True:
# try:
# except zmq.ZMQError:
# time.sleep(1e-3)
# continue
# else:
# break
return msg['content']
class AsynClient(Client):
"""An Asynchronous client, using the Tornado Event Loop"""
io_loop = None
queue_stream = None
notifier_stream = None
def __init__(self, addr, context=None, username=None, debug=False, io_loop=None):
Client.__init__(self, addr, context, username, debug)
if io_loop is None:
io_loop = ioloop.IOLoop.instance()
self.io_loop = io_loop
self.queue_stream = zmqstream.ZMQStream(self.queue_socket, io_loop)
self.control_stream = zmqstream.ZMQStream(self.control_socket, io_loop)
self.task_stream = zmqstream.ZMQStream(self.task_socket, io_loop)
self.notification_stream = zmqstream.ZMQStream(self.notification_socket, io_loop)
def spin(self):
for stream in (self.queue_stream, self.notifier_stream,
self.task_stream, self.control_stream):
stream.flush()