[WIP] feat: GPU passthrough for actors via gVisor nvproxy + cuda-checkpoint

cmd/ateapi/internal/controlapi/gpu.go

@@ -0,0 +1,27 @@
+//  Copyright 2026 Google LLC
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+
+package controlapi
+
+import (
+	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
+	"github.com/agent-substrate/substrate/pkg/api/v1alpha1"
+)
+
+func toAteletGpuSpec(r *v1alpha1.ContainerResources) *ateletpb.GpuSpec {
+	if r == nil || r.GPU == nil {
+		return nil
+	}
+	g := r.GPU
+	return &ateletpb.GpuSpec{
+		Count:              g.Count,
+		Device:             g.Device,
+		DriverCapabilities: g.DriverCapabilities,
+		DriverVersion:      g.DriverVersion,
+	}
+}

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -179,6 +179,7 @@ func (s *CallAteletRestoreStep) Execute(ctx context.Context, input *ResumeInput,
 			Name:    ctr.Name,
 			Image:   ctr.Image,
 			Command: ctr.Command,
+			Gpu:     toAteletGpuSpec(ctr.Resources),
 		}
 		for _, env := range ctr.Env {
 			ateletEnv := &ateletpb.EnvEntry{

cmd/ateapi/internal/controlapi/workflow_suspend.go

@@ -152,6 +152,7 @@ func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput
 			Name:    ctr.Name,
 			Image:   ctr.Image,
 			Command: ctr.Command,
+			Gpu:     toAteletGpuSpec(ctr.Resources),
 		}
 		for _, env := range ctr.Env {
 			ateletEnv := &ateletpb.EnvEntry{

cmd/atelet/main.go

@@ -395,6 +395,7 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele
 				"io.kubernetes.cri.container-name": "pause",
 			},
 			netnsPath,
+			firstGpuSpec(req.GetSpec().GetContainers()),
 		); err != nil {
 			return fmt.Errorf("while creating pause OCI bundle: %w", err)
 		}
@@ -426,6 +427,7 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele
 					"io.kubernetes.cri.container-name": ctr.GetName(),
 				},
 				netnsPath,
+				ctr.GetGpu(),
 			); err != nil {
 				return fmt.Errorf("while creating %q OCI bundle: %w", ctr.GetName(), err)
 			}
@@ -456,6 +458,7 @@ func (s *AteomHerder) Run(ctx context.Context, req *ateletpb.RunRequest) (*atele
 	for _, ctr := range req.GetSpec().GetContainers() {
 		ateomCtr := &ateompb.Container{
 			Name: ctr.GetName(),
+			Gpu:  toAteomGpuSpec(ctr.GetGpu()),
 		}
 		ateomReq.GetSpec().Containers = append(ateomReq.GetSpec().Containers, ateomCtr)
 	}
@@ -511,6 +514,7 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe
 	for _, ctr := range req.GetSpec().GetContainers() {
 		ateomCtr := &ateompb.Container{
 			Name: ctr.GetName(),
+			Gpu:  toAteomGpuSpec(ctr.GetGpu()),
 		}
 		ateomReq.GetSpec().Containers = append(ateomReq.GetSpec().Containers, ateomCtr)
 	}
@@ -644,6 +648,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 				"io.kubernetes.cri.container-name": "pause",
 			},
 			netnsPath,
+			firstGpuSpec(req.GetSpec().GetContainers()),
 		); err != nil {
 			return fmt.Errorf("while creating pause OCI bundle: %w", err)
 		}
@@ -675,6 +680,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 					"io.kubernetes.cri.container-name": ctr.GetName(),
 				},
 				netnsPath,
+				ctr.GetGpu(),
 			); err != nil {
 				return fmt.Errorf("while creating %q OCI bundle: %w", ctr.GetName(), err)
 			}
@@ -705,6 +711,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 	for _, ctr := range req.GetSpec().GetContainers() {
 		ateomCtr := &ateompb.Container{
 			Name: ctr.GetName(),
+			Gpu:  toAteomGpuSpec(ctr.GetGpu()),
 		}
 		ateomReq.GetSpec().Containers = append(ateomReq.GetSpec().Containers, ateomCtr)
 	}
@@ -716,6 +723,27 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 	return &ateletpb.RestoreResponse{}, nil
 }
 
