diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index 74f0f347769..7cd12f597a6 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -33,6 +33,16 @@ __all__ = ["VirtualMemoryResource", "VirtualMemoryResourceOptions"] +# Location types whose physical backing lives in host memory. Shared by +# VirtualMemoryResource.__init__ and is_host_accessible so the two cannot drift. +_HOST_LOCATION_TYPES = frozenset( + { + VirtualMemoryLocationType.HOST, + VirtualMemoryLocationType.HOST_NUMA, + VirtualMemoryLocationType.HOST_NUMA_CURRENT, + } +) + @dataclass class VirtualMemoryResourceOptions: @@ -169,8 +179,7 @@ def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions self.config: VirtualMemoryResourceOptions = check_or_create_options( # type: ignore[assignment] VirtualMemoryResourceOptions, config, "VirtualMemoryResource options", keep_none=False ) - # Matches ("host", "host_numa", "host_numa_current") - if "host" in self.config.location_type: + if self.config.location_type in _HOST_LOCATION_TYPES: self.device = None if not self.device and self.config.location_type == "device": @@ -609,7 +618,7 @@ def is_host_accessible(self) -> bool: """ Indicates whether the allocated memory is accessible from the host. """ - return self.config.location_type == "host" + return self.config.location_type in _HOST_LOCATION_TYPES @property def device_id(self) -> int: diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 98baa521ef2..8b7605efc7d 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1854,6 +1854,23 @@ def test_vmm_options_handle_type_win32_raises(): VirtualMemoryResourceOptions._handle_type_to_driver("win32") +@pytest.mark.agent_authored(model="claude-opus-5") +@pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"]) +def test_vmm_host_location_types_report_host_accessible(location_type): + """Every host-backed location type reports is_host_accessible. + + __init__ classifies "host", "host_numa" and "host_numa_current" alike when + deciding the resource is not bound to a device, so is_host_accessible must + agree; otherwise a NUMA-located resource claims to be neither host- nor + device-accessible. + """ + device = Device() + device.set_current() + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) + assert mr.device is None + assert mr.is_host_accessible is True + + def test_device_memory_resource_peer_accessible_by_non_owned(mempool_device): """peer_accessible_by on a non-owned (default) DMR queries the driver live.""" dev = mempool_device