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

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 cuda.bindings cimport cydriver 

8from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle 

9from cuda.core._memory._memory_pool cimport _MemPool 

10from cuda.core._stream cimport Stream, Stream_accept 

11from cuda.core._resource_handles cimport ( 

12 DevicePtrHandle, 

13 create_fd_handle, 

14 create_mempool_handle_ipc, 

15 deviceptr_import_ipc, 

16 get_last_error, 

17 as_cu, 

18 as_intptr, 

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__ = ['IPCBufferDescriptor', 'IPCAllocationHandle'] (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 

41  

42cdef is_supported(): (empty)

43 return IPC_HANDLE_TYPE != cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo 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): 

49 self._ipc_descriptor = ipc_descriptor 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

50 self._is_mapped = is_mapped 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

51  

52 @property 

53 def ipc_descriptor(self): 

54 return self._ipc_descriptor 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

55  

56 @property 

57 def is_mapped(self): 

58 return self._is_mapped 22 ctx1stuvwxyzABCDEFcdGHIJqr

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): 

64 self._alloc_handle = alloc_handle 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

65 self._is_mapped = is_mapped 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

66  

67 @property 

68 def alloc_handle(self): 

69 return self._alloc_handle 

70  

71 @property 

72 def is_mapped(self): 

73 return self._is_mapped 

74  

75 @property 

76 def uuid(self): 

77 return getattr(self._alloc_handle, 'uuid', None) 6 ctx1abqref

78  

79  

80cdef class IPCBufferDescriptor: 

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

82  

83 def __init__(self, *arg, **kwargs): 

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

85  

86 @staticmethod (empty)

87 def _init(reserved: bytes, size: int): 

88 cdef IPCBufferDescriptor self = IPCBufferDescriptor.__new__(IPCBufferDescriptor) 51 ctx2a b K L g h s t u v M T N U O V P W Q X R Y w x y z A B C D E F c d G H I J q r i j k l e f Lbn Mbo p

89 self._payload = reserved 51 ctx2a b K L g h s t u v M T N U O V P W Q X R Y w x y z A B C D E F c d G H I J q r i j k l e f Lbn Mbo p

90 self._size = size 51 ctx2a b K L g h s t u v M T N U O V P W Q X R Y w x y z A B C D E F c d G H I J q r i j k l e f Lbn Mbo p

91 return self 51 ctx2a b K L g h s t u v M T N U O V P W Q X R Y w x y z A B C D E F c d G H I J q r i j k l e f Lbn Mbo p

92  

93 def __reduce__(self): (empty)

94 return IPCBufferDescriptor._init, (self._payload, self._size) 48 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefop

95  

96 @property 

97 def size(self): 

98 return self._size 1 ctx1n

99  

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

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

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

103  

104  

105cdef class IPCAllocationHandle: 

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

107  

108 def __init__(self, *arg, **kwargs): 

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

110  

111 @classmethod (empty)

112 def _init(cls, handle: int, uuid): # no-cython-lint 

113 cdef IPCAllocationHandle self = IPCAllocationHandle.__new__(cls) 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

114 assert handle >= 0 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

115 self._h_fd = create_fd_handle(handle) 92 ctx2m Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

116 self._uuid = uuid 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

117 return self 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

118  

119 cpdef close(self): (empty)

120 """Close the handle.""" 

121 self._h_fd.reset() 

122  

123 def __int__(self) -> int: 

124 if not self._h_fd or as_intptr(self._h_fd) < 0: 

125 raise ValueError( 

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

127 ) 

128 return as_py(self._h_fd) 

129  

130 @property 

131 def handle(self) -> int: 

132 return as_py(self._h_fd) 55 ctx2NbOba b K L g h s t u v PbM QbRbN SbTbO UbVbP WbXbQ YbZbR 0bw x y z A B C D E F c d G H I J q r i j k l e f S

133  

134 @property 

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

136 return self._uuid 55 ctx2NbOba b K L g h s t u v PbM QbRbN SbTbO UbVbP WbXbQ YbZbR 0bw x y z A B C D E F c d G H I J q r i j k l e f S

137  

138  

139def _reduce_allocation_handle(alloc_handle): (empty)

140 check_multiprocessing_start_method() 55 ctx2NbOba b K L g h s t u v PbM QbRbN SbTbO UbVbP WbXbQ YbZbR 0bw x y z A B C D E F c d G H I J q r i j k l e f S

