Implement Actors

Author: mrocklin

distributed/__init__.py

@@ -2,6 +2,7 @@
 
 from . import config
 from dask.config import config
+from .actor import Actor, ActorFuture
 from .core import connect, rpc
 from .deploy import LocalCluster, Adaptive
 from .diagnostics import progress

distributed/actor.py

@@ -0,0 +1,213 @@
+from tornado import gen
+import functools
+
+from .client import Future, default_client
+from .compatibility import get_thread_identity, Queue
+from .protocol import to_serialize
+from .utils import sync
+from .utils_comm import WrappedKey
+from .worker import get_worker
+
+
+class Actor(WrappedKey):
+    """ Controls an object on a remote worker
+
+    An actor allows remote control of a stateful object living on a remote
+    worker.  Method calls on this object trigger operations on the remote
+    object and return ActorFutures on which we can block to get results.
+
+    Examples
+    --------
+    >>> class Counter:
+    ...    def __init__(self):
+    ...        self.n = 0
+    ...    def increment(self):
+    ...        self.n += 1
+    ...        return self.n
+
+    >>> from dask.distributed import Client
+    >>> client = Client()
+
+    You can create an actor by submitting a class with the keyword
+    ``actor=True``.
+
+    >>> future = client.submit(Counter, actor=True)
+    >>> counter = future.result()
+    >>> counter
+    <Actor: Counter, key=Counter-1234abcd>
+
+    Calling methods on this object immediately returns deferred ``ActorFuture``
+    objects.  You can call ``.result()`` on these objects to block and get the
+    result of the function call.
+
+    >>> future = counter.increment()
+    >>> future.result()
+    1
+    >>> future = counter.increment()
+    >>> future.result()
+    2
+    """
+    def __init__(self, cls, address, key, worker=None):
+        self._cls = cls
+        self._address = address
+        self.key = key
+        self._future = None
+        if worker:
+            self._worker = worker
+            self._client = None
+        else:
+            try:
+                self._worker = get_worker()
+            except ValueError:
+                self._worker = None
+            try:
+                self._client = default_client()
+                self._future = Future(key)
+            except ValueError:
+                self._client = None
+
+    def __repr__(self):
+        return '<Actor: %s, key=%s>' % (self._cls.__name__, self.key)
+
+    def __reduce__(self):
+        return (Actor, (self._cls, self._address, self.key))
+
+    @property
+    def _io_loop(self):
+        if self._worker:
+            return self._worker.io_loop
+        else:
+            return self._client.io_loop
+
+    @property
+    def _scheduler_rpc(self):
+        if self._worker:
+            return self._worker.scheduler
+        else:
+            return self._client.scheduler
+
+    @property
+    def _worker_rpc(self):
+        if self._worker:
+            return self._worker.rpc(self._address)
+        else:
+            if self._client.direct_to_workers:
+                return self._client.rpc(self._address)
+            else:
+                return ProxyRPC(self._client.scheduler, self._address)
+
+    @property
+    def _asynchronous(self):
+        if self._client:
+            return self._client.asynchronous
+        else:
+            return get_thread_identity() == self._worker.thread_id
+
+    def _sync(self, func, *args, **kwargs):
+        if self._client:
+            return self._client.sync(func, *args, **kwargs)
+        else:
+            # TODO support sync operation by checking against thread ident of loop
+            return sync(self._worker.loop, func, *args, **kwargs)
+
+    def __dir__(self):
+        o = set(dir(type(self)))
+        o.update({attr for attr in dir(self._cls) if not attr.startswith('_')})
+        return sorted(o)
+
+    def __getattr__(self, key):
+        if not hasattr(self._cls, key):
+            raise AttributeError("%s does not have attribute %s" %
+                                 (type(self).__name__, key))
+
+        attr = getattr(self._cls, key)
+
+        if self._future and not self._future.status == 'finished':
+            raise ValueError("Worker holding Actor was lost")
+
+        if callable(attr):
+            @functools.wraps(attr)
+            def func(*args, **kwargs):
+                @gen.coroutine
+                def run_actor_function_on_worker():
+                    try:
+                        result = yield self._worker_rpc.actor_execute(
+                            function=key,
+                            actor=self.key,
+                            args=[to_serialize(arg) for arg in args],
+                            kwargs={k: to_serialize(v) for k, v in kwargs.items()},
+                        )
+                    except OSError:
+                        if self._future:
+                            yield self._future
+                        else:
+                            raise OSError("Unable to contact Actor's worker")
+                    raise gen.Return(result['result'])
+
+                if self._asynchronous:
+                    return run_actor_function_on_worker()
+                else:
+                    # TODO: this mechanism is error prone
+                    # we should endeavor to make dask's standard code work here
+                    q = Queue()
+
+                    @gen.coroutine
+                    def wait_then_add_to_queue():
+                        x = yield run_actor_function_on_worker()
+                        q.put(x)
+                    self._io_loop.add_callback(wait_then_add_to_queue)
+
+                    return ActorFuture(q, self._io_loop)
+            return func
+
+        else:
+            @gen.coroutine
+            def get_actor_attribute_from_worker():
+                x = yield self._worker_rpc.actor_attribute(attribute=key, actor=self.key)
+                raise gen.Return(x['result'])
+
+            return self._sync(get_actor_attribute_from_worker)
+
+
+class ProxyRPC(object):
+    """
+    An rpc-like object that uses the scheduler's rpc to connect to a worker
+    """
+    def __init__(self, rpc, address):
+        self.rpc = rpc
+        self._address = address
+
+    def __getattr__(self, key):
+        @gen.coroutine
+        def func(**msg):
+            msg['op'] = key
+            result = yield self.rpc.proxy(worker=self._address, msg=msg)
+            raise gen.Return(result)
+
+        return func
+
+
+class ActorFuture(object):
+    """ Future to an actor's method call
+
+    Whenever you call a method on an Actor you get an ActorFuture immediately
+    while the computation happens in the background.  You can call ``.result``
+    to block and collect the full result
+
+    See Also
+    --------
+    Actor
+    """
+    def __init__(self, q, io_loop):
+        self.q = q
+        self.io_loop = io_loop
+
+    def result(self, timeout=None):
+        try:
+            return self._cached_result
+        except AttributeError:
+            self._cached_result = self.q.get(timeout=timeout)
+            return self._cached_result
+
+    def __repr__(self):
+        return '<ActorFuture>'

distributed/client.py

@@ -401,19 +401,7 @@ def lose(self):
         self._get_event().clear()
 
     def set_error(self, exception, traceback):
-        if isinstance(exception, bytes):
-            try:
-                exception = loads(exception)
-            except TypeError:
-                exception = Exception("Undeserializable exception", exception)
-        if traceback:
-            if isinstance(traceback, bytes):
-                try:
-                    traceback = loads(traceback)
-                except (TypeError, AttributeError):
-                    traceback = None
-        else:
-            traceback = None
+        _, exception, traceback = clean_exception(exception, traceback)
 
         self.status = 'error'
         self.exception = exception
@@ -1195,13 +1183,16 @@ def submit(self, func, *args, **kwargs):
             raise TypeError("First input to submit must be a callable function")
 
         key = kwargs.pop('key', None)
-        pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
         resources = kwargs.pop('resources', None)
         retries = kwargs.pop('retries', None)
         priority = kwargs.pop('priority', 0)
         fifo_timeout = kwargs.pop('fifo_timeout', '100ms')
         allow_other_workers = kwargs.pop('allow_other_workers', False)
+        actor = kwargs.pop('actor', False)
+        actors = kwargs.pop('actors', False)
+        actor = actor or actors
+        pure = kwargs.pop('pure', not actor)
 
         if allow_other_workers not in (True, False, None):
             raise TypeError("allow_other_workers= must be True or False")
@@ -1240,7 +1231,8 @@ def submit(self, func, *args, **kwargs):
                                          user_priority=priority,
                                          resources={skey: resources} if resources else None,
                                          retries=retries,
-                                         fifo_timeout=fifo_timeout)
+                                         fifo_timeout=fifo_timeout,
+                                         actors=actor)
 
         logger.debug("Submit %s(...), %s", funcname(func), key)
 
@@ -1322,13 +1314,16 @@ def map(self, func, *iterables, **kwargs):
 
         key = kwargs.pop('key', None)
         key = key or funcname(func)
-        pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
         retries = kwargs.pop('retries', None)
         resources = kwargs.pop('resources', None)
         user_priority = kwargs.pop('priority', 0)
         allow_other_workers = kwargs.pop('allow_other_workers', False)
         fifo_timeout = kwargs.pop('fifo_timeout', '100ms')
+        actor = kwargs.pop('actor', False)
+        actors = kwargs.pop('actors', False)
+        actor = actor or actors
+        pure = kwargs.pop('pure', not actor)
 
         if allow_other_workers and workers is None:
             raise ValueError("Only use allow_other_workers= if using workers=")
@@ -1386,7 +1381,8 @@ def map(self, func, *iterables, **kwargs):
                                          resources=resources,
                                          retries=retries,
                                          user_priority=user_priority,
-                                         fifo_timeout=fifo_timeout)
+                                         fifo_timeout=fifo_timeout,
+                                         actors=actor)
         logger.debug("map(%s, ...)", funcname(func))
 
         return [futures[tokey(k)] for k in keys]
@@ -2075,7 +2071,7 @@ def run_coroutine(self, function, *args, **kwargs):
     def _graph_to_futures(self, dsk, keys, restrictions=None,
                           loose_restrictions=None, priority=None,
                           user_priority=0, resources=None, retries=None,
-                          fifo_timeout=0):
+                          fifo_timeout=0, actors=None):
         with self._lock:
             if resources:
                 resources = self._expand_resources(resources,
@@ -2139,7 +2135,8 @@ def _graph_to_futures(self, dsk, keys, restrictions=None,
                                      'resources': resources,
                                      'submitting_task': getattr(thread_state, 'key', None),
                                      'retries': retries,
-                                     'fifo_timeout': fifo_timeout})
+                                     'fifo_timeout': fifo_timeout,
+                                     'actors': actors})
             return futures
 
     def get(self, dsk, keys, restrictions=None, loose_restrictions=None,

distributed/core.py

@@ -851,9 +851,17 @@ def clean_exception(exception, traceback, **kwargs):
     error_message: create and serialize errors into message
     """
     if isinstance(exception, bytes):
-        exception = protocol.pickle.loads(exception)
+        try:
+            exception = protocol.pickle.loads(exception)
+        except Exception:
+            exception = Exception(exception)
+    elif isinstance(exception, str):
+        exception = Exception(exception)
     if isinstance(traceback, bytes):
-        traceback = protocol.pickle.loads(traceback)
+        try:
+            traceback = protocol.pickle.loads(traceback)
+        except (TypeError, AttributeError):
+            traceback = None
     elif isinstance(traceback, string_types):
         traceback = None  # happens if the traceback failed serializing
     return type(exception), exception, traceback