IPython/zmq/parallel/client.py

#!/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()