allow Reference as callable in map/apply

Assumptions were made that the first argument was a callable/function with
a __name__ attribute.  These assumptions were the only barrier to using
References, and have been removed.

Associated tests included.

IPython/parallel/client/client.py

@@ -1,1443 +1,1444 b''
 """A semi-synchronous Client for the ZMQ cluster
 
 Authors:
 
 * MinRK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import os
 import json
 import sys
 import time
 import warnings
 from datetime import datetime
 from getpass import getpass
 from pprint import pprint
 
 pjoin = os.path.join
 
 import zmq
 # from zmq.eventloop import ioloop, zmqstream
 
 from IPython.config.configurable import MultipleInstanceError
 from IPython.core.application import BaseIPythonApplication
 
 from IPython.utils.jsonutil import rekey
 from IPython.utils.localinterfaces import LOCAL_IPS
 from IPython.utils.path import get_ipython_dir
 from IPython.utils.traitlets import (HasTraits, Integer, Instance, Unicode,
                                     Dict, List, Bool, Set)
 from IPython.external.decorator import decorator
 from IPython.external.ssh import tunnel
 
+from IPython.parallel import Reference
 from IPython.parallel import error
 from IPython.parallel import util
 
 from IPython.zmq.session import Session, Message
 
 from .asyncresult import AsyncResult, AsyncHubResult
 from IPython.core.profiledir import ProfileDir, ProfileDirError
 from .view import DirectView, LoadBalancedView
 
 if sys.version_info[0] >= 3:
     # xrange is used in a couple 'isinstance' tests in py2
     # should be just 'range' in 3k
     xrange = range
 
 #--------------------------------------------------------------------------
 # Decorators for Client methods
 #--------------------------------------------------------------------------
 
 @decorator
 def spin_first(f, self, *args, **kwargs):
     """Call spin() to sync state prior to calling the method."""
     self.spin()
     return f(self, *args, **kwargs)
 
 
 #--------------------------------------------------------------------------
 # Classes
 #--------------------------------------------------------------------------
 
 class Metadata(dict):
     """Subclass of dict for initializing metadata values.
 
     Attribute access works on keys.
 
     These objects have a strict set of keys - errors will raise if you try
     to add new keys.
     """
     def __init__(self, *args, **kwargs):
         dict.__init__(self)
         md = {'msg_id' : None,
               'submitted' : None,
               'started' : None,
               'completed' : None,
               'received' : None,
               'engine_uuid' : None,
               'engine_id' : None,
               'follow' : None,
               'after' : None,
               'status' : None,
 
               'pyin' : None,
               'pyout' : None,
               'pyerr' : None,
               'stdout' : '',
               'stderr' : '',
             }
         self.update(md)
         self.update(dict(*args, **kwargs))
 
     def __getattr__(self, key):
         """getattr aliased to getitem"""
         if key in self.iterkeys():
             return self[key]
         else:
             raise AttributeError(key)
 
     def __setattr__(self, key, value):
         """setattr aliased to setitem, with strict"""
         if key in self.iterkeys():
             self[key] = value
         else:
             raise AttributeError(key)
 
     def __setitem__(self, key, value):
         """strict static key enforcement"""
         if key in self.iterkeys():
             dict.__setitem__(self, key, value)
         else:
             raise KeyError(key)
 
 
 class Client(HasTraits):
     """A semi-synchronous client to the IPython ZMQ cluster
 
     Parameters
     ----------
 
     url_or_file : bytes or unicode; zmq url or path to ipcontroller-client.json
         Connection information for the Hub's registration.  If a json connector
         file is given, then likely no further configuration is necessary.
         [Default: use profile]
     profile : bytes
         The name of the Cluster profile to be used to find connector information.
         If run from an IPython application, the default profile will be the same
         as the running application, otherwise it will be 'default'.
     context : zmq.Context
         Pass an existing zmq.Context instance, otherwise the client will create its own.
     debug : bool
         flag for lots of message printing for debug purposes
     timeout : int/float
         time (in seconds) to wait for connection replies from the Hub
         [Default: 10]
 
     #-------------- session related args ----------------
 
     config : Config object
         If specified, this will be relayed to the Session for configuration
     username : str
         set username for the session object
     packer : str (import_string) or callable
         Can be either the simple keyword 'json' or 'pickle', or an import_string to a
         function to serialize messages. Must support same input as
         JSON, and output must be bytes.
         You can pass a callable directly as `pack`
     unpacker : str (import_string) or callable
         The inverse of packer.  Only necessary if packer is specified as *not* one
         of 'json' or 'pickle'.
 
     #-------------- ssh related args ----------------
     # These are args for configuring the ssh tunnel to be used
     # credentials are used to forward connections over ssh to the Controller
     # Note that the ip given in `addr` needs to be relative to sshserver
     # The most basic case is to leave addr as pointing to localhost (127.0.0.1),
     # and set sshserver as the same machine the Controller is on. However,
     # the only requirement is that sshserver is able to see the Controller
     # (i.e. is within the same trusted network).
 
     sshserver : str
         A string of the form passed to ssh, i.e. 'server.tld' or '[e-mail]:port'
         If keyfile or password is specified, and this is not, it will default to
         the ip given in addr.
     sshkey : str; path to ssh private key file
         This specifies a key to be used in ssh login, default None.
         Regular default ssh keys will be used without specifying this argument.
     password : str
         Your ssh password to sshserver. Note that if this is left None,
         you will be prompted for it if passwordless key based login is unavailable.
     paramiko : bool
         flag for whether to use paramiko instead of shell ssh for tunneling.
         [default: True on win32, False else]
 
     ------- exec authentication args -------
     If even localhost is untrusted, you can have some protection against
     unauthorized execution by signing messages with HMAC digests.
     Messages are still sent as cleartext, so if someone can snoop your
     loopback traffic this will not protect your privacy, but will prevent
     unauthorized execution.
 
     exec_key : str
         an authentication key or file containing a key
         default: None
 
 
     Attributes
     ----------
 
     ids : list of int engine IDs
         requesting the ids attribute always synchronizes
         the registration state. To request ids without synchronization,
         use semi-private _ids attributes.
 
     history : list of msg_ids
         a list of msg_ids, keeping track of all the execution
         messages you have submitted in order.
 
     outstanding : set of msg_ids
         a set of msg_ids that have been submitted, but whose
         results have not yet 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
 
     wait
         wait on one or more msg_ids
 
     execution methods
         apply
         legacy: execute, run
 
     data movement
         push, pull, scatter, gather
 
     query methods
         queue_status, get_result, purge, result_status
 
     control methods
         abort, shutdown
 
     """
 
 
     block = Bool(False)
     outstanding = Set()
     results = Instance('collections.defaultdict', (dict,))
     metadata = Instance('collections.defaultdict', (Metadata,))
     history = List()
     debug = Bool(False)
 
     profile=Unicode()
     def _profile_default(self):
         if BaseIPythonApplication.initialized():
             # an IPython app *might* be running, try to get its profile
             try:
                 return BaseIPythonApplication.instance().profile
             except (AttributeError, MultipleInstanceError):
                 # could be a *different* subclass of config.Application,
                 # which would raise one of these two errors.
                 return u'default'
         else:
             return u'default'
 
 
     _outstanding_dict = Instance('collections.defaultdict', (set,))
     _ids = List()
     _connected=Bool(False)
     _ssh=Bool(False)
     _context = Instance('zmq.Context')
     _config = Dict()
     _engines=Instance(util.ReverseDict, (), {})
     # _hub_socket=Instance('zmq.Socket')
     _query_socket=Instance('zmq.Socket')
     _control_socket=Instance('zmq.Socket')
     _iopub_socket=Instance('zmq.Socket')
     _notification_socket=Instance('zmq.Socket')
     _mux_socket=Instance('zmq.Socket')
     _task_socket=Instance('zmq.Socket')
     _task_scheme=Unicode()
     _closed = False
     _ignored_control_replies=Integer(0)
     _ignored_hub_replies=Integer(0)
 
     def __new__(self, *args, **kw):
         # don't raise on positional args
         return HasTraits.__new__(self, **kw)
 
     def __init__(self, url_or_file=None, profile=None, profile_dir=None, ipython_dir=None,
             context=None, debug=False, exec_key=None,
             sshserver=None, sshkey=None, password=None, paramiko=None,
             timeout=10, **extra_args
             ):
         if profile:
             super(Client, self).__init__(debug=debug, profile=profile)
         else:
             super(Client, self).__init__(debug=debug)
         if context is None:
             context = zmq.Context.instance()
         self._context = context
 
         self._setup_profile_dir(self.profile, profile_dir, ipython_dir)
         if self._cd is not None:
             if url_or_file is None:
                 url_or_file = pjoin(self._cd.security_dir, 'ipcontroller-client.json')
         assert url_or_file is not None, "I can't find enough information to connect to a hub!"\
             " Please specify at least one of url_or_file or profile."
 
         if not util.is_url(url_or_file):
             # it's not a url, try for a file
             if not os.path.exists(url_or_file):
                 if self._cd:
                     url_or_file = os.path.join(self._cd.security_dir, url_or_file)
                 assert os.path.exists(url_or_file), "Not a valid connection file or url: %r"%url_or_file
             with open(url_or_file) as f:
                 cfg = json.loads(f.read())
         else:
             cfg = {'url':url_or_file}
 
         # sync defaults from args, json:
         if sshserver:
             cfg['ssh'] = sshserver
         if exec_key:
             cfg['exec_key'] = exec_key
         exec_key = cfg['exec_key']
         location = cfg.setdefault('location', None)
         cfg['url'] = util.disambiguate_url(cfg['url'], location)
         url = cfg['url']
         proto,addr,port = util.split_url(url)
         if location is not None and addr == '127.0.0.1':
             # location specified, and connection is expected to be local
             if location not in LOCAL_IPS and not sshserver:
                 # load ssh from JSON *only* if the controller is not on
                 # this machine
                 sshserver=cfg['ssh']
             if location not in LOCAL_IPS and not sshserver:
                 # warn if no ssh specified, but SSH is probably needed
                 # This is only a warning, because the most likely cause
                 # is a local Controller on a laptop whose IP is dynamic
                 warnings.warn("""
             Controller appears to be listening on localhost, but not on this machine.
             If this is true, you should specify Client(...,sshserver='you@%s')
             or instruct your controller to listen on an external IP."""%location,
                 RuntimeWarning)
         elif not sshserver:
             # otherwise sync with cfg
             sshserver = cfg['ssh']
 
         self._config = cfg
 
         self._ssh = bool(sshserver or sshkey or password)
         if self._ssh and sshserver is None:
             # default to ssh via localhost
             sshserver = url.split('://')[1].split(':')[0]
         if self._ssh and password is None:
             if tunnel.try_passwordless_ssh(sshserver, sshkey, paramiko):
                 password=False
             else:
                 password = getpass("SSH Password for %s: "%sshserver)
         ssh_kwargs = dict(keyfile=sshkey, password=password, paramiko=paramiko)
 
         # configure and construct the session
         if exec_key is not None:
             if os.path.isfile(exec_key):
                 extra_args['keyfile'] = exec_key
             else:
                 exec_key = util.asbytes(exec_key)
                 extra_args['key'] = exec_key
         self.session = Session(**extra_args)
 
         self._query_socket = self._context.socket(zmq.DEALER)
         self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
         if self._ssh:
             tunnel.tunnel_connection(self._query_socket, url, sshserver, **ssh_kwargs)
         else:
             self._query_socket.connect(url)
 
         self.session.debug = self.debug
 
         self._notification_handlers = {'registration_notification' : self._register_engine,
                                     'unregistration_notification' : self._unregister_engine,
                                     'shutdown_notification' : lambda msg: self.close(),
                                     }
         self._queue_handlers = {'execute_reply' : self._handle_execute_reply,
                                 'apply_reply' : self._handle_apply_reply}
         self._connect(sshserver, ssh_kwargs, timeout)
 
     def __del__(self):
         """cleanup sockets, but _not_ context."""
         self.close()
 
     def _setup_profile_dir(self, profile, profile_dir, ipython_dir):
         if ipython_dir is None:
             ipython_dir = get_ipython_dir()
         if profile_dir is not None:
             try:
                 self._cd = ProfileDir.find_profile_dir(profile_dir)
                 return
             except ProfileDirError:
                 pass
         elif profile is not None:
             try:
                 self._cd = ProfileDir.find_profile_dir_by_name(
                     ipython_dir, profile)
                 return
             except ProfileDirError:
                 pass
         self._cd = None
 
     def _update_engines(self, engines):
         """Update our engines dict and _ids from a dict of the form: {id:uuid}."""
         for k,v in engines.iteritems():
             eid = int(k)
             self._engines[eid] = v
             self._ids.append(eid)
         self._ids = sorted(self._ids)
         if sorted(self._engines.keys()) != range(len(self._engines)) and \
                         self._task_scheme == 'pure' and self._task_socket:
             self._stop_scheduling_tasks()
 
     def _stop_scheduling_tasks(self):
         """Stop scheduling tasks because an engine has been unregistered
         from a pure ZMQ scheduler.
         """
         self._task_socket.close()
         self._task_socket = None
         msg = "An engine has been unregistered, and we are using pure " +\
               "ZMQ task scheduling.  Task farming will be disabled."
         if self.outstanding:
             msg += " If you were running tasks when this happened, " +\
                    "some `outstanding` msg_ids may never resolve."
         warnings.warn(msg, RuntimeWarning)
 
     def _build_targets(self, targets):
         """Turn valid target IDs or 'all' into two lists:
         (int_ids, uuids).
         """
         if not self._ids:
             # flush notification socket if no engines yet, just in case
             if not self.ids:
                 raise error.NoEnginesRegistered("Can't build targets without any engines")
 
         if targets is None:
             targets = self._ids
         elif isinstance(targets, basestring):
             if targets.lower() == 'all':
                 targets = self._ids
             else:
                 raise TypeError("%r not valid str target, must be 'all'"%(targets))
         elif isinstance(targets, int):
             if targets < 0:
                 targets = self.ids[targets]
             if targets not in self._ids:
                 raise IndexError("No such engine: %i"%targets)
             targets = [targets]
 
         if isinstance(targets, slice):
             indices = range(len(self._ids))[targets]
             ids = self.ids
             targets = [ ids[i] for i in indices ]
 
         if not isinstance(targets, (tuple, list, xrange)):
             raise TypeError("targets by int/slice/collection of ints only, not %s"%(type(targets)))
 
         return [util.asbytes(self._engines[t]) for t in targets], list(targets)
 
     def _connect(self, sshserver, ssh_kwargs, timeout):
         """setup all our socket connections to the cluster. This is called from
         __init__."""
 
         # Maybe allow reconnecting?
         if self._connected:
             return
         self._connected=True
 
         def connect_socket(s, url):
             url = util.disambiguate_url(url, self._config['location'])
             if self._ssh:
                 return tunnel.tunnel_connection(s, url, sshserver, **ssh_kwargs)
             else:
                 return s.connect(url)
 
         self.session.send(self._query_socket, 'connection_request')
         # use Poller because zmq.select has wrong units in pyzmq 2.1.7
         poller = zmq.Poller()
         poller.register(self._query_socket, zmq.POLLIN)
         # poll expects milliseconds, timeout is seconds
         evts = poller.poll(timeout*1000)
         if not evts:
             raise error.TimeoutError("Hub connection request timed out")
         idents,msg = self.session.recv(self._query_socket,mode=0)
         if self.debug:
             pprint(msg)
         msg = Message(msg)
         content = msg.content
         self._config['registration'] = dict(content)
         if content.status == 'ok':
             ident = self.session.bsession
             if content.mux:
                 self._mux_socket = self._context.socket(zmq.DEALER)
                 self._mux_socket.setsockopt(zmq.IDENTITY, ident)
                 connect_socket(self._mux_socket, content.mux)
             if content.task:
                 self._task_scheme, task_addr = content.task
                 self._task_socket = self._context.socket(zmq.DEALER)
                 self._task_socket.setsockopt(zmq.IDENTITY, ident)
                 connect_socket(self._task_socket, task_addr)
             if content.notification:
                 self._notification_socket = self._context.socket(zmq.SUB)
                 connect_socket(self._notification_socket, content.notification)
                 self._notification_socket.setsockopt(zmq.SUBSCRIBE, b'')
             # if content.query:
             #     self._query_socket = self._context.socket(zmq.DEALER)
             #     self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
             #     connect_socket(self._query_socket, content.query)
             if content.control:
                 self._control_socket = self._context.socket(zmq.DEALER)
                 self._control_socket.setsockopt(zmq.IDENTITY, ident)
                 connect_socket(self._control_socket, content.control)
             if content.iopub:
                 self._iopub_socket = self._context.socket(zmq.SUB)
                 self._iopub_socket.setsockopt(zmq.SUBSCRIBE, b'')
                 self._iopub_socket.setsockopt(zmq.IDENTITY, ident)
                 connect_socket(self._iopub_socket, content.iopub)
             self._update_engines(dict(content.engines))
         else:
             self._connected = False
             raise Exception("Failed to connect!")
 
     #--------------------------------------------------------------------------
     # handlers and callbacks for incoming messages
     #--------------------------------------------------------------------------
 
     def _unwrap_exception(self, content):
         """unwrap exception, and remap engine_id to int."""
         e = error.unwrap_exception(content)
         # print e.traceback
         if e.engine_info:
             e_uuid = e.engine_info['engine_uuid']
             eid = self._engines[e_uuid]
             e.engine_info['engine_id'] = eid
         return e
 
     def _extract_metadata(self, header, parent, content):
         md = {'msg_id' : parent['msg_id'],
               'received' : datetime.now(),
               'engine_uuid' : header.get('engine', None),
               'follow' : parent.get('follow', []),
               'after' : parent.get('after', []),
               'status' : content['status'],
             }
 
         if md['engine_uuid'] is not None:
             md['engine_id'] = self._engines.get(md['engine_uuid'], None)
 
         if 'date' in parent:
             md['submitted'] = parent['date']
         if 'started' in header:
             md['started'] = header['started']
         if 'date' in header:
             md['completed'] = header['date']
         return md
 
     def _register_engine(self, msg):
         """Register a new engine, and update our connection info."""
         content = msg['content']
         eid = content['id']
         d = {eid : content['queue']}
         self._update_engines(d)
 
     def _unregister_engine(self, msg):
         """Unregister an engine that has died."""
         content = msg['content']
         eid = int(content['id'])
         if eid in self._ids:
             self._ids.remove(eid)
             uuid = self._engines.pop(eid)
 
             self._handle_stranded_msgs(eid, uuid)
 
         if self._task_socket and self._task_scheme == 'pure':
             self._stop_scheduling_tasks()
 
     def _handle_stranded_msgs(self, eid, uuid):
         """Handle messages known to be on an engine when the engine unregisters.
 
         It is possible that this will fire prematurely - that is, an engine will
         go down after completing a result, and the client will be notified
         of the unregistration and later receive the successful result.
         """
 
         outstanding = self._outstanding_dict[uuid]
 
         for msg_id in list(outstanding):
             if msg_id in self.results:
                 # we already
                 continue
             try:
                 raise error.EngineError("Engine %r died while running task %r"%(eid, msg_id))
             except:
                 content = error.wrap_exception()
             # build a fake message:
             parent = {}
             header = {}
             parent['msg_id'] = msg_id
             header['engine'] = uuid
             header['date'] = datetime.now()
             msg = dict(parent_header=parent, header=header, content=content)
             self._handle_apply_reply(msg)
 
     def _handle_execute_reply(self, msg):
         """Save the reply to an execute_request into our results.
 
         execute messages are never actually used. apply is used instead.
         """
 
         parent = msg['parent_header']
         msg_id = parent['msg_id']
         if msg_id not in self.outstanding:
             if msg_id in self.history:
                 print ("got stale result: %s"%msg_id)
             else:
                 print ("got unknown result: %s"%msg_id)
         else:
             self.outstanding.remove(msg_id)
         self.results[msg_id] = self._unwrap_exception(msg['content'])
 
     def _handle_apply_reply(self, msg):
         """Save the reply to an apply_request into our results."""
         parent = msg['parent_header']
         msg_id = parent['msg_id']
         if msg_id not in self.outstanding:
             if msg_id in self.history:
                 print ("got stale result: %s"%msg_id)
                 print self.results[msg_id]
                 print msg
             else:
                 print ("got unknown result: %s"%msg_id)
         else:
             self.outstanding.remove(msg_id)
         content = msg['content']
         header = msg['header']
 
         # construct metadata:
         md = self.metadata[msg_id]
         md.update(self._extract_metadata(header, parent, content))
         # is this redundant?
         self.metadata[msg_id] = md
 
         e_outstanding = self._outstanding_dict[md['engine_uuid']]
         if msg_id in e_outstanding:
             e_outstanding.remove(msg_id)
 
         # construct result:
         if content['status'] == 'ok':
             self.results[msg_id] = util.unserialize_object(msg['buffers'])[0]
         elif content['status'] == 'aborted':
             self.results[msg_id] = error.TaskAborted(msg_id)
         elif content['status'] == 'resubmitted':
             # TODO: handle resubmission
             pass
         else:
             self.results[msg_id] = self._unwrap_exception(content)
 
     def _flush_notifications(self):
         """Flush notifications of engine registrations waiting
         in ZMQ queue."""
         idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             msg_type = msg['header']['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)
             idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
 
     def _flush_results(self, sock):
         """Flush task or queue results waiting in ZMQ queue."""
         idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             msg_type = msg['header']['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)
             idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
 
     def _flush_control(self, sock):
         """Flush replies from the control channel waiting
         in the ZMQ queue.
 
         Currently: ignore them."""
         if self._ignored_control_replies <= 0:
             return
         idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
         while msg is not None:
             self._ignored_control_replies -= 1
             if self.debug:
                 pprint(msg)
             idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
 
     def _flush_ignored_control(self):
         """flush ignored control replies"""
         while self._ignored_control_replies > 0:
             self.session.recv(self._control_socket)
             self._ignored_control_replies -= 1
 
     def _flush_ignored_hub_replies(self):
         ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
         while msg is not None:
             ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
 
     def _flush_iopub(self, sock):
         """Flush replies from the iopub channel waiting
         in the ZMQ queue.
         """
         idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             parent = msg['parent_header']
             msg_id = parent['msg_id']
             content = msg['content']
             header = msg['header']
             msg_type = msg['header']['msg_type']
 
             # init metadata:
             md = self.metadata[msg_id]
 
             if msg_type == 'stream':
                 name = content['name']
                 s = md[name] or ''
                 md[name] = s + content['data']
             elif msg_type == 'pyerr':
                 md.update({'pyerr' : self._unwrap_exception(content)})
             elif msg_type == 'pyin':
                 md.update({'pyin' : content['code']})
             else:
                 md.update({msg_type : content.get('data', '')})
 
             # reduntant?
             self.metadata[msg_id] = md
 
             idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
 
     #--------------------------------------------------------------------------
     # len, getitem
     #--------------------------------------------------------------------------
 
     def __len__(self):
         """len(client) returns # of engines."""
         return len(self.ids)
 
     def __getitem__(self, key):
         """index access returns DirectView multiplexer objects
 
         Must be int, slice, or list/tuple/xrange of ints"""
         if not isinstance(key, (int, slice, tuple, list, xrange)):
             raise TypeError("key by int/slice/iterable of ints only, not %s"%(type(key)))
         else:
             return self.direct_view(key)
 
     #--------------------------------------------------------------------------
     # Begin public methods
     #--------------------------------------------------------------------------
 
     @property
     def ids(self):
         """Always up-to-date ids property."""
         self._flush_notifications()
         # always copy:
         return list(self._ids)
 
     def close(self):
         if self._closed:
             return
         snames = filter(lambda n: n.endswith('socket'), dir(self))
         for socket in map(lambda name: getattr(self, name), snames):
             if isinstance(socket, zmq.Socket) and not socket.closed:
                 socket.close()
         self._closed = True
 
     def spin(self):
         """Flush any registration notifications and execution results
         waiting in the ZMQ queue.
         """
         if self._notification_socket:
             self._flush_notifications()
         if self._mux_socket:
             self._flush_results(self._mux_socket)
         if self._task_socket:
             self._flush_results(self._task_socket)
         if self._control_socket:
             self._flush_control(self._control_socket)
         if self._iopub_socket:
             self._flush_iopub(self._iopub_socket)
         if self._query_socket:
             self._flush_ignored_hub_replies()
 
     def wait(self, jobs=None, timeout=-1):
         """waits on one or more `jobs`, for up to `timeout` seconds.
 
         Parameters
         ----------
 
         jobs : int, str, or list of ints and/or strs, or one or more AsyncResult objects
                 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, some msg_ids still outstanding
         """
         tic = time.time()
         if jobs is None:
             theids = self.outstanding
         else:
             if isinstance(jobs, (int, basestring, AsyncResult)):
                 jobs = [jobs]
             theids = set()
             for job in jobs:
                 if isinstance(job, int):
                     # index access
                     job = self.history[job]
                 elif isinstance(job, AsyncResult):
                     map(theids.add, job.msg_ids)
                     continue
                 theids.add(job)
         if not theids.intersection(self.outstanding):
             return True
         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
 
     #--------------------------------------------------------------------------
     # Control methods
     #--------------------------------------------------------------------------
 
     @spin_first
     def clear(self, targets=None, block=None):
         """Clear the namespace in target(s)."""
         block = self.block if block is None else block
         targets = self._build_targets(targets)[0]
         for t in targets:
             self.session.send(self._control_socket, 'clear_request', content={}, ident=t)
         error = False
         if block:
             self._flush_ignored_control()
             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 = self._unwrap_exception(msg['content'])
         else:
             self._ignored_control_replies += len(targets)
         if error:
             raise error
 
 
     @spin_first
     def abort(self, jobs=None, targets=None, block=None):
         """Abort specific jobs from the execution queues of target(s).
 
         This is a mechanism to prevent jobs that have already been submitted
         from executing.
 
         Parameters
         ----------
 
         jobs : msg_id, list of msg_ids, or AsyncResult
             The jobs to be aborted
             
             If unspecified/None: abort all outstanding jobs.
 
         """
         block = self.block if block is None else block
         jobs = jobs if jobs is not None else list(self.outstanding)
         targets = self._build_targets(targets)[0]
         
         msg_ids = []
         if isinstance(jobs, (basestring,AsyncResult)):
             jobs = [jobs]
         bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
         if bad_ids:
             raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
         for j in jobs:
             if isinstance(j, AsyncResult):
                 msg_ids.extend(j.msg_ids)
             else:
                 msg_ids.append(j)
         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 block:
             self._flush_ignored_control()
             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 = self._unwrap_exception(msg['content'])
         else:
             self._ignored_control_replies += len(targets)
         if error:
             raise error
 
     @spin_first
     def shutdown(self, targets=None, restart=False, hub=False, block=None):
         """Terminates one or more engine processes, optionally including the hub."""
         block = self.block if block is None else block
         if hub:
             targets = 'all'
         targets = self._build_targets(targets)[0]
         for t in targets:
             self.session.send(self._control_socket, 'shutdown_request',
                         content={'restart':restart},ident=t)
         error = False
         if block or hub:
             self._flush_ignored_control()
             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 = self._unwrap_exception(msg['content'])
         else:
             self._ignored_control_replies += len(targets)
 
         if hub:
             time.sleep(0.25)
             self.session.send(self._query_socket, 'shutdown_request')
             idents,msg = self.session.recv(self._query_socket, 0)
             if self.debug:
                 pprint(msg)
             if msg['content']['status'] != 'ok':
                 error = self._unwrap_exception(msg['content'])
 
         if error:
             raise error
 
     #--------------------------------------------------------------------------
     # Execution related methods
     #--------------------------------------------------------------------------
 
     def _maybe_raise(self, result):
         """wrapper for maybe raising an exception if apply failed."""
         if isinstance(result, error.RemoteError):
             raise result
 
         return result
 
     def send_apply_message(self, socket, f, args=None, kwargs=None, subheader=None, track=False,
                             ident=None):
         """construct and send an apply message via a socket.
 
         This is the principal method with which all engine execution is performed by views.
         """
 
         assert not self._closed, "cannot use me anymore, I'm closed!"
         # defaults:
         args = args if args is not None else []
         kwargs = kwargs if kwargs is not None else {}
         subheader = subheader if subheader is not None else {}
 
         # validate arguments
