diff --git a/dask_cuda/__init__.py b/dask_cuda/__init__.py index 55207d08f..103ca8ea9 100644 --- a/dask_cuda/__init__.py +++ b/dask_cuda/__init__.py @@ -15,7 +15,7 @@ from .cuda_worker import CUDAWorker from .explicit_comms.dataframe.shuffle import ( get_rearrange_by_column_wrapper, - get_default_shuffle_algorithm, + get_default_shuffle_method, ) from .local_cuda_cluster import LocalCUDACluster from .proxify_device_objects import proxify_decorator, unproxify_decorator @@ -28,11 +28,11 @@ dask.dataframe.shuffle.rearrange_by_column = get_rearrange_by_column_wrapper( dask.dataframe.shuffle.rearrange_by_column ) -# We have to replace all modules that imports Dask's `get_default_shuffle_algorithm()` +# We have to replace all modules that imports Dask's `get_default_shuffle_method()` # TODO: introduce a shuffle-algorithm dispatcher in Dask so we don't need this hack -dask.dataframe.shuffle.get_default_shuffle_algorithm = get_default_shuffle_algorithm -dask.dataframe.multi.get_default_shuffle_algorithm = get_default_shuffle_algorithm -dask.bag.core.get_default_shuffle_algorithm = get_default_shuffle_algorithm +dask.dataframe.shuffle.get_default_shuffle_method = get_default_shuffle_method +dask.dataframe.multi.get_default_shuffle_method = get_default_shuffle_method +dask.bag.core.get_default_shuffle_method = get_default_shuffle_method # Monkey patching Dask to make use of proxify and unproxify in compatibility mode diff --git a/dask_cuda/explicit_comms/dataframe/shuffle.py b/dask_cuda/explicit_comms/dataframe/shuffle.py index a444fce0b..0ca1c48ee 100644 --- a/dask_cuda/explicit_comms/dataframe/shuffle.py +++ b/dask_cuda/explicit_comms/dataframe/shuffle.py @@ -585,7 +585,7 @@ def wrapper(*args, **kwargs): return wrapper -def get_default_shuffle_algorithm() -> str: +def get_default_shuffle_method() -> str: """Return the default shuffle algorithm used by Dask This changes the default shuffle algorithm from "p2p" to "tasks" @@ -594,4 +594,4 @@ def get_default_shuffle_algorithm() -> str: ret = dask.config.get("dataframe.shuffle.algorithm", None) if ret is None and _use_explicit_comms(): return "tasks" - return dask.utils.get_default_shuffle_algorithm() + return dask.utils.get_default_shuffle_method()