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

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 101 ctx2o 6 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

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 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

50 self._is_mapped = is_mapped 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

51  

52 @property 

53 def ipc_descriptor(self) -> IPCBufferDescriptor: 

54 return self._ipc_descriptor 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

55  

56 @property 

57 def is_mapped(self) -> bool: 

58 return self._is_mapped 22 ctx1uvwxyzABCDEFGHIJKLMNst

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 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

65 self._is_mapped = is_mapped 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

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

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.") 1 ctx2.b

92  

93 @staticmethod (empty)

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

95 cdef IPCBufferDescriptor self = IPCBufferDescriptor.__new__(IPCBufferDescriptor) 58 ctx2c d P Q R S g h a YbZbu v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1b+bq ,br

96 self._payload = reserved 58 ctx2c d P Q R S g h a YbZbu v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1b+bq ,br

97 self._size = size 58 ctx2c d P Q R S g h a YbZbu v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1b+bq ,br

98 return self 58 ctx2c d P Q R S g h a YbZbu v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1b+bq ,br

99  

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

101 return IPCBufferDescriptor._init, (self._payload, self._size) 50 ctx2c d P Q R S g h u v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f +b,b

102  

103 @property 

104 def size(self) -> int: 

105 return self._size 2 ctx1ap

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

110  

111  

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

113 if self._h_fd.get() == NULL: 61 ctx24b5bc d P Q R S g h O X Y b u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

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

115 return 0 60 ctx2o 4b5bc d P Q R S g h O b u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

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.") 1 ctx2.b

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) 104 ctx26 7 c 8 d 9 ! # $ % g ' h ( O X ) Y * a + , -b2bb - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

127 if handle < 0: 104 ctx26 7 c 8 d 9 ! # $ % g ' h ( O X ) Y * a + , -b2bb - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

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

129 self._h_fd = create_fd_handle(handle) 103 ctx26 7 c 8 d 9 ! # $ % g ' h ( O X ) Y * a + , 2bb - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

130 self._uuid = uuid 103 ctx26 7 c 8 d 9 ! # $ % g ' h ( O X ) Y * a + , 2bb - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

131 return self 103 ctx26 7 c 8 d 9 ! # $ % g ' h ( O X ) Y * a + , 2bb - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

132  

133 cpdef close(self): (empty)

134 """Close the handle.""" 

135 self._h_fd.reset() 4 ctx2O X Y 2b

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 ctx22b

141  

142 def __int__(self) -> int: 

143 if self._h_fd.get() == NULL: 5 ctx2X Y a 2bb

144 raise ValueError( 1 ctx22b

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

146 ) 

147 return as_py(self._h_fd) 4 ctx1XYab

148  

149 @property 

150 def handle(self) -> int: 

151 return as_py(self._h_fd) 59 ctx24b5bc d P Q R S g h O a u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

152  

153 @property 

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

155 return self._uuid 62 ctx24b5bc d P Q R S g h O X Y a b u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

156  

157  

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

159 IPCAllocationHandle_check_open(alloc_handle) 58 ctx24b5bc d P Q R S g h O u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

160 check_multiprocessing_start_method() 58 ctx24b5bc d P Q R S g h O u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

161 df = multiprocessing.reduction.DupFd(alloc_handle.handle) 58 ctx24b5bc d P Q R S g h O u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

162 return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid) 58 ctx24b5bc d P Q R S g h O u v w x 6bT 7b8bU 9b!by #b$bV %b'bW (b)bz *bA B C D E F G H I J i j K L M N s t k l m n e f Z

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

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) 54 ctx2c d P Q R S g h a u v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1bq r

176 if not self.memory_resource.is_ipc_enabled: 54 ctx2c d P Q R S g h a u v w x T 0 U 1 y 2 V 3 W 4 z 5 A B C D E F G H I J i j K L M N s t k l m n e f 0bp 1bq r

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

178 cdef cydriver.CUmemPoolPtrExportData data 

179 with nogil: 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

180 HANDLE_RETURN( 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

181 cydriver.cuMemPoolExportPointer(&data, as_cu(self._h_ptr)) 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

182 ) 

183 cdef bytes data_b = cpython.PyBytes_FromStringAndSize( 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

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

185 ) 

