Repository: dask/distributed

Issue:
Transferring Actor handle between workers makes Actor unreleasable
When worker A calls `gather_dep` on an Actor task, it gets sent an Actor handle by worker B where the Actor is running. When that handle is deserialized on worker A, it gets a Client and creates a Future reference holding onto that Actor's key. The scheduler now notes that worker A's Client desires that key.

When the actual user's Client tries to release the Actor, the scheduler notes that worker A's Client still holds a reference to it, so it is not released.

More complex case:

A user submits a task where one of the dependencies is marked as an Actor, like:
```python
with dask.annotate(workers=workers[0]):
    counter = dask.delayed(Counter)()
with dask.annotate(workers=workers[1]):
    intermediate = dask.delayed(lambda c: None)(counter)
with dask.annotate(workers=workers[0]):
    final = dask.delayed(lambda x, c: x)(intermediate, counter)

final.compute(actors=counter, optimize_graph=False)
```

In this case, the user doesn't even hold a reference to the Actor. But when the `final` task completes and the scheduler runs `_propagate_forgotten` to release its dependencies (including `Counter`), it [sees that](https://github.com/dask/distributed/blob/bb991d1befdcb6b7b30ad984e08a23e3478c7cd2/distributed/scheduler.py#L7315) some Client holds a reference to the `Counter`, so it doesn't release it—when in fact the client holding the reference is `workers[1]`'s Actor handle.

This is what's causing test failures in #4925, now that we're more likely to schedule tasks on workers that don't hold any dependencies.
Transferring Actor handle between workers makes Actor unreleasable
When worker A calls `gather_dep` on an Actor task, it gets sent an Actor handle by worker B where the Actor is running. When that handle is deserialized on worker A, it gets a Client and creates a Future reference holding onto that Actor's key. The scheduler now notes that worker A's Client desires that key.

When the actual user's Client tries to release the Actor, the scheduler notes that worker A's Client still holds a reference to it, so it is not released.

More complex case:

A user submits a task where one of the dependencies is marked as an Actor, like:
```python
with dask.annotate(workers=workers[0]):
    counter = dask.delayed(Counter)()
with dask.annotate(workers=workers[1]):
    intermediate = dask.delayed(lambda c: None)(counter)
with dask.annotate(workers=workers[0]):
    final = dask.delayed(lambda x, c: x)(intermediate, counter)

final.compute(actors=counter, optimize_graph=False)
```

In this case, the user doesn't even hold a reference to the Actor. But when the `final` task completes and the scheduler runs `_propagate_forgotten` to release its dependencies (including `Counter`), it [sees that](https://github.com/dask/distributed/blob/bb991d1befdcb6b7b30ad984e08a23e3478c7cd2/distributed/scheduler.py#L7315) some Client holds a reference to the `Counter`, so it doesn't release it—when in fact the client holding the reference is `workers[1]`'s Actor handle.

This is what's causing test failures in #4925, now that we're more likely to schedule tasks on workers that don't hold any dependencies.


Discussion:

Resolving patch:
diff --git a/distributed/actor.py b/distributed/actor.py
--- a/distributed/actor.py
+++ b/distributed/actor.py
@@ -3,11 +3,11 @@
 import threading
 from queue import Queue
 
-from .client import Future, default_client
+from .client import Future
 from .protocol import to_serialize
 from .utils import iscoroutinefunction, sync, thread_state
 from .utils_comm import WrappedKey
-from .worker import get_worker
+from .worker import get_client, get_worker
 
 
 class Actor(WrappedKey):
@@ -59,12 +59,15 @@ def __init__(self, cls, address, key, worker=None):
             self._client = None
         else:
             try:
+                # TODO: `get_worker` may return the wrong worker instance for async local clusters (most tests)
+                # when run outside of a task (when deserializing a key pointing to an Actor, etc.)
                 self._worker = get_worker()
             except ValueError:
                 self._worker = None
             try:
-                self._client = default_client()
-                self._future = Future(key)
+                self._client = get_client()
+                self._future = Future(key, inform=self._worker is None)
+                # ^ When running on a worker, only hold a weak reference to the key, otherwise the key could become unreleasable.
             except ValueError:
                 self._client = None
 
@@ -109,7 +112,8 @@ 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
+            if self._asynchronous:
+                return func(*args, **kwargs)
             return sync(self._worker.loop, func, *args, **kwargs)
 
     def __dir__(self):
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1412,7 +1412,7 @@ async def get_data(
                 if k in self.actors:
                     from .actor import Actor
 
-                    data[k] = Actor(type(self.actors[k]), self.address, k)
+                    data[k] = Actor(type(self.actors[k]), self.address, k, worker=self)
 
         msg = {"status": "OK", "data": {k: to_serialize(v) for k, v in data.items()}}
         nbytes = {k: self.tasks[k].nbytes for k in data if k in self.tasks}


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
@@ -564,6 +564,61 @@ async def wait(self):
     await c.gather(futures)
 
 
+@gen_cluster(client=True, client_kwargs=dict(set_as_default=False))
+# ^ NOTE: without `set_as_default=False`, `get_client()` within worker would return
+# the same client instance the test is using (because it's all one process).
+# Even with this, both workers will share the same client instance.
+async def test_worker_actor_handle_is_weakref(c, s, a, b):
+    counter = c.submit(Counter, actor=True, workers=[a.address])
+
+    await c.submit(lambda _: None, counter, workers=[b.address])
+
+    del counter
+
+    start = time()
+    while a.actors or b.data:
+        await asyncio.sleep(0.1)
+        assert time() < start + 30
+
+
+def test_worker_actor_handle_is_weakref_sync(client):
+    workers = list(client.run(lambda: None))
+    counter = client.submit(Counter, actor=True, workers=[workers[0]])
+
+    client.submit(lambda _: None, counter, workers=[workers[1]]).result()
+
+    del counter
+
+    def check(dask_worker):
+        return len(dask_worker.data) + len(dask_worker.actors)
+
+    start = time()
+    while any(client.run(check).values()):
+        sleep(0.01)
+        assert time() < start + 30
+
+
+def test_worker_actor_handle_is_weakref_from_compute_sync(client):
+    workers = list(client.run(lambda: None))
+
+    with dask.annotate(workers=workers[0]):
+        counter = dask.delayed(Counter)()
+    with dask.annotate(workers=workers[1]):
+        intermediate = dask.delayed(lambda c: None)(counter)
+    with dask.annotate(workers=workers[0]):
+        final = dask.delayed(lambda x, c: x)(intermediate, counter)
+
+    final.compute(actors=counter, optimize_graph=False)
+
+    def worker_tasks_running(dask_worker):
+        return len(dask_worker.data) + len(dask_worker.actors)
+
+    start = time()
+    while any(client.run(worker_tasks_running).values()):
+        sleep(0.01)
+        assert time() < start + 30
+
+
 def test_one_thread_deadlock():
     with cluster(nworkers=2) as (cl, w):
         client = Client(cl["address"])