Repository: dask/distributed

Issue:
Make actor futures play with as_completed
Based on [this slack thread](https://app.slack.com/client/T019Q6SMM5Z/C019Q0QSD4J) , it would be nice if `as_completed` respected actor futures as well as normal ones. Now actor futures are not actually awaitable:
```
    def __await__(self):
        return self.result()
```
(which is a blocking call, that waits on a queue.Queue, `self.q.get()`) but it could be, since we can check for the appearance of an item in the queue (`q.empty()`) or that the result is already cached.


Discussion:
Also doesn't play with `client.gather`: `gather(actor_future) -> actor_future`. I can tweak ActorFuture to make it more Future-like and for this work, probably, but I don't know why they had to be separate things in the first place.
```diff
--- a/distributed/actor.py
+++ b/distributed/actor.py
@@ -241,14 +241,33 @@ class ActorFuture:
         self.io_loop = io_loop
         if result:
             self._cached_result = result
+        self.status = "pending"

     def __await__(self):
         return self.result()

+    async def _result(self, raiseit=True):
+        while not hasattr(self, "_cached_result"):
+            if self.q.empty():
+                await asyncio.sleep(0.01)
+            else:
+                out = self.q.get()
+                if out["status"] == "OK":
+                    self._cached_result = out["result"]
+                else:
+                    self._cached_result = out["exception"]
+        if isinstance(self._cached_result, Exception):
+            self.status = "error"
+            raise self._cached_result
+        self.status = "finished"
+        return self._cached_result
+
     def result(self, timeout=None):
         try:
             if isinstance(self._cached_result, Exception):
+                self.status = "error"
                 raise self._cached_result
+            self.status = "finished"
             return self._cached_result
         except AttributeError:
             out = self.q.get(timeout=timeout)
```
solves `as_completes`.
There is no intention around keeping them separate.  Actors were originally built quickly as an experiment.  Their implementation was different enough that little attempt was made to make them like normal futures.

What you propose seems fine.  However, I would prefer to avoid polling.  Perhaps we should switch from a `queue.Queue` to an `asyncio.Queue` instead.
it sounds like to do this we would want to make the async `_result` method the main way to get a result.  We would then define `result` using the client.sync method.
Do you have an immediate feeling or memory of how much effort it would take to get full Future comliance?

> we should switch from a queue.Queue to an asyncio.Queue instead.

Agreed, if this is easy. I haven't tried yet. The work may be in another thread in this process (asyncio.Queue is not tread-safe), and is accessed via __getattr__ (blocking), so maybe not straight forward.

> We would then define result using the client.sync method.

This certainly occurred to me - this is, of course, scratch code.
> The work may be in another thread in this process

Nope.  All logic in distributed happens on the main event loop thread.  You shouldn't need another thread for this.  

> so maybe not straight forward

I'm confident that if you stick to everything running on the event loop thread that things will work out.

Resolving patch:
diff --git a/distributed/actor.py b/distributed/actor.py
--- a/distributed/actor.py
+++ b/distributed/actor.py
@@ -1,7 +1,6 @@
 import asyncio
 import functools
 import threading
-from queue import Queue
 
 from .client import Future, default_client
 from .protocol import to_serialize
@@ -50,9 +49,9 @@ class Actor(WrappedKey):
     """
 
     def __init__(self, cls, address, key, worker=None):
+        super().__init__(key)
         self._cls = cls
         self._address = address
-        self.key = key
         self._future = None
         if worker:
             self._worker = worker
@@ -138,7 +137,7 @@ def __getattr__(self, key):
 
             elif callable(attr):
                 return lambda *args, **kwargs: ActorFuture(
-                    None, None, result=attr(*args, **kwargs)
+                    None, self._io_loop, result=attr(*args, **kwargs)
                 )
             else:
                 return attr
@@ -158,33 +157,21 @@ async def run_actor_function_on_worker():
                             kwargs={k: to_serialize(v) for k, v in kwargs.items()},
                         )
                     except OSError:
-                        if self._future:
+                        if self._future and not self._future.done():
                             await self._future
+                            return await run_actor_function_on_worker()
                         else:
                             raise OSError("Unable to contact Actor's worker")
                     return result
 
-                if self._asynchronous:
+                q = asyncio.Queue(loop=self._io_loop.asyncio_loop)
 
-                    async def unwrap():
-                        result = await run_actor_function_on_worker()
-                        if result["status"] == "OK":
-                            return result["result"]
-                        raise result["exception"]
+                async def wait_then_add_to_queue():
+                    x = await run_actor_function_on_worker()
+                    await q.put(x)
 
-                    return asyncio.ensure_future(unwrap())
-                else:
-                    # TODO: this mechanism is error prone
-                    # we should endeavor to make dask's standard code work here
-                    q = Queue()
-
-                    async def wait_then_add_to_queue():
-                        x = await run_actor_function_on_worker()
-                        q.put(x)
-
-                    self._io_loop.add_callback(wait_then_add_to_queue)
-
-                    return ActorFuture(q, self._io_loop)
+                self._io_loop.add_callback(wait_then_add_to_queue)
+                return ActorFuture(q, self._io_loop)
 
             return func
 
@@ -241,22 +228,29 @@ def __init__(self, q, io_loop, result=None):
         self.io_loop = io_loop
         if result:
             self._cached_result = result
+        self.status = "pending"
 
     def __await__(self):
-        return self.result()
+        return self._result().__await__()
 
-    def result(self, timeout=None):
-        try:
-            if isinstance(self._cached_result, Exception):
-                raise self._cached_result
-            return self._cached_result
-        except AttributeError:
-            out = self.q.get(timeout=timeout)
+    def done(self):
+        return self.status != "pending"
+
+    async def _result(self, raiseit=True):
+        if not hasattr(self, "_cached_result"):
+            out = await self.q.get()
             if out["status"] == "OK":
+                self.status = "finished"
                 self._cached_result = out["result"]
             else:
+                self.status = "error"
                 self._cached_result = out["exception"]
-        return self.result()
+        if self.status == "error":
+            raise self._cached_result
+        return self._cached_result
+
+    def result(self, timeout=None):
+        return sync(self.io_loop, self._result, callback_timeout=timeout)
 
     def __repr__(self):
         return "<ActorFuture>"
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -4403,9 +4403,11 @@ def update(self, futures):
         """Add multiple futures to the collection.
 
         The added futures will emit from the iterator once they finish"""
+        from .actor import ActorFuture
+
         with self.lock:
             for f in futures:
-                if not isinstance(f, Future):
+                if not isinstance(f, (Future, ActorFuture)):
                     raise TypeError("Input must be a future, got %s" % f)
                 self.futures[f] += 1
                 self.loop.add_callback(self._track_future, f)


Test patch:
diff --git a/distributed/tests/test_actor.py b/distributed/tests/test_actor.py
--- a/distributed/tests/test_actor.py
+++ b/distributed/tests/test_actor.py
@@ -6,7 +6,16 @@
 
 import dask
 
-from distributed import Actor, ActorFuture, Client, Future, Nanny, get_client, wait
+from distributed import (
+    Actor,
+    ActorFuture,
+    Client,
+    Future,
+    Nanny,
+    as_completed,
+    get_client,
+    wait,
+)
 from distributed.metrics import time
 from distributed.utils_test import cluster, gen_cluster
 
@@ -626,6 +635,32 @@ def prop(self):
         await ac.prop
 
 
+def test_as_completed(client):
+    ac = client.submit(Counter, actor=True).result()
+    futures = [ac.increment() for _ in range(10)]
+    max = 0
+
+    for future in as_completed(futures):
+        value = future.result()
+        if value > max:
+            max = value
+
+    assert all(future.done() for future in futures)
+    assert max == 10
+
+
+@gen_cluster(client=True, timeout=3)
+async def test_actor_future_awaitable(client, s, a, b):
+    ac = await client.submit(Counter, actor=True)
+    futures = [ac.increment() for _ in range(10)]
+
+    assert all([isinstance(future, ActorFuture) for future in futures])
+
+    out = await asyncio.gather(*futures)
+    assert all([future.done() for future in futures])
+    assert max(out) == 10
+
+
 @gen_cluster(client=True)
 async def test_serialize_with_pickle(c, s, a, b):
     class Foo: