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

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4  

5cimport cpython 

6  

7from libc.stddef cimport size_t 

8from cuda.bindings cimport cydriver 

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

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

11from cuda.core._stream cimport Stream, Stream_accept 

12from cuda.core._resource_handles cimport ( 

13 DevicePtrHandle, 

14 create_fd_handle, 

15 create_mempool_handle_ipc, 

16 deviceptr_import_ipc, 

17 get_last_error, 

18 as_cu, 

19 as_py, 

20) 

21  

22from cuda.core._utils.cuda_utils cimport HANDLE_RETURN 

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

24  

25import multiprocessing (empty)

26import os (empty)

27import platform (empty)

28import uuid (empty)

29import weakref (empty)

30  

31__all__ = [] (empty)

32  

33  

34cdef object registry = weakref.WeakValueDictionary() (empty)

35  

36  

37cdef cydriver.CUmemAllocationHandleType IPC_HANDLE_TYPE = \ 

38 cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR \ (empty)

39 if platform.system() == "Linux" else \ (empty)

40 cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE (empty)

41  

42cdef is_supported(): (empty)

43 return IPC_HANDLE_TYPE != cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE 98 ctx2m 3 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

44  

45  

46cdef class IPCDataForBuffer: 

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

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

49 self._ipc_descriptor = ipc_descriptor 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

50 self._is_mapped = is_mapped 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

51  

52 @property 

53 def ipc_descriptor(self) -> IPCBufferDescriptor: 

54 return self._ipc_descriptor 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

55  

56 @property 

57 def is_mapped(self) -> bool: 

58 return self._is_mapped 22 ctx1stuvwxyzABCDEFGHIJKLqr

59  

60  

61cdef class IPCDataForMR: 

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

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

64 self._alloc_handle = alloc_handle 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

65 self._is_mapped = is_mapped 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

66  

67 @property 

68 def alloc_handle(self) -> IPCAllocationHandle: 

69 return self._alloc_handle 

70  

71 @property 

72 def is_mapped(self) -> bool: 

73 return self._is_mapped 

74  

75 @property 

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

77 return getattr(self._alloc_handle, 'uuid', None) 8 ctx1abUVqrcd

78  

79  

80cdef class IPCBufferDescriptor: 

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

82  

83 Note 

84 ---- 

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

86 Receivers must treat them as untrusted and import only through 

87 :meth:`Buffer.from_ipc_descriptor`. 

88 """ 

89  

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

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

92  

93 @staticmethod (empty)

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

95 cdef IPCBufferDescriptor self = IPCBufferDescriptor.__new__(IPCBufferDescriptor) 57 ctx2a b M N O P e f VbWbs t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Yb(bo )bp

96 self._payload = reserved 57 ctx2a b M N O P e f VbWbs t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Yb(bo )bp

97 self._size = size 57 ctx2a b M N O P e f VbWbs t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Yb(bo )bp

98 return self 57 ctx2a b M N O P e f VbWbs t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Yb(bo )bp

99  

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

101 return IPCBufferDescriptor._init, (self._payload, self._size) 50 ctx2a b M N O P e f s t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d (b)b

102  

103 @property 

104 def size(self) -> int: 

105 return self._size 1 ctx1n

106  

107 cdef const void* payload_ptr(self) noexcept: (empty)

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

109 return <const void*><const char*>(self._payload) 

110  

111  

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

113 if self._h_fd.get() == NULL: 59 ctx21b2ba b M N O P e f U V s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

114 raise RuntimeError("IPCAllocationHandle has been closed") 2 ctx1UV

115 return 0 58 ctx2m 1b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

116  

117  

118cdef class IPCAllocationHandle: 

119 """Shareable handle to an IPC-enabled device memory pool.""" 

120  

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

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

123  

124 @classmethod (empty)

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

126 cdef IPCAllocationHandle self = IPCAllocationHandle.__new__(cls) 101 ctx23 4 a 5 b 6 7 8 9 ! e # f $ U % V ' ( ) *bZb* + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

127 if handle < 0: 101 ctx23 4 a 5 b 6 7 8 9 ! e # f $ U % V ' ( ) *bZb* + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

128 raise ValueError(f"Invalid allocation handle (fd) {handle}: must be non-negative") 1 ctx2*b

129 self._h_fd = create_fd_handle(handle) 100 ctx23 4 a 5 b 6 7 8 9 ! e # f $ U % V ' ( ) Zb* + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

130 self._uuid = uuid 100 ctx23 4 a 5 b 6 7 8 9 ! e # f $ U % V ' ( ) Zb* + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

131 return self 100 ctx23 4 a 5 b 6 7 8 9 ! e # f $ U % V ' ( ) Zb* + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

132  

133 cpdef close(self): (empty)

134 """Close the handle.""" 

135 self._h_fd.reset() 3 ctx2U V Zb

136  

137 @property 

138 def is_closed(self) -> bool: 

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

140 return self._h_fd.get() == NULL 1 ctx2Zb

141  

142 def __int__(self) -> int: 

143 if self._h_fd.get() == NULL: 3 ctx2U V Zb

144 raise ValueError( 1 ctx2Zb

145 f"Cannot convert IPCAllocationHandle to int: the handle (id={id(self)}) is closed." 1 ctx2Zb

146 ) 

147 return as_py(self._h_fd) 2 ctx1UV

148  

149 @property 

150 def handle(self) -> int: 

151 return as_py(self._h_fd) 57 ctx21b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

152  

153 @property 

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

155 return self._uuid 59 ctx21b2ba b M N O P e f U V s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

156  

157  

158def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: (empty)

159 IPCAllocationHandle_check_open(alloc_handle) 57 ctx21b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

160 check_multiprocessing_start_method() 57 ctx21b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

161 df = multiprocessing.reduction.DupFd(alloc_handle.handle) 57 ctx21b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

162 return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid) 57 ctx21b2ba b M N O P e f s t u v 3bQ 4b5bR 6b7bw 8b9bS !b#bT $b%bx 'by z A B C D E F G H g h I J K L q r i j k l c d W

163  

164  

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

166 return cls._init(df.detach(), uuid) 

167  

168  

169multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle) (empty)

170  

171  

172# Buffer IPC Implementation 

173# ------------------------- 

174cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self): (empty)

175 Buffer_check_open(self) 53 ctx2a b M N O P e f s t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Ybo p

176 if not self.memory_resource.is_ipc_enabled: 53 ctx2a b M N O P e f s t u v Q X R Y w Z S 0 T 1 x 2 y z A B C D E F G H g h I J K L q r i j k l c d Xbn Ybo p

177 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2XbYb

178 cdef cydriver.CUmemPoolPtrExportData data 

179 with nogil: 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

180 HANDLE_RETURN( 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

181 cydriver.cuMemPoolExportPointer(&data, as_cu(self._h_ptr)) 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

182 ) 

183 cdef bytes data_b = cpython.PyBytes_FromStringAndSize( 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

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

185 ) 

186 return IPCBufferDescriptor._init(data_b, self.size) 51 ctx1abMNOPefstuvQXRYwZS0T1x2yzABCDEFGHghIJKLqrijklcdnop

187  

188cdef Buffer Buffer_from_ipc_descriptor( (empty)

189 cls, _MemPool mr, IPCBufferDescriptor ipc_descriptor, stream 

190): 

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

192 MP_check_open(mr) 4 ctx2VbWbXbYb

193 if not mr.is_ipc_enabled: 4 ctx2VbWbXbYb

194 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2XbYb

195 cdef size_t payload_size = len(ipc_descriptor._payload) 2 ctx2VbWb

196 cdef size_t expected_size = sizeof(cydriver.CUmemPoolPtrExportData) 2 ctx2VbWb

197 if payload_size < expected_size: 2 ctx2VbWb

198 raise ValueError( 2 ctx2VbWb

199 f"IPC buffer descriptor payload is {payload_size} bytes; " 2 ctx2VbWb

200 f"expected at least {expected_size}" 2 ctx2VbWb

201 ) 

202 cdef Stream s = Stream_accept(stream) 

203 cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( 

204 mr._h_pool, 

205 ipc_descriptor.payload_ptr(), 

206 s._h_stream 

207 ) 

208 if not h_ptr: 

209 HANDLE_RETURN(get_last_error()) 

210 cdef size_t mapped_size = 0 

211 cdef size_t claimed_size = ipc_descriptor.size 

212 with nogil: 

213 HANDLE_RETURN(cydriver.cuPointerGetAttribute( 

214 &mapped_size, 

215 cydriver.CU_POINTER_ATTRIBUTE_RANGE_SIZE, 

216 as_cu(h_ptr))) 

217 if claimed_size > mapped_size: 

218 h_ptr.reset() 

219 raise ValueError( 

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

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

222 ) 

223 return Buffer_from_deviceptr_handle(h_ptr, claimed_size, mr, ipc_descriptor) 

224  

225  

226# _MemPool IPC Implementation 

227# --------------------------- 

228  

229cdef _MemPool MP_from_allocation_handle(cls, alloc_handle): (empty)

230 if isinstance(alloc_handle, IPCAllocationHandle): 2 ctx1UV

231 IPCAllocationHandle_check_open(<IPCAllocationHandle>alloc_handle) 2 ctx1UV

232  

233 # Quick exit for registry hits. 

234 uuid = getattr(alloc_handle, 'uuid', None) # no-cython-lint 

235 mr = registry.get(uuid) 

236 if mr is not None: 

237 if not isinstance(mr, cls): 

238 raise TypeError( 

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

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

241 MP_check_open(<_MemPool>mr) 

242 return mr 

243  

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

245 # necessary. 

246 if isinstance(alloc_handle, int): 

247 fd = os.dup(alloc_handle) 

248 try: 

249 alloc_handle = IPCAllocationHandle._init(fd, None) 

250 except: 

251 os.close(fd) 

252 raise 

253  

254 # Construct a new mempool. 

255 cdef _MemPool self = <_MemPool>(cls.__new__(cls)) 

256 self._mempool_owned = True 

257 cdef int ipc_fd = int(alloc_handle) 

258 self._h_pool = create_mempool_handle_ipc(ipc_fd, IPC_HANDLE_TYPE) 

259 if not self._h_pool: 

260 HANDLE_RETURN(get_last_error()) 

261 raise RuntimeError( 

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

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

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

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

266 ) 

267 self._ipc_data = IPCDataForMR(alloc_handle, True) 

268  

269 # Register it. 

270 if uuid is not None: 

271 registered = self.register(uuid) 

272 assert registered is self 

273  

274 return self 

275  

276  

277cdef _MemPool MP_from_registry(uuid): (empty)

278 cdef _MemPool mr 

279 try: 1 ctx20b

280 mr = registry[uuid] 1 ctx20b

281 MP_check_open(mr) 

282 return mr 

283 except KeyError: 1 ctx20b

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

285  

286  

287cdef _MemPool MP_register(_MemPool self, uuid): (empty)

288 MP_check_open(self) 3 ctx2U V 0b

289 existing = registry.get(uuid) 3 ctx2U V 0b

290 if existing is not None: 3 ctx2U V 0b

291 MP_check_open(<_MemPool>existing) 

292 return existing 

293 if not self.is_ipc_enabled: 3 ctx2U V 0b

294 raise RuntimeError("Memory resource is not IPC-enabled") 1 ctx20b

295 assert self.uuid is None or self.uuid == uuid 2 ctx1UV

296 registry[uuid] = self 2 ctx1UV

297 self._ipc_data._alloc_handle._uuid = uuid 2 ctx1UV

298 return self 2 ctx1UV

299  

300  

301cdef IPCAllocationHandle MP_export_mempool(_MemPool self): (empty)

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

303 MP_check_open(self) 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

304 cdef int fd 

305 with nogil: 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

306 HANDLE_RETURN(cydriver.cuMemPoolExportToShareableHandle( 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

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

308 ) 

309 try: 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

310 return IPCAllocationHandle._init(fd, uuid.uuid4()) 97 ctx23 4 a 5 b 6 7 8 9 ! e # f $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbg h ubvbwbxbybzbi j k l c d AbBbCbDbEbFbGbHbIbJbKbLbMbNbObPbn QbRbSbW TbUbo p

311 except: 

312 os.close(fd) 

313 raise