Coverage for cuda/core/_memory/_ipc.pyx: 73.76%

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  

7cimport cpython 

8  

9from libc.stddef cimport size_t 

10from cuda.bindings cimport cydriver 

11from cuda.core._memory._buffer cimport Buffer, Buffer_check_open, Buffer_from_deviceptr_handle 

12from cuda.core._memory._memory_pool cimport _MemPool, MP_check_open 

13from cuda.core._stream cimport Stream, Stream_accept 

14from cuda.core._rt cimport ( 

15 DevicePtrHandle, 

16 create_fd_handle, 

17 create_mempool_handle_ipc, 

18 deviceptr_import_ipc, 

19 get_last_error, 

20 as_cu, 

21 as_py, 

22) 

23  

24from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

25from cuda.core._utils.cuda_utils import check_multiprocessing_start_method 

26  

27import multiprocessing 

28import operator 

29import os 

30import platform 

31import uuid 

32import weakref 

33from typing import TYPE_CHECKING, Iterable 

34  

35if TYPE_CHECKING: 

36 from cuda.core.typing import VirtualMemoryHandleType 

37  

38__all__ = [] 

39  

40  

41cdef object registry = weakref.WeakValueDictionary() 

42  

43  

44cdef cydriver.CUmemAllocationHandleType IPC_HANDLE_TYPE = \ 

45 cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR \ 

46 if platform.system() == "Linux" else \ 

47 cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE 

48  

49cdef is_supported(): 

50 return IPC_HANDLE_TYPE != cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

51  

52  

53cdef class IPCDataForBuffer: 

54 """Data members related to sharing memory buffers via IPC.""" 

55 def __cinit__(self, IPCBufferDescriptor ipc_descriptor, bint is_mapped) -> None: 

56 self._ipc_descriptor = ipc_descriptor 63 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v 1 | bb} c d ~ cbabm n o p g h r s t

57 self._is_mapped = is_mapped 63 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v 1 | bb} c d ~ cbabm n o p g h r s t

58  

59 @property 

60 def ipc_descriptor(self) -> IPCBufferDescriptor: 

61 return self._ipc_descriptor 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

62  

63 @property 

64 def is_mapped(self) -> bool: 

65 return self._is_mapped 3 ctx2bb~ ab

66  

67  

68cdef class IPCDataForMR: 

69 """Data members related to sharing memory resources via IPC.""" 

70 def __cinit__(self, IPCAllocationHandle alloc_handle, bint is_mapped) -> None: 

71 self._alloc_handle = alloc_handle 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

72 self._is_mapped = is_mapped 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

73  

74 @property 

75 def alloc_handle(self) -> IPCAllocationHandle: 

76 return self._alloc_handle 

77  

78 @property 

79 def is_mapped(self) -> bool: 

80 return self._is_mapped 

81  

82 @property 

83 def uuid(self) -> uuid.UUID | None: 

84 return getattr(self._alloc_handle, 'uuid', None) 10 ctx1ef45abuvgh

85  

86  

87cdef class IPCBufferDescriptor: 

88 """Serializable object describing a buffer that can be shared between processes. 

89  

90 Note 

91 ---- 

92 The payload and ``size`` fields are controlled by the exporting peer. 

93 Receivers must treat them as untrusted and import only through 

94 :meth:`Buffer.from_ipc_descriptor`. 

95 """ 

96  

97 def __init__(self, *arg, **kwargs) -> None: 

98 raise RuntimeError("IPCBufferDescriptor objects cannot be instantiated directly. Please use MemoryResource APIs.") 1 ctx2Ec

99  

100 @staticmethod 

101 def _init(reserved: bytes, size: int) -> IPCBufferDescriptor: 

102 cdef IPCBufferDescriptor self = IPCBufferDescriptor.__new__(IPCBufferDescriptor) 62 ctx2e f x y z A i j a mbnblcB C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcBcs Cct

103 self._payload = reserved 62 ctx2e f x y z A i j a mbnblcB C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcBcs Cct