-        if not callable(f):
+        if not callable(f) and not isinstance(f, Reference):
             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))
         if not isinstance(subheader, dict):
             raise TypeError("subheader must be dict, not %s"%type(subheader))
 
         bufs = util.pack_apply_message(f,args,kwargs)
 
         msg = self.session.send(socket, "apply_request", buffers=bufs, ident=ident,
                             subheader=subheader, track=track)
 
         msg_id = msg['header']['msg_id']
         self.outstanding.add(msg_id)
         if ident:
             # possibly routed to a specific engine
             if isinstance(ident, list):
                 ident = ident[-1]
             if ident in self._engines.values():
                 # save for later, in case of engine death
                 self._outstanding_dict[ident].add(msg_id)
         self.history.append(msg_id)
         self.metadata[msg_id]['submitted'] = datetime.now()
 
         return msg
 
     #--------------------------------------------------------------------------
     # construct a View object
     #--------------------------------------------------------------------------
 
     def load_balanced_view(self, targets=None):
         """construct a DirectView object.
 
         If no arguments are specified, create a LoadBalancedView
         using all engines.
 
         Parameters
         ----------
 
         targets: list,slice,int,etc. [default: use all engines]
             The subset of engines across which to load-balance
         """
         if targets == 'all':
             targets = None
         if targets is not None:
             targets = self._build_targets(targets)[1]
         return LoadBalancedView(client=self, socket=self._task_socket, targets=targets)
 
     def direct_view(self, targets='all'):
         """construct a DirectView object.
 
         If no targets are specified, create a DirectView using all engines.
         
         rc.direct_view('all') is distinguished from rc[:] in that 'all' will
         evaluate the target engines at each execution, whereas rc[:] will connect to
         all *current* engines, and that list will not change.
         
         That is, 'all' will always use all engines, whereas rc[:] will not use
         engines added after the DirectView is constructed.
 
         Parameters
         ----------
 
         targets: list,slice,int,etc. [default: use all engines]
             The engines to use for the View
         """
         single = isinstance(targets, int)
         # allow 'all' to be lazily evaluated at each execution
         if targets != 'all':
             targets = self._build_targets(targets)[1]
         if single:
             targets = targets[0]
         return DirectView(client=self, socket=self._mux_socket, targets=targets)
 
     #--------------------------------------------------------------------------
     # Query methods
     #--------------------------------------------------------------------------
 
     @spin_first
     def get_result(self, indices_or_msg_ids=None, block=None):
         """Retrieve a result by msg_id or history index, wrapped in an AsyncResult object.
 
         If the client already has the results, no request to the Hub will be made.
 
         This is a convenient way to construct AsyncResult objects, which are wrappers
         that include metadata about execution, and allow for awaiting results that
         were not submitted by this Client.
 
         It can also be a convenient way to retrieve the metadata associated with
         blocking execution, since it always retrieves
 
         Examples
         --------
         ::
 
             In [10]: r = client.apply()
 
         Parameters
         ----------
 
         indices_or_msg_ids : integer history index, str msg_id, or list of either
             The indices or msg_ids of indices to be retrieved
 
         block : bool
             Whether to wait for the result to be done
 
         Returns
         -------
 
         AsyncResult
             A single AsyncResult object will always be returned.
 
         AsyncHubResult
             A subclass of AsyncResult that retrieves results from the Hub
 
         """
         block = self.block if block is None else block
         if indices_or_msg_ids is None:
             indices_or_msg_ids = -1
 
         if not isinstance(indices_or_msg_ids, (list,tuple)):
             indices_or_msg_ids = [indices_or_msg_ids]
 
         theids = []
         for id in indices_or_msg_ids:
             if isinstance(id, int):
                 id = self.history[id]
             if not isinstance(id, basestring):
                 raise TypeError("indices must be str or int, not %r"%id)
             theids.append(id)
 
         local_ids = filter(lambda msg_id: msg_id in self.history or msg_id in self.results, theids)
         remote_ids = filter(lambda msg_id: msg_id not in local_ids, theids)
 
         if remote_ids:
             ar = AsyncHubResult(self, msg_ids=theids)
         else:
             ar = AsyncResult(self, msg_ids=theids)
 
         if block:
             ar.wait()
 
         return ar
 
     @spin_first
     def resubmit(self, indices_or_msg_ids=None, subheader=None, block=None):
         """Resubmit one or more tasks.
 
         in-flight tasks may not be resubmitted.
 
         Parameters
         ----------
 
         indices_or_msg_ids : integer history index, str msg_id, or list of either
             The indices or msg_ids of indices to be retrieved
 
         block : bool
             Whether to wait for the result to be done
 
         Returns
         -------
 
         AsyncHubResult
             A subclass of AsyncResult that retrieves results from the Hub
 
         """
         block = self.block if block is None else block
         if indices_or_msg_ids is None:
             indices_or_msg_ids = -1
 
         if not isinstance(indices_or_msg_ids, (list,tuple)):
             indices_or_msg_ids = [indices_or_msg_ids]
 
         theids = []
         for id in indices_or_msg_ids:
             if isinstance(id, int):
                 id = self.history[id]
             if not isinstance(id, basestring):
                 raise TypeError("indices must be str or int, not %r"%id)
             theids.append(id)
 
         for msg_id in theids:
             self.outstanding.discard(msg_id)
             if msg_id in self.history:
                 self.history.remove(msg_id)
             self.results.pop(msg_id, None)
             self.metadata.pop(msg_id, None)
         content = dict(msg_ids = theids)
 
         self.session.send(self._query_socket, 'resubmit_request', content)
 
         zmq.select([self._query_socket], [], [])
         idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
         if self.debug:
             pprint(msg)
         content = msg['content']
         if content['status'] != 'ok':
             raise self._unwrap_exception(content)
 
         ar = AsyncHubResult(self, msg_ids=theids)
 
         if block:
             ar.wait()
 
         return ar
 
     @spin_first
     def result_status(self, msg_ids, status_only=True):
         """Check on the status of the result(s) of the apply request with `msg_ids`.
 
         If status_only is False, then the actual results will be retrieved, else
         only the status of the results will be checked.
 
         Parameters
         ----------
 
         msg_ids : list of msg_ids
             if int:
                 Passed as index to self.history for convenience.
         status_only : bool (default: True)
             if False:
                 Retrieve the actual results of completed tasks.
 
         Returns
         -------
 
         results : dict
             There will always be the keys 'pending' and 'completed', which will
             be lists of msg_ids that are incomplete or complete. If `status_only`
             is False, then completed results will be keyed by their `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]
             if not isinstance(msg_id, basestring):
                 raise TypeError("msg_ids must be str, not %r"%msg_id)
             theids.append(msg_id)
 
         completed = []
         local_results = {}
 
         # comment this block out to temporarily disable local shortcut:
         for msg_id in theids:
             if msg_id in self.results:
                 completed.append(msg_id)
                 local_results[msg_id] = self.results[msg_id]
                 theids.remove(msg_id)
 
         if theids: # some not locally cached
             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)
             content = msg['content']
             if content['status'] != 'ok':
                 raise self._unwrap_exception(content)
             buffers = msg['buffers']
         else:
             content = dict(completed=[],pending=[])
 
         content['completed'].extend(completed)
 
         if status_only:
             return content
 
         failures = []
         # load cached results into result:
         content.update(local_results)
 
         # update cache with results:
         for msg_id in sorted(theids):
             if msg_id in content['completed']:
                 rec = content[msg_id]
                 parent = rec['header']
                 header = rec['result_header']
                 rcontent = rec['result_content']
                 iodict = rec['io']
                 if isinstance(rcontent, str):
                     rcontent = self.session.unpack(rcontent)
 
                 md = self.metadata[msg_id]
                 md.update(self._extract_metadata(header, parent, rcontent))
                 md.update(iodict)
 
                 if rcontent['status'] == 'ok':
                     res,buffers = util.unserialize_object(buffers)
                 else:
                     print rcontent
                     res = self._unwrap_exception(rcontent)
                     failures.append(res)
 
                 self.results[msg_id] = res
                 content[msg_id] = res
 
         if len(theids) == 1 and failures:
                 raise failures[0]
 
         error.collect_exceptions(failures, "result_status")
         return content
 
     @spin_first
     def queue_status(self, targets='all', verbose=False):
         """Fetch the status of engine queues.
 
         Parameters
         ----------
 
         targets : int/str/list of ints/strs
                 the engines whose states are to be queried.
                 default : all
         verbose : bool
                 Whether to return lengths only, or lists of ids for each element
         """
         engine_ids = self._build_targets(targets)[1]
         content = dict(targets=engine_ids, verbose=verbose)
         self.session.send(self._query_socket, "queue_request", content=content)
         idents,msg = self.session.recv(self._query_socket, 0)
         if self.debug:
             pprint(msg)
         content = msg['content']
         status = content.pop('status')
         if status != 'ok':
             raise self._unwrap_exception(content)
         content = rekey(content)
         if isinstance(targets, int):
             return content[targets]
         else:
             return content
 
     @spin_first
     def purge_results(self, jobs=[], targets=[]):
         """Tell the Hub to forget results.
 
         Individual results can be purged by msg_id, or the entire
         history of specific targets can be purged.
 
         Use `purge_results('all')` to scrub everything from the Hub's db.
 
         Parameters
         ----------
 
         jobs : str or list of str or AsyncResult objects
                 the msg_ids whose results should be forgotten.
         targets : int/str/list of ints/strs
                 The targets, by int_id, whose entire history is to be purged.
 
                 default : None
         """
         if not targets and not jobs:
             raise ValueError("Must specify at least one of `targets` and `jobs`")
         if targets:
             targets = self._build_targets(targets)[1]
 
         # construct msg_ids from jobs
         if jobs == 'all':
             msg_ids = jobs
         else:
             msg_ids = []
             if isinstance(jobs, (basestring,AsyncResult)):
                 jobs = [jobs]
             bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
             if bad_ids:
                 raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
             for j in jobs:
                 if isinstance(j, AsyncResult):
                     msg_ids.extend(j.msg_ids)
                 else:
                     msg_ids.append(j)
 
         content = dict(engine_ids=targets, msg_ids=msg_ids)
         self.session.send(self._query_socket, "purge_request", content=content)
         idents, msg = self.session.recv(self._query_socket, 0)
         if self.debug:
             pprint(msg)
         content = msg['content']
         if content['status'] != 'ok':
             raise self._unwrap_exception(content)
 
     @spin_first
     def hub_history(self):
         """Get the Hub's history
 
         Just like the Client, the Hub has a history, which is a list of msg_ids.
         This will contain the history of all clients, and, depending on configuration,
         may contain history across multiple cluster sessions.
 
         Any msg_id returned here is a valid argument to `get_result`.
 
         Returns
         -------
 
         msg_ids : list of strs
                 list of all msg_ids, ordered by task submission time.
         """
 
         self.session.send(self._query_socket, "history_request", content={})
         idents, msg = self.session.recv(self._query_socket, 0)
 
         if self.debug:
             pprint(msg)
         content = msg['content']
         if content['status'] != 'ok':
             raise self._unwrap_exception(content)
         else:
             return content['history']
 
     @spin_first
     def db_query(self, query, keys=None):
         """Query the Hub's TaskRecord database
 
         This will return a list of task record dicts that match `query`
 
         Parameters
         ----------
 
         query : mongodb query dict
             The search dict. See mongodb query docs for details.
         keys : list of strs [optional]
             The subset of keys to be returned.  The default is to fetch everything but buffers.
             'msg_id' will *always* be included.
         """
         if isinstance(keys, basestring):
             keys = [keys]
         content = dict(query=query, keys=keys)
         self.session.send(self._query_socket, "db_request", content=content)
         idents, msg = self.session.recv(self._query_socket, 0)
         if self.debug:
             pprint(msg)
         content = msg['content']
         if content['status'] != 'ok':
             raise self._unwrap_exception(content)
 
         records = content['records']
 
         buffer_lens = content['buffer_lens']
         result_buffer_lens = content['result_buffer_lens']
         buffers = msg['buffers']
         has_bufs = buffer_lens is not None
         has_rbufs = result_buffer_lens is not None
         for i,rec in enumerate(records):
             # relink buffers
             if has_bufs:
                 blen = buffer_lens[i]
                 rec['buffers'], buffers = buffers[:blen],buffers[blen:]
             if has_rbufs:
                 blen = result_buffer_lens[i]
                 rec['result_buffers'], buffers = buffers[:blen],buffers[blen:]
 
         return records
 
 __all__ = [ 'Client' ]

IPython/parallel/client/remotefunction.py

@@ -1,222 +1,241 b''
 """Remote Functions and decorators for Views.
 
 Authors:
 
 * Brian Granger
 * Min RK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 from __future__ import division
 
 import sys
 import warnings
 
 from IPython.testing.skipdoctest import skip_doctest
 
 from . import map as Map
 from .asyncresult import AsyncMapResult
 
 #-----------------------------------------------------------------------------
