feat(agent): validate prequalified NVIDIA DRA hosts

api/machina/v1alpha3/machine_types.go

@@ -95,6 +95,11 @@ const (
 	// For OnDelete updates this remains True until the operator deletes
 	// the Node and the agent repaves with the desired version.
 	MachineConditionRepavePending = "RepavePending"
+
+	// MachineConditionNVIDIAIMEXPrepared indicates whether the agent prepared
+	// the host and nspawn boundary for the configured NVIDIA IMEX lifecycle.
+	// It does not assert that a DRA ComputeDomain is ready.
+	MachineConditionNVIDIAIMEXPrepared = "NVIDIAIMEXPrepared"
 )
 
 // Annotation keys.

api/machina/v1alpha3/machineconfiguration_types.go

@@ -136,14 +136,43 @@ type MachineConfigurationKubernetes struct {
 type MachineConfigurationAgent struct {
 	// Image is the OCI image reference used for provisioning the
 	// nspawn machine (e.g. "ghcr.io/org/repo:tag").
-	// +kubebuilder:validation:Required
-	Image string `json:"image"`
+	// +optional
+	Image string `json:"image,omitempty"`
 
 	// LocalDNS configures the optional CoreDNS cache for this version.
 	// +optional
 	LocalDNS *LocalDNSSpec `json:"localDNS,omitempty"`
+
+	// NVIDIA configures optional NVIDIA host integration.
+	// +optional
+	NVIDIA *MachineConfigurationNVIDIA `json:"nvidia,omitempty"`
+}
+
+// MachineConfigurationNVIDIA configures NVIDIA host integration.
+type MachineConfigurationNVIDIA struct {
+	// IMEX configures multi-node NVLink IMEX preparation.
+	// +optional
+	IMEX *MachineConfigurationNVIDIAIMEX `json:"imex,omitempty"`
 }
 
+// MachineConfigurationNVIDIAIMEX configures NVIDIA IMEX integration.
+type MachineConfigurationNVIDIAIMEX struct {
+	// Mode selects the IMEX lifecycle owner. DRA prepares the node for NVIDIA
+	// DRA ComputeDomains; an empty value disables IMEX preparation.
+	// +optional
+	// +kubebuilder:validation:Enum=DRA
+	Mode NVIDIAIMEXMode `json:"mode,omitempty"`
+}
+
+// NVIDIAIMEXMode identifies the supported NVIDIA IMEX lifecycle owner.
+type NVIDIAIMEXMode string
+
+const (
+	// NVIDIAIMEXModeDRA delegates IMEX daemon and channel lifecycle to NVIDIA
+	// DRA ComputeDomains.
+	NVIDIAIMEXModeDRA NVIDIAIMEXMode = "DRA"
+)
+
 // MachineConfigurationUpdateStrategyType defines how configuration
 // changes are applied to machines.
 // +kubebuilder:validation:Enum=OnDelete;RollingUpdate

api/machina/v1alpha3/machineoperation_types.go

@@ -179,6 +179,18 @@ type MachineOperationTargetInput struct {
 	// than reselecting intent from a Machine edited during execution.
 	// +optional
 	ProvisioningFormat ProvisioningFormat `json:"provisioningFormat,omitempty"`
+
+	// ConfigurationVersionName is the exact MachineConfigurationVersion
+	// selected with HostImage for HostReplace. Retries reuse it so host and
+	// agent configuration cannot drift across version publication.
+	// +optional
+	ConfigurationVersionName string `json:"configurationVersionName,omitempty"`
+
+	// ConfigurationTemplate is the immutable versioned agent intent selected
+	// with HostImage. Snapshotting the data keeps HostReplace retries independent
+	// of Machine edits and MachineConfigurationVersion garbage collection.
+	// +optional
+	ConfigurationTemplate *MachineConfigurationTemplate `json:"configurationTemplate,omitempty"`
 }
 
 // ProviderMachineSnapshot identifies the exact provider-owned resource

cmd/agent/internal/cmd/bootstrap.go

@@ -29,6 +29,7 @@ type agentStages struct {
 	log              *slog.Logger
 	cfg              *provision.UnboundedAgentConfig
 	gs               *goalstates.MachineGoalState
+	downloads        *goalstates.DownloadOverrides
 	archives         *goalstates.ContainerImageArchiveStaging
 	reporter         *daemon.BootstrapStatusReporter
 	credentialsReady bool
@@ -69,7 +70,13 @@ func bootstrapIdentity(cfg *provision.UnboundedAgentConfig) (bootstrap.Identity,
 		KubernetesVersion string
 		OCIImage          string
 		APIServer         string
-	}{strings.TrimPrefix(cfg.Cluster.Version, "v"), canonicalImageIdentity(cfg.OCIImage), cfg.Kubelet.ApiServer})
+		NVIDIA            *provision.NVIDIAConfig `json:",omitempty"`
+	}{
+		strings.TrimPrefix(cfg.Cluster.Version, "v"),
+		canonicalImageIdentity(cfg.OCIImage),
+		cfg.Kubelet.ApiServer,
+		cfg.NVIDIA,
+	})
 	if err != nil {
 		return bootstrap.Identity{}, err
 	}
