Add CUDA process checkpointing helpers

Author: kkraus14

cuda_core/cuda/core/system/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
@@ -11,13 +11,22 @@
 
 __all__ = [
     "CUDA_BINDINGS_NVML_IS_COMPATIBLE",
+    "CudaCheckpointGpuPair",
+    "CudaProcessState",
+    "checkpoint_cuda_process",
+    "get_cuda_process_restore_thread_id",
+    "get_cuda_process_state",
     "get_driver_version",
     "get_driver_version_full",
     "get_num_devices",
     "get_process_name",
+    "lock_cuda_process",
+    "restore_cuda_process",
+    "unlock_cuda_process",
 ]
 
 
+from ._checkpoint import *
 from ._system import *
 
 if CUDA_BINDINGS_NVML_IS_COMPATIBLE:

cuda_core/cuda/core/system/_checkpoint.py

@@ -0,0 +1,210 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Any, Iterable
+
+
+class CudaProcessState(IntEnum):
+    """
+    CUDA checkpoint state for a process.
+    """
+
+    RUNNING = 0
+    LOCKED = 1
+    CHECKPOINTED = 2
+    FAILED = 3
+
+
+@dataclass(frozen=True)
+class CudaCheckpointGpuPair:
+    """
+    GPU UUID remapping pair used while restoring a checkpointed CUDA process.
+
+    Attributes
+    ----------
+    old_uuid
+        UUID of the GPU that was checkpointed.
+    new_uuid
+        UUID of the GPU to restore onto.
+    """
+
+    old_uuid: Any
+    new_uuid: Any
+
+
+def _get_driver():
+    try:
+        from cuda.bindings import driver
+    except ImportError:
+        from cuda import cuda as driver
+
+    required = (
+        "cuCheckpointProcessCheckpoint",
+        "cuCheckpointProcessGetRestoreThreadId",
+        "cuCheckpointProcessGetState",
+        "cuCheckpointProcessLock",
+        "cuCheckpointProcessRestore",
+        "cuCheckpointProcessUnlock",
+        "CUcheckpointGpuPair",
+        "CUcheckpointLockArgs",
+        "CUcheckpointRestoreArgs",
+    )
+    missing = [name for name in required if not hasattr(driver, name)]
+    if missing:
+        raise RuntimeError(
+            f"CUDA checkpointing requires cuda.bindings with CUDA checkpoint API support. Missing: {', '.join(missing)}"
+        )
+    return driver
+
+
+def _handle_return(result):
+    from cuda.core._utils.cuda_utils import handle_return
+
+    return handle_return(result)
+
+
+def _check_pid(pid: int) -> int:
+    if not isinstance(pid, int):
+        raise TypeError("pid must be an int")
+    if pid <= 0:
+        raise ValueError("pid must be a positive int")
+    return pid
+
+
+def _check_timeout_ms(timeout_ms: int) -> int:
+    if not isinstance(timeout_ms, int):
+        raise TypeError("timeout_ms must be an int")
+    if timeout_ms < 0:
+        raise ValueError("timeout_ms must be >= 0")
+    return timeout_ms
+
+
+def _make_restore_args(driver, gpu_pairs: Iterable[CudaCheckpointGpuPair] | None):
+    if gpu_pairs is None:
+        return None
+
+    pairs = []
+    for gpu_pair in gpu_pairs:
+        if isinstance(gpu_pair, driver.CUcheckpointGpuPair):
+            pair = gpu_pair
+        elif isinstance(gpu_pair, CudaCheckpointGpuPair):
+            pair = driver.CUcheckpointGpuPair()
+            pair.oldUuid = gpu_pair.old_uuid
+            pair.newUuid = gpu_pair.new_uuid
+        else:
+            raise TypeError(
+                "gpu_pairs must contain CudaCheckpointGpuPair or cuda.bindings.driver.CUcheckpointGpuPair objects"
+            )
+        pairs.append(pair)
+
+    if not pairs:
+        return None
+
+    args = driver.CUcheckpointRestoreArgs()
+    args.gpuPairs = pairs
+    args.gpuPairsCount = len(pairs)
+    return args
+
+
+def get_cuda_process_state(pid: int) -> CudaProcessState:
+    """
+    Return the CUDA checkpoint state for a process.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    """
+    driver = _get_driver()
+    state = _handle_return(driver.cuCheckpointProcessGetState(_check_pid(pid)))
+    return CudaProcessState(int(state))
+
+
+def get_cuda_process_restore_thread_id(pid: int) -> int:
+    """
+    Return the CUDA restore thread ID for a process.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    """
+    driver = _get_driver()
+    return _handle_return(driver.cuCheckpointProcessGetRestoreThreadId(_check_pid(pid)))
+
+
+def lock_cuda_process(pid: int, timeout_ms: int = 0) -> None:
+    """
+    Lock a running CUDA process, blocking further CUDA API calls.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    timeout_ms : int, optional
+        Timeout in milliseconds. A value of 0 indicates no timeout.
+    """
+    driver = _get_driver()
+    args = driver.CUcheckpointLockArgs()
+    args.timeoutMs = _check_timeout_ms(timeout_ms)
+    _handle_return(driver.cuCheckpointProcessLock(_check_pid(pid), args))
+
+
+def checkpoint_cuda_process(pid: int) -> None:
+    """
+    Checkpoint the GPU memory contents of a locked CUDA process.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    """
+    driver = _get_driver()
+    _handle_return(driver.cuCheckpointProcessCheckpoint(_check_pid(pid), None))
+
+
+def restore_cuda_process(pid: int, gpu_pairs: Iterable[CudaCheckpointGpuPair] | None = None) -> None:
+    """
+    Restore a checkpointed CUDA process.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    gpu_pairs : iterable of CudaCheckpointGpuPair, optional
+        GPU UUID remapping pairs. If provided, the array must contain every
+        checkpointed GPU.
+    """
+    driver = _get_driver()
+    args = _make_restore_args(driver, gpu_pairs)
+    _handle_return(driver.cuCheckpointProcessRestore(_check_pid(pid), args))
+
+
+def unlock_cuda_process(pid: int) -> None:
+    """
+    Unlock a locked CUDA process so it can resume CUDA API calls.
+
+    Parameters
+    ----------
+    pid : int
+        Process ID of the CUDA process.
+    """
+    driver = _get_driver()
+    _handle_return(driver.cuCheckpointProcessUnlock(_check_pid(pid), None))
+
+
+__all__ = [
+    "CudaCheckpointGpuPair",
+    "CudaProcessState",
+    "checkpoint_cuda_process",
+    "get_cuda_process_restore_thread_id",
+    "get_cuda_process_state",
+    "lock_cuda_process",
+    "restore_cuda_process",
+    "unlock_cuda_process",
+]