104 self._size = size 62 ctx2e f x y z A i j a mbnblcB C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcBcs Cct

105 return self 62 ctx2e f x y z A i j a mbnblcB C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcBcs Cct

106  

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

108 return IPCBufferDescriptor._init, (self._payload, self._size) 52 ctx2e f x y z A i j B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h BcCc

109  

110 @property 

111 def size(self) -> int: 

112 return self._size 10 ctx23 a 1 | bbw c dbd r

113  

114 cdef const void* payload_ptr(self) noexcept: 

115 """Return the payload as a const void* for C API calls.""" 

116 return <const void*><const char*>(self._payload) 1 ctx1a

117  

118  

119cdef inline int IPCAllocationHandle_check_open(IPCAllocationHandle self) except -1: 

120 if self._h_fd.get() == NULL: 93 ctx2ncoce f x y 6 z A 3 i j 7 2 4 5 b B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w c m n ] ^ o p _ ` g h {

121 raise RuntimeError("IPCAllocationHandle has been closed") 3 ctx145c

122 return 0 91 ctx2ncoce f x y 6 z A 3 i j 7 2 b B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w c m n ] ^ o p _ ` g h {

123  

124  

125cdef class IPCAllocationHandle: 

126 """Shareable OS handle to an IPC-enabled memory pool or to a physical allocation 

127 of a :class:`VirtualMemoryResource` buffer.""" 

128  

129 def __init__(self, *arg, **kwargs) -> None: 

130 raise RuntimeError("IPCAllocationHandle objects cannot be instantiated directly. Please use MemoryResource APIs.") 1 ctx2Ec

131  

132 @classmethod 

133 def _init(cls, handle: int, uuid: uuid.UUID | None) -> IPCAllocationHandle: # no-cython-lint 

134 cdef IPCAllocationHandle self = IPCAllocationHandle.__new__(cls) 147 ctx2obpbe qbf rbsbtb6 ubvb3 i wbj xb7 2 4 yb5 zba AbBbDckcb CbDb8 EbFb9 GbHb! IbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b# $ % ' ( ) 8b9b* !b#b+ $b%b, 'b(b- )b*bk l . / +b,b: -b.b; /b:b= ? 1 | lb@ [ } w c dbd ~ cbebq abm n ] ^ o p _ ` g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

135 if handle < 0: 147 ctx2obpbe qbf rbsbtb6 ubvb3 i wbj xb7 2 4 yb5 zba AbBbDckcb CbDb8 EbFb9 GbHb! IbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b# $ % ' ( ) 8b9b* !b#b+ $b%b, 'b(b- )b*bk l . / +b,b: -b.b; /b:b= ? 1 | lb@ [ } w c dbd ~ cbebq abm n ] ^ o p _ ` g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

136 raise ValueError(f"Invalid allocation handle (fd) {handle}: must be non-negative") 1 ctx2Dc

137 self._h_fd = create_fd_handle(handle) 146 ctx2obpbe qbf rbsbtb6 ubvb3 i wbj xb7 2 4 yb5 zba AbBbkcb CbDb8 EbFb9 GbHb! IbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b# $ % ' ( ) 8b9b* !b#b+ $b%b, 'b(b- )b*bk l . / +b,b: -b.b; /b:b= ? 1 | lb@ [ } w c dbd ~ cbebq abm n ] ^ o p _ ` g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

138 self._uuid = uuid 146 ctx2obpbe qbf rbsbtb6 ubvb3 i wbj xb7 2 4 yb5 zba AbBbkcb CbDb8 EbFb9 GbHb! IbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b# $ % ' ( ) 8b9b* !b#b+ $b%b, 'b(b- )b*bk l . / +b,b: -b.b; /b:b= ? 1 | lb@ [ } w c dbd ~ cbebq abm n ] ^ o p _ ` g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

139 return self 146 ctx2obpbe qbf rbsbtb6 ubvb3 i wbj xb7 2 4 yb5 zba AbBbkcb CbDb8 EbFb9 GbHb! IbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b# $ % ' ( ) 8b9b* !b#b+ $b%b, 'b(b- )b*bk l . / +b,b: -b.b; /b:b= ? 1 | lb@ [ } w c dbd ~ cbebq abm n ] ^ o p _ ` g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

140  

141 cpdef close(self): 

142 """Close the handle.""" 

143 self._h_fd.reset() 6 ctx22 4 5 kc1 c

144  

145 @property 

146 def is_closed(self) -> bool: 

147 """Whether this allocation handle has been closed.""" 

148 return self._h_fd.get() == NULL 1 ctx2kc

149  

150 def __int__(self) -> int: 

151 if self._h_fd.get() == NULL: 17 ctx24 5 a kcb 1 | lb} w c dbd ~ cbebab

