Actor: don't hold key references on workers

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):
@@ -63,8 +63,8 @@ def __init__(self, cls, address, key, worker=None):
             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)
             except ValueError:
                 self._client = None
 

distributed/tests/test_actor.py

@@ -466,7 +466,7 @@ def f(block, ps=None):
 
 
 @pytest.mark.flaky(reruns=10, reruns_delay=5)
-@gen_cluster(client=True)
+@gen_cluster(client=True, timeout=10)
 async def test_compute(c, s, a, b):
     @dask.delayed
     def f(n, counter):
@@ -485,10 +485,8 @@ def check(counter, blanks):
     result = await c.compute(final, actors=counter)
     assert result == 0 + 1 + 2 + 3 + 4
 
-    start = time()
     while a.data or b.data:
         await asyncio.sleep(0.01)
-        assert time() < start + 30
 
 
 def test_compute_sync(client):
@@ -515,7 +513,7 @@ def check(dask_worker):
     start = time()
     while any(client.run(check).values()):
         sleep(0.01)
-        assert time() < start + 30
+        assert time() < start + 10
 
 
 @gen_cluster(
@@ -566,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 + 10
+
+
+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 + 10
+
+
+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 + 10
+
+
 def test_one_thread_deadlock():
     with cluster(nworkers=2) as (cl, w):
         client = Client(cl["address"])

distributed/tests/test_worker.py

@@ -1111,6 +1111,7 @@ async def test_robust_to_bad_sizeof_estimates(c, s, a):
     np = pytest.importorskip("numpy")
     memory = psutil.Process().memory_info().rss
     a.memory_limit = memory / 0.7 + 400e6
+    print("memory limit:", format_bytes(a.memory_limit))
 
     class BadAccounting:
         def __init__(self, data):
@@ -1128,6 +1129,12 @@ def f(n):
 
     start = time()
     while not a.data.disk:
+        print(
+            "RSS:",
+            format_bytes(psutil.Process().memory_info().rss),
+            "disk:",
+            list(a.data.disk),
+        )
         await asyncio.sleep(0.1)
         assert time() < start + 5
 

distributed/worker.py

@@ -1417,7 +1417,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}