diff --git a/cmd/agent/internal/cmd/cmd.go b/cmd/agent/internal/cmd/cmd.go index c3cd57e3..a6e3e5b8 100644 --- a/cmd/agent/internal/cmd/cmd.go +++ b/cmd/agent/internal/cmd/cmd.go @@ -31,6 +31,7 @@ func Run() { newCmdReset(cmdCtx), newCmdVersion(), newCmdRecordAgentUpgradeFailureSignal(), + newCmdReconcileNVIDIA(cmdCtx), ) if err := root.Execute(); err != nil { diff --git a/cmd/agent/internal/cmd/reconcile_nvidia.go b/cmd/agent/internal/cmd/reconcile_nvidia.go new file mode 100644 index 00000000..bbac7115 --- /dev/null +++ b/cmd/agent/internal/cmd/reconcile_nvidia.go @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +package cmd + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "os" + + "github.com/spf13/cobra" + + "github.com/Azure/unbounded/internal/provision" + "github.com/Azure/unbounded/pkg/agent/goalstates" + "github.com/Azure/unbounded/pkg/agent/phases" + "github.com/Azure/unbounded/pkg/agent/phases/nodestart" +) + +var reconcileNVIDIAOnMachineStart = runNVIDIAReconciliation + +func newCmdReconcileNVIDIA(cmdCtx *CommandContext) *cobra.Command { + cmd := &cobra.Command{ + Use: "reconcile-nvidia MACHINE", + Short: "Reconcile NVIDIA state after an nspawn machine starts", + Hidden: true, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + machine := args[0] + if machine != goalstates.NSpawnMachineKube1 && machine != goalstates.NSpawnMachineKube2 { + return fmt.Errorf("unknown nspawn machine %q", machine) + } + + cmdCtx.Setup() + + return reconcileNVIDIAOnMachineStart(cmd.Context(), cmdCtx.Logger, machine) + }, + } + + return cmd +} + +func runNVIDIAReconciliation(ctx context.Context, log *slog.Logger, machine string) error { + path := goalstates.AppliedConfigPath(machine) + + data, err := os.ReadFile(path) + if errors.Is(err, os.ErrNotExist) { + log.Info("applied config not available; managed node start will reconcile NVIDIA state", + "machine", machine) + + return nil + } + + if err != nil { + return fmt.Errorf("read applied config %s: %w", path, err) + } + + if err := goalstates.VerifyChecksum(data, goalstates.AppliedConfigChecksumPath(machine)); err != nil { + return fmt.Errorf("verify applied config checksum for %s: %w", machine, err) + } + + var cfg provision.AgentConfig + if err := json.Unmarshal(data, &cfg); err != nil { + return fmt.Errorf("decode applied config %s: %w", path, err) + } + + goalState, err := goalstates.ResolveMachine(log, &cfg, machine, nil) + if err != nil { + return fmt.Errorf("resolve machine goal state: %w", err) + } + + if !goalState.NodeStart.Containerd.NvidiaRuntime.Enabled || len(goalState.NodeStart.Nvidia.LibMappings) == 0 { + return fmt.Errorf("NVIDIA setup state is unavailable for machine %s", machine) + } + + if err := nodestart.WaitForMachine(ctx, log, machine); err != nil { + return fmt.Errorf("wait for machine %s: %w", machine, err) + } + + return phases.ExecuteTask(ctx, log, nodestart.SetupNVIDIA(log, goalState.NodeStart)) +} diff --git a/cmd/agent/internal/cmd/reconcile_nvidia_test.go b/cmd/agent/internal/cmd/reconcile_nvidia_test.go new file mode 100644 index 00000000..71eac178 --- /dev/null +++ b/cmd/agent/internal/cmd/reconcile_nvidia_test.go @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +package cmd + +import ( + "context" + "log/slog" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestReconcileNVIDIACommand(t *testing.T) { + original := reconcileNVIDIAOnMachineStart + + t.Cleanup(func() { + reconcileNVIDIAOnMachineStart = original + }) + + var gotMachine string + + reconcileNVIDIAOnMachineStart = func(_ context.Context, _ *slog.Logger, machine string) error { + gotMachine = machine + + return nil + } + + commandContext := &CommandContext{LogFormat: "text"} + command := newCmdReconcileNVIDIA(commandContext) + command.SetArgs([]string{"kube1"}) + + require.NoError(t, command.ExecuteContext(context.Background())) + require.Equal(t, "kube1", gotMachine) +} + +func TestReconcileNVIDIACommandRejectsUnknownMachine(t *testing.T) { + command := newCmdReconcileNVIDIA(&CommandContext{LogFormat: "text"}) + command.SetArgs([]string{"other"}) + + require.ErrorContains(t, command.ExecuteContext(context.Background()), "unknown nspawn machine") +} diff --git a/pkg/agent/goalstates/constants.go b/pkg/agent/goalstates/constants.go index 80e0ffb0..be3ba7cc 100644 --- a/pkg/agent/goalstates/constants.go +++ b/pkg/agent/goalstates/constants.go @@ -163,6 +163,8 @@ const ( DefaultAzureLinux3NvidiaOCIImage = "ghcr.io/azure/agent-azlinux3-nvidia:v20260626" SystemdUnitContainerd = "containerd.service" + SystemdUnitNVIDIAReady = "unbounded-nvidia-ready.service" + NVIDIAReadyPath = "/run/unbounded/nvidia-ready" ContainerdConfigPath = "/etc/containerd/config.toml" ContainerdConfDropInDir = "/etc/containerd/conf.d" ContainerdCertsDir = "/etc/containerd/certs.d" diff --git a/pkg/agent/phases/nodestart/assets/containerd.service b/pkg/agent/phases/nodestart/assets/containerd.service index f004daf3..643a67e0 100644 --- a/pkg/agent/phases/nodestart/assets/containerd.service +++ b/pkg/agent/phases/nodestart/assets/containerd.service @@ -2,6 +2,10 @@ Description=containerd container runtime Documentation=https://containerd.io After=network.target local-fs.target +{{- if .NvidiaEnabled}} +Requires=unbounded-nvidia-ready.service +After=unbounded-nvidia-ready.service +{{- end}} [Service] ExecStartPre=-/sbin/modprobe overlay diff --git a/pkg/agent/phases/nodestart/assets/kubelet.service b/pkg/agent/phases/nodestart/assets/kubelet.service index 2e0a2686..1640abf4 100644 --- a/pkg/agent/phases/nodestart/assets/kubelet.service +++ b/pkg/agent/phases/nodestart/assets/kubelet.service @@ -3,6 +3,10 @@ Description=Kubelet ConditionPathExists={{.KubeletBinPath}} Wants=network-online.target containerd.service After=network-online.target containerd.service +{{- if .NvidiaEnabled}} +Requires=unbounded-nvidia-ready.service +After=unbounded-nvidia-ready.service +{{- end}} [Service] Restart=always diff --git a/pkg/agent/phases/nodestart/assets/unbounded-nvidia-ready.service b/pkg/agent/phases/nodestart/assets/unbounded-nvidia-ready.service new file mode 100644 index 00000000..0fd87ac1 --- /dev/null +++ b/pkg/agent/phases/nodestart/assets/unbounded-nvidia-ready.service @@ -0,0 +1,9 @@ +[Unit] +Description=Wait for unbounded NVIDIA runtime setup +Before=containerd.service kubelet.service + +[Service] +Type=oneshot +ExecStart=/bin/sh -c 'until test -e /run/unbounded/nvidia-ready; do sleep 1; done' +RemainAfterExit=yes +TimeoutStartSec=infinity diff --git a/pkg/agent/phases/nodestart/cri.go b/pkg/agent/phases/nodestart/cri.go index 319a2668..414af1d2 100644 --- a/pkg/agent/phases/nodestart/cri.go +++ b/pkg/agent/phases/nodestart/cri.go @@ -64,6 +64,10 @@ func (c *configureContainerd) Do(_ context.Context) error { return fmt.Errorf("ensure containerd service unit: %w", err) } + if err := c.ensureNVIDIAReadyServiceUnit(); err != nil { + return fmt.Errorf("ensure NVIDIA ready service unit: %w", err) + } + if err := c.ensureGPUDropInConfigs(); err != nil { return fmt.Errorf("ensure GPU drop-in configs: %w", err) } @@ -131,6 +135,7 @@ func (c *configureContainerd) ensureContainerdServiceUnit() error { buf := &bytes.Buffer{} if err := assetsTemplate.ExecuteTemplate(buf, "containerd.service", map[string]any{ "ContainerdBinPath": spec.ContainerdBinPath, + "NvidiaEnabled": nvidiaSetupEnabled(c.goalState), }); err != nil { return err } @@ -140,6 +145,28 @@ func (c *configureContainerd) ensureContainerdServiceUnit() error { return utilio.WriteFile(dest, buf.Bytes(), 0o644) } +func (c *configureContainerd) ensureNVIDIAReadyServiceUnit() error { + dest := filepath.Join(c.goalState.MachineDir, goalstates.SystemdSystemDir, goalstates.SystemdUnitNVIDIAReady) + if !nvidiaSetupEnabled(c.goalState) { + if err := os.Remove(dest); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + + return nil + } + + data, err := assets.ReadFile("assets/unbounded-nvidia-ready.service") + if err != nil { + return err + } + + return utilio.WriteFile(dest, data, 0o644) +} + +func nvidiaSetupEnabled(goalState *goalstates.NodeStart) bool { + return goalState.Containerd.NvidiaRuntime.Enabled && len(goalState.Nvidia.LibMappings) > 0 +} + // ensureGPUDropInConfigs manages GPU-related containerd drop-in configs. // When the nvidia runtime is enabled the drop-in is written; otherwise it is // removed. diff --git a/pkg/agent/phases/nodestart/cri_test.go b/pkg/agent/phases/nodestart/cri_test.go index 2e26259d..1d934321 100644 --- a/pkg/agent/phases/nodestart/cri_test.go +++ b/pkg/agent/phases/nodestart/cri_test.go @@ -55,6 +55,56 @@ func TestConfigureContainerdSetsObservabilityDefaults(t *testing.T) { "image_pull_progress_timeout = \"15m\"") } +func TestConfigureContainerdGatesNVIDIAStartup(t *testing.T) { + t.Parallel() + + machineDir := t.TempDir() + goalState := &goalstates.NodeStart{ + MachineDir: machineDir, + Containerd: goalstates.ResolveContainerd(""), + Nvidia: goalstates.NvidiaHost{ + GPUDevicePaths: []string{"/dev/nvidia0"}, + LibMappings: []goalstates.NvidiaLibMapping{{HostPath: "/usr/lib/libcuda.so.1"}}, + }, + } + goalState.Containerd.NvidiaRuntime.Enabled = true + + require.NoError(t, ConfigureContainerd(goalState).Do(context.Background())) + + service, err := os.ReadFile(filepath.Join(machineDir, goalstates.SystemdSystemDir, goalstates.SystemdUnitContainerd)) + require.NoError(t, err) + require.Contains(t, string(service), "Requires=unbounded-nvidia-ready.service") + require.Contains(t, string(service), "After=unbounded-nvidia-ready.service") + + readyService, err := os.ReadFile(filepath.Join( + machineDir, + goalstates.SystemdSystemDir, + goalstates.SystemdUnitNVIDIAReady, + )) + require.NoError(t, err) + require.Contains(t, string(readyService), + "ExecStart=/bin/sh -c 'until test -e /run/unbounded/nvidia-ready; do sleep 1; done'") +} + +func TestConfigureContainerdDoesNotGateCPUNodes(t *testing.T) { + t.Parallel() + + machineDir := t.TempDir() + goalState := &goalstates.NodeStart{ + MachineDir: machineDir, + Containerd: goalstates.ResolveContainerd(""), + } + + require.NoError(t, ConfigureContainerd(goalState).Do(context.Background())) + + service, err := os.ReadFile(filepath.Join(machineDir, goalstates.SystemdSystemDir, goalstates.SystemdUnitContainerd)) + require.NoError(t, err) + require.NotContains(t, string(service), "unbounded-nvidia-ready.service") + + _, err = os.Stat(filepath.Join(machineDir, goalstates.SystemdSystemDir, goalstates.SystemdUnitNVIDIAReady)) + require.ErrorIs(t, err, os.ErrNotExist) +} + func TestConfigureContainerdUpdatesManagedGantryHostsConfig(t *testing.T) { t.Parallel() diff --git a/pkg/agent/phases/nodestart/kubelet.go b/pkg/agent/phases/nodestart/kubelet.go index 3e52e6b2..e4052f53 100644 --- a/pkg/agent/phases/nodestart/kubelet.go +++ b/pkg/agent/phases/nodestart/kubelet.go @@ -201,6 +201,7 @@ func (c *configureKubelet) ensureKubeletServiceUnit() error { if err := assetsTemplate.ExecuteTemplate(buf, "kubelet.service", map[string]any{ "KubeletBinPath": spec.KubeletBinPath, "KubeletConfigurationPath": goalstates.KubeletConfigurationPath, + "NvidiaEnabled": nvidiaSetupEnabled(c.goalState), }); err != nil { return err } diff --git a/pkg/agent/phases/nodestart/kubelet_test.go b/pkg/agent/phases/nodestart/kubelet_test.go index b01f4c1d..2bed8442 100644 --- a/pkg/agent/phases/nodestart/kubelet_test.go +++ b/pkg/agent/phases/nodestart/kubelet_test.go @@ -33,7 +33,6 @@ func TestConfigureKubeletWritesHostnameOverride(t *testing.T) { ClusterDNS: "10.0.0.10", }, } - require.NoError(t, ConfigureKubelet(goalState).Do(context.Background())) data, err := os.ReadFile(filepath.Join( @@ -62,7 +61,6 @@ func TestConfigureKubeletOmitsNodeIPWhenEmpty(t *testing.T) { ClusterDNS: "10.0.0.10", }, } - require.NoError(t, ConfigureKubelet(goalState).Do(context.Background())) data, err := os.ReadFile(filepath.Join( @@ -187,6 +185,35 @@ func TestConfigureKubeletWritesConfiguration(t *testing.T) { require.NotContains(t, string(service), "--cluster-dns=") } +func TestConfigureKubeletGatesNVIDIAStartup(t *testing.T) { + t.Parallel() + + machineDir := t.TempDir() + goalState := &goalstates.NodeStart{ + MachineDir: machineDir, + NodeName: "worker-1", + Nvidia: goalstates.NvidiaHost{ + GPUDevicePaths: []string{"/dev/nvidia0"}, + LibMappings: []goalstates.NvidiaLibMapping{{HostPath: "/usr/lib/libcuda.so.1"}}, + }, + Kubelet: goalstates.Kubelet{ + CACertData: []byte("ca"), + ClusterDNS: "10.0.0.10", + KubeletAuthInfo: config.KubeletAuthInfo{ + BootstrapToken: "token", + }, + }, + } + goalState.Containerd.NvidiaRuntime.Enabled = true + + require.NoError(t, ConfigureKubelet(goalState).Do(context.Background())) + + service, err := os.ReadFile(filepath.Join(machineDir, goalstates.SystemdSystemDir, goalstates.SystemdUnitKubelet)) + require.NoError(t, err) + require.Contains(t, string(service), "Requires=unbounded-nvidia-ready.service") + require.Contains(t, string(service), "After=unbounded-nvidia-ready.service") +} + func TestConfigureKubeletWritesImageCredentialProviderFlags(t *testing.T) { t.Parallel() diff --git a/pkg/agent/phases/nodestart/nspawn.go b/pkg/agent/phases/nodestart/nspawn.go index 5430e9c0..15023c11 100644 --- a/pkg/agent/phases/nodestart/nspawn.go +++ b/pkg/agent/phases/nodestart/nspawn.go @@ -84,7 +84,7 @@ func (s *startNSpawnMachine) Do(ctx context.Context) error { return err } - if err := waitForMachine(ctx, s.log, name); err != nil { + if err := WaitForMachine(ctx, s.log, name); err != nil { return fmt.Errorf("wait for machine %s: %w", name, err) } @@ -172,10 +172,10 @@ func isAlreadyExistsErr(err error) bool { return strings.Contains(msg, "already exists") || strings.Contains(msg, "file exists") } -// waitForMachine polls the machine until it is responsive to systemd-run +// WaitForMachine polls the machine until it is responsive to systemd-run // commands. machinectl start returns before D-Bus is ready, so phases that use // executil.MachineRun() would fail without this gate. -func waitForMachine(ctx context.Context, log *slog.Logger, machine string) error { +func WaitForMachine(ctx context.Context, log *slog.Logger, machine string) error { const ( pollInterval = 500 * time.Millisecond timeout = 30 * time.Second diff --git a/pkg/agent/phases/nodestart/nvidia.go b/pkg/agent/phases/nodestart/nvidia.go index 6b618d02..d8460d09 100644 --- a/pkg/agent/phases/nodestart/nvidia.go +++ b/pkg/agent/phases/nodestart/nvidia.go @@ -68,7 +68,7 @@ func (s *setupNVIDIA) Do(ctx context.Context) error { return err } - return nil + return s.markReady(ctx) } func (s *setupNVIDIA) setupLibraries(ctx context.Context) error { @@ -119,6 +119,22 @@ func (s *setupNVIDIA) setupLibraries(ctx context.Context) error { return nil } +func (s *setupNVIDIA) markReady(ctx context.Context) error { + if _, err := executil.MachineRun(ctx, s.log, s.goalState.MachineName, + "mkdir", "-p", filepath.Dir(goalstates.NVIDIAReadyPath), + ); err != nil { + return fmt.Errorf("create NVIDIA ready directory: %w", err) + } + + if _, err := executil.MachineRun(ctx, s.log, s.goalState.MachineName, + "touch", goalstates.NVIDIAReadyPath, + ); err != nil { + return fmt.Errorf("mark NVIDIA runtime ready: %w", err) + } + + return nil +} + type nvidiaDriverRootPaths struct { rootDir string libDir string diff --git a/pkg/agent/phases/rootfs/assets/service-override.conf b/pkg/agent/phases/rootfs/assets/service-override.conf index 54e46852..b427cda8 100644 --- a/pkg/agent/phases/rootfs/assets/service-override.conf +++ b/pkg/agent/phases/rootfs/assets/service-override.conf @@ -40,6 +40,9 @@ RestartSec=10s ExecStartPre=-/usr/bin/machinectl terminate {{.MachineName}} ExecStartPre=/usr/bin/mkdir -p {{.BPFFSMountPath}} ExecStartPre=/bin/sh -c '/usr/bin/mountpoint -q {{.BPFFSMountPath}} || /usr/bin/mount -t bpf bpf {{.BPFFSMountPath}}' +{{- if .NvidiaEnabled}} +ExecStartPost=/bin/sh -c 'if [ ! -x /usr/local/bin/unbounded-agent-current ]; then exit 0; fi; exec /usr/local/bin/unbounded-agent-current reconcile-nvidia {{.MachineName}}' +{{- end}} Environment=SYSTEMD_NSPAWN_UNIFIED_HIERARCHY=1 Environment=SYSTEMD_NSPAWN_API_VFS_WRITABLE=network {{- if or .HostDevicePaths .HostDeviceGroupSpecifiers}} diff --git a/pkg/agent/phases/rootfs/nspawn.go b/pkg/agent/phases/rootfs/nspawn.go index 06c6d06e..1701408a 100644 --- a/pkg/agent/phases/rootfs/nspawn.go +++ b/pkg/agent/phases/rootfs/nspawn.go @@ -101,6 +101,7 @@ type nspawnTemplateData struct { NvidiaLibDirMounts []goalstates.NvidiaLibDirMount NvidiaI386LibDirMounts []goalstates.NvidiaLibDirMount NvidiaBinDir string + NvidiaEnabled bool AMDGPUDevicePaths []string AMDSysFSPaths []string } @@ -136,6 +137,7 @@ func (e *ensureNSpawnWorkspace) writeNSpawnConfigs() error { NvidiaLibDirMounts: e.goalState.Nvidia.LibDirMounts, NvidiaI386LibDirMounts: e.goalState.Nvidia.I386LibDirMounts, NvidiaBinDir: nvidiaHostBinDir(e.goalState.Nvidia), + NvidiaEnabled: nvidiaSetupEnabled(e.goalState.Nvidia), AMDGPUDevicePaths: amdGPUDevicePaths, AMDSysFSPaths: e.goalState.AMD.SysFSPaths, } @@ -188,6 +190,10 @@ func (e *ensureNSpawnWorkspace) writeNSpawnConfigs() error { return nil } +func nvidiaSetupEnabled(nvidia goalstates.NvidiaHost) bool { + return len(nvidia.GPUDevicePaths) > 0 && len(nvidia.LibMappings) > 0 +} + func nvidiaHostBinDir(nvidia goalstates.NvidiaHost) string { for _, path := range []string{nvidia.NvidiaSMIPath, nvidia.NvidiaIMEXPath, nvidia.NvidiaIMEXCtlPath} { if path != "" { diff --git a/pkg/agent/phases/rootfs/nspawn_render_test.go b/pkg/agent/phases/rootfs/nspawn_render_test.go index b630148e..a5e8b5c2 100644 --- a/pkg/agent/phases/rootfs/nspawn_render_test.go +++ b/pkg/agent/phases/rootfs/nspawn_render_test.go @@ -319,6 +319,32 @@ func TestServiceOverride_NoHostDevicesNoDeviceAllow(t *testing.T) { require.NotContains(t, buf.String(), "DeviceAllow=") } +func TestServiceOverride_NVIDIAReconcilesOnEveryStart(t *testing.T) { + t.Parallel() + + var buf bytes.Buffer + require.NoError(t, nspawnTemplates.ExecuteTemplate(&buf, "service-override.conf", nspawnTemplateData{ + MachineName: "kube1", + BPFFSMountPath: goalstates.BPFFSMountPath("kube1"), + NvidiaEnabled: true, + })) + + require.Contains(t, buf.String(), + "exec /usr/local/bin/unbounded-agent-current reconcile-nvidia kube1") +} + +func TestServiceOverride_CPUNodesDoNotReconcileNVIDIA(t *testing.T) { + t.Parallel() + + var buf bytes.Buffer + require.NoError(t, nspawnTemplates.ExecuteTemplate(&buf, "service-override.conf", nspawnTemplateData{ + MachineName: "kube1", + BPFFSMountPath: goalstates.BPFFSMountPath("kube1"), + })) + + require.NotContains(t, buf.String(), "reconcile-nvidia") +} + func nspawnRenderScenarioData() nspawnTemplateData { return nspawnTemplateData{ MachineName: "kube1",