Lazy evaluation

distarray/globalapi/context.py

@@ -17,6 +17,7 @@
 from abc import ABCMeta, abstractmethod
 
 from functools import wraps
+from contextlib import contextmanager
 
 import numpy
 
@@ -28,7 +29,7 @@
 
 from distarray.globalapi.ipython_utils import IPythonClient
 from distarray.utils import uid, nonce, has_exactly_one
-from distarray.localapi.proxyize import Proxy
+from distarray.localapi.proxyize import Proxy, lazy_proxyize, lazy_name
 
 # mpi context
 from distarray.mpionly_utils import (make_targets_comm,
@@ -71,7 +72,8 @@ def make_subcomm(self, new_targets):
         pass
 
     @abstractmethod
-    def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
+    def apply(self, func, args=None, kwargs=None, targets=None, nresults=None,
+              autoproxyize=False):
         pass
 
     @abstractmethod
@@ -206,7 +208,7 @@ def local_allclose(la, lb, rtol, atol):
             from numpy import allclose
             return allclose(la.ndarray, lb.ndarray, rtol, atol)
 
-        local_results = self.apply(local_allclose, 
+        local_results = self.apply(local_allclose,
                                   (a.key, b.key, rtol, atol),
                                   targets=a.targets)
         return all(local_results)
@@ -580,7 +582,7 @@ def is_NoneType(pxy):
             return pxy.type_str == str(type(None))
 
         def is_LocalArray(pxy):
-            return (isinstance(pxy, Proxy) and 
+            return (isinstance(pxy, Proxy) and
                     pxy.type_str == "<class 'distarray.localapi.localarray.LocalArray'>")
 
         if all(is_LocalArray(r) for r in results):
@@ -743,7 +745,8 @@ def _execute(self, lines, targets):
     def _push(self, d, targets):
         return self.view.push(d, targets=targets, block=True)
 
-    def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
+    def apply(self, func, args=None, kwargs=None, targets=None,
+              autoproxyize=False, nresults=1):
         """
         Analogous to IPython.parallel.view.apply_sync
 
@@ -758,6 +761,8 @@ def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
             engines func is to be run on.
         autoproxyize: bool, default False
             If True, implicitly return a Proxy object from the function.
+        nresults: int, default 1
+            Number of return values.  Only implemented for MPIContext.
 
         Returns
         -------
@@ -860,7 +865,7 @@ def delete_key(self, key, targets=None):
         if MPIContext.INTERCOMM:
             self._send_msg(msg, targets=targets)
 
-    def __init__(self, targets=None):
+    def __init__(self, targets=None, lazy=False):
 
         if MPIContext.INTERCOMM is None:
             MPIContext.INTERCOMM = initial_comm_setup()
@@ -870,6 +875,16 @@ def __init__(self, targets=None):
         self.all_targets = list(range(self.nengines))
         self.targets = self.all_targets if targets is None else sorted(targets)
 
+        self.lazy = lazy  # is the context in lazy-communication mode?
+
+        # message queues used for lazy mode
+        # mapping: target -> queue of messages for that target
+
+        # _sendq: batches up messages to send upon sync()
+        self._sendq = dict([(t, []) for t in self.targets])
+        # _recvq: stores proxy objects for expected return values
+        self._recvq = dict([(t, []) for t in self.targets])
+
         # make/get comms
         # this is the object we want to use with push, pull, etc'
         self._comm_from_targets = {}
@@ -912,16 +927,45 @@ def close(self):
 
     # End of key management routines.
 
+    @contextmanager
+    def lazy_eval(self):
+        """Context manager that enables lazy evaluation.
+
+        On exit, call `sync()` and set `self.lazy` to False
+        """
+        self.lazy = True
+        yield self
+        self.sync()
+        self.lazy = False
+
     def _send_msg(self, msg, targets=None):
         targets = self.targets if targets is None else targets
-        for t in targets:
-            MPIContext.INTERCOMM.send(msg, dest=t)
+        if self.lazy and msg[0] != 'process_message_queue':
+            for t in targets:
+                self._sendq[t].append(msg)
+        else:
+            for t in targets:
+                MPIContext.INTERCOMM.send(msg, dest=t)
 
-    def _recv_msg(self, targets=None):
+    def _recv_msg(self, targets=None, nresults=1, sync=False):
         res = []
         targets = self.targets if targets is None else targets
-        for t in targets:
-            res.append(MPIContext.INTERCOMM.recv(source=t))
+        if self.lazy and not sync:
+            result_names = [lazy_name() for n in range(nresults)]
+            for t in targets:
+                if nresults == 0:
+                    res.append(None)
+                elif nresults == 1:
+                    res.append(lazy_proxyize(name=result_names[0]))
+                else:
+                    target_results = []
+                    for name in result_names:
+                        target_results.append(lazy_proxyize(name))
+                    res.append(target_results)
+                self._recvq[t].append(res[-1])
+        else:
+            for t in targets:
+                res.append(MPIContext.INTERCOMM.recv(source=t))
         return res
 
     def make_subcomm(self, targets):
@@ -952,7 +996,8 @@ def _push(self, d, targets=None):
         msg = ('push', d)
         return self._send_msg(msg, targets=targets)
 
-    def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
+    def apply(self, func, args=None, kwargs=None, targets=None,
+              autoproxyize=False, nresults=1):
         """
         Analogous to IPython.parallel.view.apply_sync
 
@@ -967,6 +1012,8 @@ def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
             engines func is to be run on.
         autoproxyize: bool, default False
             If True, implicitly return a Proxy object from the function.
+        nresults: int, default 1
+            Number of return values.  Only needed for lazy evaluation.
 
         Returns
         -------
@@ -1000,11 +1047,47 @@ def apply(self, func, args=None, kwargs=None, targets=None, autoproxyize=False):
             msg = ('builtin_call', func, args, kwargs, autoproxyize)
 
         self._send_msg(msg, targets=targets)
-        return self._recv_msg(targets=targets)
+        return self._recv_msg(targets=targets, nresults=nresults)
 
     def push_function(self, key, func, targets=None):
         push_function(self, key, func, targets=targets)
 
+    def sync(self, targets=None):
+        """Send queued messages, fill in expected result values."""
+        targets = self.targets if targets is None else targets
+        for t in targets:
+            msg = ('process_message_queue', self._recvq[t], self._sendq[t])
+            self._send_msg(msg, targets=[t])
+            self._sendq[t] = []  # empty the send queue
+
+        for t in targets:
+            results = self._recv_msg(targets=[t], sync=True)[0]
+            lresults = self._recvq[t]
+            for lres, res in zip(lresults, results):
+                # multiple return values
+                if isinstance(res, collections.Sequence):
+                    if len(lres) != len(res):
+                        msg = ("Reserved lazy result object isn't the same"
+                               " size as the actual result object: {} != {}")
+                        raise TypeError(msg.format(len(lres), len(res)))
+                    for sublres, subres in zip(lres, res):
+                        if isinstance(subres, Proxy):
+                            sublres.__dict__ = subres.__dict__
+                        else:
+                            msg = ("Only DistArray return values are "
+                                   "supported in lazy mode.  Type is: {}"
+                                   "".format(type(subres)))
+                            raise TypeError(msg)
+                # single return value
+                else:
+                    if isinstance(res, Proxy):
+                        lres.__dict__ = res.__dict__
+                    else:
+                        msg = ("Only DistArray return values are "
+                               "currently supported in lazy mode.")
+                        raise TypeError(msg)
+                self._recvq[t] = []  # empty the recv queue
+
 
 class ContextCreationError(RuntimeError):
     pass

distarray/globalapi/distarray.py

@@ -77,7 +77,6 @@ def from_localarrays(cls, key, context=None, targets=None, distribution=None,
 
         If `dtype` is not provided, it will be fetched from the engines.
         """
-
         def get_dim_datas_and_dtype(arr):
             return (arr.dim_data, arr.dtype)
 

distarray/globalapi/functions.py

@@ -35,6 +35,25 @@
     __all__.append(func_name)
 
 
+def unary_output_dtype(ufunc, val):
+    """Determine the output dtype of a unary ufunc with input `val`.
+
+    Use the ufunc.types attribute and the input dtype.
+    """
+    input_dtype = numpy.result_type(val)  # find out dtype of scalars
+    # Look at the built-in implementations to find an output type.
+    for input_type, _, _, output_type in ufunc.types:
+        if input_dtype.char == input_type:
+            return numpy.dtype(output_type)
+    # Nothing found.  Try coercion.
+    for input_type, _, _, output_type in ufunc.types:
+        if numpy.can_cast(input_dtype, input_type):
+            return numpy.dtype(output_type)
+    else:  # Can't even coerce to a known input type.  Give up.
+        msg = "Unary ufunc doesn't have a mapping for this type: {}."
+        raise TypeError(msg.format(input_dtype))
+
+
 def unary_proxy(name):
     def proxy_func(a, *args, **kwargs):
         context = determine_context(a)
@@ -44,18 +63,41 @@ def func_call(func_name, arr_name, args, kwargs):
             dotted_name = 'distarray.localapi.%s' % (func_name,)
             func = get_from_dotted_name(dotted_name)
             res = func(arr_name, *args, **kwargs)
-            return proxyize(res), res.dtype  # noqa
+            return proxyize(res)
 
         res = context.apply(func_call, args=(name, a.key, args, kwargs),
                             targets=a.targets)
-        new_key = res[0][0]
-        dtype = res[0][1]
+        new_key = res[0]
+        dtype = unary_output_dtype(getattr(numpy, name), a)
         return DistArray.from_localarrays(new_key,
                                           distribution=a.distribution,
                                           dtype=dtype)
     return proxy_func
 
 
+def binary_output_dtype(ufunc, val0, val1):
+    """Determine the output dtype of a binary ufunc, given input values.
+
+    Use the ufunc.types attribute and the input dtypes.
+    """
+    # find out dtype of scalars
+    input_dtype_0, input_dtype_1 = map(numpy.result_type, (val0, val1))
+    # Look at the built-in implementations to find an output type.
+    for input_type_0, input_type_1, _, _, output_type in ufunc.types:
+        if ((input_dtype_0.char == input_type_0) and
+                (input_dtype_1.char == input_type_1)):
+            return numpy.dtype(output_type)
+    # Nothing found.  Try coercion.
+    for input_type_0, input_type_1, _, _, output_type in ufunc.types:
+        if (numpy.can_cast(input_dtype_0, input_type_0) and
+                numpy.can_cast(input_dtype_1, input_type_1)):
+            return numpy.dtype(output_type)
+    else:  # Can't even coerce to a known input type.  Give up.
+        msg = ("Binary ufunc doesn't have a mapping for these input types: "
+               "{}, {}")
+        raise TypeError(msg.format(input_dtype_0, input_dtype_1))
+
+
 def binary_proxy(name):
     def proxy_func(a, b, *args, **kwargs):
         context = determine_context(a, b)
@@ -83,12 +125,12 @@ def func_call(func_name, a, b, args, kwargs):
             dotted_name = 'distarray.localapi.%s' % (func_name,)
             func = get_from_dotted_name(dotted_name)
             res = func(a, b, *args, **kwargs)
-            return proxyize(res), res.dtype  # noqa
+            return proxyize(res)
 
         res = context.apply(func_call, args=(name, a_key, b_key, args, kwargs),
-                            targets=distribution.targets)
-        new_key = res[0][0]
-        dtype = res[0][1]
+                            targets=distribution.targets, nresults=1)
+        new_key = res[0]
+        dtype = binary_output_dtype(getattr(numpy, name), a, b)
         return DistArray.from_localarrays(new_key,
                                           distribution=distribution,
                                           dtype=dtype)

distarray/globalapi/maps.py

@@ -271,7 +271,7 @@ def view(self, new_dimsize):
 
     def is_compatible(self, other):
         return (isinstance(other, (NoDistMap, BlockMap, BlockCyclicMap)) and
-                other.grid_size == self.grid_size and 
+                other.grid_size == self.grid_size and
                 other.size == self.size)