Coverage for cuda / core / _memory / _device_memory_resource.pyx: 70.00%

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

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 _MemPool, _MemPoolOptions 

9from cuda.core._memory cimport _ipc 

10from cuda.core._memory._ipc cimport IPCAllocationHandle 

11from cuda.core._utils.cuda_utils cimport ( 

12 check_or_create_options, 

13 HANDLE_RETURN, 

14) 

15  

16from dataclasses import dataclass (empty)

17import multiprocessing (empty)

18import platform # no-cython-lint (empty)

19import uuid (empty)

20  

21from cuda.core._utils.cuda_utils import check_multiprocessing_start_method (empty)

22from cuda.core._resource_handles cimport as_cu 

23  

24__all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions'] (empty)

25  

26  

27@dataclass (empty)

28cdef class DeviceMemoryResourceOptions: 

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

30  

31 Attributes 

32 ---------- 

33 ipc_enabled : bool, optional 

34 Specifies whether to create an IPC-enabled memory pool. When set to 

35 True, the memory pool and its allocations can be shared with other 

36 processes. (Default to False) 

37  

38 max_size : int, optional 

39 Maximum pool size. When set to 0, defaults to a system-dependent value. 

40 (Default to 0) 

41 """ 

42 ipc_enabled : bool = False (empty)

43 max_size : int = 0 

44  

45  

46cdef class DeviceMemoryResource(_MemPool): 

47 """ 

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

49  

50 Parameters 

51 ---------- 

52 device_id : Device | int 

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

54  

55 options : DeviceMemoryResourceOptions 

56 Memory resource creation options. 

57  

58 If set to `None`, the memory resource uses the driver's current 

59 stream-ordered memory pool for the specified `device_id`. If no memory 

60 pool is set as current, the driver's default memory pool for the device 

61 is used. 

62  

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

64 the memory resource. 

65  

66 When using an existing (current or default) memory pool, the returned 

67 device memory resource does not own the pool (`is_handle_owned` is 

68 `False`), and closing the resource has no effect. 

69  

70 Notes 

71 ----- 

72 To create an IPC-Enabled memory resource (MR) that is capable of sharing 

73 allocations between processes, specify ``ipc_enabled=True`` in the initializer 

74 option. Sharing an allocation is a two-step procedure that involves 

75 mapping a memory resource and then mapping buffers owned by that resource. 

76 These steps can be accomplished in several ways. 

77  

78 An IPC-enabled memory resource can allocate memory buffers but cannot 

79 receive shared buffers. Mapping an MR to another process creates a "mapped 

80 memory resource" (MMR). An MMR cannot allocate memory buffers and can only 

81 receive shared buffers. MRs and MMRs are both of type 

82 :class:`DeviceMemoryResource` and can be distinguished via 

83 :attr:`DeviceMemoryResource.is_mapped`. 

84  

85 An MR is shared via an allocation handle obtained by calling 

86 :meth:`DeviceMemoryResource.get_allocation_handle`. The allocation handle 

87 has a platform-specific interpretation; however, memory IPC is currently 

88 only supported for Linux, and in that case allocation handles are file 

89 descriptors. After sending an allocation handle to another process, it can 

90 be used to create an MMR by invoking 

91 :meth:`DeviceMemoryResource.from_allocation_handle`. 

92  

93 Buffers can be shared as serializable descriptors obtained by calling 

94 :meth:`Buffer.get_ipc_descriptor`. In a receiving process, a shared buffer is 

95 created by invoking :meth:`Buffer.from_ipc_descriptor` with an MMR and 

96 buffer descriptor, where the MMR corresponds to the MR that created the 

97 described buffer. 

98  

99 To help manage the association between memory resources and buffers, a 

100 registry is provided. Every MR has a unique identifier (UUID). MMRs can be 

101 registered by calling :meth:`DeviceMemoryResource.register` with the UUID 

102 of the corresponding MR. Registered MMRs can be looked up via 

103 :meth:`DeviceMemoryResource.from_registry`. When registering MMRs in this 

