use prefix kwarg - comming soon to dask-expr anyway

rapidsai · Mar 21, 2024 · 69421df · 69421df
1 parent fcc0e04
commit 69421df
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git 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
@@ -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."
     ),
 )