dask / dask/distributed

Use Client.as_current more

Open
#7,377 10 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

I've been playing around with a workflow that uses many Clients in many threads. Each client looks at a remote Dask cluster running a large XGBoost computation. I run them all in the same process to use Optuna for HPO.

I find that when run naively I run into issues with different threads picking up the wrong client at times. There are various parts in the codebase where we call `get_client` or `default_client` and it's not always foolproof which one gets selected.

My approach to this was to use `Client.current()` everywhere in `default_client` and `get_client()` and then to wrap my code in a `client.as_current` context manager

```python
def f():
with Cluster(...) as cluster:
with Client(cluster) as client:
with client.as_current:
XGBoost stuff

with concurrent.futures.ThreadPoolExecutor() as e:
futures = [e.submit(f) for _ in range(10)]
results = [future.result() for future in futures]
```

We have too many ways to get "the client" clearly. It seems like some attention was paid to using `ContextVars` with `Client.current`. This seems like a sensible approach to me and seems to be decently well done. We might now want to leverage it more fully.

For example, there's this in dask/dask

```diff
diff --git a/dask/base.py b/dask/base.py
index 43b0b1adf..bb7ac9d3d 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1370,9 +1370,12 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
)
return named_schedulers[scheduler]
elif scheduler in ("dask.distributed", "distributed"):
- from distributed.worker import get_client
+ from dask.distributed import default_client, get_client

- return get_client().get
+ try:
+ return default_client().get
+ except ValueError:
+ return get_client().get
else:
raise ValueError(
"Expected one of [distributed, %s]"
```

Also, maybe `Client.__enter__` and `Client.__aenter__` should also call `client.as_current.__enter__`?

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.