Coverage for cuda/core/checkpoint.py: 76.61%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import ctypes as _ctypes (empty)

6from collections.abc import Mapping (empty)

7from typing import Any (empty)

8 

9from cuda.bindings import driver as _driver (empty)

10from cuda.core._utils.cuda_utils import handle_return as _handle_cuda_return (empty)

11from cuda.core._utils.version import binding_version as _binding_version (empty)

12from cuda.core._utils.version import driver_version as _driver_version (empty)

13from cuda.core.typing import ProcessStateType as _ProcessStateType (empty)

14 

15_PROCESS_STATE_NAME_ATTRS: tuple[tuple[str, _ProcessStateType], ...] = ( (empty)

16 ("CU_PROCESS_STATE_RUNNING", "running"), 

17 ("CU_PROCESS_STATE_LOCKED", "locked"), 

18 ("CU_PROCESS_STATE_CHECKPOINTED", "checkpointed"), 

19 ("CU_PROCESS_STATE_FAILED", "failed"), 

20) 

21 

22_REQUIRED_BINDING_ATTRS = ( (empty)

23 "cuCheckpointProcessCheckpoint", 

24 "cuCheckpointProcessGetRestoreThreadId", 

25 "cuCheckpointProcessGetState", 

26 "cuCheckpointProcessLock", 

27 "cuCheckpointProcessRestore", 

28 "cuCheckpointProcessUnlock", 

29 "CUcheckpointGpuPair", 

30 "CUcheckpointLockArgs", 

31 "CUprocessState", 

32 "CUcheckpointRestoreArgs", 

33) 

34_REQUIRED_DRIVER_VERSION = (12, 8, 0) (empty)

35_driver_capability_checked = False (empty)

36 

37 

38class Process: (empty)

39 """ 

40 CUDA process that can be locked, checkpointed, restored, and unlocked. 

41 

42 Parameters 

43 ---------- 

44 pid : int 

45 Process ID of the CUDA process. 

46 """ 

47 

48 __slots__ = ("_pid",) (empty)

49 

50 def __init__(self, pid: int): (empty)

51 self._pid = _check_pid(pid) 5 ctx1jopmn

52 

53 @property (empty)

54 def pid(self) -> int: (empty)

55 """ 

56 Process ID of the CUDA process. 

57 """ 

58 return self._pid 1 ctx1j

59 

60 @property (empty)

61 def state(self) -> _ProcessStateType: (empty)

62 """ 

63 CUDA checkpoint state for this process. 

64 """ 

65 driver = _get_driver() 

66 state = _call_driver(driver, driver.cuCheckpointProcessGetState, self._pid) 

67 state_names = _get_process_state_names(driver) 

68 try: 

69 return state_names[state] 

70 except KeyError as e: 

71 state_value = int(state) 

72 raise RuntimeError(f"Unknown CUDA checkpoint process state: {state_value}") from e 

73 

74 @property (empty)

75 def restore_thread_id(self) -> int: (empty)

76 """ 

77 CUDA restore thread ID for this process. 

78 """ 

79 driver = _get_driver() 

80 return int(_call_driver(driver, driver.cuCheckpointProcessGetRestoreThreadId, self._pid)) 

81 

82 def lock(self, timeout_ms: int = 0) -> None: (empty)

83 """ 

84 Lock this process, blocking further CUDA API calls. 

85 

86 Parameters 

87 ---------- 

88 timeout_ms : int, optional 

89 Timeout in milliseconds. A value of 0 indicates no timeout. 

90 """ 

91 driver = _get_driver() 

92 args = driver.CUcheckpointLockArgs() 

93 args.timeoutMs = _check_timeout_ms(timeout_ms) 

94 _call_driver(driver, driver.cuCheckpointProcessLock, self._pid, args) 

95 

96 def checkpoint(self) -> None: (empty)

97 """ 

98 Checkpoint the GPU memory contents of this locked process. 

99 """ 

100 driver = _get_driver() 