+func toAteomGpuSpec(g *ateletpb.GpuSpec) *ateompb.GpuSpec {
+	if g == nil {
+		return nil
+	}
+	return &ateompb.GpuSpec{
+		Count:              g.GetCount(),
+		Device:             g.GetDevice(),
+		DriverCapabilities: g.GetDriverCapabilities(),
+		DriverVersion:      g.GetDriverVersion(),
+	}
+}
+
+func firstGpuSpec(containers []*ateletpb.Container) *ateletpb.GpuSpec {
+	for _, c := range containers {
+		if g := c.GetGpu(); g != nil {
+			return g
+		}
+	}
+	return nil
+}
+
 type AteomDialer struct {
 	conns *lru.Cache
 }

cmd/atelet/oci.go

@@ -28,12 +28,14 @@ import (
 
 	"github.com/agent-substrate/substrate/internal/ateompath"
 	"github.com/agent-substrate/substrate/internal/memorypullcache"
+	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
 	"github.com/opencontainers/runtime-spec/specs-go"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
+	"golang.org/x/sys/unix"
 )
 
-func prepareOCIDirectory(ctx context.Context, pullCache *memorypullcache.MemoryPullCache, actorTemplateNamespace, actorTemplateName, actorID, containerName, ref string, args []string, env []string, annotations map[string]string, netns string) error {
+func prepareOCIDirectory(ctx context.Context, pullCache *memorypullcache.MemoryPullCache, actorTemplateNamespace, actorTemplateName, actorID, containerName, ref string, args []string, env []string, annotations map[string]string, netns string, gpu *ateletpb.GpuSpec) error {
 	tracer := otel.Tracer("prepareOCIDirectory")
 
 	ctx, span := tracer.Start(ctx, "prepareOCIDirectory")
@@ -162,6 +164,15 @@ func prepareOCIDirectory(ctx context.Context, pullCache *memorypullcache.MemoryP
 		},
 		Annotations: annotations,
 	}
+
+	if gpu != nil {
+		if err := addGPUToOCISpec(ociSpec); err != nil {
+			return fmt.Errorf("while adding GPU passthrough to OCI spec: %w", err)
+		}
+		if err := injectNVIDIAAssetsIntoRootfs(ctx, rootPath); err != nil {
+			return fmt.Errorf("while injecting NVIDIA driver assets into rootfs: %w", err)
+		}
+	}
 	ociSpecBytes, err := json.MarshalIndent(ociSpec, "", "  ")
 	if err != nil {
 		return fmt.Errorf("while marshaling OCI spec: %w", err)
@@ -279,3 +290,168 @@ func untar(ctx context.Context, tarData io.Reader, rootPath string) error {
 
 	return nil
 }
+
+var gpuDevicePaths = []string{
+	"/dev/nvidiactl",
+	"/dev/nvidia-uvm",
+	"/dev/nvidia-uvm-tools",
+	"/dev/nvidia-modeset",
+}
+
+func addGPUToOCISpec(spec *specs.Spec) error {
+	paths := append([]string{}, gpuDevicePaths...)
+	entries, err := os.ReadDir("/dev")
+	if err != nil {
+		return fmt.Errorf("reading /dev: %w", err)
+	}
+	for _, e := range entries {
+		n := e.Name()
+		if len(n) > 6 && n[:6] == "nvidia" && n[6] >= '0' && n[6] <= '9' {
+			paths = append(paths, "/dev/"+n)
+		}
+	}
+	for _, p := range paths {
+		var st unix.Stat_t
+		if err := unix.Stat(p, &st); err != nil {
+			if errors.Is(err, os.ErrNotExist) {
+				continue
+			}
+			return fmt.Errorf("stat %s: %w", p, err)
+		}
+		major := int64(unix.Major(uint64(st.Rdev))) //nolint:gosec
+		minor := int64(unix.Minor(uint64(st.Rdev))) //nolint:gosec
+		mode := os.FileMode(st.Mode & 0o777)        //nolint:gosec
+		uid := st.Uid
+		gid := st.Gid
+		spec.Linux.Devices = append(spec.Linux.Devices, specs.LinuxDevice{
+			Path: p, Type: "c", Major: major, Minor: minor,
+			FileMode: &mode, UID: &uid, GID: &gid,
+		})
+		if spec.Linux.Resources == nil {
+			spec.Linux.Resources = &specs.LinuxResources{}
+		}
+		allow := true
+		access := "rwm"
+		spec.Linux.Resources.Devices = append(spec.Linux.Resources.Devices,
+			specs.LinuxDeviceCgroup{Allow: allow, Type: "c", Major: &major, Minor: &minor, Access: access},
+		)
+	}
+	for _, name := range []string{"cuda-checkpoint", "cuda-checkpoint-wrapper.sh"} {
+		dest := "/usr/local/bin/" + name
+		// Source paths: prefer the shared /run/ateom-gvisor/static-files
+		// drop (visible inside both atelet and kind-node), fall back to
+		// /usr/local/bin if the operator installed it system-wide.
+		candidates := []string{"/run/ateom-gvisor/static-files/" + name, dest}
+		for _, src := range candidates {
+			if _, err := os.Stat(src); err != nil {
+				continue
+			}
+			spec.Mounts = append(spec.Mounts, specs.Mount{
+				Destination: dest, Type: "bind", Source: src,
+				Options: []string{"ro", "bind"},
+			})
+			break
+		}
+	}
+	return nil
+}
+
+// nvidiaLibsStagingDir is where setup-host.sh stages the host's NVIDIA
+// driver libs (libcuda.so.<v>, libnvidia-ml.so.<v>, …) plus their
+// SONAME / dev symlinks. atelet copies them into each actor's rootfs at
+// sandbox-create time so the workload image doesn't have to bake them in.
+//
+// This is the substrate-side equivalent of what
+// `nvidia-container-cli configure --compute --utility --device=all` does
+// in the standard docker+nvidia-container-runtime flow. We replicate the
+// effect in Go rather than exec'ing nvidia-container-cli because atelet
+// runs on `distroless/static-debian13` and has no dynamic linker for the
+// `nvidia-container-cli` binary's libnvidia-container.so.1 dep.
+const nvidiaLibsStagingDir = "/run/ateom-gvisor/static-files/nvidia-libs"
+
+// rootfsNVIDIALibDest is where libcuda.so.<v> et al. need to land inside
+// the sandbox rootfs. /etc/ld.so.cache on every glibc distro searches
+// this path, so dlopen("libcuda.so.1") just works.
+const rootfsNVIDIALibDest = "/usr/lib/x86_64-linux-gnu"
+
+// injectNVIDIAAssetsIntoRootfs walks nvidiaLibsStagingDir and mirrors
+// every entry into <rootfs>/usr/lib/x86_64-linux-gnu — preserving
+// symlinks as symlinks and copying real files byte-for-byte. Hard-fails
+// if the staging dir is missing or empty so an operator misconfiguration
+// surfaces immediately instead of crashing inside the sandbox.
+func injectNVIDIAAssetsIntoRootfs(ctx context.Context, rootfsPath string) error {
+	tracer := otel.Tracer("prepareOCIDirectory")
+	_, span := tracer.Start(ctx, "injectNVIDIAAssetsIntoRootfs")
+	defer span.End()
+
+	entries, err := os.ReadDir(nvidiaLibsStagingDir)
+	if err != nil {
+		return fmt.Errorf("reading NVIDIA libs staging dir %q (run setup-host.sh): %w", nvidiaLibsStagingDir, err)
+	}
+	if len(entries) == 0 {
+		return fmt.Errorf("NVIDIA libs staging dir %q is empty — re-run setup-host.sh", nvidiaLibsStagingDir)
+	}
+
+	destDir := filepath.Join(rootfsPath, rootfsNVIDIALibDest)
+	if err := os.MkdirAll(destDir, 0o755); err != nil {
+		return fmt.Errorf("creating dest dir %q: %w", destDir, err)
+	}
+
+	var copied, linked int
+	for _, e := range entries {
+		name := e.Name()
+		src := filepath.Join(nvidiaLibsStagingDir, name)
+		dst := filepath.Join(destDir, name)
+
+		info, err := os.Lstat(src)
+		if err != nil {
+			return fmt.Errorf("lstat %q: %w", src, err)
+		}
+
+		_ = os.Remove(dst)
+
+		switch {
+		case info.Mode()&os.ModeSymlink != 0:
+			target, err := os.Readlink(src)
+			if err != nil {
+				return fmt.Errorf("readlink %q: %w", src, err)
+			}
+			if err := os.Symlink(target, dst); err != nil {
+				return fmt.Errorf("symlink %q -> %q: %w", dst, target, err)
+			}
+			linked++
+		case info.Mode().IsRegular():
+			if err := copyRegularFile(src, dst, info.Mode().Perm()); err != nil {
+				return fmt.Errorf("copy %q -> %q: %w", src, dst, err)
+			}
+			copied++
+		default:
+			return fmt.Errorf("unexpected file type in %q: %s", src, info.Mode())
+		}
+	}
+
+	slog.InfoContext(ctx, "Injected NVIDIA driver assets into rootfs",
+		slog.String("source", nvidiaLibsStagingDir),
+		slog.String("dest", destDir),
+		slog.Int("files_copied", copied),
+		slog.Int("symlinks_created", linked),
+	)
+	return nil
+}
+
+func copyRegularFile(src, dst string, mode os.FileMode) error {
+	in, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer in.Close()
+	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode)
+	if err != nil {
+		return err
+	}
+	if _, err := io.Copy(out, in); err != nil {
+		out.Close()
+		return err
+	}
+	return out.Close()
+}

cmd/ateom-gvisor/main.go

@@ -170,6 +170,15 @@ func initTracing(ctx context.Context) (*sdktrace.TracerProvider, error) {
 	return tp, nil
 }
 
+func firstGPUSpec(containers []*ateompb.Container) *ateompb.GpuSpec {
+	for _, c := range containers {
+		if g := c.GetGpu(); g != nil {
+			return g
+		}
+	}
+	return nil
+}
+
 // AteomService is a service for shepherding single microvm.
 type AteomService struct {
 	ateompb.UnimplementedAteomServer
@@ -291,6 +300,7 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload
 		actorTemplateNamespace: req.GetActorTemplateNamespace(),
 		actorTemplateName:      req.GetActorTemplateName(),
 		actorID:                req.GetActorId(),
+		gpu:                    firstGPUSpec(req.GetSpec().GetContainers()),
 	}
 
 	// Create and start pause container
@@ -338,13 +348,21 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 		actorTemplateNamespace: req.GetActorTemplateNamespace(),
 		actorTemplateName:      req.GetActorTemplateName(),
 		actorID:                req.GetActorId(),
+		gpu:                    firstGPUSpec(req.GetSpec().GetContainers()),
 	}
 
 	checkpointPath := ateompath.CheckpointDir(req.GetActorTemplateNamespace(), req.GetActorTemplateName(), req.GetActorId())
 	if err := os.MkdirAll(checkpointPath, 0o700); err != nil {
 		return nil, fmt.Errorf("while creating checkpoint directory: %w", err)
 	}
 
