Skip to content

Commit 55c156e

Browse files
martindurantmartindurant
committed
Apply suggestions
1 parent a519baa commit 55c156e

3 files changed

Lines changed: 16 additions & 10 deletions

File tree

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ class VirtualMemoryResourceOptions:
6565
"""
6666
allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
6767
location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE
68-
handle_type: VirtualMemoryHandleType | None = field(default_factory=lambda: None if platform.system() == 'Windows' else VirtualMemoryHandleType.POSIX_FD)
68+
handle_type: VirtualMemoryHandleType | str | None = field(default_factory=lambda: None if platform.system() == 'Windows' else VirtualMemoryHandleType.POSIX_FD)
6969
granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED
7070
gpu_direct_rdma: bool = False
7171
addr_hint: int | None = 0
@@ -93,7 +93,7 @@ class VirtualMemoryResourceOptions:
9393
@staticmethod
9494
def _location_type_to_driver(spec: VirtualMemoryLocationType) -> int: ...
9595
@staticmethod
96-
def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int: ...
96+
def _handle_type_to_driver(spec: VirtualMemoryHandleType | str | None) -> int: ...
9797
@staticmethod
9898
def _granularity_to_driver(spec: VirtualMemoryGranularityType) -> int: ...
9999

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ class VirtualMemoryResourceOptions:
134134

135135
allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
136136
location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE
137-
handle_type: VirtualMemoryHandleType | None = field(
137+
handle_type: VirtualMemoryHandleType | str | None = field(
138138
default_factory=lambda: None if platform.system() == "Windows" else VirtualMemoryHandleType.POSIX_FD
139139
)
140140
granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED
@@ -203,7 +203,7 @@ class VirtualMemoryResourceOptions:
203203
return loc_type # type: ignore[no-any-return]
204204

205205
@staticmethod
206-
def _handle_type_to_driver(spec: VirtualMemoryHandleType | None) -> int:
206+
def _handle_type_to_driver(spec: VirtualMemoryHandleType | str | None) -> int:
207207
if spec == "win32":
208208
raise NotImplementedError("win32 is currently not supported, please reach out to the CUDA Python team")
209209
handle_type = VirtualMemoryResourceOptions._handle_types.get(spec)

‎cuda_core/tests/test_memory.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1768,7 +1768,11 @@ def test_vmm_host_modify_allocation_rejects_exportable_handle_type(init_cuda):
17681768
buf = mr.allocate(4096)
17691769
try:
17701770
with pytest.raises(ValueError, match="handle_type=None"):
1771-
mr.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions(location_type="host"))
1771+
mr.modify_allocation(
1772+
buf,
1773+
2 * buf.size,
1774+
config=VirtualMemoryResourceOptions(location_type="host", handle_type=VirtualMemoryHandleType.POSIX_FD),
1775+
)
17721776
finally:
17731777
buf.close()
17741778

@@ -3384,17 +3388,19 @@ def test_vmm_host_location_types_report_host_accessible(init_cuda, location_type
33843388
deciding the resource is not bound to a device, so is_host_accessible must
33853389
agree; otherwise a NUMA-located resource claims to be neither host- nor
33863390
device-accessible. The driver rejects an exportable handle type for HOST
3387-
memory only, so the default handle type is rejected at construction for
3388-
"host" and accepted for "host_numa".
3391+
memory only, so an exportable default handle type is rejected at
3392+
construction for "host" and accepted for "host_numa". The Windows
3393+
default, None, is accepted for both location types.
33893394
"""
33903395
device = Device()
33913396
device.set_current()
33923397
numa = {"numa_id": 0} if location_type == "host_numa" else {}
3393-
if location_type == "host":
3398+
config = VirtualMemoryResourceOptions(location_type=location_type, **numa)
3399+
if location_type == "host" and config.handle_type is not None:
33943400
with pytest.raises(ValueError, match="handle_type=None"):
3395-
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
3401+
VirtualMemoryResource(device, config=config)
33963402
else:
3397-
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type, **numa))
3403+
VirtualMemoryResource(device, config=config)
33983404
mr = VirtualMemoryResource(
33993405
device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None, **numa)
34003406
)

0 commit comments

Comments
 (0)