@@ -88,11 +95,7 @@ func (s *agentStages) ResolveInputs(ctx context.Context) error {
 	}
 
 	s.archives = archives
-
-	s.gs, err = goalstates.ResolveMachine(s.log, &s.cfg.AgentConfig, goalstates.NSpawnMachineKube1, downloads)
-	if err != nil {
-		return err
-	}
+	s.downloads = downloads
 
 	return nil
 }
@@ -106,12 +109,19 @@ func (s *agentStages) PrepareHost(ctx context.Context) error {
 		return err
 	}
 
-	if err := phases.Serial(s.log, host.InstallPackages(s.log), phases.Parallel(s.log,
+	if err := phases.Serial(s.log, host.InstallPackages(s.log), host.PrepareNVIDIAIMEXDRA(s.log, &s.cfg.AgentConfig), phases.Parallel(s.log,
 		host.ConfigureOS(s.log), host.ConfigureNFTables(s.log), phases.Serial(s.log, host.DisableDocker(s.log), host.ConfigureDocker(s.log)),
 		host.DisableContainerd(s.log), host.DisableKubelet(s.log), host.DisableSwap(s.log), host.HardenAPT(s.log))).Do(ctx); err != nil {
 		return err
 	}
 
+	gs, err := goalstates.ResolveMachine(s.log, &s.cfg.AgentConfig, goalstates.NSpawnMachineKube1, s.downloads)
+	if err != nil {
+		return err
+	}
+
+	s.gs = gs
+
 	return fsutil.SyncFilesystems("/etc", hostroot.Resolve(), installstate.DefaultDirectory)
 }
 

cmd/agent/internal/cmd/bootstrap_test.go

@@ -106,6 +106,13 @@ func TestBootstrapFingerprintAllowsCredentialAndDownloadRefresh(t *testing.T) {
 
 		cfg.Cluster.Version, cfg.OCIImage, cfg.Kubelet.ApiServer = version, image, endpoint
 	}
+
+	cfg.NVIDIA = &provision.NVIDIAConfig{
+		IMEX: &provision.NVIDIAIMEXConfig{Mode: provision.NVIDIAIMEXModeDRA},
+	}
+	changed, err := bootstrapIdentity(cfg)
+	require.NoError(t, err)
+	require.NotEqual(t, original.ConfigFingerprint, changed.ConfigFingerprint)
 }
 
 func TestBootstrapFingerprintIgnoresSignedImageQuery(t *testing.T) {

cmd/agent/internal/daemon/bootstrap_status.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"slices"
 	"strings"
 
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
@@ -29,6 +30,7 @@ type BootstrapStatusReporter struct {
 	client  client.Client
 	machine string
 	ready   bool
+	imexDRA bool
 }
 
 // NewBootstrapStatusReporter builds a best-effort Machine status reporter from
@@ -38,6 +40,7 @@ func NewBootstrapStatusReporter(ctx context.Context, log *slog.Logger, cfg *prov
 	reporter := &BootstrapStatusReporter{
 		log:     log,
 		machine: cfg.MachineName,
+		imexDRA: cfg.IMEXDRAEnabled(),
 	}
 
 	if strings.TrimSpace(cfg.Kubelet.Auth.BootstrapToken) == "" {
@@ -71,6 +74,7 @@ func NewBootstrapStatusReporter(ctx context.Context, log *slog.Logger, cfg *prov
 // Running reports that initial bootstrap is in progress.
 func (r *BootstrapStatusReporter) Running(ctx context.Context) {
 	r.set(ctx, metav1.ConditionFalse, "Running", "unbounded-agent bootstrap is running")
+	r.setIMEX(ctx, metav1.ConditionFalse, "Preparing", "NVIDIA IMEX DRA host preparation is in progress")
 }
 
 // Failed reports that initial bootstrap failed.
@@ -80,11 +84,78 @@ func (r *BootstrapStatusReporter) Failed(ctx context.Context, reason string, err
 	}
 
 	r.set(ctx, metav1.ConditionFalse, reason, truncateBootstrapConditionMessage(err.Error()))
+	r.setIMEX(ctx, metav1.ConditionFalse, reason, truncateBootstrapConditionMessage(err.Error()))
 }
 
 // Succeeded reports that initial bootstrap completed successfully.
 func (r *BootstrapStatusReporter) Succeeded(ctx context.Context) {
 	r.set(ctx, metav1.ConditionTrue, "Succeeded", "unbounded-agent bootstrap completed successfully")
+	r.setIMEX(
+		ctx,
+		metav1.ConditionTrue,
+		"Prepared",
+		"host and nspawn are prepared for NVIDIA DRA IMEX; this does not assert ComputeDomain readiness",
+	)
+}
+
+func (r *BootstrapStatusReporter) setIMEX(
+	ctx context.Context,
+	status metav1.ConditionStatus,
+	reason string,
+	message string,
+) {
+	if r == nil || !r.ready {
+		return
+	}
+
+	if !r.imexDRA {
+		r.clearIMEX(ctx)
+
+		return
+	}
+
+	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
+		var machine v1alpha3.Machine
+		if err := r.client.Get(ctx, client.ObjectKey{Name: r.machine}, &machine); err != nil {
+			return fmt.Errorf("get Machine %q: %w", r.machine, err)
+		}
+
+		apimeta.SetStatusCondition(&machine.Status.Conditions, metav1.Condition{
+			Type:               v1alpha3.MachineConditionNVIDIAIMEXPrepared,
+			Status:             status,
+			Reason:             reason,
+			Message:            message,
+			ObservedGeneration: machine.Generation,
+		})
+
+		return r.client.Status().Update(ctx, &machine)
+	})
+	if err != nil {
+		r.log.Warn("failed to update NVIDIAIMEXPrepared condition", "machine", r.machine, "reason", reason, "error", err)
+	}
+}
+
+func (r *BootstrapStatusReporter) clearIMEX(ctx context.Context) {
+	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
+		var machine v1alpha3.Machine
+		if err := r.client.Get(ctx, client.ObjectKey{Name: r.machine}, &machine); err != nil {
+			return fmt.Errorf("get Machine %q: %w", r.machine, err)
+		}
+
+		conditions := slices.DeleteFunc(machine.Status.Conditions, func(condition metav1.Condition) bool {
+			return condition.Type == v1alpha3.MachineConditionNVIDIAIMEXPrepared
+		})
+		if len(conditions) == len(machine.Status.Conditions) {
+			return nil
+		}
+
+		machine.Status.Conditions = conditions
+
+		return r.client.Status().Update(ctx, &machine)
+	})
+	if err != nil {
+		r.log.Warn("failed to clear NVIDIAIMEXPrepared condition", "machine", r.machine, "error", err)
+	}
 }
 
 func (r *BootstrapStatusReporter) set(ctx context.Context, status metav1.ConditionStatus, reason, message string) {