dask / dask/distributed

Nested `scatter` calls lead to `KeyError`

Open
#3,703 14 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

Hi All,

I am currently working on improving the `joblib-dask` integration.
It turns out that nested `Parallel` calls in `joblib` using the `dask` backend tend to error out with either `KeyError` or `CancelledError`.

I narrowed it down using only `dask` and `numpy`, and it seems that the issue comes from nested `scatter` calls.

Here is a reproducer: it consists of submitting functions that rely on scattered arrays. Each of these functions submit small arithmetic operations to be computed on scattered slices of their original input.
```python
import logging

import numpy as np

from distributed import LocalCluster, Client, get_client, secede, rejoin

NUM_INNER_TASKS = 10
NUM_OUTER_TASKS = 10

def my_sum(x, i, j):
print(f"running inner task {j} of outer task {i}")
return np.sum(x)

def outer_function(array, i):
print(f"running outer task {i}")
client = get_client()
slices = [array[i + j :] for j in range(NUM_INNER_TASKS)]

# commenting this line makes the code run successfully
slices = client.scatter(slices, broadcast=True)

futures = client.map(my_sum, slices, [i] * NUM_INNER_TASKS, range(NUM_INNER_TASKS))

secede()
results = client.gather(futures)
rejoin()
return sum(results)

if __name__ == "__main__":
my_arrays = [np.ones(100000) for _ in range(10)]

cluster = LocalCluster(
n_workers=1, threads_per_worker=1, silence_logs=logging.WARNING
)
client = Client(cluster)

future_arrays = client.scatter(my_arrays, direct=False)

# using .map() instead of .submit() makes the code run successfully.
# futures = client.map(outer_function, future_arrays, range(10))

futures = []
for i, arr in enumerate(future_arrays):
future = client.submit(outer_function, arr, i)
futures.append(future)

results = client.gather(futures)
print(results)
```

2 Remarks:

- as said in the code, using `client.map` makes the code run successfully.
- not scattering the `slices` in the `outer` functions makes the code run successfully.

My guess as of now is that dynamically creating new compute resources through `secede/rejoin` calls might interact badly with the data locality logic of `distributed`. I'm investigating this own my own, but I'm not familiar enough with the `dask/distributed` codebase to trace this back efficiently.

Is this behavior supported? Is there a clear anti-pattern that I'm missing? Any pointer would be helpful.

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.