cuda_core/docs/source/api.rst

@@ -191,6 +191,12 @@ Basic functions
    system.get_process_name
    system.get_topology_common_ancestor
    system.get_p2p_status
+   system.get_cuda_process_state
+   system.get_cuda_process_restore_thread_id
+   system.lock_cuda_process
+   system.checkpoint_cuda_process
+   system.restore_cuda_process
+   system.unlock_cuda_process
 
 Events
 ``````
@@ -227,13 +233,18 @@ Enums
    system.TemperatureThresholds
    system.ThermalController
    system.ThermalTarget
+   system.CudaProcessState
 
 Types
 `````
 
 .. autosummary::
    :toctree: generated/
 
+   :template: dataclass.rst
+
+   system.CudaCheckpointGpuPair
+
    :template: autosummary/cyclass.rst
 
    system.Device

cuda_core/docs/source/release/1.0.0-notes.rst

@@ -16,7 +16,9 @@ Highlights
 New features
 ------------
 
-- TBD
+- Added CUDA process checkpointing helpers to :mod:`cuda.core.system`, including
+  process state queries, lock/checkpoint/restore/unlock operations, and GPU UUID
+  remapping support for restore. (`#1343 <https://github.com/NVIDIA/cuda-python/issues/1343>`__)
 
 
 Fixes and enhancements