diff --git a/dask_cuda/device_host_file.py b/dask_cuda/device_host_file.py index e8d8bc08b..197ffcc65 100644 --- a/dask_cuda/device_host_file.py +++ b/dask_cuda/device_host_file.py @@ -17,7 +17,7 @@ serialize_bytelist, ) from distributed.sizeof import safe_sizeof -from distributed.utils import has_arg, nbytes +from distributed.utils import nbytes from .is_device_object import is_device_object from .is_spillable_object import is_spillable_object @@ -27,12 +27,8 @@ def _serialize_bytelist(x, **kwargs): kwargs["on_error"] = "raise" - if has_arg(serialize_bytelist, "compression"): - compression = dask.config.get("distributed.worker.memory.spill-compression") - return serialize_bytelist(x, compression=compression, **kwargs) - else: - # For Distributed < 2023.5.0 compatibility - return serialize_bytelist(x, **kwargs) + compression = dask.config.get("distributed.worker.memory.spill-compression") + return serialize_bytelist(x, compression=compression, **kwargs) class LoggedBuffer(Buffer): diff --git a/dask_cuda/tests/test_device_host_file.py b/dask_cuda/tests/test_device_host_file.py index 17d4055c9..4a4807941 100644 --- a/dask_cuda/tests/test_device_host_file.py +++ b/dask_cuda/tests/test_device_host_file.py @@ -2,10 +2,8 @@ import numpy as np import pytest -from packaging import version import dask.array -import distributed from distributed.protocol import ( deserialize, deserialize_bytes, @@ -53,16 +51,7 @@ def test_device_host_file_short( random.shuffle(full) for k, v in full: - try: - dhf[k] = v - except TypeError as e: - # TODO: Remove when pinning to distributed>=2023.5.1 . - # See https://github.com/rapidsai/dask-cuda/issues/1174 and - # https://github.com/dask/distributed/pull/7836 . - if version.parse(distributed.__version__) <= version.parse("2023.5.0"): - dhf[k] = v - else: - raise e + dhf[k] = v random.shuffle(full)