tmp

dask_cuda/proxify_device_objects.py

@@ -81,6 +81,7 @@ def proxify_device_object_proxy_object(obj, proxied_id_to_proxy, found_proxies):
             _id = id(obj._obj_pxy["obj"])
             proxied_id_to_proxy[_id] = obj
 
+    obj._obj_pxy["not-external"] = True
     found_proxies.append(obj)
     return obj
 

dask_cuda/proxify_host_file.py

@@ -189,6 +189,8 @@ def add_external(self, obj):
         last_access = time.time()
         self_weakref = weakref.ref(self)
         for p in found_proxies:
+            if "not-external" in p._obj_pxy:
+                continue
             name = id(p)
             finalize = weakref.finalize(p, self.del_external, name)
             external = weakref.proxy(p)

dask_cuda/tests/test_proxify_host_file.py

@@ -155,7 +155,7 @@ def test_externals():
 
 def test_externals_setitem():
     dhf = ProxifyHostFile(device_memory_limit=itemsize)
-    k1 = dhf.add_external(cupy.arange(1) + 1)
+    k1 = dhf.add_external(cupy.arange(1))
     assert type(k1) is dask_cuda.proxy_object.ProxyObject
     assert len(dhf) == 0
     assert "external" in k1._obj_pxy
@@ -167,11 +167,16 @@ def test_externals_setitem():
     assert "external" not in k1._obj_pxy
     assert "external_finalize" not in k1._obj_pxy
 
-    k1 = dhf.add_external(cupy.arange(1) + 1)
+    k1 = dhf.add_external(cupy.arange(1))
     k1._obj_pxy_serialize(serializers=("dask", "pickle"))
     dhf["k1"] = k1
     k1 = dhf["k1"]
     assert type(k1) is dask_cuda.proxy_object.ProxyObject
     assert len(dhf) == 1
     assert "external" not in k1._obj_pxy
     assert "external_finalize" not in k1._obj_pxy
+
+    dhf["k1"] = cupy.arange(1)
+    assert len(dhf.proxies_tally.proxy_id_to_proxy) == 1
+    k1 = dhf.add_external(k1)
+    assert len(dhf.proxies_tally.proxy_id_to_proxy) == 1