Coverage for cuda / core / _memory / _device_memory_resource.pyx: 68.06%
72 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from __future__ import annotations
7from cuda.bindings cimport cydriver
8from cuda.core._memory._memory_pool cimport (
9 _MemPool, MP_init_create_pool, MP_raise_release_threshold,
10)
11from cuda.core._memory cimport _ipc
12from cuda.core._memory._ipc cimport IPCAllocationHandle
13from cuda.core._resource_handles cimport (
14 as_cu,
15 get_device_mempool,
16 get_last_error,
17)
18from cuda.core._utils.cuda_utils cimport (
19 check_or_create_options,
20 HANDLE_RETURN,
21)
22from dataclasses import dataclass
23import multiprocessing
24import platform # no-cython-lint
25import uuid
27from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy, replace_peer_accessible_by
28from cuda.core._utils.cuda_utils import check_multiprocessing_start_method
30__all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions']
33@dataclass
34cdef class DeviceMemoryResourceOptions:
35 """Customizable :obj:`~_memory.DeviceMemoryResource` options.
37 Attributes
38 ----------
39 ipc_enabled : bool, optional
40 Specifies whether to create an IPC-enabled memory pool. When set to
41 True, the memory pool and its allocations can be shared with other
42 processes. (Default to False)
44 max_size : int, optional
45 Maximum pool size. When set to 0, defaults to a system-dependent value.
46 (Default to 0)
47 """
48 ipc_enabled : bool = False
49 max_size : int = 0
52cdef class DeviceMemoryResource(_MemPool):
53 """
54 A device memory resource managing a stream-ordered memory pool.
56 Parameters
57 ----------
58 device_id : Device | int
59 Device or Device ordinal for which a memory resource is constructed.
61 options : DeviceMemoryResourceOptions
62 Memory resource creation options.
64 If set to `None`, the memory resource uses the driver's current
65 stream-ordered memory pool for the specified `device_id`. If no memory
66 pool is set as current, the driver's default memory pool for the device
67 is used.
69 If not set to `None`, a new memory pool is created, which is owned by
70 the memory resource.
72 When using an existing (current or default) memory pool, the returned
73 device memory resource does not own the pool (`is_handle_owned` is
74 `False`), and closing the resource has no effect.
76 Notes
77 -----
78 To create an IPC-Enabled memory resource (MR) that is capable of sharing
79 allocations between processes, specify ``ipc_enabled=True`` in the initializer
80 option. Sharing an allocation is a two-step procedure that involves
81 mapping a memory resource and then mapping buffers owned by that resource.
82 These steps can be accomplished in several ways.
84 An IPC-enabled memory resource can allocate memory buffers but cannot
85 receive shared buffers. Mapping an MR to another process creates a "mapped
86 memory resource" (MMR). An MMR cannot allocate memory buffers and can only
87 receive shared buffers. MRs and MMRs are both of type
88 :class:`DeviceMemoryResource` and can be distinguished via
89 :attr:`DeviceMemoryResource.is_mapped`.
91 An MR is shared via an allocation handle accessed through the
92 :attr:`DeviceMemoryResource.allocation_handle` property. The allocation
93 handle has a platform-specific interpretation; however, memory IPC is
94 currently only supported for Linux, and in that case allocation handles
95 are file descriptors. After sending an allocation handle to another
96 process, it can be used to create an MMR by invoking
97 :meth:`DeviceMemoryResource.from_allocation_handle`.
99 Buffers can be shared as serializable descriptors accessed through the
100 :attr:`Buffer.ipc_descriptor` property. In a receiving process, a shared
101 buffer is created by invoking :meth:`Buffer.from_ipc_descriptor` with an
102 MMR and buffer descriptor, where the MMR corresponds to the MR that
103 created the described buffer.
105 To help manage the association between memory resources and buffers, a
106 registry is provided. Every MR has a unique identifier (UUID). MMRs can be
107 registered by calling :meth:`DeviceMemoryResource.register` with the UUID
108 of the corresponding MR. Registered MMRs can be looked up via
109 :meth:`DeviceMemoryResource.from_registry`. When registering MMRs in this
110 way, the use of buffer descriptors can be avoided. Instead, buffer objects
111 can themselves be serialized and transferred directly. Serialization embeds
112 the UUID, which is used to locate the correct MMR during reconstruction.
114 IPC-enabled memory resources interoperate with the :mod:`multiprocessing`
115 module to provide a simplified interface. This approach can avoid direct
116 use of allocation handles, buffer descriptors, MMRs, and the registry. When
117 using :mod:`multiprocessing` to spawn processes or send objects through
118 communication channels such as :class:`multiprocessing.Queue`,
119 :class:`multiprocessing.Pipe`, or :class:`multiprocessing.Connection`,
120 :class:`Buffer` objects may be sent directly, and in such cases the process
121 for creating MMRs and mapping buffers will be handled automatically.
123 For greater efficiency when transferring many buffers, one may also send
124 MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
125 an MMR is created and registered in the receiving process. Subsequently,
126 buffers may be serialized and transferred using ordinary :mod:`pickle`
127 methods. The reconstruction procedure uses the registry to find the
128 associated MMR.
129 """
131 def __cinit__(self, *args, **kwargs):
132 self._dev_id = cydriver.CU_DEVICE_INVALID 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
134 def __init__(self, device_id: Device | int, options=None):
135 _DMR_init(self, device_id, options) 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
137 def __reduce__(self):
138 return DeviceMemoryResource.from_registry, (self.uuid,) 1adbc
140 @staticmethod
141 def from_registry(uuid: uuid.UUID) -> DeviceMemoryResource: # no-cython-lint
142 """
143 Obtain a registered mapped memory resource.
145 Raises
146 ------
147 RuntimeError
148 If no mapped memory resource is found in the registry.
149 """
150 return <DeviceMemoryResource>(_ipc.MP_from_registry(uuid))
152 def register(self, uuid: uuid.UUID) -> DeviceMemoryResource: # no-cython-lint
153 """
154 Register a mapped memory resource.
156 Returns
157 -------
158 The registered mapped memory resource. If one was previously registered
159 with the given key, it is returned.
160 """
161 return <DeviceMemoryResource>(_ipc.MP_register(self, uuid))
163 @classmethod
164 def from_allocation_handle(
165 cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle
166 ) -> DeviceMemoryResource:
167 """Create a device memory resource from an allocation handle.
169 Construct a new `DeviceMemoryResource` instance that imports a memory
170 pool from a shareable handle. The memory pool is marked as owned, and
171 the resource is associated with the specified `device_id`.
173 Parameters
174 ----------
175 device_id : int | Device
176 The ID of the device or a Device object for which the memory
177 resource is created.
179 alloc_handle : int | IPCAllocationHandle
180 The shareable handle of the device memory resource to import. If an
181 integer is supplied, it must represent a valid platform-specific
182 handle. It is the caller's responsibility to close that handle.
184 Returns
185 -------
186 A new device memory resource instance with the imported handle.
187 """
188 cdef DeviceMemoryResource mr = <DeviceMemoryResource>(
189 _ipc.MP_from_allocation_handle(cls, alloc_handle))
190 from .._device import Device
191 mr._dev_id = Device(device_id).device_id
192 return mr
194 @property
195 def allocation_handle(self) -> IPCAllocationHandle:
196 """Shareable handle for this memory pool (requires IPC).
198 The handle can be used to share the memory pool with other processes.
199 The handle is cached in this `MemoryResource` and owned by it.
200 """
201 if not self.is_ipc_enabled: 2{ba d `be |b}bmcnc~bacocbcccpcdcecfcqcgchcicf g jckclch i j k b c s l t m n
202 raise RuntimeError("Memory resource is not IPC-enabled") 2a d `bs
203 return self._ipc_data._alloc_handle 2{ba `be |b}bmcnc~bacocbcccpcdcecfcqcgchcicf g jckclch i j k b c l t m n
205 @property
206 def device_id(self) -> int:
207 """The associated device ordinal."""
208 return self._dev_id 2O P Q R S T U V {ba d `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c ; = q r p ? l m n u v rcsctcw x y ucvcwcxcyczcz A B C D AcBcCcE F G H I J K DcEcFcGcL M N
210 @property
211 def peer_accessible_by(self):
212 """
213 Get or set the devices that can access allocations from this memory
214 pool. Access can be modified at any time and affects all allocations
215 from this memory pool.
217 Returns a set-like proxy of :obj:`~_device.Device` objects that manages
218 peer access. Inputs are accepted as either :obj:`~_device.Device`
219 objects or device-ordinal :class:`int` values.
221 Examples
222 --------
223 >>> dmr = DeviceMemoryResource(0)
224 >>> dmr.peer_accessible_by = {1} # grant access to device 1
225 >>> assert 1 in dmr.peer_accessible_by
226 >>> dmr.peer_accessible_by.add(2) # update access to include device 2
227 >>> dmr.peer_accessible_by = [] # revoke peer access
228 """
229 return PeerAccessibleBySetProxy(self)
231 @peer_accessible_by.setter
232 def peer_accessible_by(self, devices):
233 replace_peer_accessible_by(self, devices)
235 @property
236 def is_device_accessible(self) -> bool:
237 """Return True. This memory resource provides device-accessible buffers."""
238 return True 2q r p u v rcsctcw x y ucvcwcxcyczcz A B C D AcBcCcE F G H I J K DcEcFcGcL M N
240 @property
241 def is_host_accessible(self) -> bool:
242 """Return False. This memory resource does not provide host-accessible buffers."""
243 return False 1qrp
246cdef inline _DMR_init(DeviceMemoryResource self, device_id, options):
247 from .._device import Device 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
248 cdef int dev_id = Device(device_id).device_id 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
249 cdef DeviceMemoryResourceOptions opts = check_or_create_options( 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
250 DeviceMemoryResourceOptions, options, "DeviceMemoryResource options",
251 keep_none=True
252 )
253 cdef bint ipc_enabled = False 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
254 cdef size_t max_size = 0 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
256 self._dev_id = dev_id 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
258 if opts is not None: 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
259 ipc_enabled = opts.ipc_enabled 2W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c p @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs l t m n ibjb
260 if ipc_enabled and not _ipc.is_supported(): 2W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c p @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs l t m n ibjb
261 raise RuntimeError(f"IPC is not available on {platform.system()}")
262 max_size = opts.max_size 2W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c p @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs l t m n ibjb
264 if opts is None: 2kblbmbnbobO P Q R S T U V W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c ; = pbqbrbsbtbq r p ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs ubvbl t m n wbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.bibjb/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
265 self._h_pool = get_device_mempool(dev_id) 2kblbmbnbobO P Q R S T U V ; = pbqbrbsbtbq r ? ubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
266 if not self._h_pool: 2kblbmbnbobO P Q R S T U V ; = pbqbrbsbtbq r ? ubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
267 HANDLE_RETURN(get_last_error())
268 raise RuntimeError(
269 f"Failed to initialize DeviceMemoryResource for device {dev_id}: "
270 "cuda-core returned an empty memory pool handle without recording a CUDA error. "
271 "This is an internal cuda-core error; please report it with your CUDA driver, "
272 "CUDA Toolkit, and cuda-python versions."
273 )
274 self._mempool_owned = False 2kblbmbnbobO P Q R S T U V ; = pbqbrbsbtbq r ? ubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
275 MP_raise_release_threshold(self) 2kblbmbnbobO P Q R S T U V ; = pbqbrbsbtbq r ? ubvbwbxbybzbAbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*b+b,b-b.b/b:b;b=b?b@b[b]bu v w x y z A B C D E F G H I J K ^b_bL M N
276 else:
277 MP_init_create_pool( 2W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c p @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs l t m n ibjb
278 self,
279 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
280 dev_id,
281 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED,
282 ipc_enabled,
283 max_size, 2W a X d Y e Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - f g . / : h i j k b c p @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbs l t m n ibjb
284 )
287# Note: this is referenced in instructions to debug nvbug 5698116.
288cpdef DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):
289 """
290 Probes peer access from the given device using cuMemPoolGetAccess.
292 Parameters
293 ----------
294 device_id : int or Device
295 The device to query access for.
297 Returns
298 -------
299 str
300 Access permissions: "rw" for read-write, "r" for read-only, "" for no access.
301 """
302 from .._device import Device
304 cdef int dev_id = Device(device_id).device_id
305 cdef cydriver.CUmemAccess_flags flags
306 cdef cydriver.CUmemLocation location
308 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
309 location.id = dev_id
311 with nogil:
312 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location))
314 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE:
315 return "rw"
316 elif flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ:
317 return "r"
318 else:
319 return ""
322def _deep_reduce_device_memory_resource(mr):
323 check_multiprocessing_start_method() 2{ba d `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c l m n
324 from .._device import Device 2{ba d `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c l m n
325 device = Device(mr.device_id) 2{ba d `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c l m n
326 alloc_handle = mr.allocation_handle 2{ba d `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c l m n
327 return mr.from_allocation_handle, (device, alloc_handle) 2{ba `be |b}b~bacbcccdcecfcgchcicf g jckclch i j k b c l m n
330multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)