Fix DeviceHostFile mapping views

dask_cuda/device_host_file.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import time
+from contextlib import suppress
 
 import numpy
 from zict import Buffer, Func
@@ -262,9 +263,7 @@ def __init__(
         self.others = {}
 
     def __setitem__(self, key, value):
-        if key in self.device_buffer:
-            # Make sure we register the removal of an existing key
-            del self[key]
+        self._discard_key(key)
 
         if is_spillable_object(value):
             self.others[key] = value
@@ -284,22 +283,34 @@ def __getitem__(self, key):
         raise KeyError(key)
 
     def __len__(self):
-        return len(self.device_buffer) + len(self.others)
+        return sum(1 for _ in self)
 
     def __iter__(self):
-        return itertools.chain(self.device_buffer, self.others)
-
-    def __delitem__(self, key):
+        seen = set()
+        for key in itertools.chain(self.device_buffer, self.host_buffer, self.others):
+            if key not in seen:
+                seen.add(key)
+                yield key
+
+    def _discard_key(self, key):
+        removed = key in self.others
         self.device_keys.discard(key)
-        if key in self.others:
-            del self.others[key]
-        else:
-            if isinstance(self.device_buffer, dict) and key not in self.device_buffer:
-                # If `self.device_buffer` is a dictionary, host `key`s are inserted
-                # directly into `self.host_buffer`.
+        self.others.pop(key, None)
+
+        with suppress(KeyError):
+            del self.device_buffer[key]
+            removed = True
+
+        if self.host_buffer is not self.device_buffer:
+            with suppress(KeyError):
                 del self.host_buffer[key]
-            else:
-                del self.device_buffer[key]
+                removed = True
+
+        return removed
+
+    def __delitem__(self, key):
+        if not self._discard_key(key):
+            raise KeyError(key)
 
     def evict(self):
         """Evicts least recently used host buffer (aka, CPU or system memory)

dask_cuda/tests/test_device_host_file_mapping.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from dask_cuda.device_host_file import DeviceHostFile
+from dask_cuda.is_spillable_object import is_spillable_object
+
+
+class HostValue:
+    pass
+
+
+class SpillableValue:
+    pass
+
+
+class DeviceValue:
+    __cuda_array_interface__ = {
+        "shape": (),
+        "typestr": "<f8",
+        "data": (0, False),
+        "version": 3,
+    }
+
+
+@is_spillable_object.register(SpillableValue)
+def is_spillable_test_value(_):
+    return True
+
+
+def _make_value(kind):
+    if kind == "host":
+        return HostValue()
+    if kind == "others":
+        return SpillableValue()
+    if kind == "device":
+        return DeviceValue()
+    raise ValueError(kind)
+
+
+def _assert_key_location(dhf, key, location):
+    assert (key in dhf.others) is (location == "others")
+    assert (key in dhf.device_buffer) is (location == "device")
+    assert (key in dhf.host_buffer) is (location == "host")
+    assert (key in dhf.device_keys) is (location == "device")
+
+
+@pytest.mark.parametrize("device_memory_limit", [None, 1])
+def test_device_host_file_mapping_views_include_host_only_keys(
+    tmp_path, device_memory_limit
+):
+    dhf = DeviceHostFile(
+        device_memory_limit=device_memory_limit,
+        memory_limit=None,
+        worker_local_directory=tmp_path,
+    )
+    values = {
+        "numpy": np.arange(3),
+        "pandas": pd.DataFrame({"x": [1, 2, 3]}),
+    }
+
+    for key, value in values.items():
+        dhf[key] = value
+
+    assert len(dhf) == len(values)
+    assert set(dhf) == set(values)
+    assert list(dhf).count("numpy") == 1
+    assert list(dhf).count("pandas") == 1
+    np.testing.assert_array_equal(dhf["numpy"], values["numpy"])
+    pd.testing.assert_frame_equal(dhf["pandas"], values["pandas"])
+
+
+@pytest.mark.parametrize(
+    "first_kind, second_kind",
+    [
+        ("others", "host"),
+        ("others", "device"),
+        ("host", "others"),
+        ("device", "others"),
+        ("device", "host"),
+        ("host", "device"),
+    ],
+)
+def test_device_host_file_overwrites_discard_previous_storage(
+    tmp_path, first_kind, second_kind
+):
+    dhf = DeviceHostFile(
+        device_memory_limit=None,
+        memory_limit=None,
+        worker_local_directory=tmp_path,
+    )
+    key = "x"
+    first = _make_value(first_kind)
+    second = _make_value(second_kind)
+
+    dhf[key] = first
+    assert dhf[key] is first
+    _assert_key_location(dhf, key, first_kind)
+
+    dhf[key] = second
+
+    assert dhf[key] is second
+    assert len(dhf) == 1
+    assert list(dhf) == [key]
+    _assert_key_location(dhf, key, second_kind)
+
+
+@pytest.mark.parametrize("kind", ["host", "device", "others"])
+def test_device_host_file_delete_discards_storage(tmp_path, kind):
+    dhf = DeviceHostFile(
+        device_memory_limit=None,
+        memory_limit=None,
+        worker_local_directory=tmp_path,
+    )
+    key = "x"
+
+    dhf[key] = _make_value(kind)
+    del dhf[key]
+
+    assert len(dhf) == 0
+    assert list(dhf) == []
+    _assert_key_location(dhf, key, None)
+    with pytest.raises(KeyError):
+        dhf[key]