From 62fb56a47c8043c870127e945e6ac4911eb5910f Mon Sep 17 00:00:00 2001 From: "Richard (Rick) Zamora" Date: Thu, 17 Aug 2023 08:53:04 -0500 Subject: [PATCH] Avoid importing `loads_function` from distributed (#1220) Closes https://github.com/rapidsai/dask-cuda/issues/1219 Not sure if we need the caching behavior originally included in `loads_function`, or if `pickle.loads` is sufficient. Any idea @madsbk ? Authors: - Richard (Rick) Zamora (https://github.com/rjzamora) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) - Peter Andreas Entschev (https://github.com/pentschev) URL: https://github.com/rapidsai/dask-cuda/pull/1220 --- dask_cuda/proxy_object.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/dask_cuda/proxy_object.py b/dask_cuda/proxy_object.py index 2f9c774d..ddb7f329 100644 --- a/dask_cuda/proxy_object.py +++ b/dask_cuda/proxy_object.py @@ -19,7 +19,6 @@ import distributed.utils from dask.sizeof import sizeof from distributed.protocol.compression import decompress -from distributed.worker import dumps_function, loads_function from dask_cuda.disk_io import disk_read @@ -85,7 +84,7 @@ def asproxy( subclass = ProxyObject subclass_serialized = None else: - subclass_serialized = dumps_function(subclass) + subclass_serialized = pickle.dumps(subclass) ret = subclass( ProxyDetail( @@ -440,7 +439,7 @@ def __reduce__(self): pxy = self._pxy_get(copy=True) pxy.serialize(serializers=("pickle",)) if pxy.subclass: - subclass = loads_function(pxy.subclass) + subclass = pickle.loads(pxy.subclass) else: subclass = ProxyObject @@ -882,7 +881,7 @@ def obj_pxy_dask_deserialize(header, frames): if args["subclass"] is None: subclass = ProxyObject else: - subclass = loads_function(args["subclass"]) + subclass = pickle.loads(args["subclass"]) pxy = ProxyDetail(obj=(header["proxied-header"], frames), **args) if pxy.serializer == "disk": header, _ = pxy.obj