104 way, the use of buffer descriptors can be avoided. Instead, buffer objects 

105 can themselves be serialized and transferred directly. Serialization embeds 

106 the UUID, which is used to locate the correct MMR during reconstruction. 

107  

108 IPC-enabled memory resources interoperate with the :mod:`multiprocessing` 

109 module to provide a simplified interface. This approach can avoid direct 

110 use of allocation handles, buffer descriptors, MMRs, and the registry. When 

111 using :mod:`multiprocessing` to spawn processes or send objects through 

112 communication channels such as :class:`multiprocessing.Queue`, 

113 :class:`multiprocessing.Pipe`, or :class:`multiprocessing.Connection`, 

114 :class:`Buffer` objects may be sent directly, and in such cases the process 

115 for creating MMRs and mapping buffers will be handled automatically. 

116  

117 For greater efficiency when transferring many buffers, one may also send 

118 MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`, 

119 an MMR is created and registered in the receiving process. Subsequently, 

120 buffers may be serialized and transferred using ordinary :mod:`pickle` 

121 methods. The reconstruction procedure uses the registry to find the 

122 associated MMR. 

123 """ 

124  

125 def __init__(self, device_id: Device | int, options=None): 

126 from .._device import Device 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

127 cdef int dev_id = Device(device_id).device_id 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

128 cdef DeviceMemoryResourceOptions opts = check_or_create_options( 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

129 DeviceMemoryResourceOptions, options, "DeviceMemoryResource options", 

130 keep_none=True 

131 ) 

132 cdef _MemPoolOptions opts_base = _MemPoolOptions() 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

133  

134 cdef bint ipc_enabled = False 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

135 if opts: 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

136 ipc_enabled = opts.ipc_enabled 50 ctx1satbuevwxyzABCDfgEFGhijkcdpHIJKLMNOPQRSTUVWXYqlrmn

137 if ipc_enabled and not _ipc.is_supported(): 50 ctx1satbuevwxyzABCDfgEFGhijkcdpHIJKLMNOPQRSTUVWXYqlrmn

138 raise RuntimeError("IPC is not available on {platform.system()}") 

139 opts_base._max_size = opts.max_size 50 ctx1satbuevwxyzABCDfgEFGhijkcdpHIJKLMNOPQRSTUVWXYqlrmn

140 opts_base._use_current = False 50 ctx1satbuevwxyzABCDfgEFGhijkcdpHIJKLMNOPQRSTUVWXYqlrmn

141 opts_base._ipc_enabled = ipc_enabled 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

142 opts_base._location = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

143 opts_base._type = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

144  

145 super().__init__(dev_id, opts_base) 119 ctx2. / : ; = ? @ [ ] ^ _ s a t b u e v w x y z A B C D f g E F G h i j k c d ` { | } Z 0 p ~ H I J K L M N O P Q R S T U V W X Y q abl r m n bbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbxbybzbAb1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + BbCbDb, Eb-

146  

147 def __reduce__(self): (empty)

148 return DeviceMemoryResource.from_registry, (self.uuid,) 4 ctx1abcd

149  

150 @staticmethod (empty)

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

152 """ 

153 Obtain a registered mapped memory resource. 

154  

155 Raises 

156 ------ 

157 RuntimeError 

158 If no mapped memory resource is found in the registry. 

159 """ 

160 return <DeviceMemoryResource>(_ipc.MP_from_registry(uuid)) 

161  

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

163 """ 

164 Register a mapped memory resource. 

165  

166 Returns 

167 ------- 

168 The registered mapped memory resource. If one was previously registered 

169 with the given key, it is returned. 

170 """ 

171 return <DeviceMemoryResource>(_ipc.MP_register(self, uuid)) 

172  

173 @classmethod (empty)

174 def from_allocation_handle( 

175 cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle 

176 ) -> DeviceMemoryResource: 

