dask / dask/distributed

Client objects never collected in default usage unless `client.close` is called explicitly.

Open
#7,770 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

**Describe the issue**:

[This comes about as part of investigation of #7726.]

Consider
```python
from distributed import Client, LocalCluster
import gc

def which(cls):
return [x for x in gc.get_objects() if isinstance(x, cls)]

def howmany(cls):
return len(which(cls))

if __name__ == "__main__":
cluster = LocalCluster(protocol="tcp", n_workers=1, threads_per_worker=1)
for _ in range(100):
client = Client(cluster, set_as_default=False)
del cluster, client
gc.collect()
print(howmany(Client))
```
This prints `100` for me on e1944ec7.

Spelunking a bit

```python
import objgraph

objgraph.show_chain(objgraph.find_backref_chain(which(Client)[0], objgraph.is_proper_module))
```

![objgraph](https://user-images.githubusercontent.com/1126981/231224772-97110243-6ea3-4805-8625-39b293ab534d.png)

The periodic callbacks live on global state that persists, and reference the `Client` object via bound methods. This pattern also appears in (at least) `LocalCluster` setup, I have yet to check scheduler and so forth.

One approach round this would be something like the following, which at least removes the refcycles, but is not ready for prime-time since I get errors that must be related (though I do not yet see how):
```diff
diff --git a/distributed/client.py b/distributed/client.py
index 264fe232..a61538b1 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -960,12 +960,13 @@ class Client(SyncMethodMixin):
dask.config.get("distributed.client.scheduler-info-interval", default="ms")
)

+ pxy = weakref.proxy(self)
self._periodic_callbacks = dict()
self._periodic_callbacks["scheduler-info"] = PeriodicCallback(
- self._update_scheduler_info, scheduler_info_interval * 1000
+ partial(Client._update_scheduler_info, pxy), scheduler_info_interval * 1000
)
self._periodic_callbacks["heartbeat"] = PeriodicCallback(
- self._heartbeat, heartbeat_interval * 1000
+ partial(Client._heartbeat, pxy), heartbeat_interval * 1000
)

self._start_arg = address
@@ -1299,7 +1300,7 @@ class Client(SyncMethodMixin):
for preload in self.preloads:
await preload.start()

- self._handle_report_task = asyncio.create_task(self._handle_report())
+ self._handle_report_task = asyncio.create_task(Client._handle_report(weakref.proxy(self)))

return self

@@ -1353,9 +1354,9 @@ class Client(SyncMethodMixin):
)
comm.name = "Client->Scheduler"
if timeout is not None:
- await wait_for(self._update_scheduler_info(), timeout)
+ await wait_for(self._update_scheduler_info(self), timeout)
else:
- await self._update_scheduler_info()
+ await self._update_scheduler_info(self)
await comm.write(
{
"op": "register-client",
@@ -1396,13 +1397,17 @@ class Client(SyncMethodMixin):

logger.debug("Started scheduling coroutines. Synchronized")

+ @staticmethod
async def _update_scheduler_info(self):
- if self.status not in ("running", "connecting") or self.scheduler is None:
- return
try:
- self._scheduler_identity = SchedulerInfo(await self.scheduler.identity())
- except OSError:
- logger.debug("Not able to query scheduler for identity")
+ if self.status not in ("running", "connecting") or self.scheduler is None:
+ return
+ try:
+ self._scheduler_identity = SchedulerInfo(await self.scheduler.identity())
+ except OSError:
+ logger.debug("Not able to query scheduler for identity")
+ except ReferenceError:
+ return

async def _wait_for_workers(
self, n_workers: int, timeout: float | None = None
@@ -1460,13 +1465,17 @@ class Client(SyncMethodMixin):
)
return self.sync(self._wait_for_workers, n_workers, timeout=timeout)

+ @staticmethod
def _heartbeat(self):
- # Don't send heartbeat if scheduler comm or cluster are already closed
- if self.scheduler_comm is not None and not (
- self.scheduler_comm.comm.closed()
- or (self.cluster and self.cluster.status in (Status.closed, Status.closing))
- ):
- self.scheduler_comm.send({"op": "heartbeat-client"})
+ try:
+ # Don't send heartbeat if scheduler comm or cluster are already closed
+ if self.scheduler_comm is not None and not (
+ self.scheduler_comm.comm.closed()
+ or (self.cluster and self.cluster.status in (Status.closed, Status.closing))
+ ):
+ self.scheduler_comm.send({"op": "heartbeat-client"})
+ except ReferenceError:
+ return

def __enter__(self):
if not self._loop_runner.is_started():
@@ -1517,6 +1526,7 @@ class Client(SyncMethodMixin):
)

@log_errors
+ @staticmethod
async def _handle_report(self):
"""Listen to scheduler"""
try:
@@ -1572,6 +1582,8 @@ class Client(SyncMethodMixin):
break
except (CancelledError, asyncio.CancelledError):
pass
+ except ReferenceError:
+ return

def _handle_key_in_memory(self, key=None, type=None, workers=None):
state = self.futures.get(key)
```

Typical errors on exit:

```
2023-04-11 18:06:53,694 - distributed.comm.tcp - WARNING - Closing dangling stream in
/home/wence/Documents/src/rapids/third-party/distributed/distributed/client.py:1504: RuntimeWarning: coroutine 'wait_for' was never awaited
self.close()
/home/wence/Documents/src/rapids/third-party/distributed/distributed/client.py:1504: RuntimeWarning: coroutine 'Client._close' was never awaited
self.close()
2023-04-11 18:06:53,759 - distributed.comm.tcp - WARNING - Closing dangling stream in
2023-04-11 18:06:53,759 - distributed.comm.tcp - WARNING - Closing dangling stream in
2023-04-11 18:06:53,760 - distributed.comm.tcp - WARNING - Closing dangling stream in
0
2023-04-11 18:06:53,775 - tornado.application - ERROR - Exception in callback functools.partial(>, )
Traceback (most recent call last):
File "/home/wence/Documents/apps/mambaforge/envs/test-weakref/lib/python3.10/site-packages/tornado/ioloop.py", line 740, in _run_callback
ret = callback()
File "/home/wence/Documents/apps/mambaforge/envs/test-weakref/lib/python3.10/site-packages/tornado/ioloop.py", line 764, in _discard_future_result
future.result()
File "/home/wence/Documents/apps/mambaforge/envs/test-weakref/lib/python3.10/site-packages/tornado/gen.py", line 776, in run
yielded = self.gen.throw(*exc_info) # type: ignore
File "/home/wence/Documents/apps/mambaforge/envs/test-weakref/lib/python3.10/site-packages/tornado/gen.py", line 769, in run
value = future.result()
asyncio.exceptions.TimeoutError: Timeout
```

But notice that there are at least now zero live clients.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.