141 df = multiprocessing.reduction.DupFd(alloc_handle.handle) 55 ctx2NbOba b K L g h s t u v PbM QbRbN SbTbO UbVbP WbXbQ YbZbR 0bw x y z A B C D E F c d G H I J q r i j k l e f S

142 return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid) 55 ctx2NbOba b K L g h s t u v PbM QbRbN SbTbO UbVbP WbXbQ YbZbR 0bw x y z A B C D E F c d G H I J q r i j k l e f S

143  

144  

145def _reconstruct_allocation_handle(cls, df, uuid): # no-cython-lint (empty)

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

147  

148  

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

150  

151  

152# Buffer IPC Implementation 

153# ------------------------- 

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

155 if not self.memory_resource.is_ipc_enabled: 51 ctx2a b K L g h s t u v M T N U O V P W Q X R Y w x y z A B C D E F c d G H I J q r i j k l e f Lbn Mbo p

156 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2LbMb

157 cdef cydriver.CUmemPoolPtrExportData data 

158 with nogil: 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

159 HANDLE_RETURN( 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

160 cydriver.cuMemPoolExportPointer(&data, as_cu(self._h_ptr)) 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

161 ) 

162 cdef bytes data_b = cpython.PyBytes_FromStringAndSize( 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

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

164 ) 

165 return IPCBufferDescriptor._init(data_b, self.size) 49 ctx1abKLghstuvMTNUOVPWQXRYwxyzABCDEFcdGHIJqrijklefnop

166  

167cdef Buffer Buffer_from_ipc_descriptor( (empty)

168 cls, _MemPool mr, IPCBufferDescriptor ipc_descriptor, stream 

169): 

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

171 if not mr.is_ipc_enabled: 2 ctx2LbMb

172 raise RuntimeError("Memory resource is not IPC-enabled") 2 ctx2LbMb

173 cdef Stream s = Stream_accept(stream) 

174 cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( 

175 mr._h_pool, 

176 ipc_descriptor.payload_ptr(), 

177 s._h_stream 

178 ) 

179 if not h_ptr: 

180 HANDLE_RETURN(get_last_error()) 

181 return Buffer_from_deviceptr_handle(h_ptr, ipc_descriptor.size, mr, ipc_descriptor) 

182  

183  

184# _MemPool IPC Implementation 

185# --------------------------- 

186  

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

188 # Quick exit for registry hits. 

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

190 mr = registry.get(uuid) 

191 if mr is not None: 

192 if not isinstance(mr, cls): 

193 raise TypeError( 

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

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

196 return mr 

197  

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

199 # necessary. 

200 if isinstance(alloc_handle, int): 

201 fd = os.dup(alloc_handle) 

202 try: 

203 alloc_handle = IPCAllocationHandle._init(fd, None) 

204 except: 

205 os.close(fd) 

206 raise 

207  

208 # Construct a new mempool. 

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

210 self._mempool_owned = True 

211 cdef int ipc_fd = int(alloc_handle) 

212 self._h_pool = create_mempool_handle_ipc(ipc_fd, IPC_HANDLE_TYPE) 

213 if not self._h_pool: 

214 HANDLE_RETURN(get_last_error()) 

215 raise RuntimeError( 

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

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

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

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

220 ) 

221 self._ipc_data = IPCDataForMR(alloc_handle, True) 

222  

223 # Register it. 

224 if uuid is not None: 

225 registered = self.register(uuid) 

226 assert registered is self 

227  

228 return self 

229  

230  

231cdef _MemPool MP_from_registry(uuid): (empty)

232 try: 

233 return registry[uuid] 

234 except KeyError: 

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

236  

237  

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

239 existing = registry.get(uuid) 

240 if existing is not None: 

241 return existing 

242 assert self.uuid is None or self.uuid == uuid 

243 registry[uuid] = self 

244 self._ipc_data._alloc_handle._uuid = uuid 

245 return self 

246  

247  

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

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

250 cdef int fd 

251 with nogil: 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

252 HANDLE_RETURN(cydriver.cuMemPoolExportToShareableHandle( 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

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

254 ) 

255 try: 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

256 return IPCAllocationHandle._init(fd, uuid.uuid4()) 91 ctx2Z 0 a 1 b 2 3 4 g 5 h 6 7 8 9 ! # $ % ' ( ) * + , - . / : ; = ? @ [ ] ^ _ ` { | } ~ abbbcbdbebfbgbhbibjbc d kblbmbnbobpbi j k l e f qbrbsbtbubvbwbxbybzbAbBbCbDbEbFbn GbHbIbS JbKbo p

257 except: 

258 os.close(fd) 

259 raise