diff --git a/dask_cuda/proxy_object.py b/dask_cuda/proxy_object.py index 15f2ea994..36592526c 100644 --- a/dask_cuda/proxy_object.py +++ b/dask_cuda/proxy_object.py @@ -18,6 +18,11 @@ from dask.sizeof import sizeof from distributed.worker import dumps_function, loads_function +try: + from dask.dataframe.backends import concat_pandas +except ImportError: + from dask.dataframe.methods import concat_pandas + from .get_device_memory_objects import get_device_memory_objects from .is_device_object import is_device_object @@ -726,10 +731,10 @@ def wrapper(*args, **kwargs): # Register dispatch of ProxyObject on all known dispatch objects for dispatch in ( - dask.dataframe.utils.hash_object_dispatch, + dask.dataframe.core.hash_object_dispatch, dask.dataframe.utils.make_meta, dask.dataframe.utils.make_scalar, - dask.dataframe.utils.group_split_dispatch, + dask.dataframe.core.group_split_dispatch, dask.array.core.tensordot_lookup, dask.array.core.einsum_lookup, dask.array.core.concatenate_lookup, @@ -745,5 +750,5 @@ def wrapper(*args, **kwargs): # deserialize all ProxyObjects before concatenating dask.dataframe.methods.concat_dispatch.register( (pandas.DataFrame, pandas.Series, pandas.Index), - unproxify_input_wrapper(dask.dataframe.methods.concat_pandas), + unproxify_input_wrapper(concat_pandas), ) diff --git a/dask_cuda/tests/test_ucx_options.py b/dask_cuda/tests/test_ucx_options.py index 99a7ed714..77af9357b 100644 --- a/dask_cuda/tests/test_ucx_options.py +++ b/dask_cuda/tests/test_ucx_options.py @@ -50,6 +50,7 @@ def _test_global_option(seg_size): assert conf["SEG_SIZE"] == seg_size +@pytest.mark.xfail(reason="https://github.com/rapidsai/dask-cuda/issues/627") def test_global_option(): for seg_size in ["2K", "1M", "2M"]: p = mp.Process(target=_test_global_option, args=(seg_size,))