Reconcile NVIDIA state on every nspawn machine start

cmd/agent/internal/cmd/cmd.go

@@ -31,6 +31,7 @@ func Run() {
 		newCmdReset(cmdCtx),
 		newCmdVersion(),
 		newCmdRecordAgentUpgradeFailureSignal(),
+		newCmdReconcileNVIDIA(cmdCtx),
 	)
 
 	if err := root.Execute(); err != nil {

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))
+}

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")
+}

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"

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

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

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

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.

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()
 

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
 	}

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()
 

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

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