`Nanny` assumes `self.scheduler` is `Scheduler`, hangs when used in `async with`
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
CC: @mfeurer
**What happened**:
Using a `Nanny` worker with `async with` context manager, inside its own process, failed to connect to a `LocalCluster`'s scheduler.
**What you expected to happen**:
For the `Nanny` worker to have connected to the `LocalCluster`
**Reason**
This is due to a change [here](https://github.com/dask/distributed/commit/50fd3ff34e1a66e2fe0b27bce1457e8fd4b00d7d#diff-cada6df8cd92eda4c9748bdc445353abe8b8fdef8723b888a53d39e102ed2b19R311) which occured in `dask 2021.07.0`. Most of the context for this change is found in that change.
* [L333](https://github.com/dask/distributed/blob/main/distributed/nanny.py#L333) `await self.scheduler.register_nanny()` - Seems like it's expecting a `Scheduler` object, [specifically this](https://github.com/dask/distributed/blob/cc3a6dfca71e1304f1e87ae996be87c615f297f6/distributed/scheduler.py#L3862).
* [L244](https://github.com/dask/distributed/blob/cc3a6dfca71e1304f1e87ae996be87c615f297f6/distributed/scheduler.py#L3862) However, where this is set, `self.scheduler = self.rpc(self.scheduler_addr)` this returns an instance of [`PooledRPCCall`](https://github.com/dask/distributed/blob/cc3a6dfca71e1304f1e87ae996be87c615f297f6/distributed/core.py#L860)
* There is nothing I could find in that class hierarchy that would change it to a `Scheduler`
**Context**
When upgrading from `distributed 2021.06.0` to `distributed 2021.12.0`, one of our [examples](https://automl.github.io/auto-sklearn/master/examples/60_search/example_parallel_manual_spawning_python.html#sphx-glr-examples-60-search-example-parallel-manual-spawning-python-py) for [`auto-sklearn`](https://github.com/automl/auto-sklearn) fails.
**Minimal Complete Verifiable Example**:
Apologies, it could probably be made more minimal but I'm not the most familiar with dask or Python's async.
```python
import asyncio
import multiprocessing
import time
import dask
from dask.distributed import Client, LocalCluster, Nanny
def start_python_worker(ip, port):
dask.config.set({'distributed.worker.daemon': False})
async def do_work():
async with dask.distributed.Nanny(
scheduler_ip=ip,
scheduler_port=port,
nthreads=1,
) as worker: # <--- Hangs here
print("I won't print")
await worker.finished()
# asyncio.run(do_work())
asyncio.get_event_loop().run_until_complete(do_work())
if __name__ == "__main__":
with LocalCluster(
n_workers=0,
processes=True,
threads_per_worker=1,
) as cluster:
scheduler = cluster.scheduler
process = multiprocessing.Process(
target=start_python_worker,
args=(scheduler.ip, scheduler.port)
)
process.start()
time.sleep(2) # Buffer for process to spin up
with dask.distributed.Client(cluster) as client:
future = client.submit(lambda x: x * 2, 10)
# Should be async but here we are, looping
while not future.done():
time.sleep(1)
result = future.result()
print(result)
process.join()
```
**Environment**:
- Dask version: `2021.12.0`
- Python version: `3.8.5`
- Operating System: `Manjaro Linux`
- Install method (conda, pip, source): `pip`
Contributor guide
Assessment
This issue has not been assessed yet.