-# Decorators
+# Functions and Decorators
 #-----------------------------------------------------------------------------
 
 @skip_doctest
 def remote(view, block=None, **flags):
     """Turn a function into a remote function.
 
     This method can be used for map:
 
     In [1]: @remote(view,block=True)
        ...: def func(a):
        ...:    pass
     """
 
     def remote_function(f):
         return RemoteFunction(view, f, block=block, **flags)
     return remote_function
 
 @skip_doctest
 def parallel(view, dist='b', block=None, ordered=True, **flags):
     """Turn a function into a parallel remote function.
 
     This method can be used for map:
 
     In [1]: @parallel(view, block=True)
        ...: def func(a):
        ...:    pass
     """
 
     def parallel_function(f):
         return ParallelFunction(view, f, dist=dist, block=block, ordered=ordered, **flags)
     return parallel_function
 
+def getname(f):
+    """Get the name of an object.
+    
+    For use in case of callables that are not functions, and
+    thus may not have __name__ defined.
+    
+    Order: f.__name__ >  f.name > str(f)
+    """
+    try:
+        return f.__name__
+    except:
+        pass
+    try:
+        return f.name
+    except:
+        pass
+    
+    return str(f)
+
 #--------------------------------------------------------------------------
 # Classes
 #--------------------------------------------------------------------------
 
 class RemoteFunction(object):
     """Turn an existing function into a remote function.
 
     Parameters
     ----------
 
     view : View instance
         The view to be used for execution
     f : callable
         The function to be wrapped into a remote function
     block : bool [default: None]
         Whether to wait for results or not.  The default behavior is
         to use the current `block` attribute of `view`
 
     **flags : remaining kwargs are passed to View.temp_flags
     """
 
     view = None # the remote connection
     func = None # the wrapped function
     block = None # whether to block
     flags = None # dict of extra kwargs for temp_flags
 
     def __init__(self, view, f, block=None, **flags):
         self.view = view
         self.func = f
         self.block=block
         self.flags=flags
 
     def __call__(self, *args, **kwargs):
         block = self.view.block if self.block is None else self.block
         with self.view.temp_flags(block=block, **self.flags):
             return self.view.apply(self.func, *args, **kwargs)
 
 
 class ParallelFunction(RemoteFunction):
     """Class for mapping a function to sequences.
 
     This will distribute the sequences according the a mapper, and call
     the function on each sub-sequence.  If called via map, then the function
     will be called once on each element, rather that each sub-sequence.
 
     Parameters
     ----------
 
     view : View instance
         The view to be used for execution
     f : callable
         The function to be wrapped into a remote function
     dist : str [default: 'b']
         The key for which mapObject to use to distribute sequences
         options are:
           * 'b' : use contiguous chunks in order
           * 'r' : use round-robin striping
     block : bool [default: None]
         Whether to wait for results or not.  The default behavior is
         to use the current `block` attribute of `view`
     chunksize : int or None
         The size of chunk to use when breaking up sequences in a load-balanced manner
     ordered : bool [default: True]
         Whether 
     **flags : remaining kwargs are passed to View.temp_flags
     """
 
     chunksize=None
     ordered=None
     mapObject=None
 
     def __init__(self, view, f, dist='b', block=None, chunksize=None, ordered=True, **flags):
         super(ParallelFunction, self).__init__(view, f, block=block, **flags)
         self.chunksize = chunksize
         self.ordered = ordered
 
         mapClass = Map.dists[dist]
         self.mapObject = mapClass()
 
     def __call__(self, *sequences):
         client = self.view.client
         
         # check that the length of sequences match
         len_0 = len(sequences[0])
         for s in sequences:
             if len(s)!=len_0:
                 msg = 'all sequences must have equal length, but %i!=%i'%(len_0,len(s))
                 raise ValueError(msg)
         balanced = 'Balanced' in self.view.__class__.__name__
         if balanced:
             if self.chunksize:
                 nparts = len_0//self.chunksize + int(len_0%self.chunksize > 0)
             else:
                 nparts = len_0
             targets = [None]*nparts
         else:
             if self.chunksize:
                 warnings.warn("`chunksize` is ignored unless load balancing", UserWarning)
             # multiplexed:
             targets = self.view.targets
             # 'all' is lazily evaluated at execution time, which is now:
             if targets == 'all':
                 targets = client._build_targets(targets)[1]
             nparts = len(targets)
 
         msg_ids = []
         for index, t in enumerate(targets):
             args = []
             for seq in sequences:
                 part = self.mapObject.getPartition(seq, index, nparts)
                 if len(part) == 0:
                     continue
                 else:
                     args.append(part)
             if not args:
                 continue
 
             # print (args)
             if hasattr(self, '_map'):
                 if sys.version_info[0] >= 3:
                     f = lambda f, *sequences: list(map(f, *sequences))
                 else:
                     f = map
                 args = [self.func]+args
             else:
                 f=self.func
 
             view = self.view if balanced else client[t]
             with view.temp_flags(block=False, **self.flags):
                 ar = view.apply(f, *args)
 
             msg_ids.append(ar.msg_ids[0])
 
         r = AsyncMapResult(self.view.client, msg_ids, self.mapObject, 
-                            fname=self.func.__name__,
+                            fname=getname(self.func),
                             ordered=self.ordered
                         )
 
         if self.block:
             try:
                 return r.get()
             except KeyboardInterrupt:
                 return r
         else:
             return r
 
     def map(self, *sequences):
         """call a function on each element of a sequence remotely.
         This should behave very much like the builtin map, but return an AsyncMapResult
         if self.block is False.
         """
         # set _map as a flag for use inside self.__call__
         self._map = True
         try:
             ret = self.__call__(*sequences)
         finally:
             del self._map
         return ret
 
 __all__ = ['remote', 'parallel', 'RemoteFunction', 'ParallelFunction']

