@@ -15,12 +15,8 @@ from cuda.core._utils.cuda_utils cimport (
1515
1616from dataclasses import dataclass
1717import multiprocessing
18- import os
1918import platform # no-cython-lint
20- import subprocess
21- import threading
2219import uuid
23- import warnings
2420
2521from 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
4958cdef 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
174186cdef 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
223249multiprocessing.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
0 commit comments