nvproxy: implement GPU checkpoint/restore for single and multi-GPU containers

[lokashrinav]

Problem

gVisor's nvproxy proxies GPU ioctls between container processes and the host NVIDIA driver. Every save/restore method in nvproxy is panic("not implemented") — any checkpoint attempt on a container with live GPU state crashes the sentry.

Three things block checkpoint/restore:

1. Host file descriptors become stale after restore. When a container opens /dev/nvidia0, nvproxy opens the real device on the host and stores that FD. After checkpoint, the serialized FD number points at nothing — the device needs to be reopened on the new host, and every system that references that FD (the event notification layer, memory-mapped file handles) needs to be updated to use the new one.

2. GPU-backed memory regions can't be serialized. GPU device memory creates PMAs (physical memory areas) backed by frontendFDMemmapFile. Stateify's PMA serializer only handles *pgalloc.MemoryFile — it doesn't know how to save GPU device memory. These PMAs need to be dropped before save and re-created after restore.

3. Missing stateify annotations. Two nvproxy structs (miscObject, osDescMem) lack // +stateify savable, causing nil pointer panics during serialization. osDescMem also holds host memory addresses (pinnedRanges, m, len) that are invalid after restore.

For multi-GPU containers, there's an additional coordination problem: if multiple GPUs are communicating via NCCL and you try to lock them sequentially, you get dependency cycles — GPU 0 is frozen mid-transfer to GPU 1, so GPU 1 blocks waiting for that transfer and can never be locked. All GPUs must be locked atomically.

Solution

FD lifecycle (save_restore_impl.go): Replace all 6 panic methods with real implementations. afterLoadImpl reopens the host device files (via device gofer or direct Openat), re-registers with fdnotifier, and updates memmapFile. A restoreContext wrapper bridges goContext.Context to gVisor's extended context.Context (which requires Blocker). The waiter queue Entry is already in the Queue from checkpoint state — afterLoadImpl calls Init() to refresh the callback but does not call EventRegister, which would double-insert and create a linked list cycle.

PMA invalidation (save_restore.go): InvalidateUnsavable() now walks all PMAs and drops any not backed by *pgalloc.MemoryFile — unmaps from address space, removes RSS tracking, releases file references, removes the segment. After restore, access to these regions triggers a page fault → frontendFD.Translate() → mmap against the reopened device → PMA re-created on demand.

Stateify annotations (object.go): Add // +stateify savable to miscObject and osDescMem. Mark pinnedRanges, m, and len in osDescMem as state:"nosave".

SaveRestoreExec binary (gvisor-gpu-ckpt): Calls NVIDIA's cuda-checkpoint API — cuCheckpointProcessLock → cuCheckpointProcessCheckpoint on save, cuCheckpointProcessRestore → cuCheckpointProcessUnlock on restore. Uses dlopen("libcuda.so.1") at runtime. The binary runs inside the sandbox and targets PID 1 (container init), because the cuda-checkpoint API must route through nvproxy — from the host, cuCheckpointProcessLock(sentryPID) returns CUDA_ERROR_NOT_INITIALIZED since the sentry creates GPU contexts via raw ioctls without loading libcuda.so.

Multi-GPU: All GPUs in one container share one sentry PID. cuCheckpointProcessLock locks every GPU context owned by that PID atomically — one call, all GPUs, no sequential locking, no NCCL deadlock window.

Files changed

File	Change
pkg/sentry/devices/nvproxy/save_restore_impl.go	Replace 6 panics with FD lifecycle (reopen, fdnotifier, memmapFile)
pkg/sentry/mm/save_restore.go	Drop non-MemoryFile PMAs in InvalidateUnsavable()
pkg/sentry/devices/nvproxy/object.go	Add stateify annotations, nosave host-specific fields
cmd/gvisor-gpu-ckpt/main.go	SaveRestoreExec entry point, MODE dispatch, targets PID 1
cmd/gvisor-gpu-ckpt/cuda.go	cgo wrappers for cuCheckpointProcess* via dlopen

Test results

Tested on Lambda Labs bare-metal A10, driver 570.148.08, gVisor with nvproxy.

gVisor checkpoint: exit 0, 604KB kernel state + 9.2MB process memory serialized.

gVisor restore: state deserialized in 31ms, tasks resumed, GPU memory page faults handled via Translate path.

cuda-checkpoint API (from inside container, targeting PID 1):

cuCheckpointProcessLock(1)       = 0
cuCheckpointProcessCheckpoint(1) = 0
cuCheckpointProcessRestore(1)    = 0
cuCheckpointProcessUnlock(1)     = 0

CUDA process survived with 228 MiB GPU memory intact.

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.