IPython/parallel/client/view.py

@@ -1,1065 +1,1065 b''
 """Views of remote engines.
 
 Authors:
 
 * Min RK
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import imp
 import sys
 import warnings
 from contextlib import contextmanager
 from types import ModuleType
 
 import zmq
 
 from IPython.testing.skipdoctest import skip_doctest
 from IPython.utils.traitlets import (
     HasTraits, Any, Bool, List, Dict, Set, Instance, CFloat, Integer
 )
 from IPython.external.decorator import decorator
 
 from IPython.parallel import util
 from IPython.parallel.controller.dependency import Dependency, dependent
 
 from . import map as Map
 from .asyncresult import AsyncResult, AsyncMapResult
-from .remotefunction import ParallelFunction, parallel, remote
+from .remotefunction import ParallelFunction, parallel, remote, getname
 
 #-----------------------------------------------------------------------------
 # Decorators
 #-----------------------------------------------------------------------------
 
 @decorator
 def save_ids(f, self, *args, **kwargs):
     """Keep our history and outstanding attributes up to date after a method call."""
     n_previous = len(self.client.history)
     try:
         ret = f(self, *args, **kwargs)
     finally:
         nmsgs = len(self.client.history) - n_previous
         msg_ids = self.client.history[-nmsgs:]
         self.history.extend(msg_ids)
         map(self.outstanding.add, msg_ids)
     return ret
 
 @decorator
 def sync_results(f, self, *args, **kwargs):
     """sync relevant results from self.client to our results attribute."""
     ret = f(self, *args, **kwargs)
     delta = self.outstanding.difference(self.client.outstanding)
     completed = self.outstanding.intersection(delta)
     self.outstanding = self.outstanding.difference(completed)
     for msg_id in completed:
         self.results[msg_id] = self.client.results[msg_id]
     return ret
 
 @decorator
 def spin_after(f, self, *args, **kwargs):
     """call spin after the method."""
     ret = f(self, *args, **kwargs)
     self.spin()
     return ret
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 @skip_doctest
 class View(HasTraits):
     """Base View class for more convenint apply(f,*args,**kwargs) syntax via attributes.
 
     Don't use this class, use subclasses.
 
     Methods
     -------
 
     spin
         flushes incoming results and registration state changes
         control methods spin, and requesting `ids` also ensures up to date
 
     wait
         wait on one or more msg_ids
 
     execution methods
         apply
         legacy: execute, run
 
     data movement
         push, pull, scatter, gather
 
     query methods
         get_result, queue_status, purge_results, result_status
 
     control methods
         abort, shutdown
 
     """
     # flags
     block=Bool(False)
     track=Bool(True)
     targets = Any()
 
     history=List()
     outstanding = Set()
     results = Dict()
     client = Instance('IPython.parallel.Client')
 
     _socket = Instance('zmq.Socket')
     _flag_names = List(['targets', 'block', 'track'])
     _targets = Any()
     _idents = Any()
 
     def __init__(self, client=None, socket=None, **flags):
         super(View, self).__init__(client=client, _socket=socket)
         self.block = client.block
 
         self.set_flags(**flags)
 
         assert not self.__class__ is View, "Don't use base View objects, use subclasses"
 
 
     def __repr__(self):
         strtargets = str(self.targets)
         if len(strtargets) > 16:
             strtargets = strtargets[:12]+'...]'
         return "<%s %s>"%(self.__class__.__name__, strtargets)
 
     def set_flags(self, **kwargs):
         """set my attribute flags by keyword.
 
         Views determine behavior with a few attributes (`block`, `track`, etc.).
         These attributes can be set all at once by name with this method.
 
         Parameters
         ----------
 
         block : bool
             whether to wait for results
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
         """
         for name, value in kwargs.iteritems():
             if name not in self._flag_names:
                 raise KeyError("Invalid name: %r"%name)
             else:
                 setattr(self, name, value)
 
     @contextmanager
     def temp_flags(self, **kwargs):
         """temporarily set flags, for use in `with` statements.
 
         See set_flags for permanent setting of flags
 
         Examples
         --------
 
         >>> view.track=False
         ...
         >>> with view.temp_flags(track=True):
         ...    ar = view.apply(dostuff, my_big_array)
         ...    ar.tracker.wait() # wait for send to finish
         >>> view.track
         False
 
         """
         # preflight: save flags, and set temporaries
         saved_flags = {}
         for f in self._flag_names:
             saved_flags[f] = getattr(self, f)
         self.set_flags(**kwargs)
         # yield to the with-statement block
         try:
             yield
         finally:
             # postflight: restore saved flags
             self.set_flags(**saved_flags)
 
 
     #----------------------------------------------------------------
     # apply
     #----------------------------------------------------------------
 
     @sync_results
     @save_ids
     def _really_apply(self, f, args, kwargs, block=None, **options):
         """wrapper for client.send_apply_message"""
         raise NotImplementedError("Implement in subclasses")
 
     def apply(self, f, *args, **kwargs):
         """calls f(*args, **kwargs) on remote engines, returning the result.
 
         This method sets all apply flags via this View's attributes.
 
         if self.block is False:
             returns AsyncResult
         else:
             returns actual result of f(*args, **kwargs)
         """
         return self._really_apply(f, args, kwargs)
 
     def apply_async(self, f, *args, **kwargs):
         """calls f(*args, **kwargs) on remote engines in a nonblocking manner.
 
         returns AsyncResult
         """
         return self._really_apply(f, args, kwargs, block=False)
 
     @spin_after
     def apply_sync(self, f, *args, **kwargs):
         """calls f(*args, **kwargs) on remote engines in a blocking manner,
          returning the result.
 
         returns: actual result of f(*args, **kwargs)
         """
         return self._really_apply(f, args, kwargs, block=True)
 
     #----------------------------------------------------------------
     # wrappers for client and control methods
     #----------------------------------------------------------------
     @sync_results
     def spin(self):
         """spin the client, and sync"""
         self.client.spin()
 
     @sync_results
     def wait(self, jobs=None, timeout=-1):
         """waits on one or more `jobs`, for up to `timeout` seconds.
 
         Parameters
         ----------
 
         jobs : int, str, or list of ints and/or strs, or one or more AsyncResult objects
                 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, some msg_ids still outstanding
         """
         if jobs is None:
             jobs = self.history
         return self.client.wait(jobs, timeout)
 
     def abort(self, jobs=None, targets=None, block=None):
         """Abort jobs on my engines.
 
         Parameters
         ----------
 
         jobs : None, str, list of strs, optional
             if None: abort all jobs.
             else: abort specific msg_id(s).
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         jobs = jobs if jobs is not None else list(self.outstanding)
         
         return self.client.abort(jobs=jobs, targets=targets, block=block)
 
     def queue_status(self, targets=None, verbose=False):
         """Fetch the Queue status of my engines"""
         targets = targets if targets is not None else self.targets
         return self.client.queue_status(targets=targets, verbose=verbose)
 
     def purge_results(self, jobs=[], targets=[]):
         """Instruct the controller to forget specific results."""
         if targets is None or targets == 'all':
             targets = self.targets
         return self.client.purge_results(jobs=jobs, targets=targets)
 
     def shutdown(self, targets=None, restart=False, hub=False, block=None):
         """Terminates one or more engine processes, optionally including the hub.
         """
         block = self.block if block is None else block
         if targets is None or targets == 'all':
             targets = self.targets
         return self.client.shutdown(targets=targets, restart=restart, hub=hub, block=block)
 
     @spin_after
     def get_result(self, indices_or_msg_ids=None):
         """return one or more results, specified by history index or msg_id.
 
         See client.get_result for details.
 
         """
 
         if indices_or_msg_ids is None:
             indices_or_msg_ids = -1
         if isinstance(indices_or_msg_ids, int):
             indices_or_msg_ids = self.history[indices_or_msg_ids]
         elif isinstance(indices_or_msg_ids, (list,tuple,set)):
             indices_or_msg_ids = list(indices_or_msg_ids)
             for i,index in enumerate(indices_or_msg_ids):
                 if isinstance(index, int):
                     indices_or_msg_ids[i] = self.history[index]
         return self.client.get_result(indices_or_msg_ids)
 
     #-------------------------------------------------------------------
     # Map
     #-------------------------------------------------------------------
 
     def map(self, f, *sequences, **kwargs):
         """override in subclasses"""
         raise NotImplementedError
 
     def map_async(self, f, *sequences, **kwargs):
         """Parallel version of builtin `map`, using this view's engines.
 
         This is equivalent to map(...block=False)
 
         See `self.map` for details.
         """
         if 'block' in kwargs:
             raise TypeError("map_async doesn't take a `block` keyword argument.")
         kwargs['block'] = False
         return self.map(f,*sequences,**kwargs)
 
     def map_sync(self, f, *sequences, **kwargs):
         """Parallel version of builtin `map`, using this view's engines.
 
         This is equivalent to map(...block=True)
 
         See `self.map` for details.
         """
         if 'block' in kwargs:
             raise TypeError("map_sync doesn't take a `block` keyword argument.")
         kwargs['block'] = True
         return self.map(f,*sequences,**kwargs)
 
     def imap(self, f, *sequences, **kwargs):
         """Parallel version of `itertools.imap`.
 
         See `self.map` for details.
 
         """
 
         return iter(self.map_async(f,*sequences, **kwargs))
 
     #-------------------------------------------------------------------
     # Decorators
     #-------------------------------------------------------------------
 
     def remote(self, block=True, **flags):
         """Decorator for making a RemoteFunction"""
         block = self.block if block is None else block
         return remote(self, block=block, **flags)
 
     def parallel(self, dist='b', block=None, **flags):
         """Decorator for making a ParallelFunction"""
         block = self.block if block is None else block
         return parallel(self, dist=dist, block=block, **flags)
 
 @skip_doctest
 class DirectView(View):
     """Direct Multiplexer View of one or more engines.
 
     These are created via indexed access to a client:
 
     >>> dv_1 = client[1]
     >>> dv_all = client[:]
     >>> dv_even = client[::2]
     >>> dv_some = client[1:3]
 
     This object provides dictionary access to engine namespaces:
 
     # push a=5:
     >>> dv['a'] = 5
     # pull 'foo':
     >>> db['foo']
 
     """
 
     def __init__(self, client=None, socket=None, targets=None):
         super(DirectView, self).__init__(client=client, socket=socket, targets=targets)
 
     @property
     def importer(self):
         """sync_imports(local=True) as a property.
 
         See sync_imports for details.
 
         """
         return self.sync_imports(True)
 
     @contextmanager
     def sync_imports(self, local=True):
         """Context Manager for performing simultaneous local and remote imports.
 
         'import x as y' will *not* work.  The 'as y' part will simply be ignored.
         
         If `local=True`, then the package will also be imported locally.
         
         Note that remote-only (`local=False`) imports have not been implemented.
         
         >>> with view.sync_imports():
         ...    from numpy import recarray
         importing recarray from numpy on engine(s)
 
         """
         import __builtin__
         local_import = __builtin__.__import__
         modules = set()
         results = []
         @util.interactive
         def remote_import(name, fromlist, level):
             """the function to be passed to apply, that actually performs the import
             on the engine, and loads up the user namespace.
             """
             import sys
             user_ns = globals()
             mod = __import__(name, fromlist=fromlist, level=level)
             if fromlist:
                 for key in fromlist:
                     user_ns[key] = getattr(mod, key)
             else:
                 user_ns[name] = sys.modules[name]
 
         def view_import(name, globals={}, locals={}, fromlist=[], level=-1):
             """the drop-in replacement for __import__, that optionally imports
             locally as well.
             """
             # don't override nested imports
             save_import = __builtin__.__import__
             __builtin__.__import__ = local_import
 
             if imp.lock_held():
                 # this is a side-effect import, don't do it remotely, or even
                 # ignore the local effects
                 return local_import(name, globals, locals, fromlist, level)
 
             imp.acquire_lock()
             if local:
                 mod = local_import(name, globals, locals, fromlist, level)
             else:
                 raise NotImplementedError("remote-only imports not yet implemented")
             imp.release_lock()
 
             key = name+':'+','.join(fromlist or [])
             if level == -1 and key not in modules:
                 modules.add(key)
                 if fromlist:
                     print "importing %s from %s on engine(s)"%(','.join(fromlist), name)
                 else:
                     print "importing %s on engine(s)"%name
                 results.append(self.apply_async(remote_import, name, fromlist, level))
             # restore override
             __builtin__.__import__ = save_import
 
             return mod
 
         # override __import__
         __builtin__.__import__ = view_import
         try:
             # enter the block
             yield
         except ImportError:
             if local:
                 raise
             else:
                 # ignore import errors if not doing local imports
                 pass
         finally:
             # always restore __import__
             __builtin__.__import__ = local_import
 
         for r in results:
             # raise possible remote ImportErrors here
             r.get()
 
 
     @sync_results
     @save_ids
     def _really_apply(self, f, args=None, kwargs=None, targets=None, block=None, track=None):
         """calls f(*args, **kwargs) on remote engines, returning the result.
 
         This method sets all of `apply`'s flags via this View's attributes.
 
         Parameters
         ----------
 
         f : callable
 
         args : list [default: empty]
 
         kwargs : dict [default: empty]
 
         targets : target list [default: self.targets]
             where to run
         block : bool [default: self.block]
             whether to block
         track : bool [default: self.track]
             whether to ask zmq to track the message, for safe non-copying sends
 
         Returns
         -------
 
         if self.block is False:
             returns AsyncResult
         else:
             returns actual result of f(*args, **kwargs) on the engine(s)
             This will be a list of self.targets is also a list (even length 1), or
             the single result if self.targets is an integer engine id
         """
         args = [] if args is None else args
         kwargs = {} if kwargs is None else kwargs
         block = self.block if block is None else block
         track = self.track if track is None else track
         targets = self.targets if targets is None else targets
 
         _idents = self.client._build_targets(targets)[0]
         msg_ids = []
         trackers = []
         for ident in _idents:
             msg = self.client.send_apply_message(self._socket, f, args, kwargs, track=track,
                                     ident=ident)
             if track:
                 trackers.append(msg['tracker'])
             msg_ids.append(msg['header']['msg_id'])
         tracker = None if track is False else zmq.MessageTracker(*trackers)
