From a194afded26939d789fac6bde076c6ec2cd82f67 Mon Sep 17 00:00:00 2001 From: Benjamin Elder Date: Tue, 1 Sep 2026 17:43:25 -0700 Subject: [PATCH 1/2] cdi: move the CDI handling out of ateom-gvisor containerd applies CDI to a pod's containers, and actor containers sit underneath one, so whatever gives an actor a device reads the spec itself. Split by how far it travels: internal/cdi parses a spec and resolves what a named set of devices asks for, which any sandbox can use, while cmd/ateom-gvisor/internal/cdiinject applies the result to a bundle on disk and stats the host for device numbers, which only suits a sandbox on the host kernel. The NVIDIA specifics -- which devices, which hooks may run, which library directories -- are caller options rather than package globals. No behavior change, and the tests run on any host rather than Linux only. --- cmd/ateom-gvisor/gpu.go | 377 ++---------------- cmd/ateom-gvisor/gpu_test.go | 282 ------------- cmd/ateom-gvisor/internal/cdiinject/inject.go | 256 ++++++++++++ .../internal/cdiinject/inject_test.go | 357 +++++++++++++++++ cmd/ateom-gvisor/internal/cdiinject/soname.go | 111 ++++++ internal/cdi/cdi.go | 125 ++++++ internal/cdi/cdi_test.go | 135 +++++++ 7 files changed, 1009 insertions(+), 634 deletions(-) create mode 100644 cmd/ateom-gvisor/internal/cdiinject/inject.go create mode 100644 cmd/ateom-gvisor/internal/cdiinject/inject_test.go create mode 100644 cmd/ateom-gvisor/internal/cdiinject/soname.go create mode 100644 internal/cdi/cdi.go create mode 100644 internal/cdi/cdi_test.go diff --git a/cmd/ateom-gvisor/gpu.go b/cmd/ateom-gvisor/gpu.go index 9925d9fa9f..db5abc5fae 100644 --- a/cmd/ateom-gvisor/gpu.go +++ b/cmd/ateom-gvisor/gpu.go @@ -19,20 +19,16 @@ package main import ( "cmp" "context" - "debug/elf" - "encoding/json" "fmt" "log/slog" "os" "os/exec" "path/filepath" - "strings" "sync" - specs "github.com/opencontainers/runtime-spec/specs-go" - "golang.org/x/sys/unix" - + "github.com/agent-substrate/substrate/cmd/ateom-gvisor/internal/cdiinject" "github.com/agent-substrate/substrate/internal/ateompath" + "github.com/agent-substrate/substrate/internal/cdi" ) // toolkitDir is where the host's NVIDIA container toolkit (nvidia-ctk, @@ -64,6 +60,11 @@ var ( const cdiOutputDir = "/run/ate-cdi" +// cdiAllDevice is the CDI device carrying every GPU assigned to the pod. +// nvidia-ctk repeats the same nodes across per-index and per-UUID devices, so +// only one of them may be applied. +const cdiAllDevice = "all" + // enabledCDIHooks is the set of CDI createContainer hooks the actor runs. It is an // allowlist because the toolkit is mounted from the host, so its version is the // cluster's choice: a newer one can emit hooks that have never been reviewed @@ -89,18 +90,6 @@ func toolkitBinary(name string) string { func fileExists(p string) bool { _, err := os.Stat(p); return err == nil } -// dropEnvVar returns env with every "KEY=..." entry for the given key removed. -func dropEnvVar(env []string, key string) []string { - prefix := key + "=" - out := make([]string, 0, len(env)) - for _, e := range env { - if !strings.HasPrefix(e, prefix) { - out = append(out, e) - } - } - return out -} - // gpuPresent reports whether any GPU is assigned to this worker pod, matching any // device index (not just /dev/nvidia0 — the device plugin can assign 2,3 etc.). func gpuPresent() bool { @@ -171,344 +160,28 @@ func maybeInjectGPU(ctx context.Context, actorUID, containerName string) error { return nil } -// cdiSpec is the minimal shape of the JSON CDI spec (nvidia-ctk --format=json) that -// we consume: the device nodes, driver-library mounts, and env. -type cdiSpec struct { - Devices []struct { - Name string `json:"name"` - ContainerEdits cdiEdits `json:"containerEdits"` - } `json:"devices"` - ContainerEdits cdiEdits `json:"containerEdits"` -} - -// cdiAllDevice is the CDI device that carries every GPU assigned to the pod. -// nvidia-ctk also emits per-index ("0") and per-UUID devices that repeat the same -// nodes, so we apply only this one (plus the spec-level edits) to avoid injecting -// each device node several times. -const cdiAllDevice = "all" - -type cdiEdits struct { - Env []string `json:"env,omitempty"` - DeviceNodes []cdiDev `json:"deviceNodes,omitempty"` - Mounts []cdiMount `json:"mounts,omitempty"` - Hooks []cdiHook `json:"hooks,omitempty"` -} - -type cdiDev struct { - Path string `json:"path"` - Type string `json:"type,omitempty"` - Major int64 `json:"major,omitempty"` - Minor int64 `json:"minor,omitempty"` -} - -type cdiMount struct { - HostPath string `json:"hostPath"` - ContainerPath string `json:"containerPath"` - Type string `json:"type,omitempty"` - Options []string `json:"options,omitempty"` -} - -type cdiHook struct { - HookName string `json:"hookName"` - Path string `json:"path"` - Args []string `json:"args,omitempty"` - Env []string `json:"env,omitempty"` -} - -// resolveDevNumbers fills a device node's major/minor from the host when the CDI spec -// omitted them. nvidia-ctk emits deviceNodes carrying only a path (CDI delegates -// number resolution to the OCI runtime, which stats the host); we merge into runsc's -// spec ourselves, so we stat here too, otherwise the actor gets bogus 0,0 char devices -// and NVML can't reach the driver. An nvidia device major is never 0. -func resolveDevNumbers(path string, major, minor int64) (int64, int64, error) { - if major != 0 { - return major, minor, nil - } - var st unix.Stat_t - if err := unix.Stat(path, &st); err != nil { - return 0, 0, fmt.Errorf("stat device %s: %w", path, err) - } - rdev := uint64(st.Rdev) - return int64(unix.Major(rdev)), int64(unix.Minor(rdev)), nil -} - -// injectGPUIntoBundle merges the CDI spec generated in cdiSpecDir into the actor's OCI -// config.json in bundleDir: device nodes (major/minor resolved from the host), the -// driver-library mounts, and env. It does NOT run the CDI hooks; instead it stages the -// SONAME symlinks (libcuda.so.1 -> libcuda.so.580.x) into the actor rootfs, which is -// what lets the GPU worker keep the plain unprivileged posture (no user namespace, no -// unmasked /proc). The CDI spec is plain JSON, so no CDI library is needed. +// injectGPUIntoBundle merges the CDI spec generated in cdiSpecDir into the actor's +// OCI bundle. It does NOT run the CDI cache-updating hook: that hook's ldconfig +// needs a private /proc, which the pod's masked /proc rejects, so the SONAME +// symlinks it would create are staged into the rootfs instead. That is what lets +// a GPU worker keep the plain unprivileged posture. func injectGPUIntoBundle(ctx context.Context, bundleDir, cdiSpecDir string) error { - cdiData, err := os.ReadFile(filepath.Join(cdiSpecDir, "nvidia.json")) + data, err := os.ReadFile(filepath.Join(cdiSpecDir, "nvidia.json")) if err != nil { return fmt.Errorf("reading CDI spec: %w", err) } - var cdi cdiSpec - if err := json.Unmarshal(cdiData, &cdi); err != nil { - return fmt.Errorf("parsing CDI spec: %w", err) - } - // Spec-level edits (driver libs, common device nodes, env) plus only the "all" - // device's edits — applying every device would inject each GPU node several times - // (nvidia-ctk repeats nodes across its per-index, per-UUID, and "all" devices). - edits := cdi.ContainerEdits - var foundAll bool - for _, d := range cdi.Devices { - if d.Name != cdiAllDevice { - continue - } - foundAll = true - edits.Env = append(edits.Env, d.ContainerEdits.Env...) - edits.DeviceNodes = append(edits.DeviceNodes, d.ContainerEdits.DeviceNodes...) - edits.Mounts = append(edits.Mounts, d.ContainerEdits.Mounts...) - edits.Hooks = append(edits.Hooks, d.ContainerEdits.Hooks...) - } - if !foundAll { - return fmt.Errorf("CDI spec in %s has no %q device", cdiSpecDir, cdiAllDevice) - } - if len(edits.DeviceNodes) == 0 { - return fmt.Errorf("CDI spec in %s resolved no devices", cdiSpecDir) - } - - cfgPath := filepath.Join(bundleDir, "config.json") - specData, err := os.ReadFile(cfgPath) + spec, err := cdi.Parse(data) 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) - } - - // The edits below append to the bundle's config.json, so injecting twice would - // double every device, mount, env entry and hook. atelet re-unpacks the bundle - // before each Run and Restore, so this normally runs once per bundle — but that - // invariant lives in another component, so enforce it here rather than rely on it. - if hasGPUDevice(spec.Linux) { - slog.InfoContext(ctx, "Bundle already has GPU devices; skipping injection", - slog.String("bundle", bundleDir)) - return nil - } - - 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, err := resolveDevNumbers(dn.Path, dn.Major, dn.Minor) - if err != nil { - return err - } - devType := dn.Type - if devType == "" { - devType = "c" // nvidia-ctk omits type for char devices; runsc needs it. - } - spec.Linux.Devices = append(spec.Linux.Devices, specs.LinuxDevice{ - Path: dn.Path, Type: devType, Major: major, Minor: minor, - }) - 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 { - mType := m.Type - if mType == "" { - mType = "bind" // CDI omits type for its bind mounts; runsc's gofer needs it. - } - spec.Mounts = append(spec.Mounts, specs.Mount{ - Source: m.HostPath, Destination: m.ContainerPath, Type: mType, Options: m.Options, - }) - } - - if spec.Process != nil { - spec.Process.Env = append(spec.Process.Env, edits.Env...) - // runsc's nvproxy runs nvidia-container-cli when it sees NVIDIA_VISIBLE_DEVICES - // (independent of --nvproxy); we set up the GPU via CDI, so strip it. - spec.Process.Env = dropEnvVar(spec.Process.Env, "NVIDIA_VISIBLE_DEVICES") - spec.Process.Env = prependLibraryPath(spec.Process.Env, []string{driverLibDir}) - } - - // Run the allowlisted CDI createContainer hooks (see enabledCDIHooks) from the - // mounted host toolkit. Anything else the toolkit emits is skipped and logged, - // so a toolkit upgrade that adds a hook is visible rather than silent. - if spec.Hooks == nil { - spec.Hooks = &specs.Hooks{} - } - for _, h := range edits.Hooks { - if h.HookName != "createContainer" || len(h.Args) < 2 { - continue - } - if !enabledCDIHooks[h.Args[1]] { - slog.InfoContext(ctx, "Skipping CDI hook outside the allowlist", - slog.String("hook", h.Args[1])) - continue - } - binary := toolkitBinary("nvidia-cdi-hook") - spec.Hooks.CreateContainer = append(spec.Hooks.CreateContainer, specs.Hook{ - Path: binary, - Args: append([]string{binary}, h.Args[1:]...), - Env: h.Env, - }) - } - - // Create the driver SONAME symlinks in the container's rootfs (spec.Root.Path, - // relative to the bundle). - rootfs := "rootfs" - if spec.Root != nil && spec.Root.Path != "" { - rootfs = spec.Root.Path - } - if !filepath.IsAbs(rootfs) { - rootfs = filepath.Join(bundleDir, rootfs) - } - if err := stageSonameSymlinks(ctx, rootfs, spec.Mounts); err != nil { - return fmt.Errorf("staging SONAME symlinks: %w", err) - } - - out, err := json.Marshal(&spec) - if err != nil { - return fmt.Errorf("serializing OCI spec: %w", err) - } - return os.WriteFile(cfgPath, out, 0o644) -} - -// hasGPUDevice reports whether the OCI spec already carries injected NVIDIA device -// nodes, which is how an already-injected bundle is recognized. -func hasGPUDevice(l *specs.Linux) bool { - if l == nil { - return false - } - for _, d := range l.Devices { - if strings.HasPrefix(d.Path, "/dev/nvidia") { - return true - } - } - return false -} - -// prependLibraryPath puts the driver library directories at the front of -// LD_LIBRARY_PATH, keeping whatever the image or ActorTemplate already set. The -// actor needs this to find libcuda.so.1: the update-ldcache hook that would -// normally add the directory to the loader cache is not run (see enabledCDIHooks), -// so an image that does not set LD_LIBRARY_PATH itself gets a CUDA runtime that -// reports zero devices even though the GPU is fully injected. Directories already -// on the path are left alone, so an NVIDIA base image keeps its own ordering. -// -// This is weaker than the ldcache the hook would have written — LD_LIBRARY_PATH is -// inherited by child processes and takes precedence over an executable's own -// DT_RUNPATH — so it carries only driverLibDir rather than every directory the CDI -// mounts touch, which would also sweep in the driver's X server modules. -func prependLibraryPath(env, dirs []string) []string { - if len(dirs) == 0 { - return env - } - existing := "" - for _, e := range env { - if v, ok := strings.CutPrefix(e, "LD_LIBRARY_PATH="); ok { - existing = v // OCI semantics: a later entry wins. - } - } - have := map[string]bool{} - for _, d := range strings.Split(existing, ":") { - have[d] = true - } - var add []string - for _, d := range dirs { - if !have[d] { - add = append(add, d) - } - } - if len(add) == 0 { - return env - } - val := strings.Join(add, ":") - if existing != "" { - val += ":" + existing - } - return append(dropEnvVar(env, "LD_LIBRARY_PATH"), "LD_LIBRARY_PATH="+val) -} - -// stageSonameSymlinks writes each driver library's SONAME symlink (e.g. -// libcuda.so.1 -> libcuda.so.580.65.06) into rootfs, so programs that link against -// the SONAME resolve it. For each CDI library mount it reads the library's ELF -// DT_SONAME and, when that differs from the mounted filename, writes a relative -// symlink alongside the mount destination; the real library file arrives at runtime -// via the CDI bind-mount into the same directory. -// -// A library whose SONAME cannot be read is skipped rather than failing the actor, so -// one odd file cannot cost us the whole GPU. That is logged: the missing symlink -// surfaces much later, as a loader error in the workload, with nothing pointing back -// here. -// Every path component under rootfs comes from the actor image, so the writes go -// through os.Root: an image that ships a driver-mount's parent directory as a -// symlink out of the rootfs would otherwise have the kernel resolve it in ateom's -// mount namespace, where the shared image cache and other actors' bundles are -// mounted. Same treatment as createExtraDirs in internal/imagecache. -func stageSonameSymlinks(ctx context.Context, rootfs string, mounts []specs.Mount) error { - root, err := os.OpenRoot(rootfs) - if err != nil { - return fmt.Errorf("while opening rootfs %q: %w", rootfs, err) - } - defer root.Close() - - for _, m := range mounts { - base := filepath.Base(m.Destination) - // Only shared libraries (…/lib…/foo.so.). - if !strings.Contains(base, ".so.") || m.Source == "" { - continue - } - soname, err := elfSonameFn(m.Source) - if err != nil { - slog.WarnContext(ctx, "Skipping SONAME symlink for driver library", - slog.String("library", m.Source), slog.Any("err", err)) - continue - } - if soname == "" || soname == base { - continue - } - // The SONAME is read out of the library and joined into a path, so require a - // bare filename rather than trusting the file's contents. - if soname != filepath.Base(soname) || !filepath.IsLocal(soname) { - slog.WarnContext(ctx, "Skipping driver library whose SONAME is not a filename", - slog.String("library", m.Source), slog.String("soname", soname)) - continue - } - dir := strings.TrimPrefix(filepath.Dir(m.Destination), "/") - if dir != "" && !filepath.IsLocal(dir) { - return fmt.Errorf("driver mount %q escapes the rootfs", m.Destination) - } - if err := root.MkdirAll(dir, 0o755); err != nil { - return fmt.Errorf("mkdir %s: %w", dir, err) - } - link := filepath.Join(dir, soname) - _ = root.Remove(link) - if err := root.Symlink(base, link); err != nil { - return fmt.Errorf("symlink %s -> %s: %w", link, base, err) - } - } - return nil -} - -// elfSonameFn is elfSoname, as a var so tests can supply a SONAME without needing a -// real ELF file on disk. -var elfSonameFn = elfSoname - -// elfSoname returns a shared library's DT_SONAME. It returns "" with a nil error for -// a library that simply carries no DT_SONAME entry, and an error when the file could -// not be opened or parsed as ELF. -func elfSoname(path string) (string, error) { - f, err := elf.Open(path) - if err != nil { - return "", fmt.Errorf("opening ELF %s: %w", path, err) - } - defer f.Close() - names, err := f.DynString(elf.DT_SONAME) - if err != nil { - return "", fmt.Errorf("reading DT_SONAME from %s: %w", path, err) - } - if len(names) == 0 { - return "", nil + return err } - return names[0], nil + return cdiinject.IntoBundle(ctx, bundleDir, spec, cdiinject.Options{ + Devices: []string{cdiAllDevice}, + AllowedHooks: enabledCDIHooks, + HookBinary: toolkitBinary("nvidia-cdi-hook"), + LibraryDirs: []string{driverLibDir}, + // runsc's nvproxy runs nvidia-container-cli when it sees + // NVIDIA_VISIBLE_DEVICES (independent of --nvproxy); the GPU is set up via + // CDI here, so strip it. + DropEnv: []string{"NVIDIA_VISIBLE_DEVICES"}, + }) } diff --git a/cmd/ateom-gvisor/gpu_test.go b/cmd/ateom-gvisor/gpu_test.go index a5d6753bf2..4a94d12be1 100644 --- a/cmd/ateom-gvisor/gpu_test.go +++ b/cmd/ateom-gvisor/gpu_test.go @@ -18,13 +18,10 @@ 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) { @@ -93,282 +90,3 @@ func TestGenerateCDISpec_NonZeroFails(t *testing.T) { t.Fatal("expected error on non-zero exit") } } - -func TestInjectGPUIntoBundle(t *testing.T) { - dir := t.TempDir() - - // Minimal CDI spec: one device node with major/minor OMITTED (as nvidia-ctk - // emits) plus an env var. injectGPUIntoBundle must resolve major/minor from the - // host — /dev/null is a char device present everywhere as 1,3 — so we assert on it. - // The hooks cover all three cases: create-symlinks is allowlisted and kept, - // update-ldcache is excluded, and an unrecognized hook (as a newer host toolkit - // could emit) is dropped rather than run unreviewed. - specJSON := `{ - "cdiVersion": "0.6.0", - "kind": "nvidia.com/gpu", - "devices": [ - { - "name": "all", - "containerEdits": { - "deviceNodes": [{"path": "/dev/null"}], - "env": ["NVIDIA_TEST=1"], - "hooks": [ - {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", - "args": ["nvidia-cdi-hook", "create-symlinks", "--link", "a::b"]}, - {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", - "args": ["nvidia-cdi-hook", "update-ldcache", "--folder", "/usr/local/nvidia/lib64"]}, - {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", - "args": ["nvidia-cdi-hook", "some-future-hook", "--flag", "v"]} - ] - } - } - ] -}` - specDir := filepath.Join(dir, "cdi") - os.MkdirAll(specDir, 0o755) - os.WriteFile(filepath.Join(specDir, "nvidia.json"), []byte(specJSON), 0o644) - - bundle := filepath.Join(dir, "bundle") - // The rootfs is already mounted by SetupBundleRootfs before injection runs. - os.MkdirAll(filepath.Join(bundle, "rootfs"), 0o755) - // The CUDA base image sets NVIDIA_VISIBLE_DEVICES; injection must strip it. - base := &specs.Spec{Version: "1.0.0", Process: &specs.Process{ - Args: []string{"true"}, - Env: []string{"NVIDIA_VISIBLE_DEVICES=all"}, - }} - data, _ := json.Marshal(base) - os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) - - if err := injectGPUIntoBundle(context.Background(), bundle, specDir); err != nil { - t.Fatalf("inject: %v", err) - } - - out, _ := os.ReadFile(filepath.Join(bundle, "config.json")) - var got specs.Spec - json.Unmarshal(out, &got) - - var dev *specs.LinuxDevice - for i := range got.Linux.Devices { - if got.Linux.Devices[i].Path == "/dev/null" { - dev = &got.Linux.Devices[i] - } - } - if dev == nil { - t.Fatalf("expected /dev/null device injected, spec=%s", out) - } - // The library resolved major/minor and type from the host (1,3,"c"). - if dev.Type != "c" || dev.Major != 1 || dev.Minor != 3 { - t.Fatalf("device not resolved from host: type=%q major=%d minor=%d", dev.Type, dev.Major, dev.Minor) - } - var hasEnv, hasVisibleDevices bool - for _, e := range got.Process.Env { - if e == "NVIDIA_TEST=1" { - hasEnv = true - } - if strings.HasPrefix(e, "NVIDIA_VISIBLE_DEVICES=") { - hasVisibleDevices = true - } - } - if !hasEnv { - t.Fatalf("expected NVIDIA_TEST env injected, spec=%s", out) - } - // NVIDIA_VISIBLE_DEVICES must be stripped so runsc's nvproxy does not invoke - // nvidia-container-cli (we set up the GPU via CDI instead). - if hasVisibleDevices { - t.Fatalf("expected NVIDIA_VISIBLE_DEVICES stripped, spec=%s", out) - } - // Only allowlisted hooks run. update-ldcache is excluded because its ldconfig - // needs a private /proc mount (its SONAME symlinks are staged directly instead), - // and an unrecognized hook is dropped rather than run unreviewed. - var kept []string - if got.Hooks != nil { - for _, h := range got.Hooks.CreateContainer { - if len(h.Args) > 1 { - kept = append(kept, h.Args[1]) - } - if !strings.HasPrefix(h.Path, toolkitDir) { - t.Fatalf("hook path %q should point at the mounted toolkit", h.Path) - } - } - } - if len(kept) != 1 || kept[0] != "create-symlinks" { - t.Fatalf("hooks = %v, want only [create-symlinks]", kept) - } -} - -// TestPrependLibraryPath covers the two cases that decide whether a CUDA program -// can load libcuda.so.1: an image that sets no LD_LIBRARY_PATH (must get one) and -// an NVIDIA image that already lists the driver directory (must be left alone). -func TestPrependLibraryPath(t *testing.T) { - dirs := []string{"/usr/local/nvidia/lib64"} - for _, tc := range []struct { - name string - env []string - want string - }{ - {"no existing value", []string{"PATH=/bin"}, "LD_LIBRARY_PATH=/usr/local/nvidia/lib64"}, - {"template value is kept after the driver dir", []string{"LD_LIBRARY_PATH=/opt/app/lib"}, - "LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/opt/app/lib"}, - {"already present is left alone", []string{"LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/x"}, - "LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/x"}, - } { - t.Run(tc.name, func(t *testing.T) { - got := prependLibraryPath(tc.env, dirs) - var found string - for _, e := range got { - if strings.HasPrefix(e, "LD_LIBRARY_PATH=") { - found = e - } - } - if found != tc.want { - t.Fatalf("got %q, want %q", found, tc.want) - } - if n := strings.Count(strings.Join(got, " "), "LD_LIBRARY_PATH="); n != 1 { - t.Fatalf("want exactly one LD_LIBRARY_PATH entry, got %d in %v", n, got) - } - }) - } -} - -// TestInjectGPUIntoBundle_Idempotent guards the invariant that atelet re-unpacks the -// bundle before each Run/Restore. The edits are appends against config.json on disk, -// so without the guard a second injection doubles every device, mount, env entry and -// hook — silently, since nothing errors. -func TestInjectGPUIntoBundle_Idempotent(t *testing.T) { - dir := t.TempDir() - // Explicit major/minor so the device need not exist on the test host. - specJSON := `{ - "cdiVersion": "0.6.0", - "kind": "nvidia.com/gpu", - "devices": [ - { - "name": "all", - "containerEdits": { - "deviceNodes": [{"path": "/dev/nvidia0", "major": 195, "minor": 0, "type": "c"}], - "env": ["NVIDIA_TEST=1"], - "hooks": [ - {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", - "args": ["nvidia-cdi-hook", "create-symlinks", "--link", "a::b"]} - ] - } - } - ] -}` - specDir := filepath.Join(dir, "cdi") - os.MkdirAll(specDir, 0o755) - os.WriteFile(filepath.Join(specDir, "nvidia.json"), []byte(specJSON), 0o644) - - bundle := filepath.Join(dir, "bundle") - os.MkdirAll(filepath.Join(bundle, "rootfs"), 0o755) - data, _ := json.Marshal(&specs.Spec{Version: "1.0.0", Process: &specs.Process{Args: []string{"true"}}}) - os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) - - // devices, mounts, env, hooks. - shape := func() [4]int { - out, _ := os.ReadFile(filepath.Join(bundle, "config.json")) - var s specs.Spec - if err := json.Unmarshal(out, &s); err != nil { - t.Fatal(err) - } - hooks := 0 - if s.Hooks != nil { - hooks = len(s.Hooks.CreateContainer) - } - return [4]int{len(s.Linux.Devices), len(s.Mounts), len(s.Process.Env), hooks} - } - - if err := injectGPUIntoBundle(context.Background(), bundle, specDir); err != nil { - t.Fatalf("first inject: %v", err) - } - first := shape() - if first[0] == 0 { - t.Fatalf("first inject added no devices: %v", first) - } - if err := injectGPUIntoBundle(context.Background(), bundle, specDir); err != nil { - t.Fatalf("second inject: %v", err) - } - if second := shape(); second != first { - t.Fatalf("second injection changed the bundle: %v -> %v (devices, mounts, env, hooks)", first, second) - } -} - -// TestStageSonameSymlinks_ConfinedToRootfs covers the two ways the staging writes -// could escape into ateom's mount namespace, where the shared image cache and other -// actors' bundles live: a directory the image redirects out of the rootfs, and a -// SONAME read out of a library that is not a bare filename. -func TestStageSonameSymlinks_ConfinedToRootfs(t *testing.T) { - for _, tc := range []struct { - name string - soname string - // dest is the CDI mount destination inside the actor. - dest string - // escape, when set, is planted in the rootfs as a symlink to the outside dir. - escape string - }{ - { - name: "image redirects the driver dir out of the rootfs", - soname: "libcuda.so.1", - dest: "/usr/lib/gpu/libcuda.so.580.65.06", - escape: "usr/lib/gpu", - }, - { - name: "SONAME climbs out with ../", - soname: "../../../../escaped.so.1", - dest: "/usr/lib/gpu/libcuda.so.580.65.06", - }, - } { - t.Run(tc.name, func(t *testing.T) { - dir := t.TempDir() - rootfs := filepath.Join(dir, "rootfs") - outside := filepath.Join(dir, "outside") - if err := os.MkdirAll(filepath.Join(rootfs, "usr/lib"), 0o755); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(outside, 0o755); err != nil { - t.Fatal(err) - } - victim := filepath.Join(outside, "libcuda.so.1") - if err := os.WriteFile(victim, []byte("do not delete"), 0o644); err != nil { - t.Fatal(err) - } - if tc.escape != "" { - if err := os.Symlink(outside, filepath.Join(rootfs, tc.escape)); err != nil { - t.Fatal(err) - } - } - - old := elfSonameFn - elfSonameFn = func(string) (string, error) { return tc.soname, nil } - defer func() { elfSonameFn = old }() - - // Refusing outright and skipping are both acceptable; writing outside is not. - _ = stageSonameSymlinks(context.Background(), rootfs, []specs.Mount{{ - Source: "/host/libcuda.so.580.65.06", - Destination: tc.dest, - }}) - - if b, err := os.ReadFile(victim); err != nil || string(b) != "do not delete" { - t.Fatalf("a file outside the rootfs was modified: err=%v content=%q", err, b) - } - if entries, err := os.ReadDir(outside); err == nil && len(entries) != 1 { - t.Fatalf("something was created outside the rootfs: %v", entries) - } - }) - } -} - -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) - - // An empty spec dir has no nvidia.json, so injection fails to read the spec. - emptyDir := filepath.Join(dir, "cdi-empty") - os.MkdirAll(emptyDir, 0o755) - if err := injectGPUIntoBundle(context.Background(), bundle, emptyDir); err == nil { - t.Fatal("expected error when the CDI spec is missing") - } -} diff --git a/cmd/ateom-gvisor/internal/cdiinject/inject.go b/cmd/ateom-gvisor/internal/cdiinject/inject.go new file mode 100644 index 0000000000..fadc35a7b9 --- /dev/null +++ b/cmd/ateom-gvisor/internal/cdiinject/inject.go @@ -0,0 +1,256 @@ +// 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 cdiinject applies CDI container edits to an actor's OCI bundle. +// +// This is the gVisor shape and does not generalize: it writes a bundle on disk, +// and it resolves device numbers by stat'ing the host, which only means anything +// for a sandbox on the host kernel. A micro-VM ships its OCI spec to the guest +// agent and needs the device passed through by VFIO first, so it would read the +// same spec (see internal/cdi) and apply it its own way. +package cdiinject + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + + specs "github.com/opencontainers/runtime-spec/specs-go" + "golang.org/x/sys/unix" + + "github.com/agent-substrate/substrate/internal/cdi" +) + +// Options tune what a merge is allowed to do. Everything vendor-specific lives +// here rather than in the merge itself. +type Options struct { + // Devices names the CDI devices to apply. + Devices []string + + // AllowedHooks are the createContainer hooks that may run, by their first + // argument. An allowlist because the generator is usually the cluster's own + // toolkit, so its version is not ours to choose and a newer one can emit + // hooks nobody has reviewed against this sandbox's posture. + AllowedHooks map[string]bool + + // HookBinary replaces the hook path the spec carries, so hooks run from a + // binary the caller trusts rather than wherever the generator pointed. + HookBinary string + + // LibraryDirs are prepended to the container's LD_LIBRARY_PATH. + LibraryDirs []string + + // DropEnv are environment variables to remove from the container. + DropEnv []string +} + +// IntoBundle merges a CDI spec into the OCI config.json in bundleDir: device +// nodes with their numbers resolved from the host, mounts, env, and the +// allowlisted hooks. +// +// Injecting twice would double every entry, so a bundle that already carries +// injected device nodes is left alone. +func IntoBundle(ctx context.Context, bundleDir string, spec *cdi.Spec, opts Options) error { + edits, err := spec.EditsFor(opts.Devices) + if err != nil { + return err + } + if len(edits.DeviceNodes) == 0 { + return fmt.Errorf("CDI spec resolved no devices") + } + + cfgPath := filepath.Join(bundleDir, "config.json") + specData, err := os.ReadFile(cfgPath) + if err != nil { + return fmt.Errorf("reading %s: %w", cfgPath, err) + } + var ociSpec specs.Spec + if err := json.Unmarshal(specData, &ociSpec); err != nil { + return fmt.Errorf("parsing OCI spec: %w", err) + } + if hasInjectedDevice(ociSpec.Linux, edits.DeviceNodes) { + slog.InfoContext(ctx, "Bundle already has CDI devices; skipping injection", + slog.String("bundle", bundleDir)) + return nil + } + + if ociSpec.Linux == nil { + ociSpec.Linux = &specs.Linux{} + } + if ociSpec.Linux.Resources == nil { + ociSpec.Linux.Resources = &specs.LinuxResources{} + } + for _, dn := range edits.DeviceNodes { + major, minor, err := resolveDevNumbers(dn.Path, dn.Major, dn.Minor) + if err != nil { + return err + } + devType := dn.Type + if devType == "" { + devType = "c" // Generators omit type for char devices; runsc needs it. + } + ociSpec.Linux.Devices = append(ociSpec.Linux.Devices, specs.LinuxDevice{ + Path: dn.Path, Type: devType, Major: major, Minor: minor, + }) + ociSpec.Linux.Resources.Devices = append(ociSpec.Linux.Resources.Devices, specs.LinuxDeviceCgroup{ + Allow: true, Type: devType, Major: &major, Minor: &minor, Access: "rwm", + }) + } + + for _, m := range edits.Mounts { + mType := m.Type + if mType == "" { + mType = "bind" // CDI omits type for its bind mounts; runsc's gofer needs it. + } + ociSpec.Mounts = append(ociSpec.Mounts, specs.Mount{ + Source: m.HostPath, Destination: m.ContainerPath, Type: mType, Options: m.Options, + }) + } + + if ociSpec.Process != nil { + ociSpec.Process.Env = append(ociSpec.Process.Env, edits.Env...) + for _, key := range opts.DropEnv { + ociSpec.Process.Env = dropEnvVar(ociSpec.Process.Env, key) + } + ociSpec.Process.Env = prependLibraryPath(ociSpec.Process.Env, opts.LibraryDirs) + } + + if ociSpec.Hooks == nil { + ociSpec.Hooks = &specs.Hooks{} + } + for _, h := range edits.Hooks { + if h.HookName != "createContainer" || len(h.Args) < 2 { + continue + } + if !opts.AllowedHooks[h.Args[1]] { + slog.InfoContext(ctx, "Skipping CDI hook outside the allowlist", + slog.String("hook", h.Args[1])) + continue + } + binary := h.Path + if opts.HookBinary != "" { + binary = opts.HookBinary + } + ociSpec.Hooks.CreateContainer = append(ociSpec.Hooks.CreateContainer, specs.Hook{ + Path: binary, + Args: append([]string{binary}, h.Args[1:]...), + Env: h.Env, + }) + } + + rootfs := "rootfs" + if ociSpec.Root != nil && ociSpec.Root.Path != "" { + rootfs = ociSpec.Root.Path + } + if !filepath.IsAbs(rootfs) { + rootfs = filepath.Join(bundleDir, rootfs) + } + if err := StageSonameSymlinks(ctx, rootfs, ociSpec.Mounts); err != nil { + return fmt.Errorf("staging SONAME symlinks: %w", err) + } + + out, err := json.Marshal(&ociSpec) + if err != nil { + return fmt.Errorf("serializing OCI spec: %w", err) + } + return os.WriteFile(cfgPath, out, 0o644) +} + +// hasInjectedDevice reports whether the OCI spec already carries any of the +// device nodes about to be injected, which is how an already-injected bundle is +// recognized. +func hasInjectedDevice(l *specs.Linux, nodes []cdi.Dev) bool { + if l == nil { + return false + } + for _, d := range l.Devices { + for _, n := range nodes { + if d.Path == n.Path { + return true + } + } + } + return false +} + +// resolveDevNumbers fills a device node's major/minor from the host when the CDI +// spec omitted them. CDI delegates that to the OCI runtime, which stats the +// host; the merge happens here instead, so the stat happens here too. Otherwise +// the container gets bogus 0,0 char devices that reach no driver. +func resolveDevNumbers(path string, major, minor int64) (int64, int64, error) { + if major != 0 { + return major, minor, nil + } + var st unix.Stat_t + if err := unix.Stat(path, &st); err != nil { + return 0, 0, fmt.Errorf("stat device %s: %w", path, err) + } + rdev := uint64(st.Rdev) + return int64(unix.Major(rdev)), int64(unix.Minor(rdev)), nil +} + +// dropEnvVar returns env with every "KEY=..." entry for the given key removed. +func dropEnvVar(env []string, key string) []string { + prefix := key + "=" + out := make([]string, 0, len(env)) + for _, e := range env { + if !strings.HasPrefix(e, prefix) { + out = append(out, e) + } + } + return out +} + +// prependLibraryPath puts dirs at the front of LD_LIBRARY_PATH, keeping whatever +// the image already set. This is the fallback for not running the CDI cache- +// updating hook: without it an image that sets no LD_LIBRARY_PATH cannot find +// the injected driver libraries, and the runtime reports zero devices even +// though they are all present. Directories already on the path are left alone. +// +// It is weaker than a loader cache -- LD_LIBRARY_PATH is inherited by children +// and outranks an executable's own DT_RUNPATH -- so callers should pass only the +// directory holding the libraries, not every directory the mounts touch. +func prependLibraryPath(env, dirs []string) []string { + if len(dirs) == 0 { + return env + } + existing := "" + for _, e := range env { + if v, ok := strings.CutPrefix(e, "LD_LIBRARY_PATH="); ok { + existing = v // OCI semantics: a later entry wins. + } + } + have := map[string]bool{} + for _, d := range strings.Split(existing, ":") { + have[d] = true + } + var add []string + for _, d := range dirs { + if !have[d] { + add = append(add, d) + } + } + if len(add) == 0 { + return env + } + val := strings.Join(add, ":") + if existing != "" { + val += ":" + existing + } + return append(dropEnvVar(env, "LD_LIBRARY_PATH"), "LD_LIBRARY_PATH="+val) +} diff --git a/cmd/ateom-gvisor/internal/cdiinject/inject_test.go b/cmd/ateom-gvisor/internal/cdiinject/inject_test.go new file mode 100644 index 0000000000..038d46fefc --- /dev/null +++ b/cmd/ateom-gvisor/internal/cdiinject/inject_test.go @@ -0,0 +1,357 @@ +// 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 cdiinject + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + specs "github.com/opencontainers/runtime-spec/specs-go" + "golang.org/x/sys/unix" + + "github.com/agent-substrate/substrate/internal/cdi" +) + +// statDev returns a device node's numbers, so the expectation matches whatever +// host the test runs on. +func statDev(t *testing.T, path string) (int64, int64) { + t.Helper() + var st unix.Stat_t + if err := unix.Stat(path, &st); err != nil { + t.Fatalf("stat %s: %v", path, err) + } + rdev := uint64(st.Rdev) + return int64(unix.Major(rdev)), int64(unix.Minor(rdev)) +} + +// testHookBinary stands in for the trusted hook binary a caller supplies. +const testHookBinary = "/opt/toolkit/cdi-hook" + +// injectFixture reads a spec from disk and applies the "all" device with the +// options the gVisor ateom uses, so these tests exercise the real combination. +func injectFixture(t *testing.T, bundleDir, specDir string) error { + t.Helper() + data, err := os.ReadFile(filepath.Join(specDir, "nvidia.json")) + if err != nil { + return err + } + spec, err := cdi.Parse(data) + if err != nil { + return err + } + return IntoBundle(context.Background(), bundleDir, spec, Options{ + Devices: []string{"all"}, + AllowedHooks: map[string]bool{"create-symlinks": true, "enable-cuda-compat": true}, + HookBinary: testHookBinary, + LibraryDirs: []string{"/usr/local/nvidia/lib64"}, + DropEnv: []string{"NVIDIA_VISIBLE_DEVICES"}, + }) +} + +func TestInjectGPUIntoBundle(t *testing.T) { + dir := t.TempDir() + + // Minimal CDI spec: one device node with major/minor OMITTED (as nvidia-ctk + // emits) plus an env var. injectGPUIntoBundle must resolve major/minor from the + // host — /dev/null is a char device present everywhere as 1,3 — so we assert on it. + // The hooks cover all three cases: create-symlinks is allowlisted and kept, + // update-ldcache is excluded, and an unrecognized hook (as a newer host toolkit + // could emit) is dropped rather than run unreviewed. + specJSON := `{ + "cdiVersion": "0.6.0", + "kind": "nvidia.com/gpu", + "devices": [ + { + "name": "all", + "containerEdits": { + "deviceNodes": [{"path": "/dev/null"}], + "env": ["NVIDIA_TEST=1"], + "hooks": [ + {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", + "args": ["nvidia-cdi-hook", "create-symlinks", "--link", "a::b"]}, + {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", + "args": ["nvidia-cdi-hook", "update-ldcache", "--folder", "/usr/local/nvidia/lib64"]}, + {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", + "args": ["nvidia-cdi-hook", "some-future-hook", "--flag", "v"]} + ] + } + } + ] +}` + specDir := filepath.Join(dir, "cdi") + os.MkdirAll(specDir, 0o755) + os.WriteFile(filepath.Join(specDir, "nvidia.json"), []byte(specJSON), 0o644) + + bundle := filepath.Join(dir, "bundle") + // The rootfs is already mounted by SetupBundleRootfs before injection runs. + os.MkdirAll(filepath.Join(bundle, "rootfs"), 0o755) + // The CUDA base image sets NVIDIA_VISIBLE_DEVICES; injection must strip it. + base := &specs.Spec{Version: "1.0.0", Process: &specs.Process{ + Args: []string{"true"}, + Env: []string{"NVIDIA_VISIBLE_DEVICES=all"}, + }} + data, _ := json.Marshal(base) + os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) + + if err := injectFixture(t, bundle, specDir); err != nil { + t.Fatalf("inject: %v", err) + } + + out, _ := os.ReadFile(filepath.Join(bundle, "config.json")) + var got specs.Spec + json.Unmarshal(out, &got) + + var dev *specs.LinuxDevice + for i := range got.Linux.Devices { + if got.Linux.Devices[i].Path == "/dev/null" { + dev = &got.Linux.Devices[i] + } + } + if dev == nil { + t.Fatalf("expected /dev/null device injected, spec=%s", out) + } + // The spec omitted major/minor, so they must come from stat'ing the host. + wantMajor, wantMinor := statDev(t, "/dev/null") + if dev.Type != "c" || dev.Major != wantMajor || dev.Minor != wantMinor { + t.Fatalf("device not resolved from host: type=%q major=%d minor=%d, want c %d %d", + dev.Type, dev.Major, dev.Minor, wantMajor, wantMinor) + } + var hasEnv, hasVisibleDevices bool + for _, e := range got.Process.Env { + if e == "NVIDIA_TEST=1" { + hasEnv = true + } + if strings.HasPrefix(e, "NVIDIA_VISIBLE_DEVICES=") { + hasVisibleDevices = true + } + } + if !hasEnv { + t.Fatalf("expected NVIDIA_TEST env injected, spec=%s", out) + } + // NVIDIA_VISIBLE_DEVICES must be stripped so runsc's nvproxy does not invoke + // nvidia-container-cli (we set up the GPU via CDI instead). + if hasVisibleDevices { + t.Fatalf("expected NVIDIA_VISIBLE_DEVICES stripped, spec=%s", out) + } + // Only allowlisted hooks run. update-ldcache is excluded because its ldconfig + // needs a private /proc mount (its SONAME symlinks are staged directly instead), + // and an unrecognized hook is dropped rather than run unreviewed. + var kept []string + if got.Hooks != nil { + for _, h := range got.Hooks.CreateContainer { + if len(h.Args) > 1 { + kept = append(kept, h.Args[1]) + } + if !strings.HasPrefix(h.Path, testHookBinary) { + t.Fatalf("hook path %q should point at the mounted toolkit", h.Path) + } + } + } + if len(kept) != 1 || kept[0] != "create-symlinks" { + t.Fatalf("hooks = %v, want only [create-symlinks]", kept) + } +} + +// TestPrependLibraryPath covers the two cases that decide whether a CUDA program +// can load libcuda.so.1: an image that sets no LD_LIBRARY_PATH (must get one) and +// an NVIDIA image that already lists the driver directory (must be left alone). +// TestPrependLibraryPath covers the two cases that decide whether a CUDA program +// can load libcuda.so.1: an image that sets no LD_LIBRARY_PATH (must get one) and +// an NVIDIA image that already lists the driver directory (must be left alone). +func TestPrependLibraryPath(t *testing.T) { + dirs := []string{"/usr/local/nvidia/lib64"} + for _, tc := range []struct { + name string + env []string + want string + }{ + {"no existing value", []string{"PATH=/bin"}, "LD_LIBRARY_PATH=/usr/local/nvidia/lib64"}, + {"template value is kept after the driver dir", []string{"LD_LIBRARY_PATH=/opt/app/lib"}, + "LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/opt/app/lib"}, + {"already present is left alone", []string{"LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/x"}, + "LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/x"}, + } { + t.Run(tc.name, func(t *testing.T) { + got := prependLibraryPath(tc.env, dirs) + var found string + for _, e := range got { + if strings.HasPrefix(e, "LD_LIBRARY_PATH=") { + found = e + } + } + if found != tc.want { + t.Fatalf("got %q, want %q", found, tc.want) + } + if n := strings.Count(strings.Join(got, " "), "LD_LIBRARY_PATH="); n != 1 { + t.Fatalf("want exactly one LD_LIBRARY_PATH entry, got %d in %v", n, got) + } + }) + } +} + +// TestInjectGPUIntoBundle_Idempotent guards the invariant that atelet re-unpacks the +// bundle before each Run/Restore. The edits are appends against config.json on disk, +// so without the guard a second injection doubles every device, mount, env entry and +// hook — silently, since nothing errors. +// TestInjectGPUIntoBundle_Idempotent guards the invariant that atelet re-unpacks the +// bundle before each Run/Restore. The edits are appends against config.json on disk, +// so without the guard a second injection doubles every device, mount, env entry and +// hook — silently, since nothing errors. +func TestInjectGPUIntoBundle_Idempotent(t *testing.T) { + dir := t.TempDir() + // Explicit major/minor so the device need not exist on the test host. + specJSON := `{ + "cdiVersion": "0.6.0", + "kind": "nvidia.com/gpu", + "devices": [ + { + "name": "all", + "containerEdits": { + "deviceNodes": [{"path": "/dev/nvidia0", "major": 195, "minor": 0, "type": "c"}], + "env": ["NVIDIA_TEST=1"], + "hooks": [ + {"hookName": "createContainer", "path": "/x/nvidia-cdi-hook", + "args": ["nvidia-cdi-hook", "create-symlinks", "--link", "a::b"]} + ] + } + } + ] +}` + specDir := filepath.Join(dir, "cdi") + os.MkdirAll(specDir, 0o755) + os.WriteFile(filepath.Join(specDir, "nvidia.json"), []byte(specJSON), 0o644) + + bundle := filepath.Join(dir, "bundle") + os.MkdirAll(filepath.Join(bundle, "rootfs"), 0o755) + data, _ := json.Marshal(&specs.Spec{Version: "1.0.0", Process: &specs.Process{Args: []string{"true"}}}) + os.WriteFile(filepath.Join(bundle, "config.json"), data, 0o644) + + // devices, mounts, env, hooks. + shape := func() [4]int { + out, _ := os.ReadFile(filepath.Join(bundle, "config.json")) + var s specs.Spec + if err := json.Unmarshal(out, &s); err != nil { + t.Fatal(err) + } + hooks := 0 + if s.Hooks != nil { + hooks = len(s.Hooks.CreateContainer) + } + return [4]int{len(s.Linux.Devices), len(s.Mounts), len(s.Process.Env), hooks} + } + + if err := injectFixture(t, bundle, specDir); err != nil { + t.Fatalf("first inject: %v", err) + } + first := shape() + if first[0] == 0 { + t.Fatalf("first inject added no devices: %v", first) + } + if err := injectFixture(t, bundle, specDir); err != nil { + t.Fatalf("second inject: %v", err) + } + if second := shape(); second != first { + t.Fatalf("second injection changed the bundle: %v -> %v (devices, mounts, env, hooks)", first, second) + } +} + +// TestStageSonameSymlinks_ConfinedToRootfs covers the two ways the staging writes +// could escape into ateom's mount namespace, where the shared image cache and other +// actors' bundles live: a directory the image redirects out of the rootfs, and a +// SONAME read out of a library that is not a bare filename. +// TestStageSonameSymlinks_ConfinedToRootfs covers the two ways the staging writes +// could escape into ateom's mount namespace, where the shared image cache and other +// actors' bundles live: a directory the image redirects out of the rootfs, and a +// SONAME read out of a library that is not a bare filename. +func TestStageSonameSymlinks_ConfinedToRootfs(t *testing.T) { + for _, tc := range []struct { + name string + soname string + // dest is the CDI mount destination inside the actor. + dest string + // escape, when set, is planted in the rootfs as a symlink to the outside dir. + escape string + }{ + { + name: "image redirects the driver dir out of the rootfs", + soname: "libcuda.so.1", + dest: "/usr/lib/gpu/libcuda.so.580.65.06", + escape: "usr/lib/gpu", + }, + { + name: "SONAME climbs out with ../", + soname: "../../../../escaped.so.1", + dest: "/usr/lib/gpu/libcuda.so.580.65.06", + }, + } { + t.Run(tc.name, func(t *testing.T) { + dir := t.TempDir() + rootfs := filepath.Join(dir, "rootfs") + outside := filepath.Join(dir, "outside") + if err := os.MkdirAll(filepath.Join(rootfs, "usr/lib"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(outside, 0o755); err != nil { + t.Fatal(err) + } + victim := filepath.Join(outside, "libcuda.so.1") + if err := os.WriteFile(victim, []byte("do not delete"), 0o644); err != nil { + t.Fatal(err) + } + if tc.escape != "" { + if err := os.Symlink(outside, filepath.Join(rootfs, tc.escape)); err != nil { + t.Fatal(err) + } + } + + old := elfSonameFn + elfSonameFn = func(string) (string, error) { return tc.soname, nil } + defer func() { elfSonameFn = old }() + + // Refusing outright and skipping are both acceptable; writing outside is not. + _ = StageSonameSymlinks(context.Background(), rootfs, []specs.Mount{{ + Source: "/host/libcuda.so.580.65.06", + Destination: tc.dest, + }}) + + if b, err := os.ReadFile(victim); err != nil || string(b) != "do not delete" { + t.Fatalf("a file outside the rootfs was modified: err=%v content=%q", err, b) + } + if entries, err := os.ReadDir(outside); err == nil && len(entries) != 1 { + t.Fatalf("something was created outside the rootfs: %v", entries) + } + }) + } +} + +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) + + // An empty spec dir has no nvidia.json, so injection fails to read the spec. + emptyDir := filepath.Join(dir, "cdi-empty") + os.MkdirAll(emptyDir, 0o755) + if err := injectFixture(t, bundle, emptyDir); err == nil { + t.Fatal("expected error when the CDI spec is missing") + } +} diff --git a/cmd/ateom-gvisor/internal/cdiinject/soname.go b/cmd/ateom-gvisor/internal/cdiinject/soname.go new file mode 100644 index 0000000000..281f9976d6 --- /dev/null +++ b/cmd/ateom-gvisor/internal/cdiinject/soname.go @@ -0,0 +1,111 @@ +// 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 cdiinject + +import ( + "context" + "debug/elf" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + + specs "github.com/opencontainers/runtime-spec/specs-go" +) + +// StageSonameSymlinks writes each driver library's SONAME symlink (e.g. +// libcuda.so.1 -> libcuda.so.580.65.06) into rootfs, so programs that link against +// the SONAME resolve it. For each CDI library mount it reads the library's ELF +// DT_SONAME and, when that differs from the mounted filename, writes a relative +// symlink alongside the mount destination; the real library file arrives at runtime +// via the CDI bind-mount into the same directory. +// +// A library whose SONAME cannot be read is skipped rather than failing the actor, so +// one odd file cannot cost us the whole GPU. That is logged: the missing symlink +// surfaces much later, as a loader error in the workload, with nothing pointing back +// here. +// Every path component under rootfs comes from the actor image, so the writes go +// through os.Root: an image that ships a driver-mount's parent directory as a +// symlink out of the rootfs would otherwise have the kernel resolve it in ateom's +// mount namespace, where the shared image cache and other actors' bundles are +// mounted. Same treatment as createExtraDirs in internal/imagecache. +func StageSonameSymlinks(ctx context.Context, rootfs string, mounts []specs.Mount) error { + root, err := os.OpenRoot(rootfs) + if err != nil { + return fmt.Errorf("while opening rootfs %q: %w", rootfs, err) + } + defer root.Close() + + for _, m := range mounts { + base := filepath.Base(m.Destination) + // Only shared libraries (…/lib…/foo.so.). + if !strings.Contains(base, ".so.") || m.Source == "" { + continue + } + soname, err := elfSonameFn(m.Source) + if err != nil { + slog.WarnContext(ctx, "Skipping SONAME symlink for driver library", + slog.String("library", m.Source), slog.Any("err", err)) + continue + } + if soname == "" || soname == base { + continue + } + // The SONAME is read out of the library and joined into a path, so require a + // bare filename rather than trusting the file's contents. + if soname != filepath.Base(soname) || !filepath.IsLocal(soname) { + slog.WarnContext(ctx, "Skipping driver library whose SONAME is not a filename", + slog.String("library", m.Source), slog.String("soname", soname)) + continue + } + dir := strings.TrimPrefix(filepath.Dir(m.Destination), "/") + if dir != "" && !filepath.IsLocal(dir) { + return fmt.Errorf("driver mount %q escapes the rootfs", m.Destination) + } + if err := root.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("mkdir %s: %w", dir, err) + } + link := filepath.Join(dir, soname) + _ = root.Remove(link) + if err := root.Symlink(base, link); err != nil { + return fmt.Errorf("symlink %s -> %s: %w", link, base, err) + } + } + return nil +} + +// elfSonameFn is elfSoname, as a var so tests can supply a SONAME without needing a +// real ELF file on disk. +var elfSonameFn = elfSoname + +// elfSoname returns a shared library's DT_SONAME. It returns "" with a nil error for +// a library that simply carries no DT_SONAME entry, and an error when the file could +// not be opened or parsed as ELF. +func elfSoname(path string) (string, error) { + f, err := elf.Open(path) + if err != nil { + return "", fmt.Errorf("opening ELF %s: %w", path, err) + } + defer f.Close() + names, err := f.DynString(elf.DT_SONAME) + if err != nil { + return "", fmt.Errorf("reading DT_SONAME from %s: %w", path, err) + } + if len(names) == 0 { + return "", nil + } + return names[0], nil +} diff --git a/internal/cdi/cdi.go b/internal/cdi/cdi.go new file mode 100644 index 0000000000..973f2827dd --- /dev/null +++ b/internal/cdi/cdi.go @@ -0,0 +1,125 @@ +// 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 cdi reads a Container Device Interface spec and resolves what a named +// set of devices asks a container runtime to do. +// +// An ateom cannot leave this to the container runtime. containerd applies CDI to +// the pod's containers, and an actor's containers are created by the ateom +// underneath one of them, so whatever hands devices to actors reads the spec +// itself. That holds for any sandbox, and under DRA too, which names allocated +// devices as CDI device IDs. +// +// Applying the edits is a separate step and is not generic: it depends on how the +// sandbox gets an OCI spec and whether it shares the host's device nodes. See +// cmd/ateom-gvisor/internal/cdiinject for the gVisor one. +package cdi + +import ( + "encoding/json" + "fmt" + "slices" +) + +// Spec is the part of a CDI spec that is consumed here: the devices and the +// edits applied to every container. +type Spec struct { + Devices []Device `json:"devices"` + // ContainerEdits apply to any container the spec touches, whichever devices + // were named (driver libraries, control device nodes, env). + ContainerEdits Edits `json:"containerEdits"` +} + +// Device is one addressable device in a spec. Generators name the same +// underlying hardware several ways -- nvidia-ctk emits per-index ("0"), per-UUID +// and an "all" device that repeats every node -- so applying every device in a +// spec injects each node several times over. +type Device struct { + Name string `json:"name"` + ContainerEdits Edits `json:"containerEdits"` +} + +// Edits are the changes a device makes to a container. +type Edits struct { + Env []string `json:"env,omitempty"` + DeviceNodes []Dev `json:"deviceNodes,omitempty"` + Mounts []Mount `json:"mounts,omitempty"` + Hooks []Hook `json:"hooks,omitempty"` +} + +// Dev is a device node to create in the container. Major and Minor are often +// absent: CDI leaves resolving them to the OCI runtime, which stats the host. +type Dev struct { + Path string `json:"path"` + Type string `json:"type,omitempty"` + Major int64 `json:"major,omitempty"` + Minor int64 `json:"minor,omitempty"` +} + +// Mount is a host path to mount into the container. +type Mount struct { + HostPath string `json:"hostPath"` + ContainerPath string `json:"containerPath"` + Type string `json:"type,omitempty"` + Options []string `json:"options,omitempty"` +} + +// Hook is a lifecycle hook the spec asks the runtime to run. +type Hook struct { + HookName string `json:"hookName"` + Path string `json:"path"` + Args []string `json:"args,omitempty"` + Env []string `json:"env,omitempty"` +} + +// Parse reads a CDI spec in JSON form. +func Parse(data []byte) (*Spec, error) { + var spec Spec + if err := json.Unmarshal(data, &spec); err != nil { + return nil, fmt.Errorf("parsing CDI spec: %w", err) + } + return &spec, nil +} + +// EditsFor collects the spec-level edits plus those of the named devices. A name +// the spec does not carry is an error rather than a silent no-op: it means the +// caller and the generator disagree about what the host has. +func (s *Spec) EditsFor(devices []string) (Edits, error) { + // Cloned, not copied: a struct copy shares the spec's backing arrays, so + // appending here could write into them. + edits := Edits{ + Env: slices.Clone(s.ContainerEdits.Env), + DeviceNodes: slices.Clone(s.ContainerEdits.DeviceNodes), + Mounts: slices.Clone(s.ContainerEdits.Mounts), + Hooks: slices.Clone(s.ContainerEdits.Hooks), + } + for _, want := range devices { + i := -1 + for j := range s.Devices { + if s.Devices[j].Name == want { + i = j + break + } + } + if i < 0 { + return Edits{}, fmt.Errorf("CDI spec has no %q device", want) + } + d := s.Devices[i].ContainerEdits + edits.Env = append(edits.Env, d.Env...) + edits.DeviceNodes = append(edits.DeviceNodes, d.DeviceNodes...) + edits.Mounts = append(edits.Mounts, d.Mounts...) + edits.Hooks = append(edits.Hooks, d.Hooks...) + } + return edits, nil +} diff --git a/internal/cdi/cdi_test.go b/internal/cdi/cdi_test.go new file mode 100644 index 0000000000..ab8014827a --- /dev/null +++ b/internal/cdi/cdi_test.go @@ -0,0 +1,135 @@ +// 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 cdi + +import "testing" + +// A spec shaped like nvidia-ctk's: an "all" device carrying both nodes, a device +// per index, and a UUID alias for the same hardware. +const twoGPUSpec = `{ + "cdiVersion": "0.6.0", + "kind": "nvidia.com/gpu", + "containerEdits": {"env": ["COMMON=1"], "deviceNodes": [{"path": "/dev/nvidiactl"}]}, + "devices": [ + {"name": "all", "containerEdits": {"deviceNodes": [{"path": "/dev/nvidia0"}, {"path": "/dev/nvidia1"}]}}, + {"name": "0", "containerEdits": {"deviceNodes": [{"path": "/dev/nvidia0"}]}}, + {"name": "1", "containerEdits": {"deviceNodes": [{"path": "/dev/nvidia1"}]}}, + {"name": "GPU-abcdef", "containerEdits": {"deviceNodes": [{"path": "/dev/nvidia0"}]}} + ] +}` + +func mustParse(t *testing.T) *Spec { + t.Helper() + spec, err := Parse([]byte(twoGPUSpec)) + if err != nil { + t.Fatalf("Parse() failed: %v", err) + } + return spec +} + +func nodePaths(e Edits) []string { + var paths []string + for _, n := range e.DeviceNodes { + paths = append(paths, n.Path) + } + return paths +} + +func TestParseRejectsGarbage(t *testing.T) { + if _, err := Parse([]byte("not json")); err == nil { + t.Fatal("Parse() = nil error, want one") + } +} + +func TestEditsFor(t *testing.T) { + tests := []struct { + name string + devices []string + want []string + wantErr bool + }{{ + // Spec-level edits apply whether or not a device was named. + name: "no devices", + devices: nil, + want: []string{"/dev/nvidiactl"}, + }, { + name: "one device", + devices: []string{"1"}, + want: []string{"/dev/nvidiactl", "/dev/nvidia1"}, + }, { + name: "several devices", + devices: []string{"0", "1"}, + want: []string{"/dev/nvidiactl", "/dev/nvidia0", "/dev/nvidia1"}, + }, { + // Generators name the same hardware several ways; naming two of them is + // the caller's error to make, and the duplicate must be visible rather + // than quietly collapsed. + name: "aliases of one device repeat its nodes", + devices: []string{"0", "GPU-abcdef"}, + want: []string{"/dev/nvidiactl", "/dev/nvidia0", "/dev/nvidia0"}, + }, { + name: "unknown device", + devices: []string{"7"}, + wantErr: true, + }} + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := mustParse(t).EditsFor(tc.devices) + if tc.wantErr { + if err == nil { + t.Fatal("EditsFor() = nil error, want one") + } + return + } + if err != nil { + t.Fatalf("EditsFor() failed: %v", err) + } + paths := nodePaths(got) + if len(paths) != len(tc.want) { + t.Fatalf("device nodes = %v, want %v", paths, tc.want) + } + for i := range tc.want { + if paths[i] != tc.want[i] { + t.Fatalf("device nodes = %v, want %v", paths, tc.want) + } + } + }) + } +} + +// EditsFor must not mutate the spec: the spec-level edits are the starting point +// for every call, so appending to them in place would leak one call's devices +// into the next. +func TestEditsForDoesNotAccumulate(t *testing.T) { + spec := mustParse(t) + if _, err := spec.EditsFor([]string{"0"}); err != nil { + t.Fatalf("EditsFor() failed: %v", err) + } + got, err := spec.EditsFor([]string{"1"}) + if err != nil { + t.Fatalf("EditsFor() failed: %v", err) + } + want := []string{"/dev/nvidiactl", "/dev/nvidia1"} + paths := nodePaths(got) + if len(paths) != len(want) { + t.Fatalf("second call = %v, want %v", paths, want) + } + for i := range want { + if paths[i] != want[i] { + t.Fatalf("second call = %v, want %v", paths, want) + } + } +} From e0209e8d89e56765fdc7b8fbdf698fe1435ce45b Mon Sep 17 00:00:00 2001 From: Benjamin Elder Date: Tue, 1 Sep 2026 20:31:45 -0700 Subject: [PATCH 2/2] Drop device passthrough until the API is designed A GPU worker pool injected its devices into every container of every actor on it, which only held up while a worker ran one actor. A resource name and a count cannot replace it: that shape cannot express which container gets a device, sharing one between containers, or resuming onto hardware compatible with the snapshot. A nvidia.com/gpu pool limit is now an ordinary extended resource -- it places the pod on a GPU node and nothing else reads it. internal/cdi stays, unused: whatever the API becomes, an ateom still has to merge CDI edits into an actor's bundle itself. --- .../internal/controllers/workerpool_apply.go | 88 --------- .../controllers/workerpool_apply_test.go | 166 ---------------- cmd/ateom-gvisor/gpu.go | 187 ------------------ cmd/ateom-gvisor/gpu_test.go | 92 --------- cmd/ateom-gvisor/main.go | 10 - cmd/ateom-gvisor/runsc.go | 19 -- cmd/ateom-gvisor/runsc_test.go | 24 --- docs/api-guide.md | 83 +------- .../generated/ate.dev_workerpools.yaml | 11 -- pkg/api/v1alpha1/workerpool_types.go | 2 - .../v1alpha1/workerpool_validation_test.go | 45 +---- 11 files changed, 9 insertions(+), 718 deletions(-) delete mode 100644 cmd/ateom-gvisor/gpu.go delete mode 100644 cmd/ateom-gvisor/gpu_test.go diff --git a/cmd/atecontroller/internal/controllers/workerpool_apply.go b/cmd/atecontroller/internal/controllers/workerpool_apply.go index 06b0853442..c2fdf57773 100644 --- a/cmd/atecontroller/internal/controllers/workerpool_apply.go +++ b/cmd/atecontroller/internal/controllers/workerpool_apply.go @@ -15,7 +15,6 @@ package controllers import ( - "os" "slices" corev1 "k8s.io/api/core/v1" @@ -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) @@ -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) { diff --git a/cmd/atecontroller/internal/controllers/workerpool_apply_test.go b/cmd/atecontroller/internal/controllers/workerpool_apply_test.go index 69f1c20cf4..85bff190f4 100644 --- a/cmd/atecontroller/internal/controllers/workerpool_apply_test.go +++ b/cmd/atecontroller/internal/controllers/workerpool_apply_test.go @@ -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"}, diff --git a/cmd/ateom-gvisor/gpu.go b/cmd/ateom-gvisor/gpu.go deleted file mode 100644 index db5abc5fae..0000000000 --- a/cmd/ateom-gvisor/gpu.go +++ /dev/null @@ -1,187 +0,0 @@ -//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 ( - "cmp" - "context" - "fmt" - "log/slog" - "os" - "os/exec" - "path/filepath" - "sync" - - "github.com/agent-substrate/substrate/cmd/ateom-gvisor/internal/cdiinject" - "github.com/agent-substrate/substrate/internal/ateompath" - "github.com/agent-substrate/substrate/internal/cdi" -) - -// toolkitDir is where the host's NVIDIA container toolkit (nvidia-ctk, -// nvidia-cdi-hook) is mounted into the worker pod — read-only from the node, so -// the binaries match whatever toolkit/driver the cluster installed. A var (not -// const) so tests can point it at a fixture directory. -var toolkitDir = "/opt/nvidia-toolkit" - -// gpuDeviceGlob matches the per-GPU device nodes. The device plugin can assign any -// indices (a worker sharing a multi-GPU node may get /dev/nvidia2,3 with no -// /dev/nvidia0), so detection must not assume index 0. The [0-9] excludes the control -// nodes (/dev/nvidiactl, /dev/nvidia-uvm). A var so tests can point it at a fixture. -var gpuDeviceGlob = "/dev/nvidia[0-9]*" - -// nvidiaDriverRoot is where the GPU device plugin mounts the driver into the pod. -// GKE and gpu-operator both use /usr/local/nvidia, but that is a convention rather -// than a contract, so it is overridable via ATE_NVIDIA_DRIVER_ROOT (propagated onto -// GPU worker pods by the controller). -var nvidiaDriverRoot = cmp.Or(os.Getenv("ATE_NVIDIA_DRIVER_ROOT"), "/usr/local/nvidia") - -// Both directories are load-bearing for CDI generation, not just for completeness: -// without the library path nvidia-ctk cannot load libnvidia-ml.so.1 to enumerate the -// GPUs and generation fails outright, and without the bin path (which it discovers -// via PATH) the generated spec carries libraries but no nvidia-smi. -var ( - driverLibDir = filepath.Join(nvidiaDriverRoot, "lib64") - driverBinDir = filepath.Join(nvidiaDriverRoot, "bin") -) - -const cdiOutputDir = "/run/ate-cdi" - -// cdiAllDevice is the CDI device carrying every GPU assigned to the pod. -// nvidia-ctk repeats the same nodes across per-index and per-UUID devices, so -// only one of them may be applied. -const cdiAllDevice = "all" - -// enabledCDIHooks is the set of CDI createContainer hooks the actor runs. It is an -// allowlist because the toolkit is mounted from the host, so its version is the -// cluster's choice: a newer one can emit hooks that have never been reviewed -// against this worker's unprivileged posture, and those must not run by default. -// update-ldcache is absent deliberately — its ldconfig needs a private /proc -// mount, which the pod's masked /proc rejects, so the SONAME symlinks it would -// create are staged directly into the rootfs instead. -var enabledCDIHooks = map[string]bool{ - "create-symlinks": true, - "enable-cuda-compat": true, -} - -// toolkitBinary resolves a toolkit command to an executable path, preferring the -// unwrapped ".real" binary the NVIDIA toolkit ships: the plain name is often a -// /bin/sh wrapper. The ateom image is glibc-based (debian), so the glibc-dynamic -// toolkit binaries run directly — no ld-linux loader shim is needed. -func toolkitBinary(name string) string { - if real := filepath.Join(toolkitDir, name+".real"); fileExists(real) { - return real - } - return filepath.Join(toolkitDir, name) -} - -func fileExists(p string) bool { _, err := os.Stat(p); return err == nil } - -// gpuPresent reports whether any GPU is assigned to this worker pod, matching any -// device index (not just /dev/nvidia0 — the device plugin can assign 2,3 etc.). -func gpuPresent() bool { - matches, _ := filepath.Glob(gpuDeviceGlob) - return len(matches) > 0 -} - -var ( - generateMu sync.Mutex - cdiGenerated bool -) - -// generateCDISpec runs nvidia-ctk (from the host toolkit mounted into the pod) to -// produce a CDI spec scoped to this pod's assigned GPU. The glibc-based ateom image -// runs the glibc-dynamic toolkit binary directly. Runs through reaper like every -// other synchronous subprocess here. -func generateCDISpec(ctx context.Context, outDir string) error { - if err := os.MkdirAll(outDir, 0o755); err != nil { - return fmt.Errorf("creating CDI output dir %s: %w", outDir, err) - } - // No --nvidia-cdi-hook-path: we discard the CDI hooks (staging SONAME symlinks - // ourselves), so the hook paths nvidia-ctk writes into the spec are never used. - cmd := exec.CommandContext(ctx, toolkitBinary("nvidia-ctk"), - "cdi", "generate", - "--format=json", - "--library-search-path="+driverLibDir, - "--output="+filepath.Join(outDir, "nvidia.json"), - ) - // nvidia-ctk finds the driver binaries (nvidia-smi, ...) via PATH. - cmd.Env = append(os.Environ(), "PATH="+driverBinDir+":"+os.Getenv("PATH")) - if out, err := reaper.CombinedOutput(cmd); err != nil { - return fmt.Errorf("nvidia-ctk cdi generate failed: %w: %s", err, out) - } - return nil -} - -// ensureCDISpec generates the per-pod CDI spec once, on the first actor. A failure is -// not memoized: a transient error (e.g. the toolkit mount not yet ready) is retried on -// the next actor rather than bricking GPU for the pod's lifetime. -func ensureCDISpec(ctx context.Context) error { - generateMu.Lock() - defer generateMu.Unlock() - if cdiGenerated { - return nil - } - if err := generateCDISpec(ctx, cdiOutputDir); err != nil { - return err - } - cdiGenerated = true - return nil -} - -// maybeInjectGPU is a no-op unless the worker pod has a GPU. When it does, it -// 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() { - return nil - } - slog.InfoContext(ctx, "Injecting GPU into actor container", slog.String("container", containerName)) - if err := ensureCDISpec(ctx); err != nil { - return err - } - bundleDir := ateompath.OCIBundlePath(actorUID, containerName) - if err := injectGPUIntoBundle(ctx, bundleDir, cdiOutputDir); err != nil { - return fmt.Errorf("injecting GPU into %q bundle: %w", containerName, err) - } - return nil -} - -// injectGPUIntoBundle merges the CDI spec generated in cdiSpecDir into the actor's -// OCI bundle. It does NOT run the CDI cache-updating hook: that hook's ldconfig -// needs a private /proc, which the pod's masked /proc rejects, so the SONAME -// symlinks it would create are staged into the rootfs instead. That is what lets -// a GPU worker keep the plain unprivileged posture. -func injectGPUIntoBundle(ctx context.Context, bundleDir, cdiSpecDir string) error { - data, err := os.ReadFile(filepath.Join(cdiSpecDir, "nvidia.json")) - if err != nil { - return fmt.Errorf("reading CDI spec: %w", err) - } - spec, err := cdi.Parse(data) - if err != nil { - return err - } - return cdiinject.IntoBundle(ctx, bundleDir, spec, cdiinject.Options{ - Devices: []string{cdiAllDevice}, - AllowedHooks: enabledCDIHooks, - HookBinary: toolkitBinary("nvidia-cdi-hook"), - LibraryDirs: []string{driverLibDir}, - // runsc's nvproxy runs nvidia-container-cli when it sees - // NVIDIA_VISIBLE_DEVICES (independent of --nvproxy); the GPU is set up via - // CDI here, so strip it. - DropEnv: []string{"NVIDIA_VISIBLE_DEVICES"}, - }) -} diff --git a/cmd/ateom-gvisor/gpu_test.go b/cmd/ateom-gvisor/gpu_test.go deleted file mode 100644 index 4a94d12be1..0000000000 --- a/cmd/ateom-gvisor/gpu_test.go +++ /dev/null @@ -1,92 +0,0 @@ -//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" - "os" - "path/filepath" - "strings" - "testing" -) - -func TestMaybeInjectGPU_NoGPUIsNoop(t *testing.T) { - dir := t.TempDir() - old := gpuDeviceGlob - gpuDeviceGlob = filepath.Join(dir, "nvidia[0-9]*") // matches nothing - defer func() { gpuDeviceGlob = old }() - if err := maybeInjectGPU(context.Background(), "actor_uid", "c1"); err != nil { - t.Fatalf("expected no-op nil when no GPU is present, got %v", err) - } -} - -// TestGPUPresent checks detection matches any GPU index, not just nvidia0. -func TestGPUPresent(t *testing.T) { - dir := t.TempDir() - old := gpuDeviceGlob - gpuDeviceGlob = filepath.Join(dir, "nvidia[0-9]*") - defer func() { gpuDeviceGlob = old }() - - if gpuPresent() { - t.Fatal("expected absent before creation") - } - // A worker sharing a multi-GPU node can be assigned nvidia2, not nvidia0. - if err := os.WriteFile(filepath.Join(dir, "nvidia2"), nil, 0o644); err != nil { - t.Fatal(err) - } - if !gpuPresent() { - t.Fatal("expected present after creating nvidia2") - } -} - -func TestGenerateCDISpec_InvokesCtk(t *testing.T) { - dir := t.TempDir() - oldT := toolkitDir - toolkitDir = dir - defer func() { toolkitDir = oldT }() - - // Fake nvidia-ctk (run directly by the glibc ateom) that writes a minimal JSON - // spec to the --output= path. - 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" -` - os.WriteFile(filepath.Join(dir, "nvidia-ctk"), []byte(script), 0o755) - - out := filepath.Join(dir, "cdi") - if err := generateCDISpec(context.Background(), 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() - oldT := toolkitDir - toolkitDir = dir - defer func() { toolkitDir = oldT }() - os.WriteFile(filepath.Join(dir, "nvidia-ctk"), []byte("#!/bin/sh\nexit 3\n"), 0o755) - if err := generateCDISpec(context.Background(), filepath.Join(dir, "cdi")); err == nil { - t.Fatal("expected error on non-zero exit") - } -} diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index dbae3a3f4c..8322b2c429 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -174,10 +174,6 @@ func do(ctx context.Context) error { return fmt.Errorf("while setting up cgroup delegation: %w", err) } - if gpuPresent() { - slog.InfoContext(ctx, "GPU detected; enabling runsc nvproxy for all sandboxes") - } - go reaper.Run(ctx) slog.InfoContext(ctx, "Child process reaper launched") @@ -694,9 +690,6 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload return nil, fmt.Errorf("while composing %q rootfs: %w", ac.GetName(), err) } containersToDelete = append(containersToDelete, ac.GetName()) - 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) } @@ -975,9 +968,6 @@ 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: containersToDelete = append(containersToDelete, ac.GetName()) diff --git a/cmd/ateom-gvisor/runsc.go b/cmd/ateom-gvisor/runsc.go index f46126cb2a..bbc0add7c2 100644 --- a/cmd/ateom-gvisor/runsc.go +++ b/cmd/ateom-gvisor/runsc.go @@ -55,23 +55,6 @@ func durableVolumeNames(spec *ateompb.WorkloadSpec) []string { return slices.Compact(names) } -// nvproxyGlobalArgs returns the runsc global flags for GPU sandboxes, enabling -// gVisor's GPU ioctl proxy when the worker has a GPU. --nvproxy must be set when -// the sandbox is created (the pause/root container) so the sentry initializes GPU -// support up front — like enabling nvproxy in the containerd runtime config on -// normal Kubernetes. Otherwise nvproxy would try to initialize late, when the app -// subcontainer joins carrying the CDI /dev/nvidia* devices, and the running sentry -// crashes (StartSubcontainer: EOF). -// -// Only create and restore need it: both boot a sentry. `runsc start` acts on a -// sandbox that already exists, so the flag has no effect there. -func nvproxyGlobalArgs() []string { - if gpuPresent() { - return []string{"--nvproxy"} - } - return nil -} - // shapeSpec loads, shapes for gVisor, and saves the container's OCI spec. func (r *runsc) shapeSpec(containerName string) error { bundle := ateompath.OCIBundlePath(r.actorUID, containerName) @@ -109,7 +92,6 @@ func (r *runsc) cmdCreate(ctx context.Context, out io.Writer, containerName stri // otherwise sizes to all host CPUs). Global flag: before the subcommand. "--cpu-num-from-quota", } - args = append(args, nvproxyGlobalArgs()...) args = append(args, "create", "-bundle", ateompath.OCIBundlePath(r.actorUID, containerName), @@ -245,7 +227,6 @@ func (r *runsc) cmdRestore(ctx context.Context, out io.Writer, containerName, ch // Match cmdCreate: size the restored sentry from the cgroup CPU quota. "--cpu-num-from-quota", } - restoreArgs = append(restoreArgs, nvproxyGlobalArgs()...) restoreArgs = append(restoreArgs, "restore", "-bundle", ateompath.OCIBundlePath(r.actorUID, containerName), diff --git a/cmd/ateom-gvisor/runsc_test.go b/cmd/ateom-gvisor/runsc_test.go index 7be788c728..ce567a58dd 100644 --- a/cmd/ateom-gvisor/runsc_test.go +++ b/cmd/ateom-gvisor/runsc_test.go @@ -17,36 +17,12 @@ package main import ( - "os" - "path/filepath" "reflect" "testing" "github.com/agent-substrate/substrate/internal/ateompath" ) -// TestNvproxyGlobalArgs checks that runsc is told to enable nvproxy exactly when the -// worker has a GPU. The flag must be present on sandbox creation so the sentry -// initializes GPU support up front; without it the GPU subcontainer crashes. -func TestNvproxyGlobalArgs(t *testing.T) { - dir := t.TempDir() - old := gpuDeviceGlob - gpuDeviceGlob = filepath.Join(dir, "nvidia[0-9]*") - defer func() { gpuDeviceGlob = old }() - - if got := nvproxyGlobalArgs(); len(got) != 0 { - t.Fatalf("no GPU: want no flags, got %v", got) - } - - if err := os.WriteFile(filepath.Join(dir, "nvidia0"), nil, 0o644); err != nil { - t.Fatal(err) - } - got := nvproxyGlobalArgs() - if len(got) != 1 || got[0] != "--nvproxy" { - t.Fatalf("GPU: want [--nvproxy], got %v", got) - } -} - func TestKillArgs(t *testing.T) { r := &runsc{ path: "/usr/bin/runsc", diff --git a/docs/api-guide.md b/docs/api-guide.md index 1d70f27b15..ea3dd8d429 100644 --- a/docs/api-guide.md +++ b/docs/api-guide.md @@ -89,85 +89,16 @@ spec: # gvisor SandboxConfig unless sandboxConfigName is set. ``` -### GPU worker pools +### Devices (GPUs) — temporarily unsupported -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. +Substrate does not pass devices through to actors. A GPU worker pod's devices were +injected into every actor container; that was removed while the device API is +designed, since a resource name and a count cannot express which container gets a +device, sharing one, or resuming onto compatible hardware. -```yaml -apiVersion: ate.dev/v1alpha1 -kind: WorkerPool -metadata: - name: gpu-pool - namespace: ate-demo -spec: - replicas: 5 - # GPU pools need a glibc ateom-gvisor build — see Requirements below. - workerImage: /ateom-gvisor-glibc@sha256:... - template: - # (1) schedule onto GPU nodes - nodeSelector: - cloud.google.com/gke-accelerator: nvidia-tesla-t4 - tolerations: - - key: nvidia.com/gpu - operator: Exists - effect: NoSchedule - priorityClassName: substrate-workers - resources: - requests: - cpu: 500m - memory: 1Gi - limits: - cpu: "1" - memory: 2Gi - # (2) claim a GPU — this request is what triggers GPU passthrough - nvidia.com/gpu: "1" -``` +A `nvidia.com/gpu` pool limit is now an ordinary extended resource: it places the +worker pod on a GPU node and reserves the device, and nothing else reads it. -`atecontroller` propagates the request onto the `ateom` container and mounts the -host NVIDIA toolkit into the pod. `ateom-gvisor` then generates a CDI spec with -`nvidia-ctk` and injects the GPU device nodes, driver libraries, and env into -each actor container's OCI spec, and runs `runsc` with `--nvproxy` so CUDA and NVML -work inside the sandbox. A worker requesting `nvidia.com/gpu: N` passes all N -through. - -Every container in the actor gets the GPU. An actor's containers share one sandbox -and the worker's whole device set, and `ActorTemplate` has no per-container resource -fields, so GPUs are shared at the actor level rather than assigned to one container, -the same as cpu and memory. This differs from a Kubernetes Pod, where the GPU goes -only to the container that requests it. If per-container resource limits are added -later, GPU assignment should follow them. - -The driver library directory is prepended to each container's `LD_LIBRARY_PATH` so an -image does not have to set it to find `libcuda.so.1`; any existing value is kept after -it rather than replaced. - -**Requirements** - -- **A glibc `ateom-gvisor` image**, set as `spec.workerImage`. The distroless default - cannot exec `nvidia-ctk`. Build one with - `KO_DEFAULTBASEIMAGE=debian:stable-slim ko build ./cmd/ateom-gvisor`. -- **`nvidia-ctk` on the node**, at the path mounted into the worker — by default - `/usr/local/nvidia/toolkit`, overridable with the controller's - `ATE_NVIDIA_TOOLKIT_HOST_PATH`. gpu-operator installs it; GKE's built-in GPU - support does not. -- **The driver mounted into the pod by the device plugin**, at `/usr/local/nvidia` - by default. `nvidia-ctk` needs its libraries to enumerate the GPUs at all, so a - cluster whose plugin uses a different layout must set the controller's - `ATE_NVIDIA_DRIVER_ROOT`. -- **`atelet` must run on the GPU nodes**, so add a matching toleration to its - DaemonSet if those nodes are tainted (for example `nvidia.com/gpu`). -- **gVisor only.** `microvm` pools would need VFIO PCI passthrough instead. - -**Known limitation: a GPU actor can only be suspended while no CUDA context is -open.** gVisor cannot serialize GPU state, so a checkpoint taken while the workload -holds a context fails with an nvproxy encoding error and terminates the -sandbox. Workloads that run CUDA and exit snapshot normally; one that keeps a -context alive (a model resident in device memory, say) cannot be suspended or have a -golden snapshot taken. ### Status (`WorkerPoolStatus`) diff --git a/manifests/ate-install/generated/ate.dev_workerpools.yaml b/manifests/ate-install/generated/ate.dev_workerpools.yaml index d10dbd846f..df0fd546c1 100644 --- a/manifests/ate-install/generated/ate.dev_workerpools.yaml +++ b/manifests/ate-install/generated/ate.dev_workerpools.yaml @@ -451,17 +451,6 @@ spec: - replicas - workerImage type: object - x-kubernetes-validations: - - message: nvidia.com/gpu is only supported when sandboxClass is 'gvisor' - rule: '!has(self.sandboxClass) || self.sandboxClass == ''gvisor'' || - !has(self.template) || !has(self.template.resources) || !((has(self.template.resources.limits) - && ''nvidia.com/gpu'' in self.template.resources.limits) || (has(self.template.resources.requests) - && ''nvidia.com/gpu'' in self.template.resources.requests))' - - message: 'nvidia.com/gpu must be set in limits: Kubernetes does not - admit a request for an extended resource without a matching limit' - rule: '!has(self.template) || !has(self.template.resources) || !has(self.template.resources.requests) - || !(''nvidia.com/gpu'' in self.template.resources.requests) || (has(self.template.resources.limits) - && ''nvidia.com/gpu'' in self.template.resources.limits)' status: description: status is the observed state of WorkerPool properties: diff --git a/pkg/api/v1alpha1/workerpool_types.go b/pkg/api/v1alpha1/workerpool_types.go index eedac459ae..aa60c1e25a 100644 --- a/pkg/api/v1alpha1/workerpool_types.go +++ b/pkg/api/v1alpha1/workerpool_types.go @@ -77,8 +77,6 @@ type WorkerPoolPodTemplate struct { Resources *corev1.ResourceRequirements `json:"resources,omitempty"` } -// +kubebuilder:validation:XValidation:rule="!has(self.sandboxClass) || self.sandboxClass == 'gvisor' || !has(self.template) || !has(self.template.resources) || !((has(self.template.resources.limits) && 'nvidia.com/gpu' in self.template.resources.limits) || (has(self.template.resources.requests) && 'nvidia.com/gpu' in self.template.resources.requests))",message="nvidia.com/gpu is only supported when sandboxClass is 'gvisor'" -// +kubebuilder:validation:XValidation:rule="!has(self.template) || !has(self.template.resources) || !has(self.template.resources.requests) || !('nvidia.com/gpu' in self.template.resources.requests) || (has(self.template.resources.limits) && 'nvidia.com/gpu' in self.template.resources.limits)",message="nvidia.com/gpu must be set in limits: Kubernetes does not admit a request for an extended resource without a matching limit" type WorkerPoolSpec struct { // Replicas is the number of worker pods to run. // +required diff --git a/pkg/api/v1alpha1/workerpool_validation_test.go b/pkg/api/v1alpha1/workerpool_validation_test.go index 6992f0eefa..617e68d787 100644 --- a/pkg/api/v1alpha1/workerpool_validation_test.go +++ b/pkg/api/v1alpha1/workerpool_validation_test.go @@ -187,53 +187,12 @@ func TestWorkerPoolValidation(t *testing.T) { wantErr: true, errMsg: "spec.template.tolerations: Too many", }, { - name: "gpu on a gvisor pool", - mutate: func(wp *WorkerPool) { - wp.Spec.SandboxClass = SandboxClassGvisor - wp.Spec.Template = gpuTemplate(corev1.ResourceList{gpuResourceName: resource.MustParse("1")}, nil) - }, - wantErr: false, - }, { - name: "gpu limit on a micro-VM pool", + // Devices are no longer interpreted: the limit only places the pod. + name: "extended resource on a pool is not interpreted", mutate: func(wp *WorkerPool) { wp.Spec.SandboxClass = SandboxClassMicroVM wp.Spec.Template = gpuTemplate(corev1.ResourceList{gpuResourceName: resource.MustParse("1")}, nil) }, - wantErr: true, - errMsg: "nvidia.com/gpu is only supported when sandboxClass is 'gvisor'", - }, { - // The pod shape keys off limits OR requests, so the rule must reject both. - name: "gpu request on a micro-VM pool", - mutate: func(wp *WorkerPool) { - wp.Spec.SandboxClass = SandboxClassMicroVM - wp.Spec.Template = gpuTemplate(nil, corev1.ResourceList{gpuResourceName: resource.MustParse("1")}) - }, - wantErr: true, - errMsg: "nvidia.com/gpu is only supported when sandboxClass is 'gvisor'", - }, { - // Kubernetes refuses a pod that requests an extended resource without a - // matching limit, so a pool like this would shape a worker the Deployment - // then rejects. Catch it on the WorkerPool the user wrote. - name: "gpu in requests only", - mutate: func(wp *WorkerPool) { - wp.Spec.Template = gpuTemplate(nil, corev1.ResourceList{gpuResourceName: resource.MustParse("1")}) - }, - wantErr: true, - errMsg: "nvidia.com/gpu must be set in limits", - }, { - name: "gpu in both limits and requests", - mutate: func(wp *WorkerPool) { - wp.Spec.Template = gpuTemplate( - corev1.ResourceList{gpuResourceName: resource.MustParse("1")}, - corev1.ResourceList{gpuResourceName: resource.MustParse("1")}) - }, - wantErr: false, - }, { - name: "micro-VM pool without a gpu", - mutate: func(wp *WorkerPool) { - wp.Spec.SandboxClass = SandboxClassMicroVM - wp.Spec.Template = gpuTemplate(corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1")}, nil) - }, wantErr: false, }}