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

1# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5from __future__ import annotations 

6  

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 

26  

27from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy, replace_peer_accessible_by 

28from cuda.core._utils.cuda_utils import check_multiprocessing_start_method 

29  

30__all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions'] 

31  

32  

33@dataclass 

34cdef class DeviceMemoryResourceOptions: 

35 """Customizable :obj:`~_memory.DeviceMemoryResource` options. 

36  

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) 

43  

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 

50  

51  

52cdef class DeviceMemoryResource(_MemPool): 

53 """ 

54 A device memory resource managing a stream-ordered memory pool. 

55  

56 Parameters 

57 ---------- 

58 device_id : Device | int 

59 Device or Device ordinal for which a memory resource is constructed. 

60  

61 options : DeviceMemoryResourceOptions 

62 Memory resource creation options. 

63  

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. 

68  

69 If not set to `None`, a new memory pool is created, which is owned by 

70 the memory resource. 

71  

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. 

75  

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. 

83  

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`. 

90  

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`. 

98  

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. 

104  

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. 

113  

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. 

122  

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 """ 

130  

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

133  

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

136  

137 def __reduce__(self): 

138 return DeviceMemoryResource.from_registry, (self.uuid,) 1adbc

139  

140 @staticmethod 

141 def from_registry(uuid: uuid.UUID) -> DeviceMemoryResource: # no-cython-lint 

142 """ 

143 Obtain a registered mapped memory resource. 

144  

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)) 

151  

152 def register(self, uuid: uuid.UUID) -> DeviceMemoryResource: # no-cython-lint 

153 """ 

154 Register a mapped memory resource. 

155  

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)) 

162  

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. 

168  

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`. 

172  

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. 

178  

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. 

183  

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 

193  

194 @property 

195 def allocation_handle(self) -> IPCAllocationHandle: 

196 """Shareable handle for this memory pool (requires IPC). 

197  

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

204  

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

209  

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. 

216  

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. 

220  

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) 

230  

231 @peer_accessible_by.setter 

232 def peer_accessible_by(self, devices): 

233 replace_peer_accessible_by(self, devices) 

234  

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

239  

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

244  

245  

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

255  

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

257  

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

263  

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 ) 

285  

286  

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. 

291  

292 Parameters 

293 ---------- 

294 device_id : int or Device 

295 The device to query access for. 

296  

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 

303  

304 cdef int dev_id = Device(device_id).device_id 

305 cdef cydriver.CUmemAccess_flags flags 

306 cdef cydriver.CUmemLocation location 

307  

308 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

309 location.id = dev_id 

310  

311 with nogil: 

312 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location)) 

313  

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 "" 

320  

321  

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

328  

329  

330multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)