101 _call_driver(driver, driver.cuCheckpointProcessCheckpoint, self._pid, None) 

102 

103 def restore(self, gpu_mapping: Mapping[Any, Any] | None = None) -> None: (empty)

104 """ 

105 Restore this checkpointed process. 

106 

107 Parameters 

108 ---------- 

109 gpu_mapping : mapping, optional 

110 GPU UUID remapping from each checkpointed GPU UUID to the GPU UUID 

111 to restore onto. For migration workflows, provide mappings for 

112 every GPU visible to the kernel-mode driver. User-space masking 

113 such as ``CUDA_VISIBLE_DEVICES`` does not reduce this mapping 

114 requirement. 

115 """ 

116 driver = _get_driver() 

117 args = _make_restore_args(driver, gpu_mapping) 

118 _call_driver(driver, driver.cuCheckpointProcessRestore, self._pid, args) 

119 

120 def unlock(self) -> None: (empty)

121 """ 

122 Unlock this locked process so it can resume CUDA API calls. 

123 """ 

124 driver = _get_driver() 

125 _call_driver(driver, driver.cuCheckpointProcessUnlock, self._pid, None) 

126 

127 

128def _get_driver() -> Any: (empty)

129 global _driver_capability_checked 

130 if _driver_capability_checked: (empty)

131 return _driver 

132 

133 binding_ver = _binding_version() (empty)

134 if not _binding_version_supports_checkpoint(binding_ver): (empty)

135 raise RuntimeError( 

136 "CUDA checkpointing requires cuda.bindings with CUDA checkpoint API support. " 

137 f"Found cuda.bindings {'.'.join(str(part) for part in binding_ver[:3])}." 

138 ) 

139 

140 missing = [name for name in _REQUIRED_BINDING_ATTRS if not hasattr(_driver, name)] (empty)

141 if missing: (empty)

142 raise RuntimeError( 

143 f"CUDA checkpointing requires cuda.bindings with CUDA checkpoint API support. Missing: {', '.join(missing)}" 

144 ) 

145 

146 driver_ver = _driver_version() (empty)

147 if driver_ver < _REQUIRED_DRIVER_VERSION: (empty)

148 raise RuntimeError( 

149 "CUDA checkpointing is not supported by the installed NVIDIA driver. " 

150 "Upgrade to a driver version with CUDA checkpoint API support." 

151 ) 

152 

153 _driver_capability_checked = True (empty)

154 return _driver (empty)

155 

156 

157def _binding_version_supports_checkpoint(version: tuple[int, ...]) -> bool: (empty)

158 major, minor, patch = version[:3] (empty)

159 return (major == 12 and (minor, patch) >= (8, 0)) or (major == 13 and (minor, patch) >= (0, 2)) or major > 13 (empty)

160 

161 

162def _get_process_state_names(driver: Any) -> dict[Any, _ProcessStateType]: (empty)

163 return {getattr(driver.CUprocessState, attr): state_name for attr, state_name in _PROCESS_STATE_NAME_ATTRS} 

164 

165 

166def _call_driver(driver: Any, func: Any, *args: Any) -> Any: (empty)

167 try: 4 ctx1klef

168 result = func(*args) 4 ctx1klef

169 except RuntimeError as e: 2 ctx1kl

170 if "cuCheckpointProcess" in str(e) and "not found" in str(e): 2 ctx1kl

171 raise RuntimeError( 1 ctx1l

172 "CUDA checkpointing is not supported by the installed NVIDIA driver. " 

173 "Upgrade to a driver version with CUDA checkpoint API support." 

174 ) from e 

175 raise 1 ctx1k

176 

177 err = result[0] 2 ctx1ef

178 not_supported_errors = ( 2 ctx1ef

179 getattr(driver.CUresult, "CUDA_ERROR_NOT_FOUND", None), 

180 getattr(driver.CUresult, "CUDA_ERROR_NOT_SUPPORTED", None), 

181 ) 

182 if err in not_supported_errors: 2 ctx1ef

