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

76 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-10 02:27 +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._location cimport cumemlocation_from_id 

9from cuda.core._memory._memory_pool cimport ( 

10 _MemPool, MP_check_open, MP_init_create_pool, MP_raise_release_threshold, 

11) 

12from cuda.core._memory cimport _ipc 

13from cuda.core._memory._ipc cimport IPCAllocationHandle 

14from cuda.core._resource_handles cimport as_cu, get_device_mempool, get_last_error 

15from cuda.core._utils.cuda_utils cimport ( 

16 check_or_create_options, 

17 HANDLE_RETURN, 

18) 

19  

20import cython 

21from dataclasses import dataclass 

22import multiprocessing 

23import platform # no-cython-lint 

24import uuid 

25  

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

27from cuda.core._utils.cuda_utils import check_multiprocessing_start_method 

28  

29from typing import TYPE_CHECKING 

30  

31if TYPE_CHECKING: 

32 from cuda.core._device import Device 

33  

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

35  

36  

37@dataclass 

38cdef class DeviceMemoryResourceOptions: 

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

40  

41 Attributes 

42 ---------- 

43 ipc_enabled : bool, optional 

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

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

46 processes. (Default to False) 

47  

48 max_size : int, optional 

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

50 (Default to 0) 

51 """ 

52 ipc_enabled : bool = False 

53 max_size : int = 0 

54  

55  

56cdef class DeviceMemoryResource(_MemPool): 

57 """ 

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

59  

60 Parameters 

61 ---------- 

62 device_id : Device | int 

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

64  

65 options : DeviceMemoryResourceOptions 

66 Memory resource creation options. 

67  

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

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

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

71 is used. 

72  

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

74 the memory resource. 

75  

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

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

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

79  

80 Notes 

81 ----- 

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

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

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

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

86 These steps can be accomplished in several ways. 

87  

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

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

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

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

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

93 :attr:`DeviceMemoryResource.is_mapped`. 

94  

95 An MR is shared via an allocation handle accessed through the 

96 :attr:`DeviceMemoryResource.allocation_handle` property. The allocation 

97 handle has a platform-specific interpretation; however, memory IPC is 

98 currently only supported for Linux, and in that case allocation handles 

99 are file descriptors. After sending an allocation handle to another 

100 process, it can be used to create an MMR by invoking 

101 :meth:`DeviceMemoryResource.from_allocation_handle`. 

102  

103 Buffers can be shared as serializable descriptors accessed through the 

104 :attr:`Buffer.ipc_descriptor` property. In a receiving process, a shared 

105 buffer is created by invoking :meth:`Buffer.from_ipc_descriptor` with an 

106 MMR and buffer descriptor, where the MMR corresponds to the MR that 

107 created the described buffer. 

108  

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

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

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

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

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

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

115 can themselves be serialized and transferred directly. Serialization embeds 

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

117  

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

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

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

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

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

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

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

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

126  

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

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

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

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

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

132 associated MMR. Unpickling a :class:`Buffer` performs an IPC import from 

133 the embedded descriptor; only unpickle buffers received from trusted peers. 

134  

135 Warning 

136 ------- 

137 IPC descriptors and pickled buffers cross a trust boundary between 

138 cooperating same-host processes. A malicious peer can supply crafted 

139 descriptor fields. Use :meth:`Buffer.from_ipc_descriptor` only with 

140 descriptors from trusted peers, and do not unpickle buffers from 

141 untrusted sources. 

142 """ 

143  

144 def __cinit__(self, *args, **kwargs) -> None: 

145 self._dev_id = cydriver.CU_DEVICE_INVALID 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

146  

147 @cython.annotation_typing(False) 

148 def __init__( 

149 self, 

150 device_id: Device | int, 

151 options: DeviceMemoryResourceOptions | None = None 

152 ) -> None: 

153 _DMR_init(self, device_id, options) 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

154  

155 def __reduce__(self) -> tuple[object, ...]: 

156 MP_check_open(self) 1apbc

157 return DeviceMemoryResource.from_registry, (self.uuid,) 1apbc

158  

159 @staticmethod 

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

161 """ 

162 Obtain a registered mapped memory resource. 

163  

164 Raises 

165 ------ 

