From 69421df293c9e14162922049f5abeb63ead02a9b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 21 Mar 2024 12:06:11 -0700 Subject: [PATCH] use prefix kwarg - comming soon to dask-expr anyway --- dask_cuda/explicit_comms/dataframe/shuffle.py | 9 ++++++++- dask_cuda/tests/test_explicit_comms.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) 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." ), )