diff --git a/dask_cuda/explicit_comms/dataframe/shuffle.py b/dask_cuda/explicit_comms/dataframe/shuffle.py index 71e47f5b..d5dda66d 100644 --- a/dask_cuda/explicit_comms/dataframe/shuffle.py +++ b/dask_cuda/explicit_comms/dataframe/shuffle.py @@ -543,7 +543,14 @@ def shuffle( # Create a distributed Dataframe from all the pieces divs = [None] * (len(futures) + 1) - ret = dd.from_delayed(futures, meta=df_meta, divisions=divs).persist() + kwargs = {"meta": df_meta, "divisions": divs, "prefix": "explicit-comms-shuffle"} + try: + ret = dd.from_delayed(futures, **kwargs).persist() + except TypeError: + # This version of dask may not support `prefix` + # See: https://github.com/dask/dask-expr/pull/991 + kwargs.pop("prefix") + ret = dd.from_delayed(futures, **kwargs).persist() wait([ret]) # Release all temporary dataframes diff --git a/dask_cuda/tests/test_explicit_comms.py b/dask_cuda/tests/test_explicit_comms.py index 49995592..b39f2f2b 100644 --- a/dask_cuda/tests/test_explicit_comms.py +++ b/dask_cuda/tests/test_explicit_comms.py @@ -31,7 +31,7 @@ dask.config.get("dataframe.query-planning", None) is not False, reason=( "The 'explicit-comms' config is not supported " - "when query planning is enabled.", + "when query planning is enabled." ), )