183 raise RuntimeError( 2 ctx1ef

184 "CUDA checkpointing is not supported by the installed NVIDIA driver. " 

185 "Upgrade to a driver version with CUDA checkpoint API support." 

186 ) 

187 

188 return _handle_cuda_return(result) 

189 

190 

191def _check_pid(pid: int) -> int: (empty)

192 if isinstance(pid, bool) or not isinstance(pid, int): 5 ctx1jopmn

193 raise TypeError("pid must be an int") 2 ctx1op

194 if pid <= 0: 3 ctx1jmn

195 raise ValueError("pid must be a positive int") 2 ctx1mn

196 return pid 1 ctx1j

197 

198 

199def _check_timeout_ms(timeout_ms: int) -> int: (empty)

200 if isinstance(timeout_ms, bool) or not isinstance(timeout_ms, int): 4 ctx1tuqv

201 raise TypeError("timeout_ms must be an int") 3 ctx1tuv

202 if timeout_ms < 0: 1 ctx1q

203 raise ValueError("timeout_ms must be >= 0") 1 ctx1q

204 return timeout_ms 

205 

206 

207def _make_restore_args(driver: Any, gpu_mapping: Mapping[Any, Any] | None) -> Any: (empty)

208 if gpu_mapping is None: 3 ctx1bgr

209 return None 

210 if not isinstance(gpu_mapping, Mapping): 3 ctx1bgr

211 raise TypeError("gpu_mapping must be a mapping from checkpointed GPU UUID to restore GPU UUID") 1 ctx1r

212 

213 pairs = [] 2 ctx1bg

214 for old_uuid, new_uuid in gpu_mapping.items(): 2 ctx1bg

215 pair = driver.CUcheckpointGpuPair() 1 ctx1b

216 buffers: list[Any] = [] # holds ctypes string-buffer keepalives for the call below 1 ctx1b

217 pair.oldUuid = _as_cuuuid(driver, old_uuid, buffers) 1 ctx1b

218 pair.newUuid = _as_cuuuid(driver, new_uuid, buffers) 1 ctx1b

219 pairs.append(pair) 1 ctx1b

220 

221 if not pairs: 2 ctx1bg

222 return None 1 ctx1g

223 

224 args = driver.CUcheckpointRestoreArgs() 1 ctx1b

225 args.gpuPairs = pairs 1 ctx1b

226 args.gpuPairsCount = len(pairs) 1 ctx1b

227 return args 1 ctx1b

228 

229 

230def _as_cuuuid(driver: Any, value: Any, buffers: list[Any]) -> Any: (empty)

231 """Convert *value* to a ``CUuuid``. 

232 

233 Accepts a ``CUuuid`` instance (returned as-is) or a UUID string in 

234 the ``"xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"`` format returned by 

235 :attr:`Device.uuid`. 

236 """ 

237 if isinstance(value, driver.CUuuid): 7 ctx1cdwhisb

238 return value 1 ctx1w

239 if isinstance(value, str): 6 ctx1cdhisb

240 try: 5 ctx1cdhib

241 raw = bytes.fromhex(value.replace("-", "")) 5 ctx1cdhib

242 except ValueError: 1 ctx1i

243 raise ValueError( 1 ctx1i

244 f"GPU UUID string must be 32 hex characters (with optional hyphens), got {value!r}" 

245 ) from None 

246 if len(raw) != 16: 4 ctx1cdhb

247 raise ValueError(f"GPU UUID string must be 32 hex characters (with optional hyphens), got {value!r}") 1 ctx1h

248 buf = _ctypes.create_string_buffer(raw, 16) 3 ctx1cdb

249 buffers.append(buf) 3 ctx1cdb

250 return driver.CUuuid(_ctypes.addressof(buf)) 3 ctx1cdb

251 raise TypeError("GPU UUID values must be CUDA UUID objects or UUID strings") 1 ctx1s

252 

253 

254__all__ = [ (empty)

255 "Process", 

256]