Skip to content

Commit 8192df6

Browse files
kkraus14kkraus14
committed
Address checkpoint review follow-ups
1 parent 7c7f0e5 commit 8192df6

6 files changed

Lines changed: 89 additions & 42 deletions

File tree

‎cuda_core/cuda/core/checkpoint.py‎

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,28 +2,27 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5+
import ctypes as _ctypes
56
from collections.abc import Mapping as _Mapping
67
from typing import Any as _Any
7-
from typing import Literal as _Literal
88

99
from cuda.core._utils.cuda_utils import handle_return as _handle_cuda_return
1010
from cuda.core._utils.version import binding_version as _binding_version
1111
from cuda.core._utils.version import driver_version as _driver_version
12+
from cuda.core.typing import ProcessStateT as _ProcessStateT
1213

1314
try:
1415
from cuda.bindings import driver as _driver
1516
except ImportError:
1617
from cuda import cuda as _driver
1718

1819

19-
ProcessStateT = _Literal["running", "locked", "checkpointed", "failed"]
20-
21-
_PROCESS_STATE_NAMES: dict[int, ProcessStateT] = {
22-
0: "running",
23-
1: "locked",
24-
2: "checkpointed",
25-
3: "failed",
26-
}
20+
_PROCESS_STATE_NAME_ATTRS: tuple[tuple[str, _ProcessStateT], ...] = (
21+
("CU_PROCESS_STATE_RUNNING", "running"),
22+
("CU_PROCESS_STATE_LOCKED", "locked"),
23+
("CU_PROCESS_STATE_CHECKPOINTED", "checkpointed"),
24+
("CU_PROCESS_STATE_FAILED", "failed"),
25+
)
2726

2827
_REQUIRED_BINDING_ATTRS = (
2928
"cuCheckpointProcessCheckpoint",
@@ -34,6 +33,7 @@
3433
"cuCheckpointProcessUnlock",
3534
"CUcheckpointGpuPair",
3635
"CUcheckpointLockArgs",
36+
"CUprocessState",
3737
"CUcheckpointRestoreArgs",
3838
)
3939
_REQUIRED_DRIVER_VERSION = (12, 8, 0)
@@ -56,16 +56,17 @@ def __init__(self, pid: int):
5656
self.pid = _check_pid(pid)
5757

5858
@property
59-
def state(self) -> ProcessStateT:
59+
def state(self) -> _ProcessStateT:
6060
"""
6161
CUDA checkpoint state for this process.
6262
"""
6363
driver = _get_driver()
6464
state = _call_driver(driver, driver.cuCheckpointProcessGetState, self.pid)
65-
state_value = int(state)
65+
state_names = _get_process_state_names(driver)
6666
try:
67-
return _PROCESS_STATE_NAMES[state_value]
67+
return state_names[state]
6868
except KeyError as e:
69+
state_value = int(state)
6970
raise RuntimeError(f"Unknown CUDA checkpoint process state: {state_value}") from e
7071

