From f066c51b75f4b0f3c605de54c243c9913d2caeef Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 3 Nov 2022 09:22:29 +0100 Subject: [PATCH] Fix #1032 --- dask_cuda/get_device_memory_objects.py | 5 +++-- dask_cuda/tests/test_proxify_host_file.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/dask_cuda/get_device_memory_objects.py b/dask_cuda/get_device_memory_objects.py index 92055c190..44dc433ff 100644 --- a/dask_cuda/get_device_memory_objects.py +++ b/dask_cuda/get_device_memory_objects.py @@ -51,8 +51,9 @@ def get_device_memory_objects_default(obj): return dispatch(obj._pxy_get().obj) if hasattr(obj, "data"): return dispatch(obj.data) - if hasattr(obj, "_owner") and obj._owner is not None: - return dispatch(obj._owner) + owner = getattr(obj, "owner", None) or getattr(obj, "_owner", None) + if owner: + return dispatch(owner) if hasattr(obj, "__cuda_array_interface__"): return [obj] return [] diff --git a/dask_cuda/tests/test_proxify_host_file.py b/dask_cuda/tests/test_proxify_host_file.py index 1edcab09d..6758385fc 100644 --- a/dask_cuda/tests/test_proxify_host_file.py +++ b/dask_cuda/tests/test_proxify_host_file.py @@ -273,7 +273,7 @@ def test_dataframes_share_dev_mem(): # Even though the two dataframe doesn't point to the same cudf.Buffer object assert view1["a"].data is not view2["a"].data # They still share the same underlying device memory - assert view1["a"].data._owner._owner is view2["a"].data._owner._owner + view1["a"].data.ptr == view2["a"].data.ptr dhf = ProxifyHostFile( local_directory=root_dir, device_memory_limit=160, memory_limit=1000