Coverage for cuda/core/checkpoint.py: 48.28%
116 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 01:23 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5"""CUDA process checkpointing on Linux.
7This module requires the CUDA 13 build of cuda.core, because the driver
8structures it uses are CUDA 13 types. It also requires a CUDA driver of version
912.8 or newer with checkpoint API support.
10"""
12import ctypes as _ctypes
13from collections.abc import Mapping
14from typing import Any
16from cuda.bindings import driver as _driver
17from cuda.core._utils.cuda_utils import handle_return as _handle_cuda_return
18from cuda.core._utils.version import BUILD_CUDA_MAJOR as _BUILD_CUDA_MAJOR
19from cuda.core._utils.version import driver_version as _driver_version
20from cuda.core.typing import ProcessStateType as _ProcessStateType
22_PROCESS_STATE_NAME_ATTRS: tuple[tuple[str, _ProcessStateType], ...] = (
23 ("CU_PROCESS_STATE_RUNNING", "running"),
24 ("CU_PROCESS_STATE_LOCKED", "locked"),
25 ("CU_PROCESS_STATE_CHECKPOINTED", "checkpointed"),
26 ("CU_PROCESS_STATE_FAILED", "failed"),
27)
29_REQUIRED_DRIVER_VERSION = (12, 8, 0)
30_driver_capability_checked = False
33class Process:
34 """
35 CUDA process that can be locked, checkpointed, restored, and unlocked.
37 Parameters
38 ----------
39 pid : int
40 Process ID of the CUDA process.
41 """
43 __slots__ = ("_pid",)
45 def __init__(self, pid: int):
46 self._pid = _check_pid(pid) 1jopmn
48 @property
49 def pid(self) -> int:
50 """
51 Process ID of the CUDA process.
52 """
53 return self._pid 1j
55 @property
56 def state(self) -> _ProcessStateType:
57 """
58 CUDA checkpoint state for this process.
59 """
60 driver = _get_driver()
61 state = _call_driver(driver, driver.cuCheckpointProcessGetState, self._pid)
62 state_names = _get_process_state_names(driver)
63 try:
64 return state_names[state]
65 except KeyError as e:
66 state_value = int(state)
67 raise RuntimeError(f"Unknown CUDA checkpoint process state: {state_value}") from e
69 @property
70 def restore_thread_id(self) -> int:
71 """
72 CUDA restore thread ID for this process.
73 """
74 driver = _get_driver()
75 return int(_call_driver(driver, driver.cuCheckpointProcessGetRestoreThreadId, self._pid))
77 def lock(self, timeout_ms: int = 0) -> None:
78 """
79 Lock this process, blocking further CUDA API calls.
81 Parameters
82 ----------
83 timeout_ms : int, optional
84 Timeout in milliseconds. A value of 0 indicates no timeout.
85 """
86 driver = _get_driver()
87 args = driver.CUcheckpointLockArgs()
88 args.timeoutMs = _check_timeout_ms(timeout_ms)
89 _call_driver(driver, driver.cuCheckpointProcessLock, self._pid, args)
91 def checkpoint(self) -> None:
92 """
93 Checkpoint the GPU memory contents of this locked process.
94 """
95 driver = _get_driver()
96 _call_driver(driver, driver.cuCheckpointProcessCheckpoint, self._pid, None)
98 def restore(self, gpu_mapping: Mapping[Any, Any] | None = None) -> None:
99 """
100 Restore this checkpointed process.
102 Parameters
103 ----------
104 gpu_mapping : mapping, optional
105 GPU UUID remapping from each checkpointed GPU UUID to the GPU UUID
106 to restore onto. For migration workflows, provide mappings for
107 every GPU visible to the kernel-mode driver. User-space masking
108 such as ``CUDA_VISIBLE_DEVICES`` does not reduce this mapping
109 requirement.
110 """
111 driver = _get_driver()
112 args = _make_restore_args(driver, gpu_mapping)
113 _call_driver(driver, driver.cuCheckpointProcessRestore, self._pid, args)
115 def unlock(self) -> None:
116 """
117 Unlock this locked process so it can resume CUDA API calls.
118 """
119 driver = _get_driver()
120 _call_driver(driver, driver.cuCheckpointProcessUnlock, self._pid, None)
123def _get_driver() -> Any:
124 global _driver_capability_checked
125 if _driver_capability_checked:
126 return _driver
128 # A restore onto other GPUs uses CUcheckpointGpuPair, a CUDA 13 type that
129 # cuda-bindings 12.x does not have.
130 if _BUILD_CUDA_MAJOR < 13:
131 raise RuntimeError("CUDA checkpointing requires the CUDA 13 build of cuda.core. Install cuda-core[cu13].")
133 driver_ver = _driver_version()
134 if driver_ver < _REQUIRED_DRIVER_VERSION:
135 raise RuntimeError(
136 "CUDA checkpointing is not supported by the installed NVIDIA driver. "
137 "Upgrade to a driver version with CUDA checkpoint API support."
138 )
140 _driver_capability_checked = True
141 return _driver
144def _get_process_state_names(driver: Any) -> dict[Any, _ProcessStateType]:
145 return {getattr(driver.CUprocessState, attr): state_name for attr, state_name in _PROCESS_STATE_NAME_ATTRS}
148def _call_driver(driver: Any, func: Any, *args: Any) -> Any:
149 try: 1klef
150 result = func(*args) 1klef
151 except RuntimeError as e: 1kl
152 if "cuCheckpointProcess" in str(e) and "not found" in str(e): 1kl
153 raise RuntimeError( 1l
154 "CUDA checkpointing is not supported by the installed NVIDIA driver. "
155 "Upgrade to a driver version with CUDA checkpoint API support."
156 ) from e
157 raise 1k
159 err = result[0] 1ef
160 not_supported_errors = ( 1ef
161 getattr(driver.CUresult, "CUDA_ERROR_NOT_FOUND", None),
162 getattr(driver.CUresult, "CUDA_ERROR_NOT_SUPPORTED", None),
163 )
164 if err in not_supported_errors: 1ef
165 raise RuntimeError( 1ef
166 "CUDA checkpointing is not supported by the installed NVIDIA driver. "
167 "Upgrade to a driver version with CUDA checkpoint API support."
168 )
170 return _handle_cuda_return(result)
173def _check_pid(pid: int) -> int:
174 if isinstance(pid, bool) or not isinstance(pid, int): 1jopmn
175 raise TypeError("pid must be an int") 1op
176 if pid <= 0: 1jmn
177 raise ValueError("pid must be a positive int") 1mn
178 return pid 1j
181def _check_timeout_ms(timeout_ms: int) -> int:
182 if isinstance(timeout_ms, bool) or not isinstance(timeout_ms, int): 1tuqv
183 raise TypeError("timeout_ms must be an int") 1tuv
184 if timeout_ms < 0: 1q
185 raise ValueError("timeout_ms must be >= 0") 1q
186 return timeout_ms
189def _make_restore_args(driver: Any, gpu_mapping: Mapping[Any, Any] | None) -> Any:
190 if gpu_mapping is None: 1agr
191 return None
192 if not isinstance(gpu_mapping, Mapping): 1agr
193 raise TypeError("gpu_mapping must be a mapping from checkpointed GPU UUID to restore GPU UUID") 1r
195 pairs = [] 1ag
196 for old_uuid, new_uuid in gpu_mapping.items(): 1ag
197 pair = driver.CUcheckpointGpuPair() 1a
198 buffers: list[Any] = [] # holds ctypes string-buffer keepalives for the call below 1a
199 pair.oldUuid = _as_cuuuid(driver, old_uuid, buffers) 1a
200 pair.newUuid = _as_cuuuid(driver, new_uuid, buffers) 1a
201 pairs.append(pair) 1a
203 if not pairs: 1ag
204 return None 1g
206 args = driver.CUcheckpointRestoreArgs() 1a
207 args.gpuPairs = pairs 1a
208 args.gpuPairsCount = len(pairs) 1a
209 return args 1a
212def _as_cuuuid(driver: Any, value: Any, buffers: list[Any]) -> Any:
213 """Convert *value* to a ``CUuuid``.
215 Accepts a ``CUuuid`` instance (returned as-is) or a UUID string in
216 the ``"xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"`` format returned by
217 :attr:`Device.uuid`.
218 """
219 if isinstance(value, driver.CUuuid): 1bcwhisa
220 return value 1w
221 if isinstance(value, str): 1bchisa
222 try: 1bchia
223 raw = bytes.fromhex(value.replace("-", "")) 1bchia
224 except ValueError: 1i
225 raise ValueError( 1i
226 f"GPU UUID string must be 32 hex characters (with optional hyphens), got {value!r}"
227 ) from None
228 if len(raw) != 16: 1bcha
229 raise ValueError(f"GPU UUID string must be 32 hex characters (with optional hyphens), got {value!r}") 1h
230 buf = _ctypes.create_string_buffer(raw, 16) 1bca
231 buffers.append(buf) 1bca
232 return driver.CUuuid(_ctypes.addressof(buf)) 1bca
233 raise TypeError("GPU UUID values must be CUDA UUID objects or UUID strings") 1s
236__all__ = [
237 "Process",
238]