+	// Drain CUDA state out of nvproxy clients via cuda-checkpoint
+	// inside the supervisor container. Without this, runsc checkpoint
+	// returns "can't save with live nvproxy clients".
+	if err := rcmd.cmdDrainCUDA(ctx); err != nil {
+		return nil, fmt.Errorf("while draining CUDA: %w", err)
+	}
+
 	// Checkpoint pause container (root of the sandbox)
 	if err := rcmd.cmdCheckpoint(ctx, "pause", checkpointPath); err != nil {
 		return nil, fmt.Errorf("while checkpointing pause: %w", err)
@@ -469,6 +487,7 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore
 		actorTemplateNamespace: req.GetActorTemplateNamespace(),
 		actorTemplateName:      req.GetActorTemplateName(),
 		actorID:                req.GetActorId(),
+		gpu:                    firstGPUSpec(req.GetSpec().GetContainers()),
 	}
 
 	checkpointDir := ateompath.CheckpointDir(req.GetActorTemplateNamespace(), req.GetActorTemplateName(), req.GetActorId())
@@ -497,6 +516,12 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore
 		}
 	}
 
+	// CUDA was drained by cmdDrainCUDA before checkpoint; toggle it
+	// back to running now that the supervisor is restored.
+	if err := rcmd.cmdUntoggleCUDA(ctx); err != nil {
+		return nil, fmt.Errorf("while untoggling CUDA: %w", err)
+	}
+
 	s.actorLogger.EmitLifecycleLog("Actor restored", req.GetActorId(), req.GetActorTemplateName(), req.GetActorTemplateNamespace())
 
 	return &ateompb.RestoreWorkloadResponse{}, nil