NVIDIA / NVIDIA/cudf

[BUG] Dask error related to task fusion and dask.Array relating to dask config `optimization.fuse.active`

Open
#18,750 2 comments 0 reactions 0 assignees View on GitHub
bug dask Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Describe the bug**
PS : It might be more suitable in dask / cuml but haven't been able to isolate the issue so sharing here as it might be related to `dask_cudf`

The issue is if you a dask dataframe with a `list[float]` column and you extract the `list[float]` column to get `dask.Array` and then further run `arr.compute_chunk_sizes()` you might run into an error between 25.02 / 25.04 unless you wrap in `dask.config.set({"optimization.fuse.active": False})`. For 25.06 this succeeds but fails when you do `ddf["centroid"] = Kmeans.fit(arr).predict(arr)` and then `ddf.to_parquet(...)`

```python
# psuedocode

ddf = dask_cudf.from_cudf(
cudf.DataFrame({"embeddings" : [[1,2],[3,4], "id" : [1,2]}),
npartitions=2
)

arr = ddf.map_partitions(get_array_from_df)

arr.persist() # this will cause failure in 25.06 unless dask config is set
arr.compute_chunk_sizes()

ddf["center"] = KMeans().fit(arr).predict(arr)

ddf.to_parquet(..)
```

The permutations that works are
1. Remove persist for 25.02 but keep with dask.config.set
- In 25.02/25.04 compute_chunk_sizes won't work without the config
2. Remove persist for 25.06, you can also remove with dask.config.set
3. Set dask.config at global level works in all versions

**Steps/Code to reproduce bug**
See comment in code
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"

from cuml.dask.cluster import KMeans
import dask
from dask_cuda import LocalCUDACluster
from distributed import Client, wait
import dask_cudf
import cudf
import cupy as cp
import dask.dataframe as dd
import tempfile

def get_array_from_df(df: cudf.DataFrame, embedding_col: str) -> cp.ndarray:
return df[embedding_col].list.leaves.values.reshape(len(df), -1) # noqa: PD011

with Client(
LocalCUDACluster(
CUDA_VISIBLE_DEVICES="4,5",
enable_cudf_spill=True,
)
) as client, tempfile.TemporaryDirectory() as tmpdir:
# Create a dask dataframe
ddf = dask_cudf.from_cudf(cudf.DataFrame({
"embeddings": [cp.random.randn(10) for _ in range(10_000)],
"id" : range(10_000)
}), npartitions=10).optimize()

# Get array from the ddf
cupy_arr = ddf.map_partitions(
get_array_from_df, "embeddings", meta=cp.ndarray([1, 1])
)

# These three lines are the tricky part
# if you just remove the config it fails on [25.02, 25.06]
with dask.config.set({"optimization.fuse.active": False}):
# If you comment out persist + wait, code works as expected
cupy_arr = cupy_arr.persist()
wait(cupy_arr)
# If you move the dask config to a global level i.e where we create cluster
# then it works as well
cupy_arr.compute_chunk_sizes()

# Train and predict on KMeans
kmeans = KMeans(
n_clusters=10,
max_iter=100,
random_state=42,
n_init=1,
).fit(cupy_arr)
predictions = kmeans.predict(cupy_arr)
# predictions.compute()

ddf["nearest_cent"] = predictions.astype(cp.int32)

ddf.to_parquet(
tmpdir,
index=False,
partition_on="nearest_cent",
write_index=False,
)
print("DONE")
```
**Expected behavior**
I should be able to persist my dask array
**Environment overview (please complete the following information)**
Rapids 25.02 / 25.04 / 25.06

**Environment details**
Please run and paste the output of the `cudf/print_env.sh` script here, to gather any other relevant environment details

**Additional context**
Add any other context about the problem here.

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.