166 RuntimeError 

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

168 """ 

169 return <DeviceMemoryResource>(_ipc.MP_from_registry(uuid)) 1dD

170  

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

172 """ 

173 Register a mapped memory resource. 

174  

175 Returns 

176 ------- 

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

178 with the given key, it is returned. 

179 """ 

180 return <DeviceMemoryResource>(_ipc.MP_register(self, uuid)) 2rdd D

181  

182 @classmethod 

183 def from_allocation_handle( 

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

185 ) -> DeviceMemoryResource: 

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

187  

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

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

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

191  

192 Parameters 

193 ---------- 

194 device_id : int | Device 

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

196 resource is created. 

197  

198 alloc_handle : int | IPCAllocationHandle 

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

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

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

202  

203 Returns 

204 ------- 

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

206 """ 

207 cdef DeviceMemoryResource mr = <DeviceMemoryResource>( 1fd

208 _ipc.MP_from_allocation_handle(cls, alloc_handle)) 2rdf d

209 from .._device import Device 1fd

210 mr._dev_id = Device(device_id).device_id 1fd

211 return mr 1fd

212  

213 @property 

214 def allocation_handle(self) -> IPCAllocationHandle: 

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

216  

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

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

219 """ 

220 MP_check_open(self) 2~ca p adbde q u rdf d cdddsdtdedfdudgdhdvdidjdkdwdldmdndg h odpdqdi j k l b c z x m y n o

221 if not self.is_ipc_enabled: 2~ca adbde u rdf d cdddsdtdedfdudgdhdvdidjdkdwdldmdndg h odpdqdi j k l b c x m y n o

222 raise RuntimeError("Memory resource is not IPC-enabled") 1x

223 return self._ipc_data._alloc_handle 2~ca adbde u rdf d cdddsdtdedfdudgdhdvdidjdkdwdldmdndg h odpdqdi j k l b c m y n o

224  

225 @property 

226 def device_id(self) -> int: 

227 """The associated device ordinal.""" 

228 return self._dev_id 26 7 8 9 ! # $ % IdJd( ) Kd* Ld+ MdNdOdPd~ca p adbde q f cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c mbnbA B v pbm n o MbNbG xdH ydI J K L M N zdO P Q AdR S BdT U V W CdX Y Z Dd0 1 2 3 EdFdGdHd4 5 C

229  

230 @property 

231 def peer_accessible_by(self) -> PeerAccessibleBySetProxy: 

232 """ 

233 Get or set the devices that can access allocations from this memory 

234 pool. Access can be modified at any time and affects all allocations 

235 from this memory pool. 

236  

237 Returns a set-like proxy of :obj:`~_device.Device` objects that manages 

238 peer access. Inputs are accepted as either :obj:`~_device.Device` 

239 objects or device-ordinal :class:`int` values. 

240  

241 Examples 

242 -------- 

243 >>> dmr = DeviceMemoryResource(0) 

244 >>> dmr.peer_accessible_by = {1} # grant access to device 1 

245 >>> assert 1 in dmr.peer_accessible_by 

246 >>> dmr.peer_accessible_by.add(2) # update access to include device 2 

247 >>> dmr.peer_accessible_by = [] # revoke peer access 

