feat: CDI-based NVIDIA GPU passthrough into gVisor actor containers

cmd/atecontroller/internal/controllers/workerpool_apply.go

@@ -15,6 +15,8 @@
 package controllers
 
 import (
+	"os"
+
 	corev1 "k8s.io/api/core/v1"
 	appsv1ac "k8s.io/client-go/applyconfigurations/apps/v1"
 	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
@@ -145,6 +147,7 @@ func buildDeploymentApplyConfig(wp *atev1alpha1.WorkerPool, otel ateomOTelSettin
 
 	applyWorkerPoolPodTemplate(podSpecAC, containerAC, wp.Spec.Template)
 	maybeApplyMicroVMPodShape(podSpecAC, containerAC, wp.Spec.SandboxClass)
+	maybeApplyGPUPodShape(podSpecAC, containerAC, wp.Spec.Template, wp.Spec.SandboxClass)
 	podSpecAC.WithContainers(containerAC)
 	podSpecAC.WithTerminationGracePeriodSeconds(workerTerminationGracePeriodSeconds)
 
@@ -221,7 +224,8 @@ func fieldRefEnv(name, fieldPath string) *corev1ac.EnvVarApplyConfiguration {
 // veth and nftables rules (NET_ADMIN/NET_RAW), and the OCI rootfs is unpacked
 // and device nodes created as root over image-owned trees
 // (DAC_OVERRIDE/FOWNER/CHOWN/MKNOD). This replaces the former privileged worker;
-// the default seccomp and AppArmor profiles are sufficient (no Unconfined).
+// seccomp stays at the runtime default, but AppArmor must be Unconfined (see
+// ateomSecurityContext) since runsc's own mounts trip the default profile.
 var ateomGvisorCapabilities = []corev1.Capability{
 	"NET_ADMIN", "SYS_ADMIN", "SYS_CHROOT", "SYS_PTRACE",
 	"SETUID", "SETGID", "SETPCAP", "DAC_OVERRIDE",
@@ -301,6 +305,92 @@ func maybeApplyMicroVMPodShape(
 		WithEffect(corev1.TaintEffectNoSchedule))
 }
 
+// nvidiaToolkitContainerPath is where the host toolkit is mounted inside the
+// worker; ateom-gvisor's toolkitDir must match this. It sits outside
+// /usr/local/nvidia because the GPU device plugin mounts that tree into the
+// container read-only, and a mount cannot create its own mountpoint there, so
+// mounting under it only works when the toolkit happens to live inside the
+// directory the plugin mounts.
+const nvidiaToolkitContainerPath = "/opt/nvidia-toolkit"
+
+// defaultNvidiaToolkitHostPath is where gpu-operator installs the toolkit
+// (toolkit.installDir defaults to /usr/local/nvidia). It is deliberately not the
+// container path above: the two are independent, since where the node keeps the
+// toolkit says nothing about where we can mount it.
+const defaultNvidiaToolkitHostPath = "/usr/local/nvidia/toolkit"
+
+// nvidiaDriverRootEnv names the directory the GPU device plugin mounts the driver
+// into a pod at. ateom derives the driver library and binary paths from it, both of
+// which nvidia-ctk needs to generate a CDI spec. Only set it when the cluster's
+// device plugin does not use the /usr/local/nvidia convention; the controller
+// forwards its own value onto GPU worker pods.
+const nvidiaDriverRootEnv = "ATE_NVIDIA_DRIVER_ROOT"
+
+// nvidiaToolkitHostPath is where the NVIDIA container toolkit lives on the node.
+// It is platform-specific: gpu-operator and EKS install it at
+// /usr/local/nvidia/toolkit, while GKE keeps NVIDIA assets under
+// /home/kubernetes/bin/nvidia, so it is overridable via the
+// ATE_NVIDIA_TOOLKIT_HOST_PATH env var on the controller. We mount it read-only
+// so nvidia-ctk / nvidia-cdi-hook match whatever toolkit/driver the cluster runs.
+var nvidiaToolkitHostPath = envOrDefault("ATE_NVIDIA_TOOLKIT_HOST_PATH", defaultNvidiaToolkitHostPath)
+
+func envOrDefault(key, def string) string {
+	if v := os.Getenv(key); v != "" {
+		return v
+	}
+	return def
+}
+
+// maybeApplyGPUPodShape shapes a gVisor worker pod that requests a GPU so ateom
+// can inject the GPU into actors via CDI. It mounts the host NVIDIA toolkit
+// (version-matched to the node) read-only, for the glibc-based ateom image to run
+// directly. The pod keeps the same security posture as any other gVisor worker.
+// No-op for non-GPU pools and non-gVisor classes; an empty class defaults to
+// gVisor (WorkerPoolSpec kubebuilder default).
+func maybeApplyGPUPodShape(
+	podSpecAC *corev1ac.PodSpecApplyConfiguration,
+	containerAC *corev1ac.ContainerApplyConfiguration,
+	tmpl *atev1alpha1.WorkerPoolPodTemplate,
+	sandboxClass atev1alpha1.SandboxClass,
+) {
+	if sandboxClass != atev1alpha1.SandboxClassGvisor && sandboxClass != "" {
+		return
+	}
+	if !templateRequestsGPU(tmpl) {
+		return
+	}
+	// Mount the host NVIDIA toolkit (version-matched to the node) read-only.
+	containerAC.WithVolumeMounts(corev1ac.VolumeMount().
+		WithName("nvidia-toolkit").
+		WithMountPath(nvidiaToolkitContainerPath).
+		WithReadOnly(true))
+	podSpecAC.WithVolumes(corev1ac.Volume().
+		WithName("nvidia-toolkit").
+		WithHostPath(corev1ac.HostPathVolumeSource().
+			WithPath(nvidiaToolkitHostPath).
+			WithType(corev1.HostPathDirectory)))
+	// Only propagated when set, so a default deployment adds no env to worker pods.
+	if root := os.Getenv(nvidiaDriverRootEnv); root != "" {
+		containerAC.WithEnv(corev1ac.EnvVar().WithName(nvidiaDriverRootEnv).WithValue(root))
+	}
+}
+
+// templateRequestsGPU reports whether the pool template requests one or more
+// nvidia.com/gpu devices (limits or requests).
+func templateRequestsGPU(tmpl *atev1alpha1.WorkerPoolPodTemplate) bool {
+	if tmpl == nil || tmpl.Resources == nil {
+		return false
+	}
+	const gpu = corev1.ResourceName("nvidia.com/gpu")
+	if q, ok := tmpl.Resources.Limits[gpu]; ok && !q.IsZero() {
+		return true
+	}
+	if q, ok := tmpl.Resources.Requests[gpu]; ok && !q.IsZero() {
+		return true
+	}
+	return false
+}
+
 func applyWorkerPoolPodTemplate(
 	podSpecAC *corev1ac.PodSpecApplyConfiguration,
 	containerAC *corev1ac.ContainerApplyConfiguration,

cmd/atecontroller/internal/controllers/workerpool_apply_test.go

@@ -509,6 +509,172 @@ func envByName(env []corev1ac.EnvVarApplyConfiguration) map[string]envInfo {
 	return m
 }
 
+func TestGPUPoolMountsToolkit(t *testing.T) {
+	gpu := resource.MustParse("1")
+	wp := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{Name: "wp", Namespace: "ns"},
+		Spec: atev1alpha1.WorkerPoolSpec{
+			AteomImage: "img",
+			Template: &atev1alpha1.WorkerPoolPodTemplate{
+				Resources: &corev1.ResourceRequirements{
+					Limits: corev1.ResourceList{"nvidia.com/gpu": gpu},
+				},
+			},
+		},
+	}
+	dep := buildDeploymentApplyConfig(wp, ateomOTelSettings{})
+	pod := dep.Spec.Template.Spec
+
+	var found bool
+	for _, v := range pod.Volumes {
+		if v.Name != nil && *v.Name == "nvidia-toolkit" {
+			found = true
+			if v.HostPath == nil || *v.HostPath.Path != defaultNvidiaToolkitHostPath {
+				t.Fatalf("nvidia-toolkit volume has wrong hostPath: %+v", v.HostPath)
+			}
+		}
+	}
+	if !found {
+		t.Fatal("expected nvidia-toolkit host mount on a GPU pool")
+	}
+
+	var mounted bool
+	for _, c := range pod.Containers {
+		for _, m := range c.VolumeMounts {
+			if m.Name != nil && *m.Name == "nvidia-toolkit" && *m.MountPath == nvidiaToolkitContainerPath {
+				mounted = true
+			}
+		}
+	}
+	if !mounted {
+		t.Fatal("expected nvidia-toolkit mount on the ateom container")
+	}
+
+	// A GPU pool keeps the same posture as any other unprivileged gVisor worker: no
+	// user namespace and no unmasked /proc, which the skipped update-ldcache hook
+	// would otherwise force.
+	if pod.HostUsers != nil {
+		t.Error("did not expect hostUsers to be set on a GPU pool")
+	}
+	for _, c := range pod.Containers {
+		if c.SecurityContext != nil && c.SecurityContext.ProcMount != nil {
+			t.Errorf("did not expect procMount to be set, got %v", *c.SecurityContext.ProcMount)
+		}
+	}
+}
+
+// TestGPUPoolDriverRootEnv covers the override reaching the worker: ateom derives the
+// driver library and binary paths from it, and nvidia-ctk cannot generate a CDI spec
+// without them. Unset, no env is added at all.
+func TestGPUPoolDriverRootEnv(t *testing.T) {
+	gpu := resource.MustParse("1")
+	newGPUPool := func() *atev1alpha1.WorkerPool {
+		return &atev1alpha1.WorkerPool{
+			ObjectMeta: metav1.ObjectMeta{Name: "wp", Namespace: "ns"},
+			Spec: atev1alpha1.WorkerPoolSpec{
+				AteomImage: "img",
+				Template: &atev1alpha1.WorkerPoolPodTemplate{
+					Resources: &corev1.ResourceRequirements{
+						Limits: corev1.ResourceList{"nvidia.com/gpu": gpu},
+					},
+				},
+			},
+		}
+	}
+	driverRootEnv := func(wp *atev1alpha1.WorkerPool) (string, bool) {
+		for _, c := range buildDeploymentApplyConfig(wp, ateomOTelSettings{}).Spec.Template.Spec.Containers {
+			for _, e := range c.Env {
+				if e.Name != nil && *e.Name == nvidiaDriverRootEnv {
+					return *e.Value, true
+				}
+			}
+		}
+		return "", false
+	}
+
+	if v, ok := driverRootEnv(newGPUPool()); ok {
+		t.Errorf("unset: expected no %s on the worker, got %q", nvidiaDriverRootEnv, v)
+	}
+
+	t.Setenv(nvidiaDriverRootEnv, "/opt/nvidia")
+	v, ok := driverRootEnv(newGPUPool())
+	if !ok || v != "/opt/nvidia" {
+		t.Errorf("set: want %s=/opt/nvidia on the worker, got %q (present=%v)", nvidiaDriverRootEnv, v, ok)
+	}
+}
+
+func TestNonGPUPoolHasNoToolkit(t *testing.T) {
+	wp := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{Name: "wp", Namespace: "ns"},
+		Spec:       atev1alpha1.WorkerPoolSpec{AteomImage: "img"},
+	}
+	dep := buildDeploymentApplyConfig(wp, ateomOTelSettings{})
+	pod := dep.Spec.Template.Spec
+	for _, v := range pod.Volumes {
+		if v.Name != nil && *v.Name == "nvidia-toolkit" {
+			t.Fatal("non-GPU pool must not mount the toolkit")
+		}
+	}
+	// Non-GPU workers keep the tighter base posture: no user namespace, no
+	// unmasked /proc.
+	if pod.HostUsers != nil {
+		t.Error("non-GPU pool must not set hostUsers")
+	}
+	for _, c := range pod.Containers {
+		if c.SecurityContext != nil && c.SecurityContext.ProcMount != nil {
+			t.Error("non-GPU pool must not set procMount")
+		}
+	}
+}
+
+// TestGPUMicroVMPoolHasNoGPUPodShape asserts none of the GPU pod shaping is applied
+// to a non-gVisor pool: no toolkit volume, no toolkit mount, no driver-root env.
+//
+// A WorkerPool like this is rejected at apply time by the CEL rule on
+// WorkerPoolSpec, so it should never reach the controller. This covers the case
+// where one already exists — the rule was added after the fact, or the object was
+// written by a path that skipped CRD validation. The controller does not strip the
+// resource request itself, so such a pod still schedules onto a GPU node and holds a
+// device no actor can use; that gap is why the combination is rejected at the API
+// rather than only here.
+func TestGPUMicroVMPoolHasNoGPUPodShape(t *testing.T) {
+	// Set so the driver-root assertion below is not vacuous: a gVisor GPU pool would
+	// carry this env, a micro-VM one must not.
+	t.Setenv(nvidiaDriverRootEnv, "/opt/nvidia")
+	gpu := resource.MustParse("1")
+	wp := &atev1alpha1.WorkerPool{
+		ObjectMeta: metav1.ObjectMeta{Name: "wp", Namespace: "ns"},
+		Spec: atev1alpha1.WorkerPoolSpec{
+			AteomImage:   "img",
+			SandboxClass: atev1alpha1.SandboxClassMicroVM,
+			Template: &atev1alpha1.WorkerPoolPodTemplate{
+				Resources: &corev1.ResourceRequirements{
+					Limits: corev1.ResourceList{"nvidia.com/gpu": gpu},
+				},
+			},
+		},
+	}
+	pod := buildDeploymentApplyConfig(wp, ateomOTelSettings{}).Spec.Template.Spec
+
+	for _, v := range pod.Volumes {
+		if v.Name != nil && *v.Name == "nvidia-toolkit" {
+			t.Error("micro-VM pool must not mount the NVIDIA toolkit even when it requests a GPU")
+		}
+	}
+	for _, c := range pod.Containers {
+		for _, m := range c.VolumeMounts {
+			if m.Name != nil && *m.Name == "nvidia-toolkit" {
+				t.Error("micro-VM pool must not get the toolkit volume mount")
+			}
+		}
+		for _, e := range c.Env {
+			if e.Name != nil && *e.Name == nvidiaDriverRootEnv {
+				t.Errorf("micro-VM pool must not get %s", nvidiaDriverRootEnv)
+			}
+		}
+	}
+}
+
 func testWorkerPoolApplyConfig(tmpl *atev1alpha1.WorkerPoolPodTemplate) *atev1alpha1.WorkerPool {
 	return &atev1alpha1.WorkerPool{
 		ObjectMeta: metav1.ObjectMeta{Name: "pool", Namespace: "default", UID: "uid"},