152 raise ValueError( 2 ctx2kc1

153 f"Cannot convert IPCAllocationHandle to int: the handle (id={id(self)}) is closed." 2 ctx2kc1

154 ) 

155 return as_py(self._h_fd) 16 ctx24 5 a b 1 | lb} w c dbd ~ cbebab

156  

157 @property 

158 def handle(self) -> int: 

159 return as_py(self._h_fd) 90 ctx2ncoce f x y 6 z A 3 i j 7 2 a B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

160  

161 @property 

162 def uuid(self) -> uuid.UUID: 

163 return self._uuid 93 ctx2ncoce f x y 6 z A 3 i j 7 2 4 5 a b B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

164  

165  

166def _reduce_allocation_handle(IPCAllocationHandle alloc_handle not None: IPCAllocationHandle) -> tuple[object, ...]: 

167 IPCAllocationHandle_check_open(alloc_handle) 89 ctx2ncoce f x y 6 z A 3 i j 7 2 B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

168 check_multiprocessing_start_method() 89 ctx2ncoce f x y 6 z A 3 i j 7 2 B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

169 df = multiprocessing.reduction.DupFd(alloc_handle.handle) 89 ctx2ncoce f x y 6 z A 3 i j 7 2 B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

170 return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid) 89 ctx2ncoce f x y 6 z A 3 i j 7 2 B C 8 D E 9 F G ! pcH qcrcI sctcJ ucvcK wcxcL yczcM Ac# $ % ' ( ) N O * P Q + R S , T U - V W k l . / X Y : Z 0 ; u v = ? @ [ w m n ] ^ o p _ ` g h {

171  

172  

173def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None) -> IPCAllocationHandle: # no-cython-lint 

174 return cls._init(df.detach(), uuid) 1 ctx12

175  

176  

177multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle) 

178  

179  

180cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): 

181 """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. 

182  

183 The descriptor holds one exported POSIX file descriptor per physical 

184 allocation that backs the buffer, in address order (``handles``, each an 

185 :class:`IPCAllocationHandle` that the descriptor owns and closes when it 

186 is released), with each allocation's size (``sizes``) and the buffer's 

187 size (``size``, which may be smaller than their sum). ``handle_type`` is 

188 ``"posix_fd"``, the only handle type cuda.core shares today; the rules 

189 below describe file descriptors. 

190  

191 A descriptor pins the physical memory while it exists, in every process 

192 that holds one, and nothing else does: the exporting buffer keeps no 

193 descriptor, and an imported buffer does not keep the one it came from. A 

194 descriptor parked in a ``multiprocessing`` ``Queue`` or sent to a ``Pool`` 

195 pins the memory in the sender until the receiver has unpickled it; a 

196 buffer passed as a ``Process`` argument pins it until the ``Process`` 

197 object is released, because the child receives the file descriptors when 

198 it is spawned. 

199  

200 Two transports exist. ``multiprocessing`` (a ``Queue``, a ``Pipe``, 

201 ``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer, 

202 and duplicates the file descriptors into the receiving process; plain 

203 ``pickle`` cannot carry them and raises. A process with its own file 

204 descriptor passing sends ``fds``, ``sizes``, and ``size`` itself, and the 

205 receiver rebuilds the descriptor with :meth:`from_fds`. 

206  

207 Note 

208 ---- 

209 The sizes are controlled by the exporting peer. Receivers must treat them 

210 as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`, 

211 which checks what it can and fails instead of mapping a size the driver 

212 rejects. 

213 """ 

214  

215 @staticmethod 

216 def _from_exports( 

217 handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple, size: int 

218 ) -> VirtualMemoryIPCBufferDescriptor: 

219 cdef VirtualMemoryIPCBufferDescriptor self = VirtualMemoryIPCBufferDescriptor.__new__( 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

220 VirtualMemoryIPCBufferDescriptor) 

221 self._payload = b"" 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

222 self._size = size 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

223 self._handle_type = handle_type 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

224 self._chunk_sizes = tuple(chunk_sizes) 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

225 self._handles = tuple(handles) 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

226 return self 42 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? 1 | bblb@ [ } w c dbd ~ cbebq ab] ^ _ `

227  

228 @classmethod 

229 def from_fds( 

230 cls, 

231 fds: Iterable[int | IPCAllocationHandle], 

232 sizes: Iterable[int], 

233 *, 

234 size: int | None = None, 

235 ) -> VirtualMemoryIPCBufferDescriptor: 

236 """Build a descriptor from POSIX file descriptors received outside ``multiprocessing``. 

237  

238 The exporter sends ``fds``, ``sizes``, and ``size`` of its descriptor 

239 through a transport of its own (a Unix socket with ``SCM_RIGHTS``, for 

240 example); the receiver rebuilds the descriptor here and imports it 

241 with :meth:`Buffer.from_ipc_descriptor`. The result's ``handle_type`` 

242 is ``"posix_fd"``, the only handle type a file descriptor can carry. 

243  

244 Parameters 

245 ---------- 

246 fds : Iterable[int | IPCAllocationHandle] 

247 One handle per physical allocation, in address order, as the 

248 exporter's ``fds`` listed them. An ``int`` is duplicated: the 

249 caller keeps ownership of the original and closes it. An 

250 :class:`IPCAllocationHandle` is shared and closes with its last 

251 holder. 

252 sizes : Iterable[int] 

253 The exporter's ``sizes``, one per handle. 

254 size : int, optional 

255 The exporter's ``size``, at most the sum of ``sizes``; the sum by 

256 default. 

257  

258 Raises 

259 ------ 

260 ValueError 

261 If the counts differ, a size is not positive, or ``size`` exceeds 

262 the sum of ``sizes``. 

263 RuntimeError 

264 If an :class:`IPCAllocationHandle` is closed. 

265 OSError 

266 If an ``int`` is not an open file descriptor. 

267 """ 

268 cdef int driver_type = <int>cydriver.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR 2 ctx1cd

269 cdef tuple sizes_t = tuple(operator.index(n) for n in sizes) 2 ctx1cd

270 cdef list given = list(fds) 2 ctx1cd

271 if len(given) != len(sizes_t): 2 ctx1cd

272 raise ValueError(f"{len(given)} handles for {len(sizes_t)} sizes; one size per handle is required") 1 ctx1c

273 for n in sizes_t: 2 ctx1cd

274 if n <= 0: 2 ctx1cd

275 raise ValueError(f"allocation sizes must be positive, got {n}") 1 ctx1c

276 total = sum(sizes_t) 2 ctx1cd

277 size = total if size is None else operator.index(size) 2 ctx1cd

278 if size < 0 or size > total: 2 ctx1cd

279 raise ValueError(f"size {size} is outside 0..{total}, the sum of the allocation sizes") 1 ctx1c

280 cdef list handles = [] 2 ctx1cd

281 cdef int dup 

282 for item in given: 2 ctx1cd

283 if isinstance(item, IPCAllocationHandle): 2 ctx1cd

284 IPCAllocationHandle_check_open(<IPCAllocationHandle>item) 1 ctx1c

285 handles.append(item) 1 ctx1c

286 continue 1 ctx1c

287 dup = os.dup(operator.index(item)) 2 ctx1cd

288 try: 1 ctx1d

289 handle = IPCAllocationHandle._init(dup, None) 1 ctx1d

290 except: # noqa: E722 rollback-then-raise: the duplicate is not owned yet 

291 os.close(dup) 

292 raise 

293 # The handle owns the duplicate now; if append fails, the handle closes it once. 

294 handles.append(handle) 1 ctx1d

295 return VirtualMemoryIPCBufferDescriptor._from_exports(driver_type, sizes_t, tuple(handles), size) 2 ctx1cd

296  

297 @property 

298 def handle_type(self) -> VirtualMemoryHandleType: 

299 """The handle type the exporter used; the importing resource must be configured with the same.""" 

300 from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions 2 ctx1cd

301  

302 for spec, value in VirtualMemoryResourceOptions._handle_types.items(): 2 ctx1cd

303 if spec is not None and int(value) == self._handle_type: 2 ctx1cd

304 return spec 2 ctx1cd

305 raise ValueError(f"the descriptor carries an unknown handle type ({self._handle_type})") 

306  

307 @property 

308 def sizes(self) -> tuple[int, ...]: 

309 """Size of each physical allocation, in address order; the sum is at least ``size``.""" 

310 return self._chunk_sizes 5 ctx2bblbc dbd

311  

312 @property 

313 def handles(self) -> tuple[IPCAllocationHandle, ...]: 

314 """The exported handles, one :class:`IPCAllocationHandle` per allocation, in address order. 

315  

316 The descriptor owns them and closes them when it is released. 

317 """ 

318 return self._handles 3 ctx2bbc d

319  

320 @property 

321 def fds(self) -> tuple[int, ...]: 

322 """The exported file descriptors as integers, for a transport of your own. 

323  

324 They belong to this descriptor: send or duplicate them while it is 

325 alive, because each one is closed when the descriptor is released. 

326 Keep the descriptor in a variable while you use them; the numbers 

327 from ``buf.ipc_descriptor.fds`` alone may already be closed. 

328 """ 

329 return tuple(int(handle) for handle in self._handles) 7 ctx2bblb} c dbd eb

330  

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

332 # multiprocessing duplicates the IPCAllocationHandle fds into the receiver; plain pickle raises. 

333 return VirtualMemoryIPCBufferDescriptor._from_exports, ( 31 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? bb@ [ w eb] ^ _ `