-        ar = AsyncResult(self.client, msg_ids, fname=f.__name__, targets=targets, tracker=tracker)
+        ar = AsyncResult(self.client, msg_ids, fname=getname(f), targets=targets, tracker=tracker)
         if block:
             try:
                 return ar.get()
             except KeyboardInterrupt:
                 pass
         return ar
 
     @spin_after
     def map(self, f, *sequences, **kwargs):
         """view.map(f, *sequences, block=self.block) => list|AsyncMapResult
 
         Parallel version of builtin `map`, using this View's `targets`.
 
         There will be one task per target, so work will be chunked
         if the sequences are longer than `targets`.
 
         Results can be iterated as they are ready, but will become available in chunks.
 
         Parameters
         ----------
 
         f : callable
             function to be mapped
         *sequences: one or more sequences of matching length
             the sequences to be distributed and passed to `f`
         block : bool
             whether to wait for the result or not [default self.block]
 
         Returns
         -------
 
         if block=False:
             AsyncMapResult
                 An object like AsyncResult, but which reassembles the sequence of results
                 into a single list. AsyncMapResults can be iterated through before all
                 results are complete.
         else:
             list
                 the result of map(f,*sequences)
         """
 
         block = kwargs.pop('block', self.block)
         for k in kwargs.keys():
             if k not in ['block', 'track']:
                 raise TypeError("invalid keyword arg, %r"%k)
 
         assert len(sequences) > 0, "must have some sequences to map onto!"
         pf = ParallelFunction(self, f, block=block, **kwargs)
         return pf.map(*sequences)
 
     def execute(self, code, targets=None, block=None):
         """Executes `code` on `targets` in blocking or nonblocking manner.
 
         ``execute`` is always `bound` (affects engine namespace)
 
         Parameters
         ----------
 
         code : str
                 the code string to be executed
         block : bool
                 whether or not to wait until done to return
                 default: self.block
         """
         return self._really_apply(util._execute, args=(code,), block=block, targets=targets)
 
     def run(self, filename, targets=None, block=None):
         """Execute contents of `filename` on my engine(s).
 
         This simply reads the contents of the file and calls `execute`.
 
         Parameters
         ----------
 
         filename : str
                 The path to the file
         targets : int/str/list of ints/strs
                 the engines on which to execute
                 default : all
         block : bool
                 whether or not to wait until done
                 default: self.block
 
         """
         with open(filename, 'r') as f:
             # add newline in case of trailing indented whitespace
             # which will cause SyntaxError
             code = f.read()+'\n'
         return self.execute(code, block=block, targets=targets)
 
     def update(self, ns):
         """update remote namespace with dict `ns`
 
         See `push` for details.
         """
         return self.push(ns, block=self.block, track=self.track)
 
     def push(self, ns, targets=None, block=None, track=None):
         """update remote namespace with dict `ns`
 
         Parameters
         ----------
 
         ns : dict
             dict of keys with which to update engine namespace(s)
         block : bool [default : self.block]
             whether to wait to be notified of engine receipt
 
         """
 
         block = block if block is not None else self.block
         track = track if track is not None else self.track
         targets = targets if targets is not None else self.targets
         # applier = self.apply_sync if block else self.apply_async
         if not isinstance(ns, dict):
             raise TypeError("Must be a dict, not %s"%type(ns))
         return self._really_apply(util._push, (ns,), block=block, track=track, targets=targets)
 
     def get(self, key_s):
         """get object(s) by `key_s` from remote namespace
 
         see `pull` for details.
         """
         # block = block if block is not None else self.block
         return self.pull(key_s, block=True)
 
     def pull(self, names, targets=None, block=None):
         """get object(s) by `name` from remote namespace
 
         will return one object if it is a key.
         can also take a list of keys, in which case it will return a list of objects.
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         applier = self.apply_sync if block else self.apply_async
         if isinstance(names, basestring):
             pass
         elif isinstance(names, (list,tuple,set)):
             for key in names:
                 if not isinstance(key, basestring):
                     raise TypeError("keys must be str, not type %r"%type(key))
         else:
             raise TypeError("names must be strs, not %r"%names)
         return self._really_apply(util._pull, (names,), block=block, targets=targets)
 
     def scatter(self, key, seq, dist='b', flatten=False, targets=None, block=None, track=None):
         """
         Partition a Python sequence and send the partitions to a set of engines.
         """
         block = block if block is not None else self.block
         track = track if track is not None else self.track
         targets = targets if targets is not None else self.targets
 
         mapObject = Map.dists[dist]()
         nparts = len(targets)
         msg_ids = []
         trackers = []
         for index, engineid in enumerate(targets):
             partition = mapObject.getPartition(seq, index, nparts)
             if flatten and len(partition) == 1:
                 ns = {key: partition[0]}
             else:
                 ns = {key: partition}
             r = self.push(ns, block=False, track=track, targets=engineid)
             msg_ids.extend(r.msg_ids)
             if track:
                 trackers.append(r._tracker)
 
         if track:
             tracker = zmq.MessageTracker(*trackers)
         else:
             tracker = None
 
         r = AsyncResult(self.client, msg_ids, fname='scatter', targets=targets, tracker=tracker)
         if block:
             r.wait()
         else:
             return r
 
     @sync_results
     @save_ids
     def gather(self, key, dist='b', targets=None, block=None):
         """
         Gather a partitioned sequence on a set of engines as a single local seq.
         """
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         mapObject = Map.dists[dist]()
         msg_ids = []
 
         for index, engineid in enumerate(targets):
             msg_ids.extend(self.pull(key, block=False, targets=engineid).msg_ids)
 
         r = AsyncMapResult(self.client, msg_ids, mapObject, fname='gather')
 
         if block:
             try:
                 return r.get()
             except KeyboardInterrupt:
                 pass
         return r
 
     def __getitem__(self, key):
         return self.get(key)
 
     def __setitem__(self,key, value):
         self.update({key:value})
 
     def clear(self, targets=None, block=False):
         """Clear the remote namespaces on my engines."""
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         return self.client.clear(targets=targets, block=block)
 
     def kill(self, targets=None, block=True):
         """Kill my engines."""
         block = block if block is not None else self.block
         targets = targets if targets is not None else self.targets
         return self.client.kill(targets=targets, block=block)
 
     #----------------------------------------
     # activate for %px,%autopx magics
     #----------------------------------------
     def activate(self):
         """Make this `View` active for parallel magic commands.
 
         IPython has a magic command syntax to work with `MultiEngineClient` objects.
         In a given IPython session there is a single active one.  While
         there can be many `Views` created and used by the user,
         there is only one active one.  The active `View` is used whenever
         the magic commands %px and %autopx are used.
 
         The activate() method is called on a given `View` to make it
         active.  Once this has been done, the magic commands can be used.
         """
 
         try:
             # This is injected into __builtins__.
             ip = get_ipython()
         except NameError:
             print "The IPython parallel magics (%result, %px, %autopx) only work within IPython."
         else:
             pmagic = ip.plugin_manager.get_plugin('parallelmagic')
             if pmagic is None:
                 ip.magic_load_ext('parallelmagic')
                 pmagic = ip.plugin_manager.get_plugin('parallelmagic')
 
             pmagic.active_view = self
 
 
 @skip_doctest
 class LoadBalancedView(View):
     """An load-balancing View that only executes via the Task scheduler.
 
     Load-balanced views can be created with the client's `view` method:
 
     >>> v = client.load_balanced_view()
 
     or targets can be specified, to restrict the potential destinations:
 
     >>> v = client.client.load_balanced_view([1,3])
 
     which would restrict loadbalancing to between engines 1 and 3.
 
     """
 
     follow=Any()
     after=Any()
     timeout=CFloat()
     retries = Integer(0)
 
     _task_scheme = Any()
     _flag_names = List(['targets', 'block', 'track', 'follow', 'after', 'timeout', 'retries'])
 
     def __init__(self, client=None, socket=None, **flags):
         super(LoadBalancedView, self).__init__(client=client, socket=socket, **flags)
         self._task_scheme=client._task_scheme
 
     def _validate_dependency(self, dep):
         """validate a dependency.
 
         For use in `set_flags`.
         """
         if dep is None or isinstance(dep, (basestring, AsyncResult, Dependency)):
             return True
         elif isinstance(dep, (list,set, tuple)):
             for d in dep:
                 if not isinstance(d, (basestring, AsyncResult)):
                     return False
         elif isinstance(dep, dict):
             if set(dep.keys()) != set(Dependency().as_dict().keys()):
                 return False
             if not isinstance(dep['msg_ids'], list):
                 return False
             for d in dep['msg_ids']:
                 if not isinstance(d, basestring):
                     return False
         else:
             return False
 
         return True
 
     def _render_dependency(self, dep):
         """helper for building jsonable dependencies from various input forms."""
         if isinstance(dep, Dependency):
             return dep.as_dict()
         elif isinstance(dep, AsyncResult):
             return dep.msg_ids
         elif dep is None:
             return []
         else:
             # pass to Dependency constructor
             return list(Dependency(dep))
 
     def set_flags(self, **kwargs):
         """set my attribute flags by keyword.
 
         A View is a wrapper for the Client's apply method, but with attributes
         that specify keyword arguments, those attributes can be set by keyword
         argument with this method.
 
         Parameters
         ----------
 
         block : bool
             whether to wait for results
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
 
         after : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a time-based dependency.
             This job will only be run *after* the dependencies
             have been met.
 
         follow : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a location-based dependency.
             This job will only be run on an engine where this dependency
             is met.
 
         timeout : float/int or None
             Only for load-balanced execution (targets=None)
             Specify an amount of time (in seconds) for the scheduler to
             wait for dependencies to be met before failing with a
             DependencyTimeout.
 
         retries : int
             Number of times a task will be retried on failure.
         """
 
         super(LoadBalancedView, self).set_flags(**kwargs)
         for name in ('follow', 'after'):
             if name in kwargs:
                 value = kwargs[name]
                 if self._validate_dependency(value):
                     setattr(self, name, value)
                 else:
                     raise ValueError("Invalid dependency: %r"%value)
         if 'timeout' in kwargs:
             t = kwargs['timeout']
             if not isinstance(t, (int, long, float, type(None))):
                 raise TypeError("Invalid type for timeout: %r"%type(t))
             if t is not None:
                 if t < 0:
                     raise ValueError("Invalid timeout: %s"%t)
             self.timeout = t
 
     @sync_results
     @save_ids
     def _really_apply(self, f, args=None, kwargs=None, block=None, track=None,
                                         after=None, follow=None, timeout=None,
                                         targets=None, retries=None):
         """calls f(*args, **kwargs) on a remote engine, returning the result.
 
         This method temporarily sets all of `apply`'s flags for a single call.
 
         Parameters
         ----------
 
         f : callable
 
         args : list [default: empty]
 
         kwargs : dict [default: empty]
 
         block : bool [default: self.block]
             whether to block
         track : bool [default: self.track]
             whether to ask zmq to track the message, for safe non-copying sends
 
         !!!!!! TODO: THE REST HERE  !!!!
 
         Returns
         -------
 
         if self.block is False:
             returns AsyncResult
         else:
             returns actual result of f(*args, **kwargs) on the engine(s)
             This will be a list of self.targets is also a list (even length 1), or
             the single result if self.targets is an integer engine id
         """
 
         # validate whether we can run
         if self._socket.closed:
             msg = "Task farming is disabled"
             if self._task_scheme == 'pure':
                 msg += " because the pure ZMQ scheduler cannot handle"
                 msg += " disappearing engines."
             raise RuntimeError(msg)
 
         if self._task_scheme == 'pure':
             # pure zmq scheme doesn't support extra features
             msg = "Pure ZMQ scheduler doesn't support the following flags:"
             "follow, after, retries, targets, timeout"
             if (follow or after or retries or targets or timeout):
                 # hard fail on Scheduler flags
                 raise RuntimeError(msg)
             if isinstance(f, dependent):
                 # soft warn on functional dependencies
                 warnings.warn(msg, RuntimeWarning)
 
         # build args
         args = [] if args is None else args
         kwargs = {} if kwargs is None else kwargs
         block = self.block if block is None else block
         track = self.track if track is None else track
         after = self.after if after is None else after
         retries = self.retries if retries is None else retries
         follow = self.follow if follow is None else follow
         timeout = self.timeout if timeout is None else timeout
         targets = self.targets if targets is None else targets
 
         if not isinstance(retries, int):
             raise TypeError('retries must be int, not %r'%type(retries))
 
         if targets is None:
             idents = []
         else:
             idents = self.client._build_targets(targets)[0]
             # ensure *not* bytes
             idents = [ ident.decode() for ident in idents ]
 
         after = self._render_dependency(after)
         follow = self._render_dependency(follow)
         subheader = dict(after=after, follow=follow, timeout=timeout, targets=idents, retries=retries)
 
         msg = self.client.send_apply_message(self._socket, f, args, kwargs, track=track,
                                 subheader=subheader)
         tracker = None if track is False else msg['tracker']
 
-        ar = AsyncResult(self.client, msg['header']['msg_id'], fname=f.__name__, targets=None, tracker=tracker)
+        ar = AsyncResult(self.client, msg['header']['msg_id'], fname=getname(f), targets=None, tracker=tracker)
 
         if block:
             try:
                 return ar.get()
             except KeyboardInterrupt:
                 pass
         return ar
 
     @spin_after
     @save_ids
     def map(self, f, *sequences, **kwargs):
         """view.map(f, *sequences, block=self.block, chunksize=1, ordered=True) => list|AsyncMapResult
 
         Parallel version of builtin `map`, load-balanced by this View.
 
         `block`, and `chunksize` can be specified by keyword only.
 
         Each `chunksize` elements will be a separate task, and will be
         load-balanced. This lets individual elements be available for iteration
         as soon as they arrive.
 
         Parameters
         ----------
 
         f : callable
             function to be mapped
         *sequences: one or more sequences of matching length
             the sequences to be distributed and passed to `f`
         block : bool [default self.block]
             whether to wait for the result or not
         track : bool
             whether to create a MessageTracker to allow the user to
             safely edit after arrays and buffers during non-copying
             sends.
         chunksize : int [default 1]
             how many elements should be in each task.
         ordered : bool [default True]
             Whether the results should be gathered as they arrive, or enforce
             the order of submission.
             
             Only applies when iterating through AsyncMapResult as results arrive.
             Has no effect when block=True.
 
         Returns
         -------
 
         if block=False:
             AsyncMapResult
                 An object like AsyncResult, but which reassembles the sequence of results
                 into a single list. AsyncMapResults can be iterated through before all
                 results are complete.
             else:
                 the result of map(f,*sequences)
 
         """
 
         # default
         block = kwargs.get('block', self.block)
         chunksize = kwargs.get('chunksize', 1)
         ordered = kwargs.get('ordered', True)
 
         keyset = set(kwargs.keys())
         extra_keys = keyset.difference_update(set(['block', 'chunksize']))
         if extra_keys:
             raise TypeError("Invalid kwargs: %s"%list(extra_keys))
 
         assert len(sequences) > 0, "must have some sequences to map onto!"
 
         pf = ParallelFunction(self, f, block=block, chunksize=chunksize, ordered=ordered)
         return pf.map(*sequences)
 
 __all__ = ['LoadBalancedView', 'DirectView']

IPython/parallel/tests/test_view.py

@@ -1,472 +1,493 b''
 # -*- coding: utf-8 -*-
 """test View objects
 
 Authors:
 
 * Min RK
 """
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 import sys
 import time
 from tempfile import mktemp
 from StringIO import StringIO
 
 import zmq
 from nose import SkipTest
 
 from IPython.testing import decorators as dec
 
 from IPython import parallel  as pmod
 from IPython.parallel import error
 from IPython.parallel import AsyncResult, AsyncHubResult, AsyncMapResult
 from IPython.parallel import DirectView
 from IPython.parallel.util import interactive
 
 from IPython.parallel.tests import add_engines
 
 from .clienttest import ClusterTestCase, crash, wait, skip_without
 
 def setup():
     add_engines(3)
 
 class TestView(ClusterTestCase):
     
     def test_z_crash_mux(self):
         """test graceful handling of engine death (direct)"""
         raise SkipTest("crash tests disabled, due to undesirable crash reports")
         # self.add_engines(1)
         eid = self.client.ids[-1]
         ar = self.client[eid].apply_async(crash)
         self.assertRaisesRemote(error.EngineError, ar.get, 10)
         eid = ar.engine_id
         tic = time.time()
         while eid in self.client.ids and time.time()-tic < 5:
             time.sleep(.01)
             self.client.spin()
         self.assertFalse(eid in self.client.ids, "Engine should have died")
     
     def test_push_pull(self):
         """test pushing and pulling"""
         data = dict(a=10, b=1.05, c=range(10), d={'e':(1,2),'f':'hi'})
         t = self.client.ids[-1]
         v = self.client[t]
         push = v.push
         pull = v.pull
         v.block=True
         nengines = len(self.client)
         push({'data':data})
         d = pull('data')
         self.assertEquals(d, data)
         self.client[:].push({'data':data})
         d = self.client[:].pull('data', block=True)
         self.assertEquals(d, nengines*[data])
         ar = push({'data':data}, block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         r = ar.get()
         ar = self.client[:].pull('data', block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         r = ar.get()
         self.assertEquals(r, nengines*[data])
         self.client[:].push(dict(a=10,b=20))
         r = self.client[:].pull(('a','b'), block=True)
         self.assertEquals(r, nengines*[[10,20]])
     
     def test_push_pull_function(self):
         "test pushing and pulling functions"
         def testf(x):
             return 2.0*x
         
         t = self.client.ids[-1]
         v = self.client[t]
         v.block=True
         push = v.push
         pull = v.pull
         execute = v.execute
         push({'testf':testf})
         r = pull('testf')
         self.assertEqual(r(1.0), testf(1.0))
         execute('r = testf(10)')
         r = pull('r')
         self.assertEquals(r, testf(10))
         ar = self.client[:].push({'testf':testf}, block=False)
         ar.get()
         ar = self.client[:].pull('testf', block=False)
         rlist = ar.get()
         for r in rlist:
             self.assertEqual(r(1.0), testf(1.0))
         execute("def g(x): return x*x")
         r = pull(('testf','g'))
         self.assertEquals((r[0](10),r[1](10)), (testf(10), 100))
     
     def test_push_function_globals(self):
         """test that pushed functions have access to globals"""
         @interactive
         def geta():
             return a
         # self.add_engines(1)
         v = self.client[-1]
         v.block=True
         v['f'] = geta
         self.assertRaisesRemote(NameError, v.execute, 'b=f()')
         v.execute('a=5')
         v.execute('b=f()')
         self.assertEquals(v['b'], 5)
     
     def test_push_function_defaults(self):
         """test that pushed functions preserve default args"""
         def echo(a=10):
             return a
         v = self.client[-1]
         v.block=True
         v['f'] = echo
         v.execute('b=f()')
         self.assertEquals(v['b'], 10)
 
     def test_get_result(self):
         """test getting results from the Hub."""
         c = pmod.Client(profile='iptest')
         # self.add_engines(1)
         t = c.ids[-1]
         v = c[t]
         v2 = self.client[t]
         ar = v.apply_async(wait, 1)
         # give the monitor time to notice the message
         time.sleep(.25)
         ahr = v2.get_result(ar.msg_ids)
         self.assertTrue(isinstance(ahr, AsyncHubResult))
         self.assertEquals(ahr.get(), ar.get())
         ar2 = v2.get_result(ar.msg_ids)
         self.assertFalse(isinstance(ar2, AsyncHubResult))
         c.spin()
         c.close()
     
     def test_run_newline(self):
         """test that run appends newline to files"""
         tmpfile = mktemp()
         with open(tmpfile, 'w') as f:
             f.write("""def g():
                 return 5
                 """)
         v = self.client[-1]
         v.run(tmpfile, block=True)
         self.assertEquals(v.apply_sync(lambda f: f(), pmod.Reference('g')), 5)
 
     def test_apply_tracked(self):
         """test tracking for apply"""
         # self.add_engines(1)
         t = self.client.ids[-1]
         v = self.client[t]
         v.block=False
         def echo(n=1024*1024, **kwargs):
             with v.temp_flags(**kwargs):
                 return v.apply(lambda x: x, 'x'*n)
         ar = echo(1, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         ar = echo(track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertEquals(ar.sent, ar._tracker.done)
         ar._tracker.wait()
         self.assertTrue(ar.sent)
         
     def test_push_tracked(self):
         t = self.client.ids[-1]
         ns = dict(x='x'*1024*1024)
         v = self.client[t]
         ar = v.push(ns, block=False, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         
         ar = v.push(ns, block=False, track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         ar._tracker.wait()
         self.assertEquals(ar.sent, ar._tracker.done)
         self.assertTrue(ar.sent)
         ar.get()
         
     def test_scatter_tracked(self):
         t = self.client.ids
         x='x'*1024*1024
         ar = self.client[t].scatter('x', x, block=False, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         
         ar = self.client[t].scatter('x', x, block=False, track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertEquals(ar.sent, ar._tracker.done)
         ar._tracker.wait()
         self.assertTrue(ar.sent)
         ar.get()
     
     def test_remote_reference(self):
         v = self.client[-1]
         v['a'] = 123
         ra = pmod.Reference('a')
         b = v.apply_sync(lambda x: x, ra)
         self.assertEquals(b, 123)
 
 
     def test_scatter_gather(self):
         view = self.client[:]
         seq1 = range(16)
         view.scatter('a', seq1)
         seq2 = view.gather('a', block=True)
         self.assertEquals(seq2, seq1)
         self.assertRaisesRemote(NameError, view.gather, 'asdf', block=True)
     
     @skip_without('numpy')
     def test_scatter_gather_numpy(self):
         import numpy
         from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
         view = self.client[:]
         a = numpy.arange(64)
         view.scatter('a', a)
         b = view.gather('a', block=True)
         assert_array_equal(b, a)
     
     def test_map(self):
         view = self.client[:]
         def f(x):
             return x**2
         data = range(16)
         r = view.map_sync(f, data)
         self.assertEquals(r, map(f, data))
     
     def test_map_iterable(self):
         """test map on iterables (direct)"""
         view = self.client[:]
         # 101 is prime, so it won't be evenly distributed
         arr = range(101)
         # ensure it will be an iterator, even in Python 3
         it = iter(arr)
         r = view.map_sync(lambda x:x, arr)
         self.assertEquals(r, list(arr))
     
     def test_scatterGatherNonblocking(self):
         data = range(16)
         view = self.client[:]
         view.scatter('a', data, block=False)
         ar = view.gather('a', block=False)
         self.assertEquals(ar.get(), data)
     
     @skip_without('numpy')
     def test_scatter_gather_numpy_nonblocking(self):
         import numpy
         from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
         a = numpy.arange(64)
         view = self.client[:]
         ar = view.scatter('a', a, block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         amr = view.gather('a', block=False)
         self.assertTrue(isinstance(amr, AsyncMapResult))
         assert_array_equal(amr.get(), a)
 
     def test_execute(self):
         view = self.client[:]
         # self.client.debug=True
         execute = view.execute
         ar = execute('c=30', block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         ar = execute('d=[0,1,2]', block=False)
         self.client.wait(ar, 1)
         self.assertEquals(len(ar.get()), len(self.client))
         for c in view['c']:
             self.assertEquals(c, 30)
     
     def test_abort(self):
         view = self.client[-1]
         ar = view.execute('import time; time.sleep(1)', block=False)
         ar2 = view.apply_async(lambda : 2)
         ar3 = view.apply_async(lambda : 3)
         view.abort(ar2)
         view.abort(ar3.msg_ids)
         self.assertRaises(error.TaskAborted, ar2.get)
         self.assertRaises(error.TaskAborted, ar3.get)
     
     def test_abort_all(self):
         """view.abort() aborts all outstanding tasks"""
         view = self.client[-1]
         ars = [ view.apply_async(time.sleep, 1) for i in range(10) ]
         view.abort()
         view.wait(timeout=5)
         for ar in ars[5:]:
             self.assertRaises(error.TaskAborted, ar.get)
     
     def test_temp_flags(self):
         view = self.client[-1]
         view.block=True
         with view.temp_flags(block=False):
             self.assertFalse(view.block)
         self.assertTrue(view.block)
     
     @dec.known_failure_py3
     def test_importer(self):
         view = self.client[-1]
         view.clear(block=True)
         with view.importer:
             import re
         
         @interactive
         def findall(pat, s):
             # this globals() step isn't necessary in real code
             # only to prevent a closure in the test
             re = globals()['re']
             return re.findall(pat, s)
         
         self.assertEquals(view.apply_sync(findall, '\w+', 'hello world'), 'hello world'.split())
     
     # parallel magic tests
     
     def test_magic_px_blocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=True
 
         ip.magic_px('a=5')
         self.assertEquals(v['a'], 5)
         ip.magic_px('a=10')
         self.assertEquals(v['a'], 10)
         sio = StringIO()
         savestdout = sys.stdout
         sys.stdout = sio
         # just 'print a' worst ~99% of the time, but this ensures that
         # the stdout message has arrived when the result is finished:
         ip.magic_px('import sys,time;print (a); sys.stdout.flush();time.sleep(0.2)')
         sys.stdout = savestdout
         buf = sio.getvalue()
         self.assertTrue('[stdout:' in buf, buf)
         self.assertTrue(buf.rstrip().endswith('10'))
         self.assertRaisesRemote(ZeroDivisionError, ip.magic_px, '1/0')
 
     def test_magic_px_nonblocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=False
 
         ip.magic_px('a=5')
         self.assertEquals(v['a'], 5)
         ip.magic_px('a=10')
         self.assertEquals(v['a'], 10)
         sio = StringIO()
         savestdout = sys.stdout
         sys.stdout = sio
         ip.magic_px('print a')
         sys.stdout = savestdout
         buf = sio.getvalue()
         self.assertFalse('[stdout:%i]'%v.targets in buf)
         ip.magic_px('1/0')
         ar = v.get_result(-1)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
     
     def test_magic_autopx_blocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=True
 
         sio = StringIO()
         savestdout = sys.stdout
         sys.stdout = sio
         ip.magic_autopx()
         ip.run_cell('\n'.join(('a=5','b=10','c=0')))
         ip.run_cell('print b')
         ip.run_cell("b/c")
         ip.run_code(compile('b*=2', '', 'single'))
         ip.magic_autopx()
         sys.stdout = savestdout
         output = sio.getvalue().strip()
         self.assertTrue(output.startswith('%autopx enabled'))
         self.assertTrue(output.endswith('%autopx disabled'))
         self.assertTrue('RemoteError: ZeroDivisionError' in output)
         ar = v.get_result(-2)
         self.assertEquals(v['a'], 5)
         self.assertEquals(v['b'], 20)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
 
     def test_magic_autopx_nonblocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=False
 
         sio = StringIO()
         savestdout = sys.stdout
         sys.stdout = sio
         ip.magic_autopx()
         ip.run_cell('\n'.join(('a=5','b=10','c=0')))
         ip.run_cell('print b')
         ip.run_cell("b/c")
         ip.run_code(compile('b*=2', '', 'single'))
         ip.magic_autopx()
         sys.stdout = savestdout
         output = sio.getvalue().strip()
         self.assertTrue(output.startswith('%autopx enabled'))
         self.assertTrue(output.endswith('%autopx disabled'))
         self.assertFalse('ZeroDivisionError' in output)
         ar = v.get_result(-2)
         self.assertEquals(v['a'], 5)
         self.assertEquals(v['b'], 20)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
     
     def test_magic_result(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v['a'] = 111
         ra = v['a']
         
         ar = ip.magic_result()
         self.assertEquals(ar.msg_ids, [v.history[-1]])
         self.assertEquals(ar.get(), 111)
         ar = ip.magic_result('-2')
         self.assertEquals(ar.msg_ids, [v.history[-2]])
     
     def test_unicode_execute(self):
         """test executing unicode strings"""
         v = self.client[-1]
         v.block=True
         if sys.version_info[0] >= 3:
             code="a='é'"
         else:
             code=u"a=u'é'"
         v.execute(code)
         self.assertEquals(v['a'], u'é')
         
     def test_unicode_apply_result(self):
         """test unicode apply results"""
         v = self.client[-1]
         r = v.apply_sync(lambda : u'é')
         self.assertEquals(r, u'é')
     
     def test_unicode_apply_arg(self):
         """test passing unicode arguments to apply"""
         v = self.client[-1]
         
         @interactive
         def check_unicode(a, check):
             assert isinstance(a, unicode), "%r is not unicode"%a
             assert isinstance(check, bytes), "%r is not bytes"%check
             assert a.encode('utf8') == check, "%s != %s"%(a,check)
         
         for s in [ u'é', u'ßø®∫',u'asdf' ]:
             try:
                 v.apply_sync(check_unicode, s, s.encode('utf8'))
             except error.RemoteError as e:
                 if e.ename == 'AssertionError':
                     self.fail(e.evalue)
                 else:
                     raise e
     
+    def test_map_reference(self):
+        """view.map(<Reference>, *seqs) should work"""
+        v = self.client[:]
+        v.scatter('n', self.client.ids, flatten=True)
+        v.execute("f = lambda x,y: x*y")
+        rf = pmod.Reference('f')
+        nlist = list(range(10))
+        mlist = nlist[::-1]
+        expected = [ m*n for m,n in zip(mlist, nlist) ]
+        result = v.map_sync(rf, mlist, nlist)
+        self.assertEquals(result, expected)
+
+    def test_apply_reference(self):
+        """view.apply(<Reference>, *args) should work"""
+        v = self.client[:]
+        v.scatter('n', self.client.ids, flatten=True)
+        v.execute("f = lambda x: n*x")
+        rf = pmod.Reference('f')
+        result = v.apply_sync(rf, 5)
+        expected = [ 5*id for id in self.client.ids ]
+        self.assertEquals(result, expected)