Fix deadlock and simplify proxy tracking

dask_cuda/device_host_file.py

@@ -175,14 +175,12 @@ def __init__(
         local_directory=None,
         log_spilling=False,
     ):
-        if local_directory is None:
-            local_directory = dask.config.get("temporary-directory") or os.getcwd()
-
-        if local_directory and not os.path.exists(local_directory):
-            os.makedirs(local_directory, exist_ok=True)
-        local_directory = os.path.join(local_directory, "dask-worker-space")
-
-        self.disk_func_path = os.path.join(local_directory, "storage")
+        self.disk_func_path = os.path.join(
+            local_directory or dask.config.get("temporary-directory") or os.getcwd(),
+            "dask-worker-space",
+            "storage",
+        )
+        os.makedirs(self.disk_func_path, exist_ok=True)
 
         self.host_func = dict()
         self.disk_func = Func(

dask_cuda/explicit_comms/dataframe/shuffle.py

@@ -18,7 +18,7 @@
 from distributed import wait
 from distributed.protocol import nested_deserialize, to_serialize
 
-from ...proxify_host_file import ProxifyHostFile
+from ...proxify_host_file import ProxyManager
 from .. import comms
 
 
@@ -148,19 +148,17 @@ async def local_shuffle(
     eps = s["eps"]
 
     try:
-        hostfile = first(iter(in_parts[0].values()))._obj_pxy.get(
-            "hostfile", lambda: None
-        )()
+        manager = first(iter(in_parts[0].values()))._obj_pxy.get("manager", None)
     except AttributeError:
-        hostfile = None
+        manager = None
 
-    if isinstance(hostfile, ProxifyHostFile):
+    if isinstance(manager, ProxyManager):
 
         def concat(args, ignore_index=False):
             if len(args) < 2:
                 return args[0]
 
-            return hostfile.add_external(dd_concat(args, ignore_index=ignore_index))
+            return manager.proxify(dd_concat(args, ignore_index=ignore_index))
 
     else:
         concat = dd_concat

dask_cuda/get_device_memory_objects.py

@@ -28,10 +28,7 @@ def get_device_memory_objects(obj) -> set:
 @dispatch.register(object)
 def get_device_memory_objects_default(obj):
     if hasattr(obj, "_obj_pxy"):
-        if obj._obj_pxy["serializers"] is None:
-            return dispatch(obj._obj_pxy["obj"])
-        else:
-            return []
+        return dispatch(obj._obj_pxy["obj"])
     if hasattr(obj, "data"):
         return dispatch(obj.data)
     if hasattr(obj, "_owner") and obj._owner is not None:

dask_cuda/local_cuda_cluster.py

@@ -273,9 +273,7 @@ def __init__(
                     {
                         "device_memory_limit": self.device_memory_limit,
                         "memory_limit": self.host_memory_limit,
-                        "local_directory": local_directory
-                        or dask.config.get("temporary-directory")
-                        or os.getcwd(),
+                        "local_directory": local_directory,
                         "log_spilling": log_spilling,
                     },
                 )

dask_cuda/proxify_device_objects.py

@@ -165,14 +165,9 @@ def wrapper(*args, **kwargs):
 
 def proxify(obj, proxied_id_to_proxy, found_proxies, subclass=None):
     _id = id(obj)
-    if _id in proxied_id_to_proxy:
-        ret = proxied_id_to_proxy[_id]
-        finalize = ret._obj_pxy.get("external_finalize", None)
-        if finalize:
-            finalize()
-            proxied_id_to_proxy[_id] = ret = asproxy(obj, subclass=subclass)
-    else:
-        proxied_id_to_proxy[_id] = ret = asproxy(obj, subclass=subclass)
+    if _id not in proxied_id_to_proxy:
+        proxied_id_to_proxy[_id] = asproxy(obj, subclass=subclass)
+    ret = proxied_id_to_proxy[_id]
     found_proxies.append(ret)
     return ret
 
@@ -190,11 +185,6 @@ def proxify_device_object_default(
 def proxify_device_object_proxy_object(
     obj, proxied_id_to_proxy, found_proxies, excl_proxies
 ):
-    # We deserialize CUDA-serialized objects since it is very cheap and
-    # makes it easy to administrate device memory usage
-    if obj._obj_pxy_is_serialized() and "cuda" in obj._obj_pxy["serializers"]:
-        obj._obj_pxy_deserialize()
-
     # Check if `obj` is already known
     if not obj._obj_pxy_is_serialized():
         _id = id(obj._obj_pxy["obj"])
@@ -203,14 +193,6 @@ def proxify_device_object_proxy_object(
         else:
             proxied_id_to_proxy[_id] = obj
 
-    finalize = obj._obj_pxy.get("external_finalize", None)
-    if finalize:
-        finalize()
-        obj = obj._obj_pxy_copy()
-        if not obj._obj_pxy_is_serialized():
-            _id = id(obj._obj_pxy["obj"])
-            proxied_id_to_proxy[_id] = obj
-
     if not excl_proxies:
         found_proxies.append(obj)
     return obj