177 """Create a device memory resource from an allocation handle. 

178  

179 Construct a new `DeviceMemoryResource` instance that imports a memory 

180 pool from a shareable handle. The memory pool is marked as owned, and 

181 the resource is associated with the specified `device_id`. 

182  

183 Parameters 

184 ---------- 

185 device_id : int | Device 

186 The ID of the device or a Device object for which the memory 

187 resource is created. 

188  

189 alloc_handle : int | IPCAllocationHandle 

190 The shareable handle of the device memory resource to import. If an 

191 integer is supplied, it must represent a valid platform-specific 

192 handle. It is the caller's responsibility to close that handle. 

193  

194 Returns 

195 ------- 

196 A new device memory resource instance with the imported handle. 

197 """ 

198 cdef DeviceMemoryResource mr = <DeviceMemoryResource>( 

199 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 

200 from .._device import Device 

201 mr._dev_id = Device(device_id).device_id 

202 return mr 

203  

204 def get_allocation_handle(self) -> IPCAllocationHandle: (empty)

205 """Export the memory pool handle to be shared (requires IPC). 

206  

207 The handle can be used to share the memory pool with other processes. 

208 The handle is cached in this `MemoryResource` and owned by it. 

209  

210 Returns 

211 ------- 

212 The shareable handle for the memory pool. 

213 """ 

214 if not self.is_ipc_enabled: 28 ctx2Fba b Gbe HbIbJbQbKbLbMbf g NbObPbh i j k c d q l r m n

215 raise RuntimeError("Memory resource is not IPC-enabled") 1 ctx1q

216 return self._ipc_data._alloc_handle 27 ctx2Fba b Gbe HbIbJbQbKbLbMbf g NbObPbh i j k c d l r m n

217  

218 @property 

219 def is_device_accessible(self) -> bool: 

220 """Return True. This memory resource provides device-accessible buffers.""" 

221 return True 39 ctx2Z 0 p 1 2 Rb3 4 5 SbTb6 7 8 Ub9 ! Vb# $ % WbXbYb' Zb0b( 1b2b) 3b* 4b+ 5b, 6b-

222  

223 @property 

224 def is_host_accessible(self) -> bool: 

225 """Return False. This memory resource does not provide host-accessible buffers.""" 

226 return False 3 ctx1Z0p

227  

228  

229# Note: this is referenced in instructions to debug nvbug 5698116. 

230cpdef DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id): (empty)

231 """ 

232 Probes peer access from the given device using cuMemPoolGetAccess. 

233  

234 Parameters 

235 ---------- 

236 device_id : int or Device 

237 The device to query access for. 

238  

239 Returns 

240 ------- 

241 str 

242 Access permissions: "rw" for read-write, "r" for read-only, "" for no access. 

243 """ 

244 from .._device import Device 

245  

246 cdef int dev_id = Device(device_id).device_id 

247 cdef cydriver.CUmemAccess_flags flags 

248 cdef cydriver.CUmemLocation location 

249  

250 location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE 

251 location.id = dev_id 

252  

253 with nogil: 

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

255  

256 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: 

257 return "rw" 

258 elif flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ: 

259 return "r" 

260 else: 

261 return "" 

262  

263  

264def _deep_reduce_device_memory_resource(mr): (empty)

265 check_multiprocessing_start_method() 25 ctx2Fba b Gbe HbIbJbKbLbMbf g NbObPbh i j k c d l m n

266 from .._device import Device 25 ctx2Fba b Gbe HbIbJbKbLbMbf g NbObPbh i j k c d l m n

267 device = Device(mr.device_id) 25 ctx2Fba b Gbe HbIbJbKbLbMbf g NbObPbh i j k c d l m n

268 alloc_handle = mr.get_allocation_handle() 25 ctx2Fba b Gbe HbIbJbKbLbMbf g NbObPbh i j k c d l m n

269 return mr.from_allocation_handle, (device, alloc_handle) 25 ctx2Fba b Gbe HbIbJbKbLbMbf g NbObPbh i j k c d l m n

270  

271  

272multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource) (empty)