Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 0 additions & 88 deletions cmd/atecontroller/internal/controllers/workerpool_apply.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
package controllers

import (
"os"
"slices"

corev1 "k8s.io/api/core/v1"
Expand Down Expand Up @@ -176,7 +175,6 @@ 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)

Expand Down Expand Up @@ -401,92 +399,6 @@ const (
tunDevicePath = "/dev/net/tun"
)

// 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
}

// addDeviceResourceLimits requests one unit of each named extended resource,
// merging into whatever limits the pod template already set.
func addDeviceResourceLimits(containerAC *corev1ac.ContainerApplyConfiguration, resourceNames ...string) {
Expand Down
166 changes: 0 additions & 166 deletions cmd/atecontroller/internal/controllers/workerpool_apply_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -657,172 +657,6 @@ 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{
WorkerImage: "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{
WorkerImage: "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{WorkerImage: "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{
WorkerImage: "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"},
Expand Down
Loading
Loading