186 return IPCBufferDescriptor._init(data_b, self.size) 52 ctx1cdPQRSghauvwxT0U1y2V3W4z5ABCDEFGHIJijKLMNstklmnefpqr

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) 5 ctx2a YbZb0b1b

193 if not mr.is_ipc_enabled: 5 ctx2a YbZb0b1b

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

195 cdef size_t payload_size = len(ipc_descriptor._payload) 3 ctx2a YbZb

196 cdef size_t expected_size = sizeof(cydriver.CUmemPoolPtrExportData) 3 ctx2a YbZb

197 if payload_size < expected_size: 3 ctx2a YbZb

198 raise ValueError( 2 ctx2YbZb

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

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

201 ) 

202 cdef Stream s = Stream_accept(stream) 1 ctx1a

203 cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( 1 ctx1a

204 mr._h_pool, 

205 ipc_descriptor.payload_ptr(), 

206 s._h_stream 

207 ) 

208 if not h_ptr: 1 ctx1a

209 HANDLE_RETURN(get_last_error()) 

210 cdef size_t mapped_size = 0 1 ctx1a

211 cdef size_t claimed_size = ipc_descriptor.size 1 ctx1a

212 with nogil: 1 ctx1a

213 HANDLE_RETURN(cydriver.cuPointerGetAttribute( 1 ctx1a

214 &mapped_size, 

215 cydriver.CU_POINTER_ATTRIBUTE_RANGE_SIZE, 

216 as_cu(h_ptr))) 

217 if claimed_size > mapped_size: 1 ctx1a

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

224  

225  

226# _MemPool IPC Implementation 

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

228  

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

230 if isinstance(alloc_handle, IPCAllocationHandle): 4 ctx1XYab

231 IPCAllocationHandle_check_open(<IPCAllocationHandle>alloc_handle) 3 ctx1XYb

232  

233 # Quick exit for registry hits. 

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

235 mr = registry.get(uuid) 2 ctx1ab

236 if mr is not None: 2 ctx1ab

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

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

242 return mr 1 ctx1b

243  

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

245 # necessary. 

246 if isinstance(alloc_handle, int): 2 ctx1ab

247 fd = os.dup(alloc_handle) 1 ctx1a

248 try: 1 ctx1a

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

250 except: 

251 os.close(fd) 

252 raise 

253  

254 # Construct a new mempool. 

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

256 self._mempool_owned = True 2 ctx1ab

257 cdef int ipc_fd = int(alloc_handle) 2 ctx1ab

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

259 if not self._h_pool: 2 ctx1ab

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

268  

269 # Register it. 

270 if uuid is not None: 2 ctx1ab

271 registered = self.register(uuid) 1 ctx1b

272 assert registered is self 1 ctx1b

273  

274 return self 2 ctx1ab

275  

276  

277cdef _MemPool MP_from_registry(uuid): (empty)

278 cdef _MemPool mr 

279 try: 2 ctx2b 3b

280 mr = registry[uuid] 2 ctx2b 3b

281 MP_check_open(mr) 1 ctx1b

282 return mr 1 ctx1b

283 except KeyError: 1 ctx23b

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

285  

286  

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

288 MP_check_open(self) 4 ctx2X Y b 3b

289 existing = registry.get(uuid) 4 ctx2X Y b 3b

290 if existing is not None: 4 ctx2X Y b 3b

291 MP_check_open(<_MemPool>existing) 1 ctx1b

292 return existing 1 ctx1b

293 if not self.is_ipc_enabled: 4 ctx2X Y b 3b

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

295 assert self.uuid is None or self.uuid == uuid 3 ctx1XYb

296 registry[uuid] = self 3 ctx1XYb

297 self._ipc_data._alloc_handle._uuid = uuid 3 ctx1XYb

298 return self 3 ctx1XYb

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) 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

304 cdef int fd 

305 with nogil: 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

306 HANDLE_RETURN(cydriver.cuMemPoolExportToShareableHandle( 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

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

308 ) 

309 try: 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

310 return IPCAllocationHandle._init(fd, uuid.uuid4()) 100 ctx26 7 c 8 d 9 ! # $ % g ' h ( O ) * a + , b - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbkblbmbnbobpbqbrbsbtbubvbwbi j xbybzbAbBbCbk l m n e f DbEbFbGbHbIbJbKbLbMbNbObPbQbRbSbp TbUbVbZ WbXbq r

311 except: 

312 os.close(fd) 

313 raise