334 self._handle_type, self._chunk_sizes, self._handles, self._size) 31 ctx26 3 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? bb@ [ w eb] ^ _ `

335  

336  

337# Buffer IPC Implementation 

338# ------------------------- 

339cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self): 

340 Buffer_check_open(self) 57 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcs t

341 if not self.memory_resource.is_ipc_enabled: 57 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h icr jcs t

342 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2icjc

343 if not isinstance(self.memory_resource, _MemPool): 55 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v q m n o p g h r s t

344 # VirtualMemoryBuffer overrides ipc_descriptor; a plain Buffer wrapped 

345 # with Buffer.from_handle(mr=<VirtualMemoryResource>) lands here. 

346 raise TypeError( 1 ctx1q

347 f"a Buffer that {type(self.memory_resource).__name__} did not allocate cannot be shared") 1 ctx1q

348 cdef cydriver.CUmemPoolPtrExportData data 

349 with nogil: 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

350 HANDLE_RETURN( 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

351 cydriver.cuMemPoolExportPointer(&data, as_cu(self._h_ptr)) 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

352 ) 

353 cdef bytes data_b = cpython.PyBytes_FromStringAndSize( 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

354 <char*>(data.reserved), sizeof(data.reserved) 

355 ) 

356 return IPCBufferDescriptor._init(data_b, self.size) 54 ctx2e f x y z A i j a B C D E F G H fbI gbJ hbK ibL jbM kbN O P Q R S T U V W k l X Y Z 0 u v m n o p g h r s t

357  

358cdef Buffer Buffer_from_ipc_descriptor( 

359 cls, object mr, IPCBufferDescriptor ipc_descriptor, stream 

360): 

361 """Import a buffer that was exported from another process.""" 

362 if isinstance(ipc_descriptor, VirtualMemoryIPCBufferDescriptor): 16 ctx2a mbnblc1 | bb} c d ~ cbq abicjc

363 # Imported here rather than cimported: _virtual_memory_resource 

364 # cimports this module. 

365 from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource 10 ctx21 | bb} c d ~ cbq ab

366 if not isinstance(mr, VirtualMemoryResource): 10 ctx21 | bb} c d ~ cbq ab

367 raise TypeError( 1 ctx1q

368 "this descriptor was exported from a VirtualMemoryResource buffer and must be " 1 ctx1q

369 f"imported with a VirtualMemoryResource, not {type(mr).__name__}") 1 ctx1q

370 return mr._import_ipc_buffer(ipc_descriptor, stream) 10 ctx21 | bb} c d ~ cbq ab

371 if not isinstance(mr, _MemPool): 7 ctx2a mbnblcq icjc

372 raise TypeError( 2 ctx2lcq

373 "this descriptor was exported from a memory pool buffer and must be imported with a " 2 ctx2lcq

374 f"DeviceMemoryResource or PinnedMemoryResource, not {type(mr).__name__}") 2 ctx2lcq

375 cdef _MemPool pool = <_MemPool>mr 5 ctx2a mbnbicjc

376 MP_check_open(pool) 5 ctx2a mbnbicjc

377 if not mr.is_ipc_enabled: 5 ctx2a mbnbicjc

378 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2icjc

379 cdef size_t payload_size = len(ipc_descriptor._payload) 3 ctx2a mbnb

380 cdef size_t expected_size = sizeof(cydriver.CUmemPoolPtrExportData) 3 ctx2a mbnb

381 if payload_size < expected_size: 3 ctx2a mbnb

382 raise ValueError( 2 ctx2mbnb

383 f"IPC buffer descriptor payload is {payload_size} bytes; " 2 ctx2mbnb

384 f"expected at least {expected_size}" 2 ctx2mbnb

385 ) 

386 cdef Stream s = Stream_accept(stream) 1 ctx1a

387 cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( 1 ctx1a

388 pool._h_pool, 

389 ipc_descriptor.payload_ptr(), 

390 s._h_stream 

391 ) 

392 if not h_ptr: 1 ctx1a

393 HANDLE_RETURN(get_last_error()) 

394 cdef size_t mapped_size = 0 1 ctx1a

395 cdef size_t claimed_size = ipc_descriptor.size 1 ctx1a

396 with nogil: 1 ctx1a

397 HANDLE_RETURN(cydriver.cuPointerGetAttribute( 1 ctx1a

398 &mapped_size, 

399 cydriver.CU_POINTER_ATTRIBUTE_RANGE_SIZE, 

400 as_cu(h_ptr))) 

401 if claimed_size > mapped_size: 1 ctx1a

402 h_ptr.reset() 

403 raise ValueError( 

404 f"IPC buffer descriptor size ({claimed_size}) exceeds " 

405 f"mapped allocation extent ({mapped_size} bytes)" 

406 ) 

407 return Buffer_from_deviceptr_handle(h_ptr, claimed_size, pool, ipc_descriptor) 1 ctx1a

408  

409  

410# _MemPool IPC Implementation 

411# --------------------------- 

412  

413cdef _MemPool MP_from_allocation_handle(cls, alloc_handle): 

414 if isinstance(alloc_handle, IPCAllocationHandle): 4 ctx145ab

415 IPCAllocationHandle_check_open(<IPCAllocationHandle>alloc_handle) 3 ctx145b

416  

417 # Quick exit for registry hits. 

418 uuid = getattr(alloc_handle, 'uuid', None) # no-cython-lint 2 ctx1ab

419 mr = registry.get(uuid) 2 ctx1ab

420 if mr is not None: 2 ctx1ab

421 if not isinstance(mr, cls): 1 ctx1b

422 raise TypeError( 

423 f"Registry contains a {type(mr).__name__} for uuid " 

424 f"{uuid}, but {cls.__name__} was requested") 

425 MP_check_open(<_MemPool>mr) 1 ctx1b

426 return mr 1 ctx1b

427  

428 # Ensure we have an allocation handle. Duplicate the file descriptor, if 

429 # necessary. 

430 if isinstance(alloc_handle, int): 2 ctx1ab

431 fd = os.dup(alloc_handle) 1 ctx1a

432 try: 1 ctx1a

433 alloc_handle = IPCAllocationHandle._init(fd, None) 1 ctx1a

434 except: 

435 os.close(fd) 

436 raise 

437  

438 # Construct a new mempool. 

439 cdef _MemPool self = <_MemPool>(cls.__new__(cls)) 2 ctx1ab

440 self._mempool_owned = True 2 ctx1ab

441 cdef int ipc_fd = int(alloc_handle) 2 ctx1ab

442 self._h_pool = create_mempool_handle_ipc(ipc_fd, IPC_HANDLE_TYPE) 2 ctx1ab

443 if not self._h_pool: 2 ctx1ab

444 HANDLE_RETURN(get_last_error()) 

445 raise RuntimeError( 

446 f"Failed to import {cls.__name__} from an allocation handle: " 

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

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

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

450 ) 

451 self._ipc_data = IPCDataForMR(alloc_handle, True) 2 ctx1ab

452  

453 # Register it. 

454 if uuid is not None: 2 ctx1ab

455 registered = self.register(uuid) 1 ctx1b

456 assert registered is self 1 ctx1b

457  

458 return self 2 ctx1ab

459  

460  

461cdef _MemPool MP_from_registry(uuid): 

462 cdef _MemPool mr 

463 try: 2 ctx2b mc

464 mr = registry[uuid] 2 ctx2b mc

465 MP_check_open(mr) 1 ctx1b

466 return mr 1 ctx1b

467 except KeyError: 1 ctx2mc

468 raise RuntimeError(f"Memory resource {uuid} was not found") from None 1 ctx2mc

469  

470  

471cdef _MemPool MP_register(_MemPool self, uuid): 

472 MP_check_open(self) 4 ctx24 5 b mc

473 existing = registry.get(uuid) 4 ctx24 5 b mc

474 if existing is not None: 4 ctx24 5 b mc

475 MP_check_open(<_MemPool>existing) 1 ctx1b

476 return existing 1 ctx1b

477 if not self.is_ipc_enabled: 4 ctx24 5 b mc

478 raise RuntimeError("Memory resource is not IPC-enabled") 1 ctx2mc

479 assert self.uuid is None or self.uuid == uuid 3 ctx145b

480 registry[uuid] = self 3 ctx145b

481 self._ipc_data._alloc_handle._uuid = uuid 3 ctx145b

482 return self 3 ctx145b

483  

484  

485cdef IPCAllocationHandle MP_export_mempool(_MemPool self): 

486 # Note: This is Linux only (int for file descriptor) 

487 MP_check_open(self) 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

488 cdef int fd 

489 with nogil: 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

490 HANDLE_RETURN(cydriver.cuMemPoolExportToShareableHandle( 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

491 &fd, as_cu(self._h_pool), IPC_HANDLE_TYPE, 0) 

492 ) 

493 try: 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

494 return IPCAllocationHandle._init(fd, uuid.uuid4()) 102 ctx2obpbe qbf rbsbtbubvbi wbj xb2 ybzba AbBbb CbDbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbTbUbVbWbXbYbZb0b1b2b3b4b5b6b7b8b9b!b#b$b%b'b(b)b*bk l +b,b-b.b/b:bm n o p g h ;b=b?b@b[b]b^b_b`b{b|b}b~bacbcccr dcecfc{ gchcs t

495 except: 

496 os.close(fd) 

497 raise