From df0c6fe14656bd9c4bc4ecbf2f0e75409ceeffea Mon Sep 17 00:00:00 2001 From: Anson Qian Date: Mon, 22 Jun 2026 13:07:58 -0400 Subject: [PATCH] agent: add DisableNvidia config switch --- cmd/agent/internal/daemon/nodeoperator.go | 4 ++ .../internal/daemon/nodeoperator_test.go | 7 +++ docs/content/reference/gpu/nvidia.md | 6 +++ pkg/agent/config/config.go | 4 ++ pkg/agent/config/config_test.go | 5 ++- pkg/agent/goalstates/containerd.go | 4 +- pkg/agent/goalstates/nvidia.go | 8 ++-- pkg/agent/goalstates/resolve.go | 20 ++++++--- pkg/agent/goalstates/resolve_test.go | 43 +++++++++++++++++++ 9 files changed, 87 insertions(+), 14 deletions(-) diff --git a/cmd/agent/internal/daemon/nodeoperator.go b/cmd/agent/internal/daemon/nodeoperator.go index b253e45c5..03e6c2251 100644 --- a/cmd/agent/internal/daemon/nodeoperator.go +++ b/cmd/agent/internal/daemon/nodeoperator.go @@ -136,6 +136,10 @@ func hasDrift(applied, desired *provision.AgentConfig) bool { return true } + if applied.DisableNvidia != desired.DisableNvidia { + return true + } + if applied.Kubelet.ApiServer != desired.Kubelet.ApiServer { return true } diff --git a/cmd/agent/internal/daemon/nodeoperator_test.go b/cmd/agent/internal/daemon/nodeoperator_test.go index 38243b00a..78b1c16d9 100644 --- a/cmd/agent/internal/daemon/nodeoperator_test.go +++ b/cmd/agent/internal/daemon/nodeoperator_test.go @@ -98,6 +98,13 @@ func Test_hasDrift_OciImageChange(t *testing.T) { assert.True(t, hasDrift(applied, desired)) } +func Test_hasDrift_DisableNvidiaChange(t *testing.T) { + applied := baseConfig() + desired := baseConfig() + desired.DisableNvidia = true + assert.True(t, hasDrift(applied, desired)) +} + func Test_hasDrift_LabelsChange(t *testing.T) { applied := baseConfig() desired := baseConfig() diff --git a/docs/content/reference/gpu/nvidia.md b/docs/content/reference/gpu/nvidia.md index 5dbc87839..ed3a71aaf 100644 --- a/docs/content/reference/gpu/nvidia.md +++ b/docs/content/reference/gpu/nvidia.md @@ -8,6 +8,12 @@ The unbounded-agent automatically detects NVIDIA GPUs on the host, forwards the driver's userspace libraries into the nspawn container, generates a CDI specification, and configures containerd so that GPU workloads can be scheduled on the node. +Set `DisableNvidia` to `true` in the agent config to disable NVIDIA GPU +discovery and runtime setup even when `/dev/nvidia*` devices are present on the +host. When disabled, the agent does not select the default NVIDIA OCI image, +bind-mount NVIDIA devices or libraries, generate the NVIDIA CDI spec, or +register the NVIDIA containerd runtime. + ## Prerequisites Before the agent can expose GPUs, the **host VM** must have: diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go index 7e52c357a..65ee99555 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -47,6 +47,10 @@ type AgentConfig struct { // "ghcr.io/org/repo:tag") used to bootstrap the machine rootfs. // When empty the agent falls back to debootstrap. OCIImage string `json:"OCIImage,omitempty"` + + // DisableNvidia disables NVIDIA GPU discovery and runtime setup even when + // NVIDIA devices are present on the host. + DisableNvidia bool `json:"DisableNvidia,omitempty"` } // BackfillNodeName resolves and stores the Kubernetes Node name once. An diff --git a/pkg/agent/config/config_test.go b/pkg/agent/config/config_test.go index f734050dd..c2b7a61a0 100644 --- a/pkg/agent/config/config_test.go +++ b/pkg/agent/config/config_test.go @@ -84,7 +84,8 @@ func TestCRIConfig_JSONRoundTrip(t *testing.T) { t.Parallel() cfg := AgentConfig{ - MachineName: "test", + MachineName: "test", + DisableNvidia: true, CRI: CRIConfig{ Containerd: ContainerdConfig{Version: "2.1.0"}, Runc: RuncConfig{Version: "1.2.0"}, @@ -101,6 +102,7 @@ func TestCRIConfig_JSONRoundTrip(t *testing.T) { assert.Equal(t, "2.1.0", decoded.CRI.Containerd.Version) assert.Equal(t, "1.2.0", decoded.CRI.Runc.Version) assert.Equal(t, "1.6.0", decoded.CNI.PluginVersion) + assert.True(t, decoded.DisableNvidia) } func TestCRIConfig_OmittedWhenEmpty(t *testing.T) { @@ -124,6 +126,7 @@ func TestCRIConfig_OmittedWhenEmpty(t *testing.T) { cni := parsed["CNI"].(map[string]interface{}) assert.NotContains(t, cni, "PluginVersion") + assert.NotContains(t, parsed, "DisableNvidia") } func TestAgentConfig_DeepCopy(t *testing.T) { diff --git a/pkg/agent/goalstates/containerd.go b/pkg/agent/goalstates/containerd.go index 3ee8fb953..ec90d09ba 100644 --- a/pkg/agent/goalstates/containerd.go +++ b/pkg/agent/goalstates/containerd.go @@ -17,7 +17,7 @@ type Containerd struct { } // ResolveContainerd returns the containerd configuration goal state. -func ResolveContainerd() Containerd { +func ResolveContainerd(nvidia NvidiaHost) Containerd { return Containerd{ SandboxImage: SandboxImage, ContainerdBinPath: filepath.Join("/"+BinDir, "containerd"), @@ -25,6 +25,6 @@ func ResolveContainerd() Containerd { CNIBinDir: CNIBinDir, CNIConfDir: CNIConfigDir, MetricsAddress: ContainerdMetricsAddress, - NvidiaRuntime: resolveNvidiaRuntime(), + NvidiaRuntime: resolveNvidiaRuntime(nvidia), } } diff --git a/pkg/agent/goalstates/nvidia.go b/pkg/agent/goalstates/nvidia.go index 0262050c9..86732b0ef 100644 --- a/pkg/agent/goalstates/nvidia.go +++ b/pkg/agent/goalstates/nvidia.go @@ -113,11 +113,11 @@ func ResolveNvidiaHost(arch string) (NvidiaHost, error) { } // resolveNvidiaRuntime returns the NVIDIA container runtime goal state. -// When GPU devices are present the runtime is enabled with default paths; -// otherwise it is disabled. -func resolveNvidiaRuntime() NvidiaRuntime { +// When GPU devices are present in the resolved host state the runtime is +// enabled with default paths; otherwise it is disabled. +func resolveNvidiaRuntime(nvidia NvidiaHost) NvidiaRuntime { return NvidiaRuntime{ - Enabled: len(discoverNVIDIADevices()) > 0, + Enabled: len(nvidia.GPUDevicePaths) > 0, RuntimeClassName: NvidiaRuntimeClassName, RuntimePath: NvidiaContainerRuntimePath, DisableSetAsDefaultRuntime: false, diff --git a/pkg/agent/goalstates/resolve.go b/pkg/agent/goalstates/resolve.go index b5621791e..70080d5cd 100644 --- a/pkg/agent/goalstates/resolve.go +++ b/pkg/agent/goalstates/resolve.go @@ -25,9 +25,9 @@ type MachineGoalState struct { NodeStart *NodeStart } -// ResolveMachine probes the host (kernel version, hostname, GPU hardware) and -// resolves the complete goal state for the named nspawn machine from an agent -// config. +// ResolveMachine probes the host (kernel version, hostname, GPU hardware unless +// disabled) and resolves the complete goal state for the named nspawn machine +// from an agent config. func ResolveMachine(log *slog.Logger, cfg *config.AgentConfig, machineName string, downloads *DownloadOverrides) (*MachineGoalState, error) { kernel, err := hostKernel() if err != nil { @@ -39,9 +39,15 @@ func ResolveMachine(log *slog.Logger, cfg *config.AgentConfig, machineName strin return nil, fmt.Errorf("get host hostname: %w", err) } - nvidia, err := ResolveNvidiaHost(runtime.GOARCH) - if err != nil { - return nil, fmt.Errorf("resolve nvidia host: %w", err) + var nvidia NvidiaHost + if cfg.DisableNvidia { + log.Info("NVIDIA support disabled by agent config") + } else { + var err error + nvidia, err = ResolveNvidiaHost(runtime.GOARCH) + if err != nil { + return nil, fmt.Errorf("resolve nvidia host: %w", err) + } } ociImage := ResolveOCIImage(log, cfg.OCIImage, len(nvidia.GPUDevicePaths) > 0) @@ -95,7 +101,7 @@ func ResolveMachine(log *slog.Logger, cfg *config.AgentConfig, machineName strin KubeMachineName: cfg.MachineName, NodeName: cfg.NodeName, MachineDir: filepath.Join("/var/lib/machines", machineName), - Containerd: ResolveContainerd(), + Containerd: ResolveContainerd(nvidia), Kubelet: kubelet, Nvidia: nvidia, } diff --git a/pkg/agent/goalstates/resolve_test.go b/pkg/agent/goalstates/resolve_test.go index afdddef07..1d276bd02 100644 --- a/pkg/agent/goalstates/resolve_test.go +++ b/pkg/agent/goalstates/resolve_test.go @@ -6,6 +6,7 @@ package goalstates import ( "log/slog" "os" + "runtime" "testing" "github.com/stretchr/testify/assert" @@ -102,6 +103,44 @@ func TestResolveOCIImage_DefaultWithGPU(t *testing.T) { assert.Equal(t, DefaultNvidiaOCImage, got) } +func TestResolveContainerd_NvidiaRuntimeFollowsResolvedHost(t *testing.T) { + withoutGPU := ResolveContainerd(NvidiaHost{}) + assert.False(t, withoutGPU.NvidiaRuntime.Enabled) + + withGPU := ResolveContainerd(NvidiaHost{GPUDevicePaths: []string{"/dev/nvidia0"}}) + assert.True(t, withGPU.NvidiaRuntime.Enabled) +} + +func TestResolveMachine_DisableNvidia(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("ResolveMachine requires Linux host kernel discovery") + } + + t.Setenv("AGENT_DISABLE_OCI_IMAGE", "") + t.Setenv("AGENT_OCI_IMAGE", "") + + cfg := &config.AgentConfig{ + MachineName: "machine-1", + DisableNvidia: true, + Cluster: config.AgentClusterConfig{ + CaCertBase64: "Y2EtYnl0ZXM=", + }, + Kubelet: config.AgentKubeletConfig{ + ApiServer: "https://api.example.com", + }, + } + + got, err := ResolveMachine(discardLogger(), cfg, "kube1", nil) + require.NoError(t, err) + + assert.Equal(t, DefaultOCIImage, got.RootFS.OCIImage) + assert.Empty(t, got.RootFS.Nvidia.GPUDevicePaths) + assert.Empty(t, got.RootFS.Nvidia.LibMappings) + assert.Empty(t, got.RootFS.Nvidia.LibDirMounts) + assert.Empty(t, got.NodeStart.Nvidia.GPUDevicePaths) + assert.False(t, got.NodeStart.Containerd.NvidiaRuntime.Enabled) +} + func TestResolveOCIImage_Priority(t *testing.T) { // Verify the full priority chain: config > disable > env var > default. log := discardLogger() @@ -223,6 +262,10 @@ func TestResolveKubelet_InvalidNodeIPRejected(t *testing.T) { } func TestResolveMachine_UsesConfigNodeName(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("ResolveMachine requires Linux host kernel discovery") + } + cfg := &config.AgentConfig{ MachineName: "machine-1", NodeName: "configured-node",