diff --git a/README.md b/README.md index 0f3d9d2fa..9704e3e97 100644 --- a/README.md +++ b/README.md @@ -202,7 +202,7 @@ We provide several sample applications demonstrating Agent Substrate's capabilit 4. **[Secret Agent](demos/agent-secret/README.md)**: Highlights Substrate's "Zero-Idle" self-suspension and re-animation of volatile process memory. ### Documentation & Guides -* [API Configuration Guide](docs/api-guide.md): Detailed reference for configuring WorkerPools, ActorTemplates, Secrets, and Volumes. +* [API Configuration Guide](docs/api-guide.md): Detailed reference for configuring WorkerPools, ActorTemplates, Secrets, Volumes, and [GPU worker pools](docs/api-guide.md#gpu-worker-pools). * [Full CLI Documentation](cmd/kubectl-ate/README.md): Installation and usage for `kubectl-ate`. * [Glossary](docs/glossary.md): Core terms (Actor, ActorTemplate, WorkerPool, Worker, ate-api-server, atenet, atelet, ateom) and how they relate. * [Observability Guide](docs/observability.md): Guide to actor logging, metrics, and distributed tracing. @@ -216,7 +216,7 @@ We provide several sample applications demonstrating Agent Substrate's capabilit * `cmd/atelet`: A node-level DaemonSet that supervises physical worker pods, coordinates snapshotting, and manages state transfers. * `cmd/atecontroller`: A Kubernetes controller that reconciles WorkerPool and ActorTemplate custom resources. * `cmd/atenet`: A combined networking controller providing DNS, Envoy routing, and proxy sidecars. -* `cmd/ateom-gvisor`: An interior-pod helper running inside sandboxed worker pods to execute `runsc` checkpoint and restore commands. +* `cmd/ateom-gvisor`: An interior-pod helper running inside sandboxed worker pods to execute `runsc` checkpoint and restore commands, and to inject GPU devices (CDI) into actor containers on GPU pools. * `cmd/podcertcontroller`: A "polyfill" that provides Pod Certificate signers that will eventually ship in upstream Kubernetes (with different names). * `cmd/kubectl-ate`: A CLI tool for managing Agent Substrate resources. See its [README](cmd/kubectl-ate/README.md). diff --git a/cmd/atecontroller/internal/controllers/workerpool_apply.go b/cmd/atecontroller/internal/controllers/workerpool_apply.go index 5b35a291a..7bd9f6210 100644 --- a/cmd/atecontroller/internal/controllers/workerpool_apply.go +++ b/cmd/atecontroller/internal/controllers/workerpool_apply.go @@ -61,6 +61,7 @@ func buildDeploymentApplyConfig(wp *atev1alpha1.WorkerPool) *appsv1ac.Deployment applyWorkerPoolPodTemplate(podSpecAC, containerAC, wp.Spec.Template) maybeApplyMicroVMPodShape(podSpecAC, containerAC, wp.Spec.SandboxClass) + maybeApplyGPUPodShape(podSpecAC, containerAC, wp.Spec.Template, wp.Spec.SandboxClass) podSpecAC.WithContainers(containerAC) return appsv1ac.Deployment(wp.Name, wp.Namespace). @@ -129,6 +130,58 @@ func maybeApplyMicroVMPodShape( WithEffect(corev1.TaintEffectNoSchedule)) } +// nvidiaToolkitHostPath is where the NVIDIA container toolkit is installed on GPU +// nodes and where we mount it into the worker pod. ateom-gvisor reads nvidia-ctk and +// nvidia-cdi-hook from this same path (see its toolkitDir constant) — the two must +// match. +const nvidiaToolkitHostPath = "/usr/local/nvidia/toolkit" + +// maybeApplyGPUPodShape mounts the host NVIDIA container toolkit into the worker +// pod when a gVisor pool requests a GPU. ateom-gvisor uses nvidia-ctk (generate) +// and nvidia-cdi-hook (runtime CDI hooks) from this mount to inject the GPU into +// the actor's gVisor sandbox. No-op for non-GPU pools and for non-gVisor classes: +// the toolkit is specific to the CDI/nvproxy path, and micro-VM GPU would instead +// be VFIO PCI passthrough into a guest that carries its own driver. 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 + } + containerAC.WithVolumeMounts(corev1ac.VolumeMount(). + WithName("nvidia-toolkit"). + WithMountPath(nvidiaToolkitHostPath). + WithReadOnly(true)) + podSpecAC.WithVolumes(corev1ac.Volume(). + WithName("nvidia-toolkit"). + WithHostPath(corev1ac.HostPathVolumeSource(). + WithPath(nvidiaToolkitHostPath). + WithType(corev1.HostPathDirectory))) +} + +// 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, diff --git a/cmd/atecontroller/internal/controllers/workerpool_apply_test.go b/cmd/atecontroller/internal/controllers/workerpool_apply_test.go index 19f6d44d9..dee6220ac 100644 --- a/cmd/atecontroller/internal/controllers/workerpool_apply_test.go +++ b/cmd/atecontroller/internal/controllers/workerpool_apply_test.go @@ -261,6 +261,86 @@ func TestMicroVMPodShape(t *testing.T) { } } +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) + 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 != "/usr/local/nvidia/toolkit" { + t.Fatalf("nvidia-toolkit volume has wrong hostPath: %+v", v.HostPath) + } + } + } + if !found { + t.Fatal("expected nvidia-toolkit volume 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 == "/usr/local/nvidia/toolkit" { + mounted = true + } + } + } + if !mounted { + t.Fatal("expected nvidia-toolkit mount on the ateom container") + } +} + +func TestNonGPUPoolHasNoToolkit(t *testing.T) { + wp := &atev1alpha1.WorkerPool{ + ObjectMeta: metav1.ObjectMeta{Name: "wp", Namespace: "ns"}, + Spec: atev1alpha1.WorkerPoolSpec{AteomImage: "img"}, + } + dep := buildDeploymentApplyConfig(wp) + for _, v := range dep.Spec.Template.Spec.Volumes { + if v.Name != nil && *v.Name == "nvidia-toolkit" { + t.Fatal("non-GPU pool must not mount the toolkit") + } + } +} + +// TestGPUMicroVMPoolHasNoToolkit asserts the NVIDIA toolkit mount is gVisor-only: +// a micro-VM pool that requests a GPU must not get it (kata GPU is VFIO +// passthrough into the guest, which needs no host toolkit). +func TestGPUMicroVMPoolHasNoToolkit(t *testing.T) { + 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}, + }, + }, + }, + } + dep := buildDeploymentApplyConfig(wp) + for _, v := range dep.Spec.Template.Spec.Volumes { + if v.Name != nil && *v.Name == "nvidia-toolkit" { + t.Fatal("micro-VM pool must not mount the NVIDIA toolkit even when it requests a GPU") + } + } +} + func testWorkerPoolApplyConfig(tmpl *atev1alpha1.WorkerPoolPodTemplate) *atev1alpha1.WorkerPool { return &atev1alpha1.WorkerPool{ ObjectMeta: metav1.ObjectMeta{Name: "pool", Namespace: "default", UID: "uid"}, diff --git a/cmd/ateom-gvisor/gpu.go b/cmd/ateom-gvisor/gpu.go new file mode 100644 index 000000000..60cc68f6d --- /dev/null +++ b/cmd/ateom-gvisor/gpu.go @@ -0,0 +1,258 @@ +//go:build linux + +// 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 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + + specs "github.com/opencontainers/runtime-spec/specs-go" + + "github.com/agent-substrate/substrate/internal/ateompath" +) + +const ( + gpuSentinelDev = "/dev/nvidia0" + toolkitDir = "/usr/local/nvidia/toolkit" + cdiOutputDir = "/run/ate-cdi" + procMountsPath = "/proc/mounts" +) + +// cdiSpec is the minimal CDI spec shape we need. nvidia-ctk writes JSON with +// --format=json so no YAML library is required. +type cdiSpec struct { + Devices []struct { + Name string `json:"name"` + ContainerEdits cdiEdits `json:"containerEdits"` + } `json:"devices"` + ContainerEdits cdiEdits `json:"containerEdits"` +} + +type cdiEdits struct { + Env []string `json:"env,omitempty"` + DeviceNodes []cdiDev `json:"deviceNodes,omitempty"` + Hooks []cdiHook `json:"hooks,omitempty"` + Mounts []cdiMount `json:"mounts,omitempty"` +} + +type cdiDev struct { + Path string `json:"path"` + Type string `json:"type,omitempty"` + Major int64 `json:"major,omitempty"` + Minor int64 `json:"minor,omitempty"` + FileMode *os.FileMode `json:"fileMode,omitempty"` + UID *uint32 `json:"uid,omitempty"` + GID *uint32 `json:"gid,omitempty"` +} + +type cdiHook struct { + HookName string `json:"hookName"` + Path string `json:"path"` + Args []string `json:"args,omitempty"` + Env []string `json:"env,omitempty"` + Timeout *int `json:"timeout,omitempty"` +} + +type cdiMount struct { + HostPath string `json:"hostPath"` + ContainerPath string `json:"containerPath"` + Type string `json:"type,omitempty"` + Options []string `json:"options,omitempty"` +} + +// gpuPresent reports whether a GPU is available to this worker pod. +func gpuPresent(sentinelPath string) bool { + _, err := os.Stat(sentinelPath) + return err == nil +} + +// enforceCDIMode fails fast when the cluster injected the GPU via the legacy NVIDIA +// runtime rather than CDI. The legacy runtime overmounts /proc/driver/nvidia, which +// trips the kernel's mount_too_revealing() when runsc runs the CDI update-ldcache +// hook in the deprivileged gofer; CDI mode never creates that overmount. +func enforceCDIMode(procMountsFile string) error { + data, err := os.ReadFile(procMountsFile) + if err != nil { + return fmt.Errorf("reading %s: %w", procMountsFile, err) + } + for _, line := range strings.Split(string(data), "\n") { + fields := strings.Fields(line) + if len(fields) >= 2 && fields[1] == "/proc/driver/nvidia" { + return fmt.Errorf("GPU injection requires the cluster in CDI mode; detected legacy /proc/driver/nvidia overmount") + } + } + return nil +} + +var ( + generateOnce sync.Once + generateErr error +) + +// generateCDISpec runs nvidia-ctk (from the mounted host toolkit) to produce a CDI +// spec scoped to this pod's assigned GPU. Runs under reapLock like every other +// subprocess in this process (a child reaper is running). +func generateCDISpec(ctx context.Context, ctkPath, hookPath, outDir string) error { + if err := os.MkdirAll(outDir, 0o755); err != nil { + return fmt.Errorf("creating CDI output dir %s: %w", outDir, err) + } + reapLock.RLock() + defer reapLock.RUnlock() + cmd := exec.CommandContext(ctx, ctkPath, "cdi", "generate", + "--format=json", + "--nvidia-cdi-hook-path="+hookPath, + "--output="+filepath.Join(outDir, "nvidia.json"), + ) + if out, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("nvidia-ctk cdi generate failed: %w: %s", err, out) + } + return nil +} + +// ensureCDISpec generates the per-pod CDI spec exactly once for this process. +func ensureCDISpec(ctx context.Context) error { + generateOnce.Do(func() { + generateErr = generateCDISpec(ctx, + filepath.Join(toolkitDir, "nvidia-ctk"), + filepath.Join(toolkitDir, "nvidia-cdi-hook"), + cdiOutputDir) + }) + return generateErr +} + +// maybeInjectGPU is a no-op unless the worker pod has a GPU. When it does, +// it enforces CDI mode, generates the per-pod CDI spec once, and injects the +// GPU into the actor container's OCI bundle before runsc create. +func maybeInjectGPU(ctx context.Context, actorUID, containerName string) error { + if !gpuPresent(gpuSentinelDev) { + return nil + } + slog.InfoContext(ctx, "Injecting GPU into actor container", slog.String("container", containerName)) + if err := enforceCDIMode(procMountsPath); err != nil { + return err + } + if err := ensureCDISpec(ctx); err != nil { + return err + } + bundleDir := ateompath.OCIBundlePath(actorUID, containerName) + if err := injectGPUIntoBundle(bundleDir, filepath.Join(cdiOutputDir, "nvidia.json")); err != nil { + return fmt.Errorf("injecting GPU into %q bundle: %w", containerName, err) + } + return nil +} + +// injectGPUIntoBundle reads the JSON CDI spec at cdiSpecPath and merges its +// device nodes, mounts, hooks, and env vars into the OCI config.json in bundleDir. +// The CDI spec format is JSON (nvidia-ctk --format=json) so encoding/json is +// sufficient; no CDI library is needed. +func injectGPUIntoBundle(bundleDir, cdiSpecPath string) error { + cdiData, err := os.ReadFile(cdiSpecPath) + if err != nil { + return fmt.Errorf("reading CDI spec %s: %w", cdiSpecPath, err) + } + var cdi cdiSpec + if err := json.Unmarshal(cdiData, &cdi); err != nil { + return fmt.Errorf("parsing CDI spec: %w", err) + } + + // Collect all edits: spec-level (common to all nvidia devices) + per-device. + var edits cdiEdits + edits.Env = append(edits.Env, cdi.ContainerEdits.Env...) + edits.DeviceNodes = append(edits.DeviceNodes, cdi.ContainerEdits.DeviceNodes...) + edits.Hooks = append(edits.Hooks, cdi.ContainerEdits.Hooks...) + edits.Mounts = append(edits.Mounts, cdi.ContainerEdits.Mounts...) + for _, d := range cdi.Devices { + edits.Env = append(edits.Env, d.ContainerEdits.Env...) + edits.DeviceNodes = append(edits.DeviceNodes, d.ContainerEdits.DeviceNodes...) + edits.Hooks = append(edits.Hooks, d.ContainerEdits.Hooks...) + edits.Mounts = append(edits.Mounts, d.ContainerEdits.Mounts...) + } + if len(edits.DeviceNodes) == 0 && len(edits.Mounts) == 0 && len(edits.Hooks) == 0 { + return fmt.Errorf("CDI spec %s resolved no devices", cdiSpecPath) + } + + cfgPath := filepath.Join(bundleDir, "config.json") + specData, err := os.ReadFile(cfgPath) + if err != nil { + return fmt.Errorf("reading %s: %w", cfgPath, err) + } + var spec specs.Spec + if err := json.Unmarshal(specData, &spec); err != nil { + return fmt.Errorf("parsing OCI spec: %w", err) + } + + if spec.Process != nil { + spec.Process.Env = append(spec.Process.Env, edits.Env...) + } + + if spec.Linux == nil { + spec.Linux = &specs.Linux{} + } + if spec.Linux.Resources == nil { + spec.Linux.Resources = &specs.LinuxResources{} + } + for _, dn := range edits.DeviceNodes { + major, minor := dn.Major, dn.Minor + // CDI leaves deviceNodes[].type unset for char devices (nvidia-ctk emits + // no "type" for /dev/nvidia*); default it to "c" as the CDI spec does, + // otherwise runsc rejects the OCI device with `invalid type ""`. + devType := dn.Type + if devType == "" { + devType = "c" + } + spec.Linux.Devices = append(spec.Linux.Devices, specs.LinuxDevice{ + Path: dn.Path, Type: devType, Major: major, Minor: minor, + FileMode: dn.FileMode, UID: dn.UID, GID: dn.GID, + }) + spec.Linux.Resources.Devices = append(spec.Linux.Resources.Devices, specs.LinuxDeviceCgroup{ + Allow: true, Type: devType, Major: &major, Minor: &minor, Access: "rwm", + }) + } + + for _, m := range edits.Mounts { + spec.Mounts = append(spec.Mounts, specs.Mount{ + Source: m.HostPath, Destination: m.ContainerPath, + Type: m.Type, Options: m.Options, + }) + } + + if spec.Hooks == nil { + spec.Hooks = &specs.Hooks{} + } + for _, h := range edits.Hooks { + hook := specs.Hook{Path: h.Path, Args: h.Args, Env: h.Env, Timeout: h.Timeout} + switch h.HookName { + case "createContainer": + spec.Hooks.CreateContainer = append(spec.Hooks.CreateContainer, hook) + case "startContainer": + spec.Hooks.StartContainer = append(spec.Hooks.StartContainer, hook) + } + } + + out, err := json.Marshal(&spec) + if err != nil { + return fmt.Errorf("serializing OCI spec: %w", err) + } + return os.WriteFile(cfgPath, out, 0o644) +} diff --git a/cmd/ateom-gvisor/gpu_test.go b/cmd/ateom-gvisor/gpu_test.go new file mode 100644 index 000000000..f956b96fc --- /dev/null +++ b/cmd/ateom-gvisor/gpu_test.go @@ -0,0 +1,176 @@ +//go:build linux + +// 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 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + specs "github.com/opencontainers/runtime-spec/specs-go" +) + +func TestMaybeInjectGPU_NoGPUIsNoop(t *testing.T) { + if _, err := os.Stat(gpuSentinelDev); err == nil { + t.Skip("host has a real GPU; skipping no-op assertion") + } + if err := maybeInjectGPU(context.Background(), "actor_uid", "c1"); err != nil { + t.Fatalf("expected no-op nil on non-GPU host, got %v", err) + } +} + +func TestGPUPresent(t *testing.T) { + dir := t.TempDir() + sentinel := filepath.Join(dir, "nvidia0") + if gpuPresent(sentinel) { + t.Fatal("expected absent before creation") + } + if err := os.WriteFile(sentinel, nil, 0o644); err != nil { + t.Fatal(err) + } + if !gpuPresent(sentinel) { + t.Fatal("expected present after creation") + } +} + +func TestEnforceCDIMode_LegacyOvermountFails(t *testing.T) { + dir := t.TempDir() + mounts := filepath.Join(dir, "mounts") + os.WriteFile(mounts, []byte("proc /proc proc rw 0 0\ntmpfs /proc/driver/nvidia tmpfs rw 0 0\n"), 0o644) + err := enforceCDIMode(mounts) + if err == nil || !strings.Contains(err.Error(), "CDI mode") { + t.Fatalf("expected CDI-mode error, got %v", err) + } +} + +func TestEnforceCDIMode_OK(t *testing.T) { + dir := t.TempDir() + mounts := filepath.Join(dir, "mounts") + os.WriteFile(mounts, []byte("proc /proc proc rw 0 0\ntmpfs /dev tmpfs rw 0 0\n"), 0o644) + if err := enforceCDIMode(mounts); err != nil { + t.Fatalf("expected nil, got %v", err) + } +} + +func TestGenerateCDISpec_InvokesCtk(t *testing.T) { + dir := t.TempDir() + ctk := filepath.Join(dir, "nvidia-ctk") + // fake nvidia-ctk that writes a minimal JSON spec to the path given by --output=. + const script = `#!/bin/sh +out="" +for a in "$@"; do + case "$a" in --output=*) out="${a#--output=}" ;; esac +done +printf '{"cdiVersion":"0.6.0","kind":"nvidia.com/gpu","devices":[{"name":"all","containerEdits":{"deviceNodes":[{"path":"/dev/nvidia0","type":"c","major":195,"minor":0}]}}]}' > "$out" +` + if err := os.WriteFile(ctk, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + out := filepath.Join(dir, "cdi") + if err := generateCDISpec(context.Background(), ctk, filepath.Join(dir, "nvidia-cdi-hook"), out); err != nil { + t.Fatalf("generate: %v", err) + } + data, err := os.ReadFile(filepath.Join(out, "nvidia.json")) + if err != nil || !strings.Contains(string(data), "nvidia.com/gpu") { + t.Fatalf("spec not written correctly: %q err=%v", data, err) + } +} + +func TestGenerateCDISpec_NonZeroFails(t *testing.T) { + dir := t.TempDir() + ctk := filepath.Join(dir, "nvidia-ctk") + os.WriteFile(ctk, []byte("#!/bin/sh\nexit 3\n"), 0o755) + err := generateCDISpec(context.Background(), ctk, "hook", filepath.Join(dir, "cdi")) + if err == nil { + t.Fatal("expected error on non-zero exit") + } +} + +func TestInjectGPUIntoBundle(t *testing.T) { + dir := t.TempDir() + + // Minimal JSON CDI spec with one device node and an env var. + specJSON := `{ + "cdiVersion": "0.6.0", + "kind": "nvidia.com/gpu", + "devices": [ + { + "name": "all", + "containerEdits": { + "deviceNodes": [{"path": "/dev/nvidia0", "major": 195, "minor": 0}], + "env": ["NVIDIA_TEST=1"] + } + } + ] +}` + cdiSpecPath := filepath.Join(dir, "nvidia.json") + os.WriteFile(cdiSpecPath, []byte(specJSON), 0o644) + + bundle := filepath.Join(dir, "bundle") + os.MkdirAll(bundle, 0o755) + base := &specs.Spec{Version: "1.0.0", Process: &specs.Process{Args: []string{"true"}}} + data, _ := json.Marshal(base) + os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) + + if err := injectGPUIntoBundle(bundle, cdiSpecPath); err != nil { + t.Fatalf("inject: %v", err) + } + + out, _ := os.ReadFile(filepath.Join(bundle, "config.json")) + var got specs.Spec + json.Unmarshal(out, &got) + + var hasDev bool + for _, d := range got.Linux.Devices { + if d.Path == "/dev/nvidia0" { + hasDev = true + // nvidia-ctk omits "type" for char devices; it must default to "c" + // or runsc rejects the device with `invalid type ""`. + if d.Type != "c" { + t.Fatalf("expected device type defaulted to \"c\", got %q", d.Type) + } + } + } + if !hasDev { + t.Fatalf("expected /dev/nvidia0 injected, spec=%s", out) + } + var hasEnv bool + for _, e := range got.Process.Env { + if e == "NVIDIA_TEST=1" { + hasEnv = true + } + } + if !hasEnv { + t.Fatalf("expected NVIDIA_TEST env injected, spec=%s", out) + } +} + +func TestInjectGPUIntoBundle_MissingSpecFails(t *testing.T) { + dir := t.TempDir() + bundle := filepath.Join(dir, "bundle") + os.MkdirAll(bundle, 0o755) + base := &specs.Spec{Version: "1.0.0", Process: &specs.Process{}} + data, _ := json.Marshal(base) + os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) + + if err := injectGPUIntoBundle(bundle, filepath.Join(dir, "nonexistent.json")); err == nil { + t.Fatal("expected error for missing CDI spec") + } +} diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index 4d245dc6f..81f9b1c72 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -238,6 +238,9 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload if err := imagecache.SetupBundleRootfs(ateompath.OCIBundlePath(req.GetActorUid(), ac.GetName())); err != nil { return nil, fmt.Errorf("while composing %q rootfs: %w", ac.GetName(), err) } + if err := maybeInjectGPU(ctx, req.GetActorUid(), ac.GetName()); err != nil { + return nil, fmt.Errorf("while injecting GPU for %q: %w", ac.GetName(), err) + } if err := rcmd.cmdCreate(ctx, pw, ac.GetName(), nil); err != nil { return nil, fmt.Errorf("while creating %q application container: %w", ac.GetName(), err) } @@ -451,6 +454,9 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore if err := imagecache.SetupBundleRootfs(ateompath.OCIBundlePath(req.GetActorUid(), ac.GetName())); err != nil { return nil, fmt.Errorf("while composing %q rootfs: %w", ac.GetName(), err) } + if err := maybeInjectGPU(ctx, req.GetActorUid(), ac.GetName()); err != nil { + return nil, fmt.Errorf("while injecting GPU for %q: %w", ac.GetName(), err) + } switch req.GetScope() { case ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA: if err := rcmd.cmdCreate(ctx, pw, ac.GetName(), nil); err != nil { diff --git a/docs/api-guide.md b/docs/api-guide.md index a67776c75..565efb9e9 100644 --- a/docs/api-guide.md +++ b/docs/api-guide.md @@ -43,7 +43,13 @@ spec: # gvisor SandboxConfig unless sandboxConfigName is set. ``` -### Example with GPU node scheduling +### GPU worker pools + +A GPU pool needs two things: (1) scheduling onto GPU nodes, and (2) a +`nvidia.com/gpu` request in `template.resources`. The request does double duty — +it makes the device plugin assign a GPU to the worker pod **and** triggers +Substrate to pass that GPU **through to each actor's sandbox**. No per-actor +configuration is needed. ```yaml apiVersion: ate.dev/v1alpha1 @@ -55,6 +61,7 @@ spec: replicas: 5 ateomImage: ko://github.com/agent-substrate/substrate/cmd/ateom-gvisor template: + # (1) schedule onto GPU nodes nodeSelector: cloud.google.com/gke-accelerator: nvidia-tesla-t4 tolerations: @@ -62,13 +69,6 @@ spec: operator: Exists effect: NoSchedule priorityClassName: substrate-workers - nodeAffinity: - requiredDuringSchedulingIgnoredDuringExecution: - nodeSelectorTerms: - - matchExpressions: - - key: workload - operator: In - values: [substrate] resources: requests: cpu: 500m @@ -76,8 +76,35 @@ spec: limits: cpu: "1" memory: 2Gi + # (2) claim a GPU — this request is what triggers GPU passthrough + nvidia.com/gpu: "1" ``` +What that `nvidia.com/gpu` request sets in motion: + +- **`atecontroller`** propagates the request onto the `ateom` container (so the + device plugin assigns a GPU and its `/dev/nvidia*` nodes appear in the worker + pod) and mounts the host NVIDIA container toolkit (`/usr/local/nvidia/toolkit`, + read-only) into the pod. +- **`ateom-gvisor`**, at actor create/restore, generates a per-pod CDI spec with + `nvidia-ctk` and injects the GPU (device nodes, driver libraries, env) into each + application container's OCI spec before `runsc create`. gVisor's `--nvproxy` + proxies the driver `ioctl`s to the host driver, so CUDA / NVML work inside the + sandbox. + +**Requirements & scope** + +- **CDI mode is required.** The cluster's NVIDIA stack must run in CDI mode (e.g. + gpu-operator with `cdi.default=true`). If `ateom` detects the legacy runtime's + `/proc/driver/nvidia` overmount it **fails the actor's container creation** with a + clear error rather than silently starting it CPU-only. +- **gVisor only.** Passthrough is implemented for the `gvisor` sandbox class. + `microvm` pools use a different mechanism (VFIO PCI passthrough) and do not yet + support it. +- **Non-GPU pools are unaffected** — the passthrough path is a no-op unless a GPU + is present in the worker pod, so pools without a `nvidia.com/gpu` request behave + exactly as before. + --- ## 2. ActorTemplate: The Workload Blueprint