7172
@property
@@ -105,8 +106,8 @@ def restore(self, gpu_mapping: _Mapping[_Any, _Any] | None = None) -> None:
105106
----------
106107
gpu_mapping : mapping, optional
107108
GPU UUID remapping from each checkpointed GPU UUID to the GPU UUID
108-
to restore onto. If provided, the mapping must contain every
109-
checkpointed GPU UUID.
109+
to restore onto. For migration workflows, provide mappings for
110+
every CUDA-visible GPU.
110111
"""
111112
driver = _get_driver()
112113
args = _make_restore_args(driver, gpu_mapping)
@@ -154,6 +155,10 @@ def _binding_version_supports_checkpoint(version) -> bool:
154155
return (major == 12 and (minor, patch) >= (8, 0)) or (major == 13 and (minor, patch) >= (0, 2)) or major > 13
155156

156157

158+
def _get_process_state_names(driver) -> dict[_Any, _ProcessStateT]:
159+
return {getattr(driver.CUprocessState, attr): state_name for attr, state_name in _PROCESS_STATE_NAME_ATTRS}
160+
161+
157162
def _call_driver(driver, func, *args):
158163
try:
159164
result = func(*args)
@@ -207,8 +212,9 @@ def _make_restore_args(driver, gpu_mapping: _Mapping[_Any, _Any] | None):
207212
pairs = []
208213
for old_uuid, new_uuid in gpu_mapping.items():
209214
pair = driver.CUcheckpointGpuPair()
210-
pair.oldUuid = _as_cuuuid(driver, old_uuid)
211-
pair.newUuid = _as_cuuuid(driver, new_uuid)
215+
buffers = []
216+
pair.oldUuid = _as_cuuuid(driver, old_uuid, buffers)
217+
pair.newUuid = _as_cuuuid(driver, new_uuid, buffers)
212218
pairs.append(pair)
213219

214220
if not pairs:
@@ -220,21 +226,20 @@ def _make_restore_args(driver, gpu_mapping: _Mapping[_Any, _Any] | None):
220226
return args
221227

222228

223-
def _as_cuuuid(driver, value):
229+
def _as_cuuuid(driver, value, buffers):
224230
"""Convert *value* to a ``CUuuid``.
225231
226232
Accepts a ``CUuuid`` instance (returned as-is) or a UUID string in
227233
the ``"xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"`` format returned by
228234
:attr:`Device.uuid`.
229235
"""
230236
if isinstance(value, str):
231-
import ctypes
232-
233237
raw = bytes.fromhex(value.replace("-", ""))
234238
if len(raw) != 16:
235239
raise ValueError(f"GPU UUID string must be 32 hex characters (with optional hyphens), got {value!r}")
236-
buf = ctypes.create_string_buffer(raw, 16)
237-
return driver.CUuuid(ctypes.addressof(buf))
240+
buf = _ctypes.create_string_buffer(raw, 16)
241+
buffers.append(buf)
242+
return driver.CUuuid(_ctypes.addressof(buf))
238243
return value
239244

240245

‎cuda_core/cuda/core/typing.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,15 @@
44

55
"""Public type aliases and protocols used in cuda.core API signatures."""
66

7+
from typing import Literal as _Literal
8+
79
from cuda.core._memory._buffer import DevicePointerT
810
from cuda.core._stream import IsStreamT
911

12+
ProcessStateT = _Literal["running", "locked", "checkpointed", "failed"]
13+
1014
__all__ = [
1115
"DevicePointerT",
1216
"IsStreamT",
17+
"ProcessStateT",
1318
]

‎cuda_core/docs/source/api.rst‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -182,30 +182,48 @@ checkpoint APIs. These APIs are intended for Linux process checkpoint and
182182
restore workflows, and require a CUDA driver with checkpoint API support and
183183
a ``cuda-bindings`` version that exposes those driver entry points.
184184

185-
A checkpoint workflow operates on a CUDA process by process ID. The typical
186-
sequence is to lock the process, capture its GPU memory state, restore it
187-
when needed, and then unlock it so CUDA API calls can resume:
185+
Checkpointing is typically driven by a coordinator process acting on a target
186+
CUDA process, similar to attaching a debugger or sending a signal. The target
187+
process is identified by process ID. Linux and the CUDA driver enforce process
188+
permissions; checkpointing another user's process may require elevated
189+
permissions such as ``CAP_SYS_PTRACE`` or administrator privileges.
190+
191+
The CUDA checkpoint APIs prepare CUDA-managed GPU state for process-level
192+
checkpoint and restore. They do not capture the CPU process image by
193+
themselves; full process checkpoint workflows still need a CPU-side process
194+
checkpointing tool such as CRIU. A minimal coordinator-side sequence looks like
195+
this:
188196

189197
.. code-block:: python
190198
199+
import os
200+
191201
from cuda.core import checkpoint
192202
193-
process = checkpoint.Process(pid)
203+
target_pid = os.getpid() # or the PID of another CUDA process
204+
process = checkpoint.Process(target_pid)
194205
process.lock(timeout_ms=5000)
195206
process.checkpoint()
207+
208+
# Capture or restore the CPU process image outside cuda.core.
209+
196210
process.restore()
197211
process.unlock()
198212
199213
``Process.state`` returns one of ``"running"``, ``"locked"``,
200214
``"checkpointed"``, or ``"failed"``. Restore may optionally remap GPUs by
201215
passing ``gpu_mapping`` from each checkpointed GPU UUID to the GPU UUID that
202-
should be used during restore. A successful restore returns the process to
203-
the locked state; call ``Process.unlock`` after restore to allow CUDA API
204-
calls to resume.
216+
should be used during restore. For migration workflows, provide mappings for
217+
every CUDA-visible GPU. The mapping may use ``CUuuid`` objects or the UUID
218+
strings returned by :attr:`Device.uuid`. A successful restore returns the
219+
process to the locked state; call ``Process.unlock`` after restore to allow
220+
CUDA API calls to resume.
205221

206222
The CUDA driver requires restore to run from the process restore thread.
207223
Use ``Process.restore_thread_id`` to discover that thread before calling
208-
``Process.restore`` from a checkpoint coordinator.
224+
``Process.restore`` from a checkpoint coordinator. Restore also requires
225+
persistence mode to be enabled or ``cuInit`` to have been called before
226+
execution.
209227

210228
.. autosummary::
211229
:toctree: generated/

‎cuda_core/docs/source/api_private.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ CUDA runtime
1717
:toctree: generated/
1818

1919
typing.DevicePointerT
20+
typing.ProcessStateT
2021
_memory._virtual_memory_resource.VirtualMemoryAllocationTypeT
2122
_memory._virtual_memory_resource.VirtualMemoryLocationTypeT
2223
_memory._virtual_memory_resource.VirtualMemoryGranularityT

‎cuda_core/tests/test_checkpoint.py‎

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import os
1616
import sys
17+
from contextlib import suppress
1718

1819
import pytest
1920

@@ -62,6 +63,15 @@ def _find_same_chip_pair(devices):
6263
return None
6364

6465

66+
def _run_or_skip_unsupported(func, *args, **kwargs):
67+
try:
68+
return func(*args, **kwargs)
69+
except RuntimeError as exc:
70+
if "CUDA checkpointing is not supported" in str(exc):
71+
pytest.skip(str(exc))
72+
raise
73+
74+
6575
# -- Fixtures --------------------------------------------------------------
6676

6777

@@ -78,13 +88,16 @@ def self_process(init_cuda):
7888
# Ensure the process is not left locked if the test fails mid-lifecycle.
7989
try:
8090
st = proc.state
81-
if st == "checkpointed":
91+
except Exception:
92+
st = None
93+
if st == "checkpointed":
94+
with suppress(Exception):
8295
proc.restore()
96+
with suppress(Exception):
8397
proc.unlock()
84-
elif st == "locked":
98+
elif st == "locked":
99+
with suppress(Exception):
85100
proc.unlock()
86-
except Exception: # noqa: S110 — best-effort teardown, nothing useful to log
87-
pass
88101
# Restore the original device so init_cuda's teardown pops the right context.
89102
original_device.set_current()
90103

@@ -108,6 +121,7 @@ def test_process_rejects_invalid_pid(self, args, error_type, match):
108121

109122
def test_public_symbols(self):
110123
assert checkpoint.__all__ == ["Process"]
124+
assert not hasattr(checkpoint, "ProcessStateT")
111125

112126

113127
# -- Lifecycle (single GPU, real driver) -----------------------------------
@@ -124,31 +138,31 @@ def test_restore_thread_id_is_positive(self, self_process):
124138
assert tid > 0
125139

126140
def test_lock_unlock(self, self_process):
127-
self_process.lock()
141+
_run_or_skip_unsupported(self_process.lock)
128142
assert self_process.state == "locked"
129143
self_process.unlock()
130144
assert self_process.state == "running"
131145

132146
def test_lock_default_timeout(self, self_process):
133147
"""lock() with the default timeout_ms=0 (no timeout)."""
134-
self_process.lock()
148+
_run_or_skip_unsupported(self_process.lock)
135149
assert self_process.state == "locked"
136150
self_process.unlock()
137151

138152
def test_lock_with_timeout(self, self_process):
139-
self_process.lock(timeout_ms=5000)
153+
_run_or_skip_unsupported(self_process.lock, timeout_ms=5000)
140154
assert self_process.state == "locked"
141155
self_process.unlock()
142156

143157
def test_full_cycle_no_migration(self, self_process):
144158
"""lock -> checkpoint -> restore -> unlock, verify state at each step."""
145-
self_process.lock()
159+
_run_or_skip_unsupported(self_process.lock)
146160
assert self_process.state == "locked"
147161

148-
self_process.checkpoint()
162+
_run_or_skip_unsupported(self_process.checkpoint)
149163
assert self_process.state == "checkpointed"
150164

151-
self_process.restore()
165+
_run_or_skip_unsupported(self_process.restore)
152166
assert self_process.state == "locked" # restore leaves process locked
153167

154168
self_process.unlock()
@@ -177,10 +191,10 @@ def _try_migration(proc, gpu_mapping):
177191
the migration with CUDA_ERROR_INVALID_VALUE (known limitation
178192
on some architectures / driver versions).
179193
"""
180-
proc.lock()
181-
proc.checkpoint()
194+
_run_or_skip_unsupported(proc.lock)
195+
_run_or_skip_unsupported(proc.checkpoint)
182196
try:
183-
proc.restore(gpu_mapping=gpu_mapping)
197+
_run_or_skip_unsupported(proc.restore, gpu_mapping=gpu_mapping)
184198
except (CUDAError, RuntimeError) as exc:
185199
# Recover: restore without migration, then unlock.
186200
proc.restore()

‎cuda_core/tests/test_typing_imports.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,12 @@ def test_typing_module_imports():
1010
from cuda.core.typing import (
1111
DevicePointerT,
1212
IsStreamT,
13+
ProcessStateT,
1314
)
1415

1516
assert DevicePointerT is not None
1617
assert IsStreamT is not None
18+
assert set(ProcessStateT.__args__) == {"running", "locked", "checkpointed", "failed"}
1719

1820

1921
def test_typing_matches_private_definitions():
@@ -23,7 +25,9 @@ def test_typing_matches_private_definitions():
2325
from cuda.core.typing import (
2426
DevicePointerT,
2527
IsStreamT,
28+
ProcessStateT,
2629
)
2730

2831
assert DevicePointerT is _DevicePointerT
2932
assert IsStreamT is _IsStreamT
33+
assert set(ProcessStateT.__args__) == {"running", "locked", "checkpointed", "failed"}

0 commit comments

Comments
 (0)