Use Client.as_current more
- 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
Assessment
This issue has not been assessed yet.