248 """ 

249 MP_check_open(self) 1zEw

250 return PeerAccessibleBySetProxy(self) 1zEw

251  

252 @peer_accessible_by.setter 

253 def peer_accessible_by(self, devices) -> None: 

254 replace_peer_accessible_by(self, devices) 1w

255  

256 @property 

257 def is_device_accessible(self) -> bool: 

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

259 return True 2A B v F G xdH ydI J K L M N zdO P Q AdR S BdT U V W CdX Y Z Dd0 1 2 3 EdFdGdHd4 5 C

260  

261 @property 

262 def is_host_accessible(self) -> bool: 

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

264 return False 1ABvC

265  

266  

267cdef inline _DMR_init(DeviceMemoryResource self, device_id, options): 

268 from .._device import Device 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

269 cdef int dev_id = Device(device_id).device_id 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

270 cdef DeviceMemoryResourceOptions opts = check_or_create_options( 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

271 DeviceMemoryResourceOptions, options, "DeviceMemoryResource options", 

272 keep_none=True 

273 ) 

274 cdef bint ipc_enabled = False 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

275 cdef size_t max_size = 0 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

276  

277 self._dev_id = dev_id 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

278  

279 if opts is not None: 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

280 ipc_enabled = opts.ipc_enabled 2' , a - p . / e : q u ; f = d ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c v obF t w qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbJbm y n o KbLb

281 if ipc_enabled and not _ipc.is_supported(): 2' , a - p . / e : q u ; f = d ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c v obF t w qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbJbm y n o KbLb

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

283 max_size = opts.max_size 2' , a - p . / e : q u ; f = d ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c v obF t w qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbJbm y n o KbLb

284  

285 if opts is None: 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b' $b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbc, a - p . / e : q u ; f = d D ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c mbnbccdcecfcgcz A B E v obF t r w pbqbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbhcJbicjckcm y n o lcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#cKbLb$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

286 self._h_pool = get_device_mempool(dev_id) 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b$b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcD mbnbccdcecfcgcz A B E t r pbhcicjckclcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

287 if not self._h_pool: 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b$b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcD mbnbccdcecfcgcz A B E t r pbhcicjckclcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

288 HANDLE_RETURN(get_last_error()) 

289 raise RuntimeError( 

290 f"Failed to initialize DeviceMemoryResource for device {dev_id}: " 

291 "cuda-core returned an empty memory pool handle without recording a CUDA error. " 

292 "This is an internal cuda-core error; please report it with your CUDA driver, " 

293 "CUDA Toolkit, and cuda-python versions." 

294 ) 

295 self._mempool_owned = False 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b$b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcD mbnbccdcecfcgcz A B E t r pbhcicjckclcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

296 MP_raise_release_threshold(self) 2ObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b6 7 8 9 ! # $ % !b#b$b%b'b( ) (b)b*b+b,b-b.b* /b:b+ ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcD mbnbccdcecfcgcz A B E t r pbhcicjckclcmcncocpcqcrcsctcucvcwcxcyczcAcBcCcDcEcFcGcHcIcJcKcLcMcNcOcPcQcRcScTcUcVcWcXcYcZc0c1c2c3c4c5c6c7c8c9c!c#c$c%c'c(c)c*c+c,c-c.c/c:c;c=c?c@c[c]c^c_c`c{cMbNbG H I J K L M N O P Q R S T U V W X Y Z 0 1 2 3 |c4 5 }cC

297 else: 

298 MP_init_create_pool( 2' , a - p . / e : q u ; f = d ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c v obF t w qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbJbm y n o KbLb

299 self, 

300 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 

301 dev_id, 

302 cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED, 

303 ipc_enabled, 

304 max_size, 2' , a - p . / e : q u ; f = d ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibg h jbkblbi j k l b c v obF t w qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbGbHbx IbJbm y n o KbLb

305 ) 

306  

307  

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

309cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id): 

310 """ 

311 Probes peer access from the given device using cuMemPoolGetAccess. 

312  

313 Parameters 

314 ---------- 

315 device_id : int or Device 

316 The device to query access for. 

317  

318 Returns 

319 ------- 

320 str 

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

322 """ 

323 from .._device import Device 1r

324  

325 cdef int dev_id = Device(device_id).device_id 1r

326 cdef cydriver.CUmemAccess_flags flags 

327 cdef cydriver.CUmemLocation location = cumemlocation_from_id( 1r

328 cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id) 

329  

330 with nogil: 1r

331 HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location)) 1r

332  

333 if flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE: 1r

334 return "rw" 1r

335 elif flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ: 

336 return "r" 

337 else: 

338 return "" 

339  

340  

341def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]: 

342 check_multiprocessing_start_method() 2~ca p adbde q cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c m n o

343 from .._device import Device 2~ca p adbde q cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c m n o

344 device = Device(mr.device_id) 2~ca p adbde q cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c m n o

345 alloc_handle = mr.allocation_handle 2~ca p adbde q cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c m n o

346 return mr.from_allocation_handle, (device, alloc_handle) 2~ca adbde cdddedfdgdhdidjdkdldmdndg h odpdqdi j k l b c m n o

347  

348  

349multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)