Skip to content

Commit e55a26b

Browse files
committed
Fix PinnedMemoryResource IPC to derive NUMA ID from active device (#1603)
PinnedMemoryResource(ipc_enabled=True) hardcoded host NUMA ID 0, causing failures on multi-NUMA systems where the active device is attached to a different NUMA node. Now derives the NUMA ID from the current device's host_numa_id attribute, and adds an explicit numa_id option for manual override. Removes the _check_numa_nodes warning machinery in favor of proper NUMA node selection. Made-with: Cursor
1 parent 61810bd commit e55a26b

3 files changed

Lines changed: 116 additions & 77 deletions

File tree

‎cuda_core/cuda/core/_memory/_pinned_memory_resource.pxd‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@ from cuda.core._memory._ipc cimport IPCDataForMR
77

88

99
cdef class PinnedMemoryResource(_MemPool):
10-
pass
10+
cdef int _numa_id

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

Lines changed: 44 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,8 @@ from cuda.core._utils.cuda_utils cimport (
1515

1616
from dataclasses import dataclass
1717
import multiprocessing
18-
import os
1918
import platform # no-cython-lint
20-
import subprocess
21-
import threading
2219
import uuid
23-
import warnings
2420

2521
from cuda.core._utils.cuda_utils import check_multiprocessing_start_method
2622

@@ -41,9 +37,22 @@ cdef class PinnedMemoryResourceOptions:
4137
max_size : int, optional
4238
Maximum pool size. When set to 0, defaults to a system-dependent value.
4339
(Default to 0)
40+
41+
numa_id : int or None, optional
42+
Host NUMA node ID for pool placement. When set to None (the default),
43+
the behavior depends on ``ipc_enabled``:
44+
45+
- ``ipc_enabled=False``: OS-managed placement (location type HOST).
46+
- ``ipc_enabled=True``: automatically derived from the current CUDA
47+
device's ``host_numa_id`` attribute, requiring an active CUDA
48+
context.
49+
50+
When set to a non-negative integer, that NUMA node is used explicitly
51+
regardless of ``ipc_enabled`` (location type HOST_NUMA).
4452
"""
4553
ipc_enabled : bool = False
4654
max_size : int = 0
55+
numa_id : int | None = None
4756

4857

4958
cdef class PinnedMemoryResource(_MemPool):
@@ -71,12 +80,10 @@ cdef class PinnedMemoryResource(_MemPool):
7180
-----
7281
To create an IPC-Enabled memory resource (MR) that is capable of sharing
7382
allocations between processes, specify ``ipc_enabled=True`` in the initializer
74-
option. When IPC is enabled, the location type is automatically set to
75-
CU_MEM_LOCATION_TYPE_HOST_NUMA instead of CU_MEM_LOCATION_TYPE_HOST,
76-
with location ID 0.
77-
78-
Note: IPC support for pinned memory requires a single NUMA node. A warning
79-
is issued if multiple NUMA nodes are detected.
83+
option. When IPC is enabled and ``numa_id`` is not specified, the NUMA node
84+
is automatically derived from the current CUDA device's ``host_numa_id``
85+
attribute, which requires an active CUDA context. If ``numa_id`` is
86+
explicitly set, that value is used regardless of ``ipc_enabled``.
8087
8188
See :class:`DeviceMemoryResource` for more details on IPC usage patterns.
8289
"""
@@ -160,6 +167,11 @@ cdef class PinnedMemoryResource(_MemPool):
160167
"""Return -1. Pinned memory is host memory and is not associated with a specific device."""
161168
return -1
162169

170+
@property
171+
def numa_id(self) -> int:
172+
"""The host NUMA node ID used for pool placement, or -1 for OS-managed placement."""
173+
return self._numa_id
174+
163175
@property
164176
def is_device_accessible(self) -> bool:
165177
"""Return True. This memory resource provides device-accessible buffers."""
@@ -172,42 +184,56 @@ cdef class PinnedMemoryResource(_MemPool):
172184

173185

174186
cdef inline _PMR_init(PinnedMemoryResource self, options):
187+
from .._device import Device
188+
175189
cdef PinnedMemoryResourceOptions opts = check_or_create_options(
176190
PinnedMemoryResourceOptions, options, "PinnedMemoryResource options",
177191
keep_none=True
178192
)
179193
cdef bint ipc_enabled = False
180194
cdef size_t max_size = 0
181195
cdef cydriver.CUmemLocationType loc_type
182-
cdef int location_id
196+
cdef int numa_id = -1
183197

184198
if opts is not None:
185199
ipc_enabled = opts.ipc_enabled
186200
if ipc_enabled and not _ipc.is_supported():
187201
raise RuntimeError(f"IPC is not available on {platform.system()}")
188-
if ipc_enabled:
189-
_check_numa_nodes()
190202
max_size = opts.max_size
191203

192-
if ipc_enabled:
204+
if opts.numa_id is not None:
205+
numa_id = opts.numa_id
206+
if numa_id < 0:
207+
raise ValueError(f"numa_id must be >= 0, got {numa_id}")
208+
elif ipc_enabled:
209+
dev = Device()
210+
numa_id = dev.properties.host_numa_id
211+
if numa_id < 0:
212+
raise RuntimeError(
213+
"Cannot determine host NUMA ID for IPC-enabled pinned "
214+
"memory pool. The system may not support NUMA, or no "
215+
"CUDA context is active. Set numa_id explicitly or "
216+
"call Device.set_current() first.")
217+
218+
if numa_id >= 0:
193219
loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA
194-
location_id = 0
195220
else:
196221
loc_type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
197-
location_id = -1
222+
223+
self._numa_id = numa_id
198224

199225
if opts is None:
200226
MP_init_current_pool(
201227
self,
202228
loc_type,
203-
location_id,
229+
numa_id,
204230
cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED,
205231
)
206232
else:
207233
MP_init_create_pool(
208234
self,
209235
loc_type,
210-
location_id,
236+
numa_id,
211237
cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED,
212238
ipc_enabled,
213239
max_size,
@@ -221,60 +247,3 @@ def _deep_reduce_pinned_memory_resource(mr):
221247

222248

223249
multiprocessing.reduction.register(PinnedMemoryResource, _deep_reduce_pinned_memory_resource)
224-
225-
226-
cdef bint _numa_warning_shown = False
227-
cdef object _numa_lock = threading.Lock()
228-
229-
230-
cdef inline _check_numa_nodes():
231-
"""Check if system has multiple NUMA nodes and warn if so."""
232-
global _numa_warning_shown
233-
if _numa_warning_shown:
234-
return
235-
236-
with _numa_lock:
237-
if _numa_warning_shown:
238-
return
239-
240-
if platform.system() != "Linux":
241-
_numa_warning_shown = True
242-
return
243-
244-
numa_count = None
245-
246-
# Try /sys filesystem first (most reliable and doesn't require external tools)
247-
try:
248-
node_path = "/sys/devices/system/node"
249-
if os.path.exists(node_path):
250-
nodes = [d for d in os.listdir(node_path) if d.startswith("node") and d[4:].isdigit()]
251-
numa_count = len(nodes)
252-
except (OSError, PermissionError):
253-
pass
254-
255-
# Fallback to lscpu if /sys check didn't work
256-
if numa_count is None:
257-
try:
258-
result = subprocess.run(
259-
["lscpu"],
260-
capture_output=True,
261-
text=True,
262-
timeout=1
263-
)
264-
for line in result.stdout.splitlines():
265-
if line.startswith("NUMA node(s):"):
266-
numa_count = int(line.split(":")[1].strip())
267-
break
268-
except (subprocess.SubprocessError, ValueError, FileNotFoundError):
269-
pass
270-
271-
if numa_count is not None and numa_count > 1:
272-
warnings.warn(
273-
f"System has {numa_count} NUMA nodes. IPC-enabled pinned memory "
274-
f"uses location ID 0, which may not work correctly with multiple "
275-
f"NUMA nodes.",
276-
UserWarning,
277-
stacklevel=3
278-
)
279-
280-
_numa_warning_shown = True

‎cuda_core/tests/test_memory.py‎

Lines changed: 71 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1039,6 +1039,7 @@ def test_pinned_mempool_ipc_basic():
10391039
assert mr.is_device_accessible
10401040
assert mr.is_host_accessible
10411041
assert mr.device_id == -1 # pinned memory is not device-specific
1042+
assert mr.numa_id >= 0 # IPC requires a concrete NUMA node
10421043

10431044
# Test allocation handle export
10441045
alloc_handle = mr.get_allocation_handle()
@@ -1070,7 +1071,8 @@ def test_pinned_mempool_ipc_errors():
10701071
options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=False)
10711072
mr = PinnedMemoryResource(options)
10721073
assert not mr.is_ipc_enabled
1073-
assert mr.device_id == -1 # Non-IPC uses location id -1
1074+
assert mr.device_id == -1
1075+
assert mr.numa_id == -1 # Non-IPC uses OS-managed placement
10741076

10751077
buffer = mr.allocate(64)
10761078
ipc_error_msg = "Memory resource is not IPC-enabled"
@@ -1089,6 +1091,74 @@ def test_pinned_mempool_ipc_errors():
10891091
mr.close()
10901092

10911093

1094+
def test_pinned_mr_numa_id_default_no_ipc(init_cuda):
1095+
"""numa_id defaults to -1 (OS-managed) when IPC is disabled."""
1096+
device = Device()
1097+
skip_if_pinned_memory_unsupported(device)
1098+
1099+
mr = PinnedMemoryResource(PinnedMemoryResourceOptions())
1100+
assert mr.numa_id == -1
1101+
mr.close()
1102+
1103+
mr = PinnedMemoryResource(PinnedMemoryResourceOptions(ipc_enabled=False))
1104+
assert mr.numa_id == -1
1105+
mr.close()
1106+
1107+
1108+
def test_pinned_mr_numa_id_default_with_ipc(init_cuda):
1109+
"""numa_id is derived from the current device when IPC is enabled."""
1110+
device = Device()
1111+
skip_if_pinned_memory_unsupported(device)
1112+
1113+
if platform.system() == "Windows":
1114+
pytest.skip("IPC not implemented for Windows")
1115+
if not supports_ipc_mempool(device):
1116+
pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")
1117+
1118+
expected_numa_id = device.properties.host_numa_id
1119+
if expected_numa_id < 0:
1120+
pytest.skip("System does not support NUMA")
1121+
1122+
mr = PinnedMemoryResource(PinnedMemoryResourceOptions(ipc_enabled=True, max_size=POOL_SIZE))
1123+
assert mr.numa_id == expected_numa_id
1124+
mr.close()
1125+
1126+
1127+
def test_pinned_mr_numa_id_explicit(init_cuda):
1128+
"""Explicit numa_id is used regardless of ipc_enabled."""
1129+
device = Device()
1130+
skip_if_pinned_memory_unsupported(device)
1131+
1132+
host_numa_id = device.properties.host_numa_id
1133+
if host_numa_id < 0:
1134+
pytest.skip("System does not support NUMA")
1135+
1136+
mr = PinnedMemoryResource(PinnedMemoryResourceOptions(numa_id=host_numa_id))
1137+
assert mr.numa_id == host_numa_id
1138+
mr.close()
1139+
1140+
if platform.system() == "Windows":
1141+
pytest.skip("IPC not implemented for Windows")
1142+
if not supports_ipc_mempool(device):
1143+
pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")
1144+
1145+
mr = PinnedMemoryResource(PinnedMemoryResourceOptions(ipc_enabled=True, numa_id=host_numa_id, max_size=POOL_SIZE))
1146+
assert mr.numa_id == host_numa_id
1147+
mr.close()
1148+
1149+
1150+
def test_pinned_mr_numa_id_negative_error(init_cuda):
1151+
"""Negative numa_id raises ValueError."""
1152+
device = Device()
1153+
skip_if_pinned_memory_unsupported(device)
1154+
1155+
with pytest.raises(ValueError, match="numa_id must be >= 0"):
1156+
PinnedMemoryResource(PinnedMemoryResourceOptions(numa_id=-1))
1157+
1158+
with pytest.raises(ValueError, match="numa_id must be >= 0"):
1159+
PinnedMemoryResource(PinnedMemoryResourceOptions(numa_id=-42))
1160+
1161+
10921162
@pytest.mark.parametrize("ipc_enabled", [True, False])
10931163
@pytest.mark.parametrize(
10941164
"property_name,expected_type",

0 commit comments

Comments
 (0)