add rich AsyncResult behavior

IPython/zmq/parallel/asyncresult.py

@@ -1,113 +1,186 b''
 """AsyncResult objects for the client"""
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010  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 IPython.external.decorator import decorator
 import error
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
+@decorator
+def check_ready(f, self, *args, **kwargs):
+    """Call spin() to sync state prior to calling the method."""
+    self.wait(0)
+    if not self._ready:
+        raise error.TimeoutError("result not ready")
+    return f(self, *args, **kwargs)
+
 class AsyncResult(object):
     """Class for representing results of non-blocking calls.
     
     Provides the same interface as :py:class:`multiprocessing.AsyncResult`.
     """
     def __init__(self, client, msg_ids, fname=''):
         self._client = client
         self.msg_ids = msg_ids
         self._fname=fname
         self._ready = False
         self._success = None
     
     def __repr__(self):
         if self._ready:
             return "<%s: finished>"%(self.__class__.__name__)
         else:
             return "<%s: %s>"%(self.__class__.__name__,self._fname)
     
     
     def _reconstruct_result(self, res):
         """
         Override me in subclasses for turning a list of results
         into the expected form.
         """
         if len(self.msg_ids) == 1:
             return res[0]
         else:
             return res
         
     def get(self, timeout=-1):
         """Return the result when it arrives. 
         
         If `timeout` is not ``None`` and the result does not arrive within
         `timeout` seconds then ``TimeoutError`` is raised. If the
         remote call raised an exception then that exception will be reraised
         by get().
         """
         if not self.ready():
             self.wait(timeout)
         
         if self._ready:
             if self._success:
                 return self._result
             else:
                 raise self._exception
         else:
             raise error.TimeoutError("Result not ready.")
     
     def ready(self):
         """Return whether the call has completed."""
         if not self._ready:
             self.wait(0)
         return self._ready
     
     def wait(self, timeout=-1):
         """Wait until the result is available or until `timeout` seconds pass.
         """
         if self._ready:
             return
         self._ready = self._client.barrier(self.msg_ids, timeout)
         if self._ready:
             try:
                 results = map(self._client.results.get, self.msg_ids)
+                self._result = results
                 results = error.collect_exceptions(results, self._fname)
                 self._result = self._reconstruct_result(results)
             except Exception, e:
                 self._exception = e
                 self._success = False
             else:
                 self._success = True
+            finally:
+                self._metadata = map(self._client.metadata.get, self.msg_ids)
             
     
     def successful(self):
         """Return whether the call completed without raising an exception. 
         
         Will raise ``AssertionError`` if the result is not ready.
         """
         assert self._ready
         return self._success
+    
+    #----------------------------------------------------------------
+    # Extra methods not in mp.pool.AsyncResult
+    #----------------------------------------------------------------
+    
+    def get_dict(self, timeout=-1):
+        """Get the results as a dict, keyed by engine_id."""
+        results = self.get(timeout)
+        engine_ids = [md['engine_id'] for md in self._metadata ]
+        bycount = sorted(engine_ids, key=lambda k: engine_ids.count(k))
+        maxcount = bycount.count(bycount[-1])
+        if maxcount > 1:
+            raise ValueError("Cannot build dict, %i jobs ran on engine #%i"%(
+                    maxcount, bycount[-1]))
+        
+        return dict(zip(engine_ids,results))
+    
+    @property
+    @check_ready
+    def result(self):
+        """result property."""
+        return self._result
+    
+    @property
+    @check_ready
+    def metadata(self):
+        """metadata property."""
+        return self._metadata
+    
+    @property
+    @check_ready
+    def result_dict(self):
+        """result property as a dict."""
+        return self.get_dict(0)
+    
+    #-------------------------------------
+    # dict-access
+    #-------------------------------------
+    
+    @check_ready
+    def __getitem__(self, key):
+        """getitem returns result value(s) if keyed by int/slice, or metadata if key is str.
+        """
+        if isinstance(key, int):
+            return error.collect_exceptions([self._result[key]], self._fname)[0]
+        elif isinstance(key, slice):
+            return error.collect_exceptions(self._result[key], self._fname)
+        elif isinstance(key, basestring):
+            return [ md[key] for md in self._metadata ]
+        else:
+            raise TypeError("Invalid key type %r, must be 'int','slice', or 'str'"%type(key))
+    
+    @check_ready
+    def __getattr__(self, key):
+        """getattr maps to getitem for convenient access to metadata."""
+        if key not in self._metadata[0].keys():
+            raise AttributeError("%r object has no attribute %r"%(
+                    self.__class__.__name__, key))
+        return self.__getitem__(key)
 
+    
 class AsyncMapResult(AsyncResult):
     """Class for representing results of non-blocking gathers.
     
     This will properly reconstruct the gather.
     """
     
     def __init__(self, client, msg_ids, mapObject, fname=''):
         self._mapObject = mapObject
         AsyncResult.__init__(self, client, msg_ids, fname=fname)
     
     def _reconstruct_result(self, res):
         """Perform the gather on the actual results."""
         return self._mapObject.joinPartitions(res)
     
         
+__all__ = ['AsyncResult', 'AsyncMapResult']
  No newline at end of file

IPython/zmq/parallel/error.py

@@ -1,283 +1,283 b''
 # encoding: utf-8
 
 """Classes and functions for kernel related errors and exceptions."""
 from __future__ import print_function
 
 __docformat__ = "restructuredtext en"
 
 # Tell nose to skip this module
 __test__ = {}
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2008  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.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Error classes
 #-------------------------------------------------------------------------------
 class IPythonError(Exception):
     """Base exception that all of our exceptions inherit from.
 
     This can be raised by code that doesn't have any more specific
     information."""
 
     pass
 
 # Exceptions associated with the controller objects
 class ControllerError(IPythonError): pass
 
 class ControllerCreationError(ControllerError): pass
 
 
 # Exceptions associated with the Engines
 class EngineError(IPythonError): pass
 
 class EngineCreationError(EngineError): pass
 
 class KernelError(IPythonError):
     pass
 
 class NotDefined(KernelError):
     def __init__(self, name):
         self.name = name
         self.args = (name,)
 
     def __repr__(self):
         return '<NotDefined: %s>' % self.name
     
     __str__ = __repr__
 
 
 class QueueCleared(KernelError):
     pass
 
 
 class IdInUse(KernelError):
     pass
 
 
 class ProtocolError(KernelError):
     pass
 
 
 class ConnectionError(KernelError):
     pass
 
 
 class InvalidEngineID(KernelError):
     pass
 
 
 class NoEnginesRegistered(KernelError):
     pass
 
 
 class InvalidClientID(KernelError):
     pass
 
 
 class InvalidDeferredID(KernelError):
     pass
 
 
 class SerializationError(KernelError):
     pass
 
 
 class MessageSizeError(KernelError):
     pass
 
 
 class PBMessageSizeError(MessageSizeError):
     pass
 
 
 class ResultNotCompleted(KernelError):
     pass
 
 
 class ResultAlreadyRetrieved(KernelError):
     pass
 
 class ClientError(KernelError):
     pass
 
 
 class TaskAborted(KernelError):
     pass
 
 
 class TaskTimeout(KernelError):
     pass
 
 
 class NotAPendingResult(KernelError):
     pass
 
 
 class UnpickleableException(KernelError):
     pass
 
 
 class AbortedPendingDeferredError(KernelError):
     pass
 
 
 class InvalidProperty(KernelError):
     pass
 
 
 class MissingBlockArgument(KernelError):
     pass
 
 
 class StopLocalExecution(KernelError):
     pass
 
 
 class SecurityError(KernelError):
     pass
 
 
 class FileTimeoutError(KernelError):
     pass
 
 class TimeoutError(KernelError):
     pass
 
 class RemoteError(KernelError):
     """Error raised elsewhere"""
     ename=None
     evalue=None
     traceback=None
     engine_info=None
     
     def __init__(self, ename, evalue, traceback, engine_info=None):
         self.ename=ename
         self.evalue=evalue
         self.traceback=traceback
         self.engine_info=engine_info or {}
         self.args=(ename, evalue)
     
     def __repr__(self):
         engineid = self.engine_info.get('engineid', ' ')
         return "<Remote[%s]:%s(%s)>"%(engineid, self.ename, self.evalue)
         
     def __str__(self):
         sig = "%s(%s)"%(self.ename, self.evalue)
         if self.traceback:
             return sig + '\n' + self.traceback
         else:
             return sig
     
 
 class TaskRejectError(KernelError):
     """Exception to raise when a task should be rejected by an engine.
     
     This exception can be used to allow a task running on an engine to test
     if the engine (or the user's namespace on the engine) has the needed
     task dependencies.  If not, the task should raise this exception.  For
     the task to be retried on another engine, the task should be created
     with the `retries` argument > 1.
     
     The advantage of this approach over our older properties system is that
     tasks have full access to the user's namespace on the engines and the
     properties don't have to be managed or tested by the controller.
     """
 
 
 class CompositeError(KernelError):
     """Error for representing possibly multiple errors on engines"""
     def __init__(self, message, elist):
         Exception.__init__(self, *(message, elist))
         # Don't use pack_exception because it will conflict with the .message
         # attribute that is being deprecated in 2.6 and beyond.
         self.msg = message
         self.elist = elist
         self.args = [ e[0] for e in elist ]
 
     def _get_engine_str(self, ei):
         if not ei:
             return '[Engine Exception]'
         else:
             return '[%i:%s]: ' % (ei['engineid'], ei['method'])
 
     def _get_traceback(self, ev):
         try:
             tb = ev._ipython_traceback_text
         except AttributeError:
             return 'No traceback available'
         else:
             return tb
 
     def __str__(self):
         s = str(self.msg)
         for en, ev, etb, ei in self.elist:
             engine_str = self._get_engine_str(ei)
             s = s + '\n' + engine_str + en + ': ' + str(ev)
         return s
     
     def __repr__(self):
         return "CompositeError(%i)"%len(self.elist)
 
     def print_tracebacks(self, excid=None):
         if excid is None:
             for (en,ev,etb,ei) in self.elist:
                 print (self._get_engine_str(ei))
                 print (etb or 'No traceback available')
                 print ()
         else:
             try:
                 en,ev,etb,ei = self.elist[excid]
             except:
                 raise IndexError("an exception with index %i does not exist"%excid)
             else:
                 print (self._get_engine_str(ei))
                 print (etb or 'No traceback available')
     
     def raise_exception(self, excid=0):
         try:
             en,ev,etb,ei = self.elist[excid]
         except:
             raise IndexError("an exception with index %i does not exist"%excid)
         else:
             try:
                 raise RemoteError(en, ev, etb, ei)
             except:
                 et,ev,tb = sys.exc_info()
 
 
-def collect_exceptions(rdict_or_list, method):
+def collect_exceptions(rdict_or_list, method='unspecified'):
     """check a result dict for errors, and raise CompositeError if any exist.
     Passthrough otherwise."""
     elist = []
     if isinstance(rdict_or_list, dict):
         rlist = rdict_or_list.values()
     else:
         rlist = rdict_or_list
     for r in rlist:
         if isinstance(r, RemoteError):
             en, ev, etb, ei = r.ename, r.evalue, r.traceback, r.engine_info
             # Sometimes we could have CompositeError in our list.  Just take
             # the errors out of them and put them in our new list.  This 
             # has the effect of flattening lists of CompositeErrors into one
             # CompositeError
             if en=='CompositeError':
                 for e in ev.elist:
                     elist.append(e)
             else:
                 elist.append((en, ev, etb, ei))
     if len(elist)==0:
         return rdict_or_list
     else:
         msg = "one or more exceptions from call to method: %s" % (method)
         # This silliness is needed so the debugger has access to the exception
         # instance (e in this case)
         try:
             raise CompositeError(msg, elist)
         except CompositeError, e:
             raise e