avoid using loads_function

rapidsai · Aug 16, 2023 · 5ba44a2 · 5ba44a2
1 parent f06d0af
commit 5ba44a2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/dask_cuda/proxy_object.py b/dask_cuda/proxy_object.py
@@ -19,7 +19,7 @@
 import distributed.utils
 from dask.sizeof import sizeof
 from distributed.protocol.compression import decompress
-from distributed.worker import dumps_function, loads_function
+from distributed.worker import dumps_function
 
 from dask_cuda.disk_io import disk_read
 
@@ -440,7 +440,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 +882,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