diff --git a/deploy/helm/charts/platform/components/operator/files/manager-role.yaml b/deploy/helm/charts/platform/components/operator/files/manager-role.yaml index 66d383706fd4..b3681d70fd7d 100644 --- a/deploy/helm/charts/platform/components/operator/files/manager-role.yaml +++ b/deploy/helm/charts/platform/components/operator/files/manager-role.yaml @@ -128,6 +128,18 @@ rules: - patch - update - watch +- apiGroups: + - disaggregatedset.x-k8s.io + resources: + - disaggregatedsets + verbs: + - create + - delete + - get + - list + - patch + - update + - watch - apiGroups: - discovery.k8s.io resources: diff --git a/deploy/operator/config/rbac/role.yaml b/deploy/operator/config/rbac/role.yaml index 66d383706fd4..b3681d70fd7d 100644 --- a/deploy/operator/config/rbac/role.yaml +++ b/deploy/operator/config/rbac/role.yaml @@ -128,6 +128,18 @@ rules: - patch - update - watch +- apiGroups: + - disaggregatedset.x-k8s.io + resources: + - disaggregatedsets + verbs: + - create + - delete + - get + - list + - patch + - update + - watch - apiGroups: - discovery.k8s.io resources: diff --git a/deploy/operator/internal/consts/consts.go b/deploy/operator/internal/consts/consts.go index 39b8f6ee8c0c..d2ea93025ff4 100644 --- a/deploy/operator/internal/consts/consts.go +++ b/deploy/operator/internal/consts/consts.go @@ -40,7 +40,8 @@ const ( KubeLabelDynamoSelector = "nvidia.com/selector" - KubeAnnotationEnableGrove = "nvidia.com/enable-grove" + KubeAnnotationEnableGrove = "nvidia.com/enable-grove" + KubeAnnotationEnableDisaggregatedSet = "nvidia.com/enable-disaggregatedset" // KubeAnnotationGroveUpdateStrategy temporarily exposes the Grove // PodCliqueSet update strategy while the long-term DGD API is settled. diff --git a/deploy/operator/internal/controller/dynamocomponentdeployment_controller.go b/deploy/operator/internal/controller/dynamocomponentdeployment_controller.go index f71ef43a6d44..bb3653525c23 100644 --- a/deploy/operator/internal/controller/dynamocomponentdeployment_controller.go +++ b/deploy/operator/internal/controller/dynamocomponentdeployment_controller.go @@ -551,61 +551,74 @@ func (r *DynamoComponentDeploymentReconciler) generateLeaderWorkerSet(ctx contex logs := log.FromContext(ctx) logs.Info("Generating LeaderWorkerSet") + leaderPodTemplateSpec, workerPodTemplateSpec, err := r.renderMultinodePodTemplateSpecs(ctx, opt) + if err != nil { + return nil, false, err + } + kubeName := leaderWorkerSetName(opt.dynamoComponentDeployment) kubeNs := opt.dynamoComponentDeployment.Namespace labels := dynamo.GetDCDKubeLabels(opt.dynamoComponentDeployment) - if labels == nil { labels = make(map[string]string) } - podLabels, err := r.getDCDWorkloadPodLabels(ctx, opt.dynamoComponentDeployment) - if err != nil { - return nil, false, err + + desiredReplicas := int32(1) + if opt.dynamoComponentDeployment.Spec.Replicas != nil { + desiredReplicas = *opt.dynamoComponentDeployment.Spec.Replicas } + groupSize := opt.dynamoComponentDeployment.GetNumberOfNodes() - leaderWorkerSet := &leaderworkersetv1.LeaderWorkerSet{ + return &leaderworkersetv1.LeaderWorkerSet{ ObjectMeta: metav1.ObjectMeta{ Name: kubeName, Namespace: kubeNs, Labels: labels, }, - } + Spec: leaderworkersetv1.LeaderWorkerSetSpec{ + Replicas: &desiredReplicas, + StartupPolicy: leaderworkersetv1.LeaderCreatedStartupPolicy, + LeaderWorkerTemplate: leaderworkersetv1.LeaderWorkerTemplate{ + LeaderTemplate: leaderPodTemplateSpec, + WorkerTemplate: *workerPodTemplateSpec, + Size: &groupSize, + }, + }, + }, false, nil +} - leaderPodLabels := make(map[string]string) - for k, v := range podLabels { - leaderPodLabels[k] = v - } - leaderPodTemplateSpec, err := r.generateLeaderPodTemplateSpec(ctx, opt, leaderPodLabels) +// renderMultinodePodTemplateSpecs is the shared multinode render path used by +// both the LWS pathway and the DisaggregatedSet pathway. Returning the two +// pod templates (rather than a fully composed LWS) keeps downstream callers +// (DS, future templating work) free to compose their own role structure. +func (r *DynamoComponentDeploymentReconciler) renderMultinodePodTemplateSpecs( + ctx context.Context, + opt generateResourceOption, +) (*corev1.PodTemplateSpec, *corev1.PodTemplateSpec, error) { + podLabels, err := r.getDCDWorkloadPodLabels(ctx, opt.dynamoComponentDeployment) if err != nil { - return nil, false, errors.Wrap(err, "generateLeaderWorkerSet: failed to generate leader pod template") + return nil, nil, err } - workerPodLabels := make(map[string]string) + leaderLabels := make(map[string]string, len(podLabels)) for k, v := range podLabels { - workerPodLabels[k] = v + leaderLabels[k] = v } - workerPodTemplateSpec, err := r.generateWorkerPodTemplateSpec(ctx, opt, workerPodLabels) + leaderPodTemplateSpec, err := r.generateLeaderPodTemplateSpec(ctx, opt, leaderLabels) if err != nil { - return nil, false, errors.Wrap(err, "generateLeaderWorkerSet: failed to generate worker pod template") + return nil, nil, err } - desiredReplicas := int32(1) - if opt.dynamoComponentDeployment.Spec.Replicas != nil { - desiredReplicas = *opt.dynamoComponentDeployment.Spec.Replicas + workerLabels := make(map[string]string, len(podLabels)) + for k, v := range podLabels { + workerLabels[k] = v } - groupSize := opt.dynamoComponentDeployment.GetNumberOfNodes() - - leaderWorkerSet.Spec = leaderworkersetv1.LeaderWorkerSetSpec{ - Replicas: &desiredReplicas, - StartupPolicy: leaderworkersetv1.LeaderCreatedStartupPolicy, - LeaderWorkerTemplate: leaderworkersetv1.LeaderWorkerTemplate{ - LeaderTemplate: leaderPodTemplateSpec, - WorkerTemplate: *workerPodTemplateSpec, - Size: &groupSize, - }, + workerPodTemplateSpec, err := r.generateWorkerPodTemplateSpec(ctx, opt, workerLabels) + if err != nil { + return nil, nil, err } - return leaderWorkerSet, false, nil + return leaderPodTemplateSpec, workerPodTemplateSpec, nil } // getDCDWorkloadPodLabels keeps LWS pod labels aligned with the workload diff --git a/deploy/operator/internal/controller/dynamographdeployment_controller.go b/deploy/operator/internal/controller/dynamographdeployment_controller.go index 62d69cc001b5..7e431b0ef390 100644 --- a/deploy/operator/internal/controller/dynamographdeployment_controller.go +++ b/deploy/operator/internal/controller/dynamographdeployment_controller.go @@ -25,6 +25,7 @@ import ( groveconstants "github.com/ai-dynamo/grove/operator/api/common/constants" grovev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1" + "github.com/go-logr/logr" "github.com/imdario/mergo" "k8s.io/apimachinery/pkg/api/equality" "k8s.io/apimachinery/pkg/api/errors" @@ -53,6 +54,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" "sigs.k8s.io/controller-runtime/pkg/predicate" + leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1" configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1" nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1" @@ -107,6 +109,7 @@ type DynamoGraphDeploymentReconciler struct { // +kubebuilder:rbac:groups=grove.io,resources=podcliquescalinggroups,verbs=get;list;watch // +kubebuilder:rbac:groups=grove.io,resources=podcliquescalinggroups/scale,verbs=get;update;patch // +kubebuilder:rbac:groups=grove.io,resources=clustertopologybindings,verbs=get;list;watch +// +kubebuilder:rbac:groups=disaggregatedset.x-k8s.io,resources=disaggregatedsets,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=scheduling.run.ai,resources=queues,verbs=get;list // +kubebuilder:rbac:groups=inference.networking.k8s.io,resources=inferencepools,verbs=get;list;watch;create;update;patch;delete // +kubebuilder:rbac:groups=networking.istio.io,resources=destinationrules,verbs=get;list;watch;create;update;patch;delete @@ -337,42 +340,13 @@ type ReconcileResult struct { RestartStatus *nvidiacomv1beta1.RestartStatus } +const resourceNotFoundReason = "resource not found" + func (r *DynamoGraphDeploymentReconciler) reconcileResources(ctx context.Context, dynamoDeployment *nvidiacomv1beta1.DynamoGraphDeployment) (ReconcileResult, error) { logger := log.FromContext(ctx) - // Ensure planner RBAC exists in cluster-wide mode - if r.Config.Namespace.Restricted == "" { - if r.RBACManager == nil { - return ReconcileResult{}, fmt.Errorf("RBAC manager not initialized in cluster-wide mode") - } - if r.Config.RBAC.PlannerClusterRoleName == "" { - return ReconcileResult{}, fmt.Errorf("planner ClusterRole name is required in cluster-wide mode") - } - if err := r.RBACManager.EnsureServiceAccountWithRBAC( - ctx, - dynamoDeployment.Namespace, - consts.PlannerServiceAccountName, - r.Config.RBAC.PlannerClusterRoleName, - ); err != nil { - logger.Error(err, "Failed to ensure planner RBAC") - return ReconcileResult{}, fmt.Errorf("failed to ensure planner RBAC: %w", err) - } - - // Ensure EPP RBAC exists in cluster-wide mode if EPP service is present - if dynamoDeployment.HasEPPComponent() { - if r.Config.RBAC.EPPClusterRoleName == "" { - return ReconcileResult{}, fmt.Errorf("EPP ClusterRole name is required in cluster-wide mode when EPP service is present") - } - if err := r.RBACManager.EnsureServiceAccountWithRBAC( - ctx, - dynamoDeployment.Namespace, - consts.EPPServiceAccountName, - r.Config.RBAC.EPPClusterRoleName, - ); err != nil { - logger.Error(err, "Failed to ensure EPP RBAC") - return ReconcileResult{}, fmt.Errorf("failed to ensure EPP RBAC: %w", err) - } - } + if err := r.ensureClusterScopedRBAC(ctx, dynamoDeployment, logger); err != nil { + return ReconcileResult{}, err } // Reconcile top-level PVCs first @@ -427,27 +401,31 @@ func (r *DynamoGraphDeploymentReconciler) reconcileResources(ctx context.Context } } - // return error early if Grove and LWS is not available for multinode - if !r.isGrovePathway(dynamoDeployment) && hasMultinode && !r.RuntimeConfig.Gate.Enabled(features.LWS) { + useDisaggregatedSet, disaggregatedSetFallbackReason := r.shouldUseDisaggregatedSet(dynamoDeployment) + + // return error early if Grove, DS, and LWS are all unavailable for multinode + if !r.isGrovePathway(dynamoDeployment) && hasMultinode && !r.RuntimeConfig.Gate.Enabled(features.LWS) && !useDisaggregatedSet { err := fmt.Errorf("no multinode orchestrator available") - logger.Error(err, err.Error(), "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS)) + logger.Error(err, err.Error(), "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS), "disaggregatedSetEnabled", r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet)) return ReconcileResult{}, fmt.Errorf("failed to reconcile Dynamo components deployments: %w", err) } restartStatus := r.computeRestartStatus(ctx, dynamoDeployment) restartState := dynamo.DetermineRestartState(dynamoDeployment, restartStatus) + restartState = r.coalesceDisaggregatedSetRestartStateForPath(dynamoDeployment, useDisaggregatedSet, restartState) - var result ReconcileResult - if r.isGrovePathway(dynamoDeployment) { - logger.Info("Reconciling Grove resources", "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS)) - result, err = r.reconcileGroveResources(ctx, dynamoDeployment, restartState, checkpointInfos) - } else { - logger.Info("Reconciling Dynamo components deployments", "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS)) - result, err = r.reconcileDynamoComponentsDeployments(ctx, dynamoDeployment, restartState, checkpointInfos) - } + result, err := r.reconcileWorkloadResources( + ctx, + dynamoDeployment, + hasMultinode, + useDisaggregatedSet, + disaggregatedSetFallbackReason, + restartState, + checkpointInfos, + ) if err != nil { - logger.Error(err, "Failed to reconcile Dynamo components deployments") - return ReconcileResult{}, fmt.Errorf("failed to reconcile Dynamo components deployments: %w", err) + logger.Error(err, "Failed to reconcile workload resources") + return ReconcileResult{}, fmt.Errorf("failed to reconcile workload resources: %w", err) } result.RestartStatus = restartStatus result = applyCheckpointStartupReadiness(result, checkpointInfos) @@ -467,6 +445,132 @@ func (r *DynamoGraphDeploymentReconciler) reconcileResources(ctx context.Context return result, nil } +func (r *DynamoGraphDeploymentReconciler) ensureClusterScopedRBAC( + ctx context.Context, + dynamoDeployment *nvidiacomv1beta1.DynamoGraphDeployment, + logger logr.Logger, +) error { + if r.Config.Namespace.Restricted != "" { + return nil + } + if r.RBACManager == nil { + return fmt.Errorf("RBAC manager not initialized in cluster-wide mode") + } + if r.Config.RBAC.PlannerClusterRoleName == "" { + return fmt.Errorf("planner ClusterRole name is required in cluster-wide mode") + } + if err := r.RBACManager.EnsureServiceAccountWithRBAC( + ctx, + dynamoDeployment.Namespace, + consts.PlannerServiceAccountName, + r.Config.RBAC.PlannerClusterRoleName, + ); err != nil { + logger.Error(err, "Failed to ensure planner RBAC") + return fmt.Errorf("failed to ensure planner RBAC: %w", err) + } + + if !dynamoDeployment.HasEPPComponent() { + return nil + } + if r.Config.RBAC.EPPClusterRoleName == "" { + return fmt.Errorf("EPP ClusterRole name is required in cluster-wide mode when EPP service is present") + } + if err := r.RBACManager.EnsureServiceAccountWithRBAC( + ctx, + dynamoDeployment.Namespace, + consts.EPPServiceAccountName, + r.Config.RBAC.EPPClusterRoleName, + ); err != nil { + logger.Error(err, "Failed to ensure EPP RBAC") + return fmt.Errorf("failed to ensure EPP RBAC: %w", err) + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) reconcileWorkloadResources( + ctx context.Context, + dynamoDeployment *nvidiacomv1beta1.DynamoGraphDeployment, + hasMultinode bool, + useDisaggregatedSet bool, + disaggregatedSetFallbackReason string, + restartState *dynamo.RestartState, + checkpointInfos map[string]*checkpoint.CheckpointInfo, +) (ReconcileResult, error) { + logger := log.FromContext(ctx) + + if r.isGrovePathway(dynamoDeployment) { + logger.Info("Reconciling Grove resources", "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS)) + return r.reconcileReplacementBeforeDisaggregatedSetCleanup(ctx, dynamoDeployment, false, func() (ReconcileResult, error) { + return r.reconcileGroveResources(ctx, dynamoDeployment, restartState, checkpointInfos) + }) + } + + if useDisaggregatedSet { + logger.Info("Reconciling DisaggregatedSet resources", "hasMultinode", hasMultinode, "disaggregatedSetEnabled", r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet)) + return r.reconcileDisaggregatedSetResources(ctx, dynamoDeployment, restartState, checkpointInfos) + } + + if r.wantsDisaggregatedSet(dynamoDeployment) && disaggregatedSetFallbackReason != "" { + logger.Info("DisaggregatedSet requested but falling back to DynamoComponentDeployments", "reason", disaggregatedSetFallbackReason) + if r.Recorder != nil { + r.Recorder.Eventf(dynamoDeployment, corev1.EventTypeWarning, "DisaggregatedSetFallback", "DisaggregatedSet requested but falling back to DynamoComponentDeployments: %s", disaggregatedSetFallbackReason) + } + } + logger.Info("Reconciling Dynamo components deployments", "hasMultinode", hasMultinode, "lwsEnabled", r.RuntimeConfig.Gate.Enabled(features.LWS)) + return r.reconcileReplacementBeforeDisaggregatedSetCleanup(ctx, dynamoDeployment, true, func() (ReconcileResult, error) { + return r.reconcileDynamoComponentsDeployments(ctx, dynamoDeployment, restartState, checkpointInfos) + }) +} + +func (r *DynamoGraphDeploymentReconciler) reconcileReplacementBeforeDisaggregatedSetCleanup( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + restoreDCDServiceOwnership bool, + reconcileReplacement func() (ReconcileResult, error), +) (ReconcileResult, error) { + result, err := reconcileReplacement() + if err != nil || result.State != nvidiacomv1beta1.DGDStateSuccessful { + return result, err + } + if restoreDCDServiceOwnership { + if err := r.restoreDisaggregatedSetServiceOwnershipToDCDs(ctx, dgd); err != nil { + return ReconcileResult{}, err + } + } + if err := r.deleteDisaggregatedSetOnLegacyPath(ctx, dgd); err != nil { + return ReconcileResult{}, err + } + return result, nil +} + +func (r *DynamoGraphDeploymentReconciler) deleteDisaggregatedSetOnLegacyPath(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment) error { + if !r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet) { + return nil + } + return r.deleteDisaggregatedSetIfExists(ctx, dgd) +} + +func (r *DynamoGraphDeploymentReconciler) deleteGrovePodCliqueSetOnDisaggregatedSetPath(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment) error { + if !r.RuntimeConfig.Gate.Enabled(features.Grove) { + return nil + } + pcs := &grovev1alpha1.PodCliqueSet{} + key := types.NamespacedName{Name: dynamo.PCSNameForDGD(dgd.Name, dgd.Spec.Components), Namespace: dgd.Namespace} + if err := r.Client.Get(ctx, key, pcs); err != nil { + if errors.IsNotFound(err) { + return nil + } + return fmt.Errorf("failed to get stale PodCliqueSet %s: %w", key, err) + } + if !isControlledByBetaDGD(pcs, dgd) { + return fmt.Errorf("refusing to delete PodCliqueSet %s because it is not controlled by DynamoGraphDeployment %s/%s", key, dgd.Namespace, dgd.Name) + } + if err := r.Client.Delete(ctx, pcs); err != nil && !errors.IsNotFound(err) { + return fmt.Errorf("failed to delete stale PodCliqueSet %s: %w", key, err) + } + return nil +} + func (r *DynamoGraphDeploymentReconciler) isGrovePathway(dgd *nvidiacomv1beta1.DynamoGraphDeployment) bool { return r.RuntimeConfig.Gate.Enabled(features.Grove) && (dgd.Annotations == nil || strings.ToLower(dgd.Annotations[consts.KubeAnnotationEnableGrove]) != consts.KubeLabelValueFalse) @@ -476,6 +580,9 @@ func (r *DynamoGraphDeploymentReconciler) getUpdatedInProgress(ctx context.Conte if r.isGrovePathway(dgd) { return r.getUpdatedInProgressForGrove(ctx, dgd, inProgress) } + if useDisaggregatedSet, _ := r.shouldUseDisaggregatedSet(dgd); useDisaggregatedSet { + return r.getUpdatedInProgressForDisaggregatedSet(ctx, dgd, inProgress) + } return r.getUpdatedInProgressForComponent(ctx, dgd, inProgress) } @@ -1424,7 +1531,7 @@ func (r *DynamoGraphDeploymentReconciler) computeSequentialRestartStatus( logger.Info("Component restart completed", "component", currentComponent) // Find the next component. - nextComponent, currentFound := getNextComponentInOrder(order, currentComponent) + nextComponent, currentFound := r.getNextSequentialRestartComponent(dgd, order, currentComponent) if !currentFound { logger.Info("Current restart component is no longer in order, restarting sequence from first component", "component", currentComponent, "firstComponent", order[0]) return &nvidiacomv1beta1.RestartStatus{ @@ -1452,6 +1559,40 @@ func (r *DynamoGraphDeploymentReconciler) computeSequentialRestartStatus( } } +// getNextSequentialRestartComponent skips selected DisaggregatedSet roles after +// one of them completes. Those roles share one DisaggregatedSet revision and +// are restarted together, even when the DGD restart strategy is sequential. +func (r *DynamoGraphDeploymentReconciler) getNextSequentialRestartComponent( + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + order []string, + currentComponent string, +) (string, bool) { + nextComponent, currentFound := getNextComponentInOrder(order, currentComponent) + if !currentFound || nextComponent == "" { + return nextComponent, currentFound + } + + useDisaggregatedSet, _ := r.shouldUseDisaggregatedSet(dgd) + if !useDisaggregatedSet || r.isGrovePathway(dgd) { + return nextComponent, currentFound + } + selection, reason := selectDisaggregatedSetComponents(dgd) + if reason != "" { + return nextComponent, currentFound + } + if _, selected := selection.componentToRole[currentComponent]; !selected { + return nextComponent, currentFound + } + + for nextComponent != "" { + if _, selected := selection.componentToRole[nextComponent]; !selected { + return nextComponent, true + } + nextComponent, _ = getNextComponentInOrder(order, nextComponent) + } + return "", true +} + // getNextComponentInOrder returns the component after the current component. // The boolean reports whether currentComponent was found in order. func getNextComponentInOrder(order []string, currentComponent string) (string, bool) { @@ -1517,7 +1658,7 @@ func (r *DynamoGraphDeploymentReconciler) checkComponentFullyUpdated(ctx context for _, hash := range r.activeWorkerHashCandidates(dgd, hashes) { resourceName := dynamo.GetDCDResourceName(dgd, componentName, hash) ready, reason := checkDCDReady(ctx, r.Client, resourceName, dgd.Namespace) - if ready || reason != "resource not found" { + if ready || reason != resourceNotFoundReason { return ready, reason } lastReason = reason @@ -1536,7 +1677,7 @@ func checkDCDReady(ctx context.Context, client client.Client, resourceName, name if err != nil { if errors.IsNotFound(err) { logger.V(2).Info("DynamoComponentDeployment not found", "resourceName", resourceName) - return false, "resource not found" + return false, resourceNotFoundReason } logger.V(1).Info("Failed to get DynamoComponentDeployment", "error", err, "resourceName", resourceName) return false, fmt.Sprintf("get error: %v", err) @@ -2769,6 +2910,30 @@ func (r *DynamoGraphDeploymentReconciler) FinalizeResource(ctx context.Context, return r.deleteAutoCheckpointsForDGD(ctx, dynamoDeployment) } +func workloadRoutingAnnotationsChanged(update event.UpdateEvent) bool { + oldDGD, oldOK := update.ObjectOld.(*nvidiacomv1beta1.DynamoGraphDeployment) + newDGD, newOK := update.ObjectNew.(*nvidiacomv1beta1.DynamoGraphDeployment) + if !oldOK || !newOK { + return false + } + annotationValue := func(dgd *nvidiacomv1beta1.DynamoGraphDeployment, key string) string { + return strings.ToLower(dgd.GetAnnotations()[key]) + } + return annotationValue(oldDGD, consts.KubeAnnotationEnableDisaggregatedSet) != annotationValue(newDGD, consts.KubeAnnotationEnableDisaggregatedSet) || + annotationValue(oldDGD, consts.KubeAnnotationEnableGrove) != annotationValue(newDGD, consts.KubeAnnotationEnableGrove) +} + +func dgdOwnedServiceEventPredicate() predicate.Predicate { + return predicate.Funcs{ + // DS-path Services are controlled directly by the DGD. Ignore creation + // because the current reconcile just created the desired object. + CreateFunc: func(ce event.CreateEvent) bool { return false }, + DeleteFunc: func(de event.DeleteEvent) bool { return true }, + UpdateFunc: func(ue event.UpdateEvent) bool { return true }, + GenericFunc: func(ge event.GenericEvent) bool { return true }, + } +} + // SetupWithManager sets up the controller with the Manager. func (r *DynamoGraphDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) error { if err := mgr.GetFieldIndexer().IndexField( @@ -2782,7 +2947,10 @@ func (r *DynamoGraphDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) err ctrlBuilder := ctrl.NewControllerManagedBy(mgr). For(&nvidiacomv1beta1.DynamoGraphDeployment{}, builder.WithPredicates( - predicate.GenerationChangedPredicate{}, + predicate.Or( + predicate.GenerationChangedPredicate{}, + predicate.Funcs{UpdateFunc: workloadRoutingAnnotationsChanged}, + ), )). Named(consts.ResourceTypeDynamoGraphDeployment). Watches( @@ -2807,6 +2975,7 @@ func (r *DynamoGraphDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) err UpdateFunc: func(de event.UpdateEvent) bool { return true }, GenericFunc: func(ge event.GenericEvent) bool { return true }, })). + Owns(&corev1.Service{}, builder.WithPredicates(dgdOwnedServiceEventPredicate())). Owns(&nvidiacomv1alpha1.DynamoGraphDeploymentScalingAdapter{}, builder.WithPredicates(predicate.Funcs{ // ignore creation cause we don't want to be called again after we create the adapter CreateFunc: func(ce event.CreateEvent) bool { return false }, @@ -2830,6 +2999,23 @@ func (r *DynamoGraphDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) err GenericFunc: func(ge event.GenericEvent) bool { return false }, })) } + if r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet) { + ctrlBuilder = ctrlBuilder.Owns(newDisaggregatedSetObject(), builder.WithPredicates(predicate.Funcs{ + CreateFunc: func(ce event.CreateEvent) bool { return false }, + DeleteFunc: func(de event.DeleteEvent) bool { return true }, + UpdateFunc: func(ue event.UpdateEvent) bool { return disaggregatedSetStatusChanged(ue.ObjectOld, ue.ObjectNew) }, + GenericFunc: func(ge event.GenericEvent) bool { return true }, + })).Watches( + &leaderworkersetv1.LeaderWorkerSet{}, + handler.EnqueueRequestsFromMapFunc(r.mapDisaggregatedSetChildLWSToDGD), + builder.WithPredicates(predicate.Funcs{ + CreateFunc: func(ce event.CreateEvent) bool { return true }, + DeleteFunc: func(de event.DeleteEvent) bool { return true }, + UpdateFunc: func(ue event.UpdateEvent) bool { return leaderWorkerSetStatusChanged(ue.ObjectOld, ue.ObjectNew) }, + GenericFunc: func(ge event.GenericEvent) bool { return false }, + }), + ) + } if r.RuntimeConfig.Gate.Enabled(features.Grove) { ctrlBuilder = ctrlBuilder.Owns(&grovev1alpha1.PodCliqueSet{}, builder.WithPredicates(predicate.Funcs{ // ignore creation cause we don't want to be called again after we create the pod gang set diff --git a/deploy/operator/internal/controller/dynamographdeployment_controller_test.go b/deploy/operator/internal/controller/dynamographdeployment_controller_test.go index 07aafd5af21b..31e5b7ed2b97 100644 --- a/deploy/operator/internal/controller/dynamographdeployment_controller_test.go +++ b/deploy/operator/internal/controller/dynamographdeployment_controller_test.go @@ -78,6 +78,37 @@ func newDynamoGraphDeploymentControllerTestScheme(t testing.TB) *runtime.Scheme return s } +func TestDeleteGrovePodCliqueSetOnDisaggregatedSetPath(t *testing.T) { + testScheme := newDynamoGraphDeploymentControllerTestScheme(t) + dgd := &v1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "dgd-uid"}, + Spec: v1beta1.DynamoGraphDeploymentSpec{Components: []v1beta1.DynamoComponentDeploymentSharedSpec{ + {ComponentName: "prefill", ComponentType: v1beta1.ComponentTypePrefill}, + }}, + } + pcs := &grovev1alpha1.PodCliqueSet{ObjectMeta: metav1.ObjectMeta{ + Name: dynamo.PCSNameForDGD(dgd.Name, dgd.Spec.Components), + Namespace: dgd.Namespace, + OwnerReferences: []metav1.OwnerReference{{ + APIVersion: v1beta1.GroupVersion.String(), + Kind: "DynamoGraphDeployment", + Name: dgd.Name, + UID: dgd.UID, + Controller: ptr.To(true), + }}, + }} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(testScheme).WithObjects(dgd, pcs).Build(), + RuntimeConfig: &controller_common.RuntimeConfig{Gate: features.Gates{Grove: true}}, + } + + require.NoError(t, reconciler.deleteGrovePodCliqueSetOnDisaggregatedSetPath(t.Context(), dgd)) + err := reconciler.Get(t.Context(), client.ObjectKeyFromObject(pcs), &grovev1alpha1.PodCliqueSet{}) + require.True(t, apierrors.IsNotFound(err)) + // Deletion races and repeated reconciliation are idempotent. + require.NoError(t, reconciler.deleteGrovePodCliqueSetOnDisaggregatedSetPath(t.Context(), dgd)) +} + func TestDynamoGraphDeploymentReconciler_preserveExistingDCDBackendFramework(t *testing.T) { ctx := context.Background() testScheme := newDynamoGraphDeploymentControllerTestScheme(t) diff --git a/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset.go b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset.go new file mode 100644 index 000000000000..e1997ac24874 --- /dev/null +++ b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset.go @@ -0,0 +1,1435 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 controller + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "maps" + "sort" + "strings" + + "github.com/ai-dynamo/dynamo/deploy/operator/internal/checkpoint" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/equality" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/apimachinery/pkg/types" + "k8s.io/utils/ptr" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/log" + disaggregatedsetv1 "sigs.k8s.io/lws/api/disaggregatedset/v1" + leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1" + disaggregatedsetutils "sigs.k8s.io/lws/pkg/utils/disaggregatedset" + + nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1" + nvidiacomv1beta1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta1" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" +) + +var disaggregatedSetGVK = schema.GroupVersionKind{ + Group: "disaggregatedset.x-k8s.io", + Version: "v1", + Kind: "DisaggregatedSet", +} + +const ( + maxDisaggregatedSetRoles = 10 + disaggregatedSetRevisionLength = 8 + // LWS names children -- without truncation. Reserve + // readable budgets for both user-derived names while keeping the result at + // the DNS label limit: 31 + 1 + 8 + 1 + 22 = 63. + maxDisaggregatedSetNameLength = 31 + maxDisaggregatedSetRoleNameLength = 63 - maxDisaggregatedSetNameLength - disaggregatedSetRevisionLength - 2 + disaggregatedSetNameHashLength = 8 + dynamoGraphDeploymentKind = "DynamoGraphDeployment" + dynamoComponentDeploymentKind = "DynamoComponentDeployment" +) + +type disaggregatedSetSelection struct { + componentToRole map[string]string + desiredReplicas map[string]int32 +} + +func newDisaggregatedSetObject() *unstructured.Unstructured { + obj := &unstructured.Unstructured{} + obj.SetGroupVersionKind(disaggregatedSetGVK) + return obj +} + +func (r *DynamoGraphDeploymentReconciler) wantsDisaggregatedSet(dgd *nvidiacomv1beta1.DynamoGraphDeployment) bool { + if dgd == nil || dgd.Annotations == nil { + return false + } + return strings.ToLower(dgd.Annotations[consts.KubeAnnotationEnableDisaggregatedSet]) == consts.KubeLabelValueTrue +} + +func (r *DynamoGraphDeploymentReconciler) shouldUseDisaggregatedSet(dgd *nvidiacomv1beta1.DynamoGraphDeployment) (bool, string) { + if !r.wantsDisaggregatedSet(dgd) { + return false, "" + } + if r.RuntimeConfig == nil { + return false, "runtime config is not initialized" + } + if !r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet) { + return false, "DisaggregatedSet API is not available" + } + selection, reason := selectDisaggregatedSetComponents(dgd) + if reason != "" { + return false, reason + } + if len(selection.componentToRole) < 2 { + return false, "DisaggregatedSet requires at least two eligible multinode worker roles" + } + if !r.RuntimeConfig.Gate.Enabled(features.LWS) { + for i := range dgd.Spec.Components { + component := &dgd.Spec.Components[i] + if component.GetNumberOfNodes() <= 1 { + continue + } + if _, selected := selection.componentToRole[component.ComponentName]; !selected { + return false, fmt.Sprintf("multinode component %q is not eligible for DisaggregatedSet and requires LeaderWorkerSet support", component.ComponentName) + } + } + } + return true, "" +} + +// coalesceDisaggregatedSetRestartState treats all selected DS roles as one +// restart unit. A DisaggregatedSet revision covers the complete role list, so +// annotating only one selected role would still roll every role and a later +// sequential step would otherwise create another whole-set revision. +func coalesceDisaggregatedSetRestartState( + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + restartState *dynamo.RestartState, +) *dynamo.RestartState { + if restartState == nil || dynamo.IsParallelRestart(dgd) { + return restartState + } + selection, reason := selectDisaggregatedSetComponents(dgd) + if reason != "" { + return restartState + } + selectedRestarting := false + for componentName := range selection.componentToRole { + if restartState.ShouldAnnotateComponent(componentName) { + selectedRestarting = true + break + } + } + if !selectedRestarting { + return restartState + } + for componentName := range selection.componentToRole { + restartState.ComponentsToAnnotate[componentName] = true + } + return restartState +} + +func (r *DynamoGraphDeploymentReconciler) coalesceDisaggregatedSetRestartStateForPath( + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + useDisaggregatedSet bool, + restartState *dynamo.RestartState, +) *dynamo.RestartState { + if !useDisaggregatedSet || r.isGrovePathway(dgd) { + return restartState + } + return coalesceDisaggregatedSetRestartState(dgd, restartState) +} + +func selectDisaggregatedSetComponents(dgd *nvidiacomv1beta1.DynamoGraphDeployment) (disaggregatedSetSelection, string) { + selection := disaggregatedSetSelection{ + componentToRole: make(map[string]string), + desiredReplicas: make(map[string]int32), + } + if dgd == nil { + return selection, "DynamoGraphDeployment is nil" + } + + usedRoles := make(map[string]struct{}) + zeroReplicas := 0 + positiveReplicas := 0 + for i := range dgd.Spec.Components { + component := &dgd.Spec.Components[i] + if !isDisaggregatedSetEligibleComponent(component) { + continue + } + if component.ScalingAdapter != nil { + return selection, fmt.Sprintf("component %q uses scalingAdapter, but DisaggregatedSet does not support scale subresource integration", component.ComponentName) + } + roleName := disaggregatedSetRoleName(component, usedRoles) + usedRoles[roleName] = struct{}{} + selection.componentToRole[component.ComponentName] = roleName + + desiredReplicas := desiredComponentReplicas(component) + selection.desiredReplicas[component.ComponentName] = desiredReplicas + if desiredReplicas == 0 { + zeroReplicas++ + } else { + positiveReplicas++ + } + } + + if len(selection.componentToRole) == 0 { + return selection, "no eligible multinode worker roles found" + } + if len(selection.componentToRole) > maxDisaggregatedSetRoles { + return selection, fmt.Sprintf("DisaggregatedSet supports at most %d roles", maxDisaggregatedSetRoles) + } + if zeroReplicas > 0 && positiveReplicas > 0 { + return selection, "DisaggregatedSet requires replicas to be zero for all selected roles or positive for all selected roles" + } + return selection, "" +} + +func isDisaggregatedSetEligibleComponent(component *nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec) bool { + return component != nil && component.GetNumberOfNodes() > 1 && dynamo.IsWorkerComponent(string(component.ComponentType)) +} + +func desiredComponentReplicas(component *nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec) int32 { + if component == nil || component.Replicas == nil { + return 1 + } + return *component.Replicas +} + +func disaggregatedSetRoleName(component *nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec, used map[string]struct{}) string { + preferred := strings.ToLower(string(component.ComponentType)) + if preferred != consts.ComponentTypePrefill && preferred != consts.ComponentTypeDecode { + preferred = "" + } + if preferred == "" || roleNameUsed(preferred, used) { + preferred = dynamo.NormalizeKubeResourceName(component.ComponentName) + } + preferred = truncateDNSLabelWithHash(preferred, maxDisaggregatedSetRoleNameLength) + roleName := preferred + for i := 2; roleNameUsed(roleName, used); i++ { + suffix := fmt.Sprintf("-%d", i) + roleName = truncateDNSLabel(preferred, maxDisaggregatedSetRoleNameLength-len(suffix)) + suffix + } + return roleName +} + +func roleNameUsed(roleName string, used map[string]struct{}) bool { + _, ok := used[roleName] + return ok +} + +func truncateDNSLabel(value string, maxLength int) string { + if len(value) <= maxLength { + return value + } + return strings.TrimRight(value[:maxLength], "-") +} + +func truncateDNSLabelWithHash(value string, maxLength int) string { + if len(value) <= maxLength { + return value + } + hash := sha256.Sum256([]byte(value)) + hashText := hex.EncodeToString(hash[:])[:disaggregatedSetNameHashLength] + if maxLength <= len(hashText) { + return hashText[:maxLength] + } + suffix := "-" + hashText + prefix := strings.TrimRight(value[:maxLength-len(suffix)], "-") + return prefix + suffix +} + +func disaggregatedSetName(dgd *nvidiacomv1beta1.DynamoGraphDeployment) string { + return truncateDNSLabelWithHash(dynamo.NormalizeKubeResourceName(dgd.Name), maxDisaggregatedSetNameLength) +} + +func (r *DynamoGraphDeploymentReconciler) reconcileDisaggregatedSetResources( + ctx context.Context, + dynamoDeployment *nvidiacomv1beta1.DynamoGraphDeployment, + restartState *dynamo.RestartState, + checkpointInfos map[string]*checkpoint.CheckpointInfo, +) (ReconcileResult, error) { + resources := []Resource{} + logger := log.FromContext(ctx) + + rollingUpdateCtx, err := r.buildRollingUpdateContext(ctx, dynamoDeployment) + if err != nil { + return ReconcileResult{}, fmt.Errorf("failed to build rolling update context: %w", err) + } + selection, reason := selectDisaggregatedSetComponents(dynamoDeployment) + if reason != "" { + return ReconcileResult{}, fmt.Errorf("failed to select DisaggregatedSet roles: %s", reason) + } + + existingRestartAnnotations, err := r.getExistingRestartAnnotationsDCD(ctx, dynamoDeployment) + if err != nil { + logger.Error(err, "failed to get existing restart annotations") + return ReconcileResult{}, fmt.Errorf("failed to get existing restart annotations: %w", err) + } + existingDSRestartAnnotations, err := r.getExistingRestartAnnotationsDisaggregatedSet(ctx, dynamoDeployment, selection) + if err != nil { + logger.Error(err, "failed to get existing DisaggregatedSet restart annotations") + return ReconcileResult{}, fmt.Errorf("failed to get existing DisaggregatedSet restart annotations: %w", err) + } + // Selected DCDs are deleted after cutover, so the current DS role templates + // become the source of truth for their restart annotations. Preserve those + // values until the sequential restart reaches the DS restart unit. + maps.Copy(existingRestartAnnotations, existingDSRestartAnnotations) + + dynamoComponentsDeployments, err := dynamo.GenerateDynamoComponentsDeployments( + dynamoDeployment, + restartState, + existingRestartAnnotations, + rollingUpdateCtx, + ) + if err != nil { + return ReconcileResult{}, fmt.Errorf("failed to generate DynamoComponentDeployments for DisaggregatedSet path: %w", err) + } + + checkpointGated, err := applyDisaggregatedSetCheckpointStartupPolicies(dynamoComponentsDeployments, checkpointInfos, selection) + if err != nil { + return ReconcileResult{}, err + } + + desiredDS, err := r.generateDisaggregatedSet(ctx, dynamoDeployment, dynamoComponentsDeployments, selection) + if err != nil { + return ReconcileResult{}, err + } + dsModified, syncedDS, err := r.syncDisaggregatedSet(ctx, dynamoDeployment, desiredDS) + if err != nil { + return ReconcileResult{}, err + } + + if err := r.reconcileDisaggregatedSetSideResources(ctx, dynamoDeployment, dynamoComponentsDeployments, selection); err != nil { + return ReconcileResult{}, err + } + dsReady, dsReason, dsStatuses, err := r.checkDisaggregatedSetReadiness(ctx, syncedDS, selection) + if err != nil { + return ReconcileResult{}, err + } + // A patched DS can still expose readiness from its previous revision. Do not + // retire the legacy DCDs until the DS controller observes the new spec. + dsReady = dsReady && !dsModified && !checkpointGated + + syncedDSResource, err := commoncontroller.NewResourceWithComponentStatuses( + syncedDS, + func() (bool, string, map[string]nvidiacomv1beta1.ComponentReplicaStatus) { + if dsModified { + return false, "DisaggregatedSet spec was updated; waiting for controller status", dsStatuses + } + if checkpointGated { + return false, "DisaggregatedSet roles are waiting for checkpoint readiness", dsStatuses + } + return dsReady, dsReason, dsStatuses + }, + ) + if err != nil { + return ReconcileResult{}, err + } + resources = append(resources, syncedDSResource) + + nonSelectedResources, err := r.reconcileDisaggregatedSetNonSelectedDCDs( + ctx, dynamoDeployment, dynamoComponentsDeployments, selection, + ) + if err != nil { + return ReconcileResult{}, err + } + resources = append(resources, nonSelectedResources...) + + if dsReady { + if err := r.deleteGrovePodCliqueSetOnDisaggregatedSetPath(ctx, dynamoDeployment); err != nil { + return ReconcileResult{}, err + } + if err := r.deleteOwnedSelectedDCDs(ctx, dynamoDeployment, selection); err != nil { + return ReconcileResult{}, err + } + } + + result := r.checkResourcesReadiness(resources) + if result.State == nvidiacomv1beta1.DGDStateSuccessful { + if err := r.deleteStaleDisaggregatedSetComponentServices(ctx, dynamoDeployment, dynamoComponentsDeployments); err != nil { + return ReconcileResult{}, err + } + } + return result, nil +} + +func (r *DynamoGraphDeploymentReconciler) getExistingRestartAnnotationsDisaggregatedSet( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + selection disaggregatedSetSelection, +) (map[string]string, error) { + ds := newDisaggregatedSetObject() + key := types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace} + if err := r.Get(ctx, key, ds); err != nil { + if apierrors.IsNotFound(err) { + return map[string]string{}, nil + } + return nil, fmt.Errorf("failed to get DisaggregatedSet %s: %w", key, err) + } + return restartAnnotationsFromDisaggregatedSet(ds, selection) +} + +func restartAnnotationsFromDisaggregatedSet( + ds *unstructured.Unstructured, + selection disaggregatedSetSelection, +) (map[string]string, error) { + restartAnnotations := make(map[string]string) + if ds == nil { + return restartAnnotations, nil + } + spec, found, err := unstructured.NestedMap(ds.Object, "spec") + if err != nil { + return nil, fmt.Errorf("failed to read DisaggregatedSet spec: %w", err) + } + if !found { + return restartAnnotations, nil + } + typedSpec := disaggregatedsetv1.DisaggregatedSetSpec{} + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(spec, &typedSpec); err != nil { + return nil, fmt.Errorf("failed to decode DisaggregatedSet spec: %w", err) + } + roleToComponent := make(map[string]string, len(selection.componentToRole)) + for componentName, roleName := range selection.componentToRole { + roleToComponent[roleName] = componentName + } + for i := range typedSpec.Roles { + role := &typedSpec.Roles[i] + componentName, selected := roleToComponent[role.Name] + if !selected { + continue + } + if role.Spec.LeaderWorkerTemplate.LeaderTemplate != nil { + if timestamp := role.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations[consts.RestartAnnotation]; timestamp != "" { + restartAnnotations[componentName] = timestamp + continue + } + } + if timestamp := role.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations[consts.RestartAnnotation]; timestamp != "" { + restartAnnotations[componentName] = timestamp + } + } + return restartAnnotations, nil +} + +func applyDisaggregatedSetCheckpointStartupPolicies( + dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment, + checkpointInfos map[string]*checkpoint.CheckpointInfo, + selection disaggregatedSetSelection, +) (bool, error) { + for _, componentName := range sortedDCDKeys(dcds) { + if err := applyDCDCheckpointStartupPolicy(dcds[componentName], checkpointInfos[componentName]); err != nil { + return false, fmt.Errorf("failed to apply checkpoint startup policy for %s: %w", componentName, err) + } + } + + // DS requires all roles to be either zero or positive. If any selected role + // is waiting for a checkpoint, gate the whole set until it is ready. + gateSelectedRoles := false + for componentName := range selection.componentToRole { + info := checkpointInfos[componentName] + if info != nil && info.Enabled && info.StartupPolicy == nvidiacomv1alpha1.CheckpointStartupPolicyWaitForCheckpoint && !info.Ready { + gateSelectedRoles = true + break + } + } + for componentName := range selection.componentToRole { + if gateSelectedRoles { + dcds[componentName].Spec.Replicas = ptr.To(int32(0)) + } + selection.desiredReplicas[componentName] = desiredComponentReplicas(&dcds[componentName].Spec.DynamoComponentDeploymentSharedSpec) + } + return gateSelectedRoles, nil +} + +func (r *DynamoGraphDeploymentReconciler) reconcileDisaggregatedSetNonSelectedDCDs( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment, + selection disaggregatedSetSelection, +) ([]Resource, error) { + resources := []Resource{} + for _, componentName := range sortedDCDKeys(dcds) { + dcd := dcds[componentName] + if _, selected := selection.componentToRole[componentName]; selected { + continue + } + if err := r.preserveExistingDCDBackendFramework(ctx, dcd); err != nil { + return nil, fmt.Errorf("failed to preserve existing DynamoComponentDeployment backendFramework: %w", err) + } + _, syncedDCD, err := commoncontroller.SyncResource(ctx, r, dgd, func(ctx context.Context) (*nvidiacomv1beta1.DynamoComponentDeployment, bool, error) { + return dcd, false, nil + }) + if err != nil { + return nil, fmt.Errorf("failed to sync non-DisaggregatedSet DynamoComponentDeployment %s: %w", dcd.Name, err) + } + resources = append(resources, syncedDCD) + } + return resources, nil +} + +func (r *DynamoGraphDeploymentReconciler) generateDisaggregatedSet( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment, + selection disaggregatedSetSelection, +) (*unstructured.Unstructured, error) { + ds := newDisaggregatedSetObject() + ds.SetName(disaggregatedSetName(dgd)) + ds.SetNamespace(dgd.Namespace) + ds.SetLabels(map[string]string{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + consts.KubeLabelDynamoSelector: disaggregatedSetName(dgd), + }) + if ownerRef := dgdControllerOwnerReference(dgd); ownerRef != nil { + ds.SetOwnerReferences([]metav1.OwnerReference{*ownerRef}) + } + + roles := make([]any, 0, len(selection.componentToRole)) + for i := range dgd.Spec.Components { + componentName := dgd.Spec.Components[i].ComponentName + roleName, ok := selection.componentToRole[componentName] + if !ok { + continue + } + dcd := dcds[componentName] + if dcd == nil { + return nil, fmt.Errorf("generated DynamoComponentDeployment missing for selected component %q", componentName) + } + renderDCD := dcd.DeepCopy() + if ownerRef := dgdControllerOwnerReference(dgd); ownerRef != nil { + renderDCD.SetOwnerReferences([]metav1.OwnerReference{*ownerRef}) + } + role, err := r.buildDisaggregatedSetRole(ctx, renderDCD) + if err != nil { + return nil, fmt.Errorf("failed to build DisaggregatedSet role %q: %w", roleName, err) + } + role["name"] = roleName + roles = append(roles, role) + } + if len(roles) < 2 { + return nil, fmt.Errorf("DisaggregatedSet requires at least two roles, got %d", len(roles)) + } + ds.Object["spec"] = map[string]any{"roles": roles} + return ds, nil +} + +// buildDisaggregatedSetRole renders a single DS role from a generated DCD by +// reusing the shared multinode render path that the LWS pathway uses. The DS +// pathway only needs the LWS spec fields as unstructured; the LWS object's +// own metadata is intentionally dropped so the DGD controller (and not the +// LWS controller) remains the visible owner of the role. +func (r *DynamoGraphDeploymentReconciler) buildDisaggregatedSetRole( + ctx context.Context, + dcd *nvidiacomv1beta1.DynamoComponentDeployment, +) (map[string]any, error) { + dcdReconciler := &DynamoComponentDeploymentReconciler{ + Client: r.Client, + Recorder: r.Recorder, + Config: r.Config, + RuntimeConfig: r.RuntimeConfig, + DockerSecretRetriever: r.DockerSecretRetriever, + } + + leaderPodTemplateSpec, workerPodTemplateSpec, err := dcdReconciler.renderMultinodePodTemplateSpecs( + ctx, + generateResourceOption{dynamoComponentDeployment: dcd}, + ) + if err != nil { + return nil, err + } + + desiredReplicas := int32(1) + if dcd.Spec.Replicas != nil { + desiredReplicas = *dcd.Spec.Replicas + } + groupSize := dcd.GetNumberOfNodes() + + lwsSpec := leaderworkersetv1.LeaderWorkerSetSpec{ + Replicas: &desiredReplicas, + StartupPolicy: leaderworkersetv1.LeaderCreatedStartupPolicy, + RolloutStrategy: leaderworkersetv1.RolloutStrategy{ + Type: leaderworkersetv1.RollingUpdateStrategyType, + }, + LeaderWorkerTemplate: leaderworkersetv1.LeaderWorkerTemplate{ + LeaderTemplate: leaderPodTemplateSpec, + WorkerTemplate: *workerPodTemplateSpec, + Size: &groupSize, + RestartPolicy: leaderworkersetv1.RecreateGroupOnPodRestart, + }, + } + lwsSpecUnstructured, err := runtime.DefaultUnstructuredConverter.ToUnstructured(&lwsSpec) + if err != nil { + return nil, fmt.Errorf("failed to convert LeaderWorkerSet spec: %w", err) + } + + return map[string]any{ + "spec": lwsSpecUnstructured, + }, nil +} + +func (r *DynamoGraphDeploymentReconciler) reconcileDisaggregatedSetSideResources( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment, + selection disaggregatedSetSelection, +) error { + if err := dynamo.ReconcileModelServicesForComponents(ctx, r, dgd, selectedComponentsByName(dgd, selection), dgd.Namespace); err != nil { + return fmt.Errorf("failed to reconcile DisaggregatedSet model services: %w", err) + } + if err := r.adoptSelectedModelServices(ctx, dgd, selection); err != nil { + return err + } + + dcdReconciler := &DynamoComponentDeploymentReconciler{ + Client: r.Client, + Recorder: r.Recorder, + Config: r.Config, + RuntimeConfig: r.RuntimeConfig, + DockerSecretRetriever: r.DockerSecretRetriever, + } + for _, componentName := range sortedSelectionComponentNames(selection) { + dcd := dcds[componentName] + if dcd == nil { + return fmt.Errorf("generated DynamoComponentDeployment missing for selected component %q", componentName) + } + _, syncedService, err := commoncontroller.SyncResource(ctx, r, dgd, func(ctx context.Context) (*corev1.Service, bool, error) { + return dcdReconciler.generateService(ctx, generateResourceOption{dynamoComponentDeployment: dcd}) + }) + if err != nil { + return fmt.Errorf("failed to reconcile DisaggregatedSet component service for %q: %w", componentName, err) + } + if syncedService != nil { + if err := r.ensureControlledByDGD(ctx, dgd, syncedService); err != nil { + return fmt.Errorf("failed to adopt DisaggregatedSet component service %s/%s: %w", syncedService.Namespace, syncedService.Name, err) + } + } + } + return nil +} + +func selectedComponentsByName(dgd *nvidiacomv1beta1.DynamoGraphDeployment, selection disaggregatedSetSelection) map[string]*nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec { + components := map[string]*nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{} + for i := range dgd.Spec.Components { + component := &dgd.Spec.Components[i] + if _, selected := selection.componentToRole[component.ComponentName]; selected { + components[component.ComponentName] = component + } + } + return components +} + +func sortedSelectionComponentNames(selection disaggregatedSetSelection) []string { + names := make([]string, 0, len(selection.componentToRole)) + for componentName := range selection.componentToRole { + names = append(names, componentName) + } + sort.Strings(names) + return names +} + +func (r *DynamoGraphDeploymentReconciler) adoptSelectedModelServices(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, selection disaggregatedSetSelection) error { + seen := map[string]struct{}{} + for i := range dgd.Spec.Components { + component := &dgd.Spec.Components[i] + if _, selected := selection.componentToRole[component.ComponentName]; !selected { + continue + } + if component.ModelRef == nil || component.ModelRef.Name == "" { + continue + } + serviceName := dynamo.GenerateServiceName(component.ModelRef.Name) + if _, ok := seen[serviceName]; ok { + continue + } + seen[serviceName] = struct{}{} + service := &corev1.Service{} + err := r.Get(ctx, types.NamespacedName{Name: serviceName, Namespace: dgd.Namespace}, service) + if apierrors.IsNotFound(err) { + continue + } + if err != nil { + return fmt.Errorf("failed to get DisaggregatedSet model service %s/%s: %w", dgd.Namespace, serviceName, err) + } + canAdopt, err := r.canAdoptModelServiceForDisaggregatedSet(ctx, dgd, service) + if err != nil { + return err + } + if !canAdopt { + continue + } + if err := r.ensureControlledByDGD(ctx, dgd, service); err != nil { + return fmt.Errorf("failed to adopt DisaggregatedSet model service %s/%s: %w", service.Namespace, service.Name, err) + } + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) canAdoptModelServiceForDisaggregatedSet( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + service *corev1.Service, +) (bool, error) { + owner := metav1.GetControllerOf(service) + if owner == nil || isControlledByBetaDGD(service, dgd) { + return true, nil + } + if owner.APIVersion != nvidiacomv1beta1.GroupVersion.String() || owner.Kind != dynamoComponentDeploymentKind { + return false, nil + } + dcd := &nvidiacomv1beta1.DynamoComponentDeployment{} + if err := r.Get(ctx, types.NamespacedName{Name: owner.Name, Namespace: service.Namespace}, dcd); err != nil { + if apierrors.IsNotFound(err) { + return false, nil + } + return false, fmt.Errorf("failed to verify model Service owner %s/%s: %w", service.Namespace, owner.Name, err) + } + return isControlledByBetaDGD(dcd, dgd), nil +} + +func sortedDCDKeys(dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment) []string { + keys := make([]string, 0, len(dcds)) + for key := range dcds { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} + +func (r *DynamoGraphDeploymentReconciler) syncDisaggregatedSet(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, desired *unstructured.Unstructured) (bool, *unstructured.Unstructured, error) { + current := newDisaggregatedSetObject() + key := types.NamespacedName{Name: desired.GetName(), Namespace: desired.GetNamespace()} + err := r.Get(ctx, key, current) + if apierrors.IsNotFound(err) { + if err := r.Create(ctx, desired); err != nil { + return false, nil, fmt.Errorf("failed to create DisaggregatedSet %s/%s: %w", desired.GetNamespace(), desired.GetName(), err) + } + return true, desired, nil + } + if err != nil { + return false, nil, fmt.Errorf("failed to get DisaggregatedSet %s/%s: %w", desired.GetNamespace(), desired.GetName(), err) + } + if !isControlledByBetaDGD(current, dgd) { + return false, nil, fmt.Errorf("refusing to update DisaggregatedSet %s/%s because it is not controlled by DynamoGraphDeployment %s/%s", desired.GetNamespace(), desired.GetName(), dgd.Namespace, dgd.Name) + } + original := current.DeepCopy() + labels := current.GetLabels() + if labels == nil { + labels = map[string]string{} + } + maps.Copy(labels, desired.GetLabels()) + current.SetLabels(labels) + annotations := current.GetAnnotations() + if annotations == nil && len(desired.GetAnnotations()) > 0 { + annotations = map[string]string{} + } + maps.Copy(annotations, desired.GetAnnotations()) + current.SetAnnotations(annotations) + setDGDControllerOwnerReference(dgd, current) + current.Object["spec"] = desired.Object["spec"] + if disaggregatedSetDesiredStateEqual(original, current) { + return false, current, nil + } + + // Ask the API server to apply structural defaults before deciding whether + // the desired state differs. This keeps reconciliation stable across the + // pinned CRD and forward-compatible CRDs that add new defaulted fields. + if err := r.Patch(ctx, current, client.MergeFrom(original), client.DryRunAll); err != nil { + return false, nil, fmt.Errorf("failed to dry-run patch DisaggregatedSet %s/%s: %w", current.GetNamespace(), current.GetName(), err) + } + if disaggregatedSetDesiredStateEqual(original, current) { + return false, current, nil + } + if err := r.Patch(ctx, current, client.MergeFrom(original)); err != nil { + return false, nil, fmt.Errorf("failed to patch DisaggregatedSet %s/%s: %w", current.GetNamespace(), current.GetName(), err) + } + return true, current, nil +} + +func disaggregatedSetDesiredStateEqual(a, b *unstructured.Unstructured) bool { + return equality.Semantic.DeepEqual(a.Object["spec"], b.Object["spec"]) && + equality.Semantic.DeepEqual(a.GetLabels(), b.GetLabels()) && + equality.Semantic.DeepEqual(a.GetAnnotations(), b.GetAnnotations()) && + equality.Semantic.DeepEqual(a.GetOwnerReferences(), b.GetOwnerReferences()) +} + +func (r *DynamoGraphDeploymentReconciler) deleteDisaggregatedSetIfExists(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment) error { + ds := newDisaggregatedSetObject() + key := types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace} + if err := r.Get(ctx, key, ds); err != nil { + if apierrors.IsNotFound(err) { + return nil + } + return fmt.Errorf("failed to get stale DisaggregatedSet %s/%s: %w", dgd.Namespace, disaggregatedSetName(dgd), err) + } + if !isControlledByBetaDGD(ds, dgd) { + return fmt.Errorf("refusing to delete DisaggregatedSet %s/%s because it is not controlled by DynamoGraphDeployment %s/%s", dgd.Namespace, disaggregatedSetName(dgd), dgd.Namespace, dgd.Name) + } + if err := r.Delete(ctx, ds); err != nil && !apierrors.IsNotFound(err) { + return fmt.Errorf("failed to delete stale DisaggregatedSet %s/%s: %w", dgd.Namespace, disaggregatedSetName(dgd), err) + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) deleteOwnedSelectedDCDs(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, selection disaggregatedSetSelection) error { + dcds, err := r.listOwnedSelectedDCDs(ctx, dgd, selection) + if err != nil { + return err + } + for i := range dcds { + dcd := &dcds[i] + if err := r.Delete(ctx, dcd); err != nil && !apierrors.IsNotFound(err) { + return fmt.Errorf("failed to delete selected DynamoComponentDeployment %s/%s: %w", dcd.Namespace, dcd.Name, err) + } + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) deleteStaleDisaggregatedSetComponentServices( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + dcds map[string]*nvidiacomv1beta1.DynamoComponentDeployment, +) error { + services := &corev1.ServiceList{} + if err := r.List(ctx, services, client.InNamespace(dgd.Namespace), client.MatchingLabels{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + }); err != nil { + return fmt.Errorf("failed to list DisaggregatedSet component Services for cleanup: %w", err) + } + + for i := range services.Items { + service := &services.Items[i] + if !isControlledByBetaDGD(service, dgd) { + continue + } + componentName := service.Labels[consts.KubeLabelDynamoComponent] + dcd := dcds[componentName] + if dcd != nil && service.Name == dynamo.NormalizeKubeResourceName(dcd.Name) { + continue + } + if componentName == "" { + continue + } + if err := r.Delete(ctx, service); err != nil && !apierrors.IsNotFound(err) { + return fmt.Errorf("failed to delete stale DisaggregatedSet component Service %s/%s: %w", service.Namespace, service.Name, err) + } + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) listOwnedSelectedDCDs(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, selection disaggregatedSetSelection) ([]nvidiacomv1beta1.DynamoComponentDeployment, error) { + dcdList := &nvidiacomv1beta1.DynamoComponentDeploymentList{} + if err := r.List(ctx, dcdList, client.InNamespace(dgd.Namespace), client.MatchingLabels{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + }); err != nil { + return nil, fmt.Errorf("failed to list DynamoComponentDeployments for DisaggregatedSet cleanup: %w", err) + } + selectedDCDs := []nvidiacomv1beta1.DynamoComponentDeployment{} + for _, dcd := range dcdList.Items { + if !isControlledByBetaDGD(&dcd, dgd) { + continue + } + componentName := dynamo.GetDCDComponentName(&dcd) + if _, selected := selection.componentToRole[componentName]; selected { + selectedDCDs = append(selectedDCDs, dcd) + } + } + return selectedDCDs, nil +} + +func (r *DynamoGraphDeploymentReconciler) ensureControlledByDGD(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, obj client.Object) error { + if dgdControllerOwnerReference(dgd) == nil || isControlledByBetaDGD(obj, dgd) { + return nil + } + controllerOwner := metav1.GetControllerOf(obj) + if controllerOwner != nil { + if controllerOwner.APIVersion != nvidiacomv1beta1.GroupVersion.String() || controllerOwner.Kind != dynamoComponentDeploymentKind { + return fmt.Errorf("resource is controlled by %s/%s %q", controllerOwner.APIVersion, controllerOwner.Kind, controllerOwner.Name) + } + dcd := &nvidiacomv1beta1.DynamoComponentDeployment{} + if err := r.Get(ctx, types.NamespacedName{Name: controllerOwner.Name, Namespace: obj.GetNamespace()}, dcd); err != nil { + return fmt.Errorf("failed to verify current DynamoComponentDeployment owner %s/%s: %w", obj.GetNamespace(), controllerOwner.Name, err) + } + if !isControlledByBetaDGD(dcd, dgd) { + return fmt.Errorf("current DynamoComponentDeployment owner %s/%s is not controlled by DynamoGraphDeployment %s/%s", dcd.Namespace, dcd.Name, dgd.Namespace, dgd.Name) + } + } + original := obj.DeepCopyObject().(client.Object) + setDGDControllerOwnerReference(dgd, obj) + if equality.Semantic.DeepEqual(original.GetOwnerReferences(), obj.GetOwnerReferences()) { + return nil + } + if err := r.Patch(ctx, obj, client.MergeFrom(original)); err != nil { + return fmt.Errorf("failed to update owner references: %w", err) + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) restoreDisaggregatedSetServiceOwnershipToDCDs( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, +) error { + if r.RuntimeConfig == nil || !r.RuntimeConfig.Gate.Enabled(features.DisaggregatedSet) { + return nil + } + ds := newDisaggregatedSetObject() + if err := r.Get(ctx, types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace}, ds); err != nil { + if apierrors.IsNotFound(err) { + return nil + } + return fmt.Errorf("failed to get DisaggregatedSet before restoring Service ownership: %w", err) + } + if !isControlledByBetaDGD(ds, dgd) { + return fmt.Errorf("refusing to restore Service ownership because DisaggregatedSet %s/%s is not controlled by DynamoGraphDeployment %s/%s", ds.GetNamespace(), ds.GetName(), dgd.Namespace, dgd.Name) + } + + dcds := &nvidiacomv1beta1.DynamoComponentDeploymentList{} + if err := r.List(ctx, dcds, client.InNamespace(dgd.Namespace), client.MatchingLabels{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + }); err != nil { + return fmt.Errorf("failed to list replacement DynamoComponentDeployments: %w", err) + } + sort.Slice(dcds.Items, func(i, j int) bool { return dcds.Items[i].Name < dcds.Items[j].Name }) + + serviceOwners := map[string]*nvidiacomv1beta1.DynamoComponentDeployment{} + modelServiceOwners := map[string]*nvidiacomv1beta1.DynamoComponentDeployment{} + componentOwners := map[string]*nvidiacomv1beta1.DynamoComponentDeployment{} + for i := range dcds.Items { + dcd := &dcds.Items[i] + if !isControlledByBetaDGD(dcd, dgd) { + continue + } + serviceOwners[dynamo.NormalizeKubeResourceName(dcd.Name)] = dcd + componentName := dynamo.GetDCDComponentName(dcd) + if componentOwners[componentName] == nil { + componentOwners[componentName] = dcd + } + } + for i := range dgd.Spec.Components { + component := &dgd.Spec.Components[i] + if component.ModelRef == nil || component.ModelRef.Name == "" { + continue + } + if owner := componentOwners[component.ComponentName]; owner != nil { + modelServiceName := dynamo.GenerateServiceName(component.ModelRef.Name) + if serviceOwners[modelServiceName] == nil { + serviceOwners[modelServiceName] = owner + modelServiceOwners[modelServiceName] = owner + } + } + } + + serviceNames := make([]string, 0, len(serviceOwners)) + for serviceName := range serviceOwners { + serviceNames = append(serviceNames, serviceName) + } + sort.Strings(serviceNames) + for _, serviceName := range serviceNames { + service := &corev1.Service{} + if err := r.Get(ctx, types.NamespacedName{Name: serviceName, Namespace: dgd.Namespace}, service); err != nil { + if apierrors.IsNotFound(err) { + // A component Service is already reconciled by its replacement + // DCD. A shared model Service can still be DGD-owned at this + // point, so only the DGD watches its deletion until handoff. + if modelServiceOwner := modelServiceOwners[serviceName]; modelServiceOwner != nil { + componentName := dynamo.GetDCDComponentName(modelServiceOwner) + if err := dynamo.ReconcileModelServicesForComponents( + ctx, + r, + modelServiceOwner, + map[string]*nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + componentName: &modelServiceOwner.Spec.DynamoComponentDeploymentSharedSpec, + }, + dgd.Namespace, + ); err != nil { + return fmt.Errorf("failed to recreate replacement model Service %s/%s: %w", dgd.Namespace, serviceName, err) + } + } + continue + } + return fmt.Errorf("failed to get replacement Service %s/%s: %w", dgd.Namespace, serviceName, err) + } + if err := r.ensureControlledByDCD(ctx, dgd, serviceOwners[serviceName], service); err != nil { + return fmt.Errorf("failed to restore Service %s/%s ownership: %w", service.Namespace, service.Name, err) + } + } + return nil +} + +func (r *DynamoGraphDeploymentReconciler) ensureControlledByDCD( + ctx context.Context, + dgd *nvidiacomv1beta1.DynamoGraphDeployment, + dcd *nvidiacomv1beta1.DynamoComponentDeployment, + obj client.Object, +) error { + ownerRef := dcdControllerOwnerReference(dcd) + if ownerRef == nil { + return fmt.Errorf("DynamoComponentDeployment %s/%s has no UID", dcd.Namespace, dcd.Name) + } + if metav1.IsControlledBy(obj, dcd) { + return nil + } + if controllerOwner := metav1.GetControllerOf(obj); controllerOwner != nil && !ownerReferenceMatchesDGD(controllerOwner, dgd) { + if controllerOwner.APIVersion == nvidiacomv1beta1.GroupVersion.String() && controllerOwner.Kind == dynamoComponentDeploymentKind { + currentOwner := &nvidiacomv1beta1.DynamoComponentDeployment{} + if err := r.Get(ctx, types.NamespacedName{Name: controllerOwner.Name, Namespace: obj.GetNamespace()}, currentOwner); err != nil { + return fmt.Errorf("failed to verify current DynamoComponentDeployment owner %s/%s: %w", obj.GetNamespace(), controllerOwner.Name, err) + } + if isControlledByBetaDGD(currentOwner, dgd) { + return nil + } + } + return fmt.Errorf("resource is controlled by %s/%s %q", controllerOwner.APIVersion, controllerOwner.Kind, controllerOwner.Name) + } + + original := obj.DeepCopyObject().(client.Object) + ownerRefs := make([]metav1.OwnerReference, 0, len(obj.GetOwnerReferences())+1) + for _, ref := range obj.GetOwnerReferences() { + if ptr.Deref(ref.Controller, false) { + continue + } + if ref.APIVersion == ownerRef.APIVersion && ref.Kind == ownerRef.Kind && ref.Name == ownerRef.Name { + continue + } + ownerRefs = append(ownerRefs, ref) + } + ownerRefs = append(ownerRefs, *ownerRef) + obj.SetOwnerReferences(ownerRefs) + if err := r.Patch(ctx, obj, client.MergeFrom(original)); err != nil { + return fmt.Errorf("failed to update owner references: %w", err) + } + return nil +} + +func dcdControllerOwnerReference(dcd *nvidiacomv1beta1.DynamoComponentDeployment) *metav1.OwnerReference { + if dcd == nil || dcd.UID == "" { + return nil + } + return &metav1.OwnerReference{ + APIVersion: nvidiacomv1beta1.GroupVersion.String(), + Kind: dynamoComponentDeploymentKind, + Name: dcd.Name, + UID: dcd.UID, + Controller: ptr.To(true), + BlockOwnerDeletion: ptr.To(true), + } +} + +func ownerReferenceMatchesDGD(owner *metav1.OwnerReference, dgd *nvidiacomv1beta1.DynamoGraphDeployment) bool { + return owner != nil && + dgd != nil && + owner.APIVersion == nvidiacomv1beta1.GroupVersion.String() && + owner.Kind == dynamoGraphDeploymentKind && + owner.Name == dgd.Name && + owner.UID != "" && + owner.UID == dgd.UID +} + +func dgdControllerOwnerReference(dgd *nvidiacomv1beta1.DynamoGraphDeployment) *metav1.OwnerReference { + if dgd == nil || dgd.UID == "" { + return nil + } + return &metav1.OwnerReference{ + APIVersion: nvidiacomv1beta1.GroupVersion.String(), + Kind: dynamoGraphDeploymentKind, + Name: dgd.Name, + UID: dgd.UID, + Controller: ptr.To(true), + BlockOwnerDeletion: ptr.To(true), + } +} + +func setDGDControllerOwnerReference(dgd *nvidiacomv1beta1.DynamoGraphDeployment, obj client.Object) { + ownerRef := dgdControllerOwnerReference(dgd) + if ownerRef == nil { + return + } + ownerRefs := make([]metav1.OwnerReference, 0, len(obj.GetOwnerReferences())+1) + for _, ref := range obj.GetOwnerReferences() { + if ptr.Deref(ref.Controller, false) { + continue + } + if ref.APIVersion == ownerRef.APIVersion && ref.Kind == ownerRef.Kind && ref.Name == ownerRef.Name { + continue + } + ownerRefs = append(ownerRefs, ref) + } + ownerRefs = append(ownerRefs, *ownerRef) + obj.SetOwnerReferences(ownerRefs) +} + +func isControlledByBetaDGD(obj client.Object, dgd *nvidiacomv1beta1.DynamoGraphDeployment) bool { + if obj == nil || dgd == nil { + return false + } + if dgd.UID != "" { + return metav1.IsControlledBy(obj, dgd) + } + controllerOwner := metav1.GetControllerOf(obj) + return controllerOwner != nil && + controllerOwner.APIVersion == nvidiacomv1beta1.GroupVersion.String() && + controllerOwner.Kind == dynamoGraphDeploymentKind && + controllerOwner.Name == dgd.Name +} + +func checkDisaggregatedSetReadiness(ds *unstructured.Unstructured, selection disaggregatedSetSelection) (bool, string, map[string]nvidiacomv1beta1.ComponentReplicaStatus) { + statuses := make(map[string]nvidiacomv1beta1.ComponentReplicaStatus, len(selection.componentToRole)) + roleStatuses := disaggregatedSetRoleStatuses(ds) + notReadyReasons := []string{} + for componentName, roleName := range selection.componentToRole { + desiredReplicas := selection.desiredReplicas[componentName] + componentStatus := nvidiacomv1beta1.ComponentReplicaStatus{ + ComponentKind: nvidiacomv1beta1.ComponentKindLeaderWorkerSet, + } + roleStatus, found := roleStatuses[roleName] + if found { + componentStatus.Replicas = nestedInt32(roleStatus, "replicas") + componentStatus.UpdatedReplicas = nestedInt32(roleStatus, "updatedReplicas") + readyReplicas := nestedInt32(roleStatus, "readyReplicas") + componentStatus.ReadyReplicas = &readyReplicas + } + statuses[componentName] = componentStatus + if !found { + notReadyReasons = append(notReadyReasons, fmt.Sprintf("%s role %q has no status yet", componentName, roleName)) + continue + } + if desiredReplicas == 0 { + if componentStatus.Replicas != 0 || componentStatus.UpdatedReplicas != 0 || ptr.Deref(componentStatus.ReadyReplicas, 0) != 0 { + notReadyReasons = append(notReadyReasons, fmt.Sprintf("%s role %q has not scaled to zero", componentName, roleName)) + } + continue + } + if componentStatus.Replicas != desiredReplicas || + componentStatus.UpdatedReplicas != desiredReplicas || + componentStatus.ReadyReplicas == nil || + *componentStatus.ReadyReplicas != desiredReplicas { + notReadyReasons = append(notReadyReasons, fmt.Sprintf( + "%s role %q replicas not ready (desired=%d replicas=%d updated=%d ready=%d)", + componentName, + roleName, + desiredReplicas, + componentStatus.Replicas, + componentStatus.UpdatedReplicas, + ptr.Deref(componentStatus.ReadyReplicas, 0), + )) + } + } + if current, reason := disaggregatedSetStatusObserved(ds); !current { + return false, reason, statuses + } + if len(notReadyReasons) > 0 { + sort.Strings(notReadyReasons) + return false, strings.Join(notReadyReasons, "; "), statuses + } + return true, "All DisaggregatedSet roles are ready", statuses +} + +// checkDisaggregatedSetReadiness falls back to the child LWS objects while +// DisaggregatedSet controllers that expose the v1 API do not publish either +// roleStatuses or generation observation. A role status without observation +// evidence may still describe the previous spec revision. +func (r *DynamoGraphDeploymentReconciler) checkDisaggregatedSetReadiness( + ctx context.Context, + ds *unstructured.Unstructured, + selection disaggregatedSetSelection, +) (bool, string, map[string]nvidiacomv1beta1.ComponentReplicaStatus, error) { + if len(disaggregatedSetRoleStatuses(ds)) > 0 && disaggregatedSetStatusHasObservation(ds) { + ready, reason, statuses := checkDisaggregatedSetReadiness(ds, selection) + return ready, reason, statuses, nil + } + + children := &leaderworkersetv1.LeaderWorkerSetList{} + if err := r.List(ctx, children, client.InNamespace(ds.GetNamespace()), client.MatchingLabels{ + disaggregatedsetv1.SetNameLabelKey: ds.GetName(), + }); err != nil { + return false, "", nil, fmt.Errorf("failed to list DisaggregatedSet child LeaderWorkerSets: %w", err) + } + typedDS := &disaggregatedsetv1.DisaggregatedSet{} + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(ds.Object, typedDS); err != nil { + return false, "", nil, fmt.Errorf("failed to decode DisaggregatedSet for child readiness: %w", err) + } + targetRevision := disaggregatedsetutils.ComputeRevision(typedDS.Spec.Roles) + targetByRole := make(map[string]*leaderworkersetv1.LeaderWorkerSet) + childrenByRole := make(map[string][]*leaderworkersetv1.LeaderWorkerSet) + for i := range children.Items { + child := &children.Items[i] + if !metav1.IsControlledBy(child, ds) { + continue + } + roleName := child.Labels[disaggregatedsetv1.RoleLabelKey] + childrenByRole[roleName] = append(childrenByRole[roleName], child) + if child.Labels[disaggregatedsetv1.RevisionLabelKey] != targetRevision { + continue + } + targetByRole[roleName] = child + } + ready, reason, statuses := checkDisaggregatedSetChildLWSReadiness(selection, targetByRole, childrenByRole) + return ready, reason, statuses, nil +} + +func checkDisaggregatedSetChildLWSReadiness( + selection disaggregatedSetSelection, + targetByRole map[string]*leaderworkersetv1.LeaderWorkerSet, + childrenByRole map[string][]*leaderworkersetv1.LeaderWorkerSet, +) (bool, string, map[string]nvidiacomv1beta1.ComponentReplicaStatus) { + statuses := make(map[string]nvidiacomv1beta1.ComponentReplicaStatus, len(selection.componentToRole)) + notReadyReasons := []string{} + for componentName, roleName := range selection.componentToRole { + desiredReplicas := selection.desiredReplicas[componentName] + child := targetByRole[roleName] + status := nvidiacomv1beta1.ComponentReplicaStatus{ComponentKind: nvidiacomv1beta1.ComponentKindLeaderWorkerSet} + children := childrenByRole[roleName] + sort.Slice(children, func(i, j int) bool { return children[i].Name < children[j].Name }) + readyReplicas := int32(0) + for _, roleChild := range children { + status.ComponentNames = append(status.ComponentNames, roleChild.Name) + status.Replicas += roleChild.Status.Replicas + readyReplicas += roleChild.Status.ReadyReplicas + } + status.ReadyReplicas = ptr.To(readyReplicas) + if child == nil { + statuses[componentName] = status + notReadyReasons = append(notReadyReasons, fmt.Sprintf("%s role %q has no child LeaderWorkerSet yet", componentName, roleName)) + continue + } + status.UpdatedReplicas = child.Status.UpdatedReplicas + statuses[componentName] = status + if child.Status.ObservedGeneration < child.Generation { + notReadyReasons = append(notReadyReasons, fmt.Sprintf("%s child LeaderWorkerSet %q has not observed generation %d", componentName, child.Name, child.Generation)) + continue + } + if child.Status.Replicas != desiredReplicas || child.Status.UpdatedReplicas != desiredReplicas || child.Status.ReadyReplicas != desiredReplicas { + notReadyReasons = append(notReadyReasons, fmt.Sprintf( + "%s child LeaderWorkerSet %q replicas not ready (desired=%d replicas=%d updated=%d ready=%d)", + componentName, child.Name, desiredReplicas, child.Status.Replicas, child.Status.UpdatedReplicas, child.Status.ReadyReplicas, + )) + } + for _, roleChild := range children { + if roleChild != child && roleChild.Status.Replicas != 0 { + notReadyReasons = append(notReadyReasons, fmt.Sprintf("%s old child LeaderWorkerSet %q still has %d replicas", componentName, roleChild.Name, roleChild.Status.Replicas)) + } + } + } + if len(notReadyReasons) > 0 { + sort.Strings(notReadyReasons) + return false, strings.Join(notReadyReasons, "; "), statuses + } + return true, "All DisaggregatedSet child LeaderWorkerSets are ready", statuses +} + +func disaggregatedSetStatusObserved(ds *unstructured.Unstructured) (bool, string) { + if ds == nil || ds.GetGeneration() == 0 { + return true, "" + } + if observedGeneration, found := nestedInt64FromObject(ds.Object, "status", "observedGeneration"); found && observedGeneration < ds.GetGeneration() { + return false, fmt.Sprintf("DisaggregatedSet status has not observed generation %d (observedGeneration=%d)", ds.GetGeneration(), observedGeneration) + } + conditions, found, _ := unstructured.NestedSlice(ds.Object, "status", "conditions") + if !found { + return true, "" + } + for _, item := range conditions { + condition, ok := item.(map[string]any) + if !ok { + continue + } + observedGeneration, ok := nestedInt64(condition, "observedGeneration") + if !ok || observedGeneration >= ds.GetGeneration() { + continue + } + conditionType, _ := condition["type"].(string) + return false, fmt.Sprintf("DisaggregatedSet condition %q has not observed generation %d (observedGeneration=%d)", conditionType, ds.GetGeneration(), observedGeneration) + } + return true, "" +} + +func disaggregatedSetStatusHasObservation(ds *unstructured.Unstructured) bool { + if ds == nil || ds.GetGeneration() == 0 { + return true + } + if _, found := nestedInt64FromObject(ds.Object, "status", "observedGeneration"); found { + return true + } + conditions, found, _ := unstructured.NestedSlice(ds.Object, "status", "conditions") + if !found { + return false + } + for _, item := range conditions { + condition, ok := item.(map[string]any) + if !ok { + continue + } + if _, found := nestedInt64(condition, "observedGeneration"); found { + return true + } + } + return false +} + +func disaggregatedSetRoleStatuses(ds *unstructured.Unstructured) map[string]map[string]any { + out := map[string]map[string]any{} + roleStatuses, found, _ := unstructured.NestedSlice(ds.Object, "status", "roleStatuses") + if !found { + return out + } + for _, item := range roleStatuses { + roleStatus, ok := item.(map[string]any) + if !ok { + continue + } + name, ok := roleStatus["name"].(string) + if !ok || name == "" { + continue + } + out[name] = roleStatus + } + return out +} + +func nestedInt32(obj map[string]any, key string) int32 { + value, _ := nestedInt64(obj, key) + return int32(value) +} + +func nestedInt64FromObject(obj map[string]any, fields ...string) (int64, bool) { + value, found, err := unstructured.NestedFieldNoCopy(obj, fields...) + if err != nil || !found { + return 0, false + } + return int64Value(value) +} + +func nestedInt64(obj map[string]any, key string) (int64, bool) { + return int64Value(obj[key]) +} + +func int64Value(value any) (int64, bool) { + switch v := value.(type) { + case int32: + return int64(v), true + case int64: + return v, true + case int: + return int64(v), true + case float64: + return int64(v), true + default: + return 0, false + } +} + +func disaggregatedSetStatusChanged(oldObj, newObj client.Object) bool { + oldDS, okOld := oldObj.(*unstructured.Unstructured) + newDS, okNew := newObj.(*unstructured.Unstructured) + if !okOld || !okNew { + return false + } + return oldDS.GetGeneration() != newDS.GetGeneration() || + !equality.Semantic.DeepEqual(oldDS.Object["status"], newDS.Object["status"]) || + !equality.Semantic.DeepEqual(oldDS.GetLabels(), newDS.GetLabels()) || + !equality.Semantic.DeepEqual(oldDS.GetOwnerReferences(), newDS.GetOwnerReferences()) +} + +func leaderWorkerSetStatusChanged(oldObj, newObj client.Object) bool { + oldLWS, okOld := oldObj.(*leaderworkersetv1.LeaderWorkerSet) + newLWS, okNew := newObj.(*leaderworkersetv1.LeaderWorkerSet) + if !okOld || !okNew { + return false + } + return oldLWS.Generation != newLWS.Generation || + !equality.Semantic.DeepEqual(oldLWS.Status, newLWS.Status) || + !equality.Semantic.DeepEqual(oldLWS.GetLabels(), newLWS.GetLabels()) || + !equality.Semantic.DeepEqual(oldLWS.GetOwnerReferences(), newLWS.GetOwnerReferences()) +} + +func (r *DynamoGraphDeploymentReconciler) mapDisaggregatedSetChildLWSToDGD(ctx context.Context, obj client.Object) []ctrl.Request { + setName := obj.GetLabels()[disaggregatedsetv1.SetNameLabelKey] + if setName == "" { + return nil + } + ds := newDisaggregatedSetObject() + if err := r.Get(ctx, types.NamespacedName{Name: setName, Namespace: obj.GetNamespace()}, ds); err != nil { + if !apierrors.IsNotFound(err) { + log.FromContext(ctx).Error(err, "failed to map DisaggregatedSet child LeaderWorkerSet", "leaderWorkerSet", obj.GetName()) + } + return nil + } + owner := metav1.GetControllerOf(ds) + if owner == nil || owner.APIVersion != nvidiacomv1beta1.GroupVersion.String() || owner.Kind != dynamoGraphDeploymentKind { + return nil + } + return []ctrl.Request{{NamespacedName: types.NamespacedName{Name: owner.Name, Namespace: ds.GetNamespace()}}} +} + +func (r *DynamoGraphDeploymentReconciler) getUpdatedInProgressForDisaggregatedSet(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment, inProgress []string) []string { + logger := log.FromContext(ctx) + selection, reason := selectDisaggregatedSetComponents(dgd) + if reason != "" { + logger.V(1).Info("failed to select DisaggregatedSet components for restart progress", "reason", reason) + return inProgress + } + + ds := newDisaggregatedSetObject() + dsErr := r.Get(ctx, types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace}, ds) + if dsErr != nil && !apierrors.IsNotFound(dsErr) { + logger.V(1).Info("failed to get DisaggregatedSet for restart progress", "error", dsErr) + } + dsReady := false + dsReason := resourceNotFoundReason + if dsErr == nil { + var err error + dsReady, dsReason, _, err = r.checkDisaggregatedSetReadiness(ctx, ds, selection) + if err != nil { + dsReason = err.Error() + } + } + + updatedInProgress := make([]string, 0, len(inProgress)) + for _, componentName := range inProgress { + if _, selected := selection.componentToRole[componentName]; !selected { + isFullyUpdated, reason := r.checkComponentFullyUpdated(ctx, dgd, componentName) + if !isFullyUpdated { + logger.V(1).Info("component not fully updated", "componentName", componentName, "reason", reason) + updatedInProgress = append(updatedInProgress, componentName) + } + continue + } + + if dsErr != nil { + reason := resourceNotFoundReason + if !apierrors.IsNotFound(dsErr) { + reason = dsErr.Error() + } + logger.V(1).Info("DisaggregatedSet component not fully updated", "componentName", componentName, "reason", reason) + updatedInProgress = append(updatedInProgress, componentName) + continue + } + + if !dsReady { + logger.V(1).Info("DisaggregatedSet component not fully updated", "componentName", componentName, "reason", dsReason) + updatedInProgress = append(updatedInProgress, componentName) + } + } + return updatedInProgress +} diff --git a/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_envtest_test.go b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_envtest_test.go new file mode 100644 index 000000000000..9ea3ff3f41eb --- /dev/null +++ b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_envtest_test.go @@ -0,0 +1,370 @@ +//go:build !clustertest + +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "context" + + configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1" + nvidiacomv1beta1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta1" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/tools/record" + "k8s.io/utils/ptr" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("DisaggregatedSet", func() { + It("creates a DisaggregatedSet with two roles", func() { + ctx := context.Background() + + dgd := newDSHappyPathDGD() + Expect(k8sClient.Create(ctx, dgd)).To(Succeed()) + DeferCleanup(func() { + _ = k8sClient.Delete(ctx, dgd) + }) + + reconciler := &DynamoGraphDeploymentReconciler{ + Client: k8sClient, + Recorder: record.NewFakeRecorder(10), + Config: &configv1alpha1.OperatorConfiguration{}, + RuntimeConfig: newTestRuntimeConfig(true), + } + + res, err := reconciler.reconcileDisaggregatedSetResources(ctx, dgd, nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(res.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + + ds := &unstructured.Unstructured{} + ds.SetGroupVersionKind(disaggregatedSetGVK) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace}, ds)).To(Succeed()) + + roles, found, err := unstructured.NestedSlice(ds.Object, "spec", "roles") + Expect(err).NotTo(HaveOccurred()) + Expect(found).To(BeTrue()) + Expect(roles).To(HaveLen(2)) + }) + + It("keeps the serving resources during DCD to DS cutover", func() { + ctx := context.Background() + By("creating a ready two-component DCD deployment with stable component Services") + dgd := newDSHappyPathDGD() + dgd.Name = "demo-ds-cutover" + dgd.UID = "" + dgd.Annotations = nil + for i := range dgd.Spec.Components { + dgd.Spec.Components[i].ModelRef = &nvidiacomv1beta1.ModelReference{Name: "shared-smoke-model"} + } + Expect(k8sClient.Create(ctx, dgd)).To(Succeed()) + DeferCleanup(func() { + _ = k8sClient.Delete(ctx, dgd) + }) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, dgd)).To(Succeed()) + + runtimeConfig := &commoncontroller.RuntimeConfig{Gate: features.Gates{LWS: true, DisaggregatedSet: true}} + operatorConfig := &configv1alpha1.OperatorConfiguration{ + Discovery: configv1alpha1.DiscoveryConfiguration{Backend: configv1alpha1.DiscoveryBackendKubernetes}, + } + reconciler := &DynamoGraphDeploymentReconciler{ + Client: k8sClient, + Recorder: record.NewFakeRecorder(100), + Config: operatorConfig, + RuntimeConfig: runtimeConfig, + } + dcdReconciler := &DynamoComponentDeploymentReconciler{ + Client: k8sClient, + Recorder: record.NewFakeRecorder(100), + Config: reconciler.Config, + RuntimeConfig: runtimeConfig, + } + + result, err := reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + legacyDCDs := ownedCutoverDCDs(ctx, dgd) + Expect(legacyDCDs).To(HaveLen(2)) + serviceUIDs := map[string]types.UID{} + for i := range legacyDCDs { + modified, err := dcdReconciler.createOrUpdateOrDeleteServices(ctx, generateResourceOption{dynamoComponentDeployment: &legacyDCDs[i]}) + Expect(err).NotTo(HaveOccurred()) + Expect(modified).To(BeTrue()) + service := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: legacyDCDs[i].Name, Namespace: dgd.Namespace}, service)).To(Succeed()) + serviceUIDs[service.Name] = service.UID + } + Expect(serviceUIDs).To(HaveLen(2)) + + By("enabling DisaggregatedSet without removing the serving DCDs") + markCutoverDCDsReady(ctx, dgd) + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStateSuccessful)) + + dgd.Annotations = map[string]string{consts.KubeAnnotationEnableDisaggregatedSet: consts.KubeLabelValueTrue} + Expect(k8sClient.Update(ctx, dgd)).To(Succeed()) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, dgd)).To(Succeed()) + useDS, reason := reconciler.shouldUseDisaggregatedSet(dgd) + Expect(useDS).To(BeTrue(), reason) + + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, useDS, reason, nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + Expect(ownedCutoverDCDs(ctx, dgd)).To(HaveLen(2), "legacy DCDs must remain until DS is ready") + Expect(cutoverServiceUIDs(ctx, dgd.Namespace, serviceUIDs)).To(Equal(serviceUIDs)) + + By("marking the DisaggregatedSet ready and completing ownership handoff") + ds := newDisaggregatedSetObject() + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace}, ds)).To(Succeed()) + ds.Object["status"] = map[string]any{ + "observedGeneration": ds.GetGeneration(), + "roleStatuses": []any{ + map[string]any{"name": "prefill", "replicas": int64(1), "updatedReplicas": int64(1), "readyReplicas": int64(1)}, + map[string]any{"name": "decode", "replicas": int64(1), "updatedReplicas": int64(1), "readyReplicas": int64(1)}, + }, + } + Expect(k8sClient.Status().Update(ctx, ds)).To(Succeed()) + + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, useDS, reason, nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStateSuccessful)) + Expect(ownedCutoverDCDs(ctx, dgd)).To(BeEmpty()) + Expect(cutoverServiceUIDs(ctx, dgd.Namespace, serviceUIDs)).To(Equal(serviceUIDs), "DS cutover must preserve component Services") + modelServiceName := dynamo.GenerateServiceName("shared-smoke-model") + modelService := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: modelServiceName, Namespace: dgd.Namespace}, modelService)).To(Succeed()) + Expect(metav1.IsControlledBy(modelService, dgd)).To(BeTrue()) + modelServiceUID := modelService.UID + + By("falling back to DCDs while preserving Services until replacements are ready") + dgd.Annotations = nil + Expect(k8sClient.Update(ctx, dgd)).To(Succeed()) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, dgd)).To(Succeed()) + + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + Expect(ownedCutoverDCDs(ctx, dgd)).To(HaveLen(2)) + markCutoverDCDsReady(ctx, dgd) + + By("restoring component and shared model Service ownership before deleting the stale DisaggregatedSet") + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStateSuccessful)) + Expect(apierrors.IsNotFound(k8sClient.Get(ctx, types.NamespacedName{Name: disaggregatedSetName(dgd), Namespace: dgd.Namespace}, newDisaggregatedSetObject()))).To(BeTrue()) + + fallbackDCDs := ownedCutoverDCDs(ctx, dgd) + fallbackDCDNames := map[string]struct{}{} + for _, dcd := range fallbackDCDs { + fallbackDCDNames[dcd.Name] = struct{}{} + service := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dynamo.NormalizeKubeResourceName(dcd.Name), Namespace: dgd.Namespace}, service)).To(Succeed()) + owner := metav1.GetControllerOf(service) + Expect(owner).NotTo(BeNil()) + Expect(owner.Kind).To(Equal("DynamoComponentDeployment")) + Expect(owner.Name).To(Equal(dcd.Name)) + } + Expect(cutoverServiceUIDs(ctx, dgd.Namespace, serviceUIDs)).To(Equal(serviceUIDs)) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: modelServiceName, Namespace: dgd.Namespace}, modelService)).To(Succeed()) + Expect(modelService.UID).To(Equal(modelServiceUID)) + modelOwner := metav1.GetControllerOf(modelService) + Expect(modelOwner).NotTo(BeNil()) + Expect(modelOwner.Kind).To(Equal("DynamoComponentDeployment")) + Expect(fallbackDCDNames).To(HaveKey(modelOwner.Name)) + }) + + It("recreates a missing shared model Service before completing DisaggregatedSet fallback", func() { + ctx := context.Background() + + By("creating a DisaggregatedSet deployment with component and shared model Services") + dgd := newDSHappyPathDGD() + dgd.Name = "demo-ds-missing-services" + dgd.UID = "" + for i := range dgd.Spec.Components { + dgd.Spec.Components[i].ModelRef = &nvidiacomv1beta1.ModelReference{Name: "shared-missing-model"} + } + Expect(k8sClient.Create(ctx, dgd)).To(Succeed()) + DeferCleanup(func() { + _ = k8sClient.Delete(ctx, dgd) + }) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, dgd)).To(Succeed()) + + runtimeConfig := &commoncontroller.RuntimeConfig{Gate: features.Gates{LWS: true, DisaggregatedSet: true}} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: k8sClient, + Recorder: record.NewFakeRecorder(100), + Config: &configv1alpha1.OperatorConfiguration{ + Discovery: configv1alpha1.DiscoveryConfiguration{Backend: configv1alpha1.DiscoveryBackendKubernetes}, + }, + RuntimeConfig: runtimeConfig, + } + + result, err := reconciler.reconcileWorkloadResources(ctx, dgd, true, true, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + + By("creating replacement DCDs and their component Services") + dgd.Annotations = nil + Expect(k8sClient.Update(ctx, dgd)).To(Succeed()) + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, dgd)).To(Succeed()) + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStatePending)) + + replacementDCDs := ownedCutoverDCDs(ctx, dgd) + Expect(replacementDCDs).To(HaveLen(2)) + dcdReconciler := &DynamoComponentDeploymentReconciler{ + Client: k8sClient, + Recorder: record.NewFakeRecorder(100), + Config: reconciler.Config, + RuntimeConfig: runtimeConfig, + } + for i := range replacementDCDs { + modified, err := dcdReconciler.createOrUpdateOrDeleteServices(ctx, generateResourceOption{ + dynamoComponentDeployment: &replacementDCDs[i], + }) + Expect(err).NotTo(HaveOccurred()) + Expect(modified).To(BeTrue()) + } + + By("deleting the DGD-owned shared model Service before ownership handoff") + modelServiceName := dynamo.GenerateServiceName("shared-missing-model") + modelService := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: modelServiceName, Namespace: dgd.Namespace}, modelService)).To(Succeed()) + Expect(k8sClient.Delete(ctx, modelService)).To(Succeed()) + markCutoverDCDsReady(ctx, dgd) + + By("recreating the missing model Service under a replacement DCD owner before deleting the DisaggregatedSet") + result, err = reconciler.reconcileWorkloadResources(ctx, dgd, true, false, "", nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(result.State).To(Equal(nvidiacomv1beta1.DGDStateSuccessful)) + Expect(apierrors.IsNotFound(k8sClient.Get(ctx, types.NamespacedName{ + Name: disaggregatedSetName(dgd), + Namespace: dgd.Namespace, + }, newDisaggregatedSetObject()))).To(BeTrue()) + + replacementDCDNames := map[string]struct{}{} + for i := range replacementDCDs { + dcd := &replacementDCDs[i] + replacementDCDNames[dcd.Name] = struct{}{} + service := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: dynamo.NormalizeKubeResourceName(dcd.Name), + Namespace: dgd.Namespace, + }, service)).To(Succeed()) + owner := metav1.GetControllerOf(service) + Expect(owner).NotTo(BeNil()) + Expect(owner.Kind).To(Equal("DynamoComponentDeployment")) + Expect(owner.Name).To(Equal(dcd.Name)) + } + Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: modelServiceName, + Namespace: dgd.Namespace, + }, modelService)).To(Succeed()) + modelOwner := metav1.GetControllerOf(modelService) + Expect(modelOwner).NotTo(BeNil()) + Expect(modelOwner.Kind).To(Equal("DynamoComponentDeployment")) + Expect(replacementDCDNames).To(HaveKey(modelOwner.Name)) + }) +}) + +func ownedCutoverDCDs(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment) []nvidiacomv1beta1.DynamoComponentDeployment { + list := &nvidiacomv1beta1.DynamoComponentDeploymentList{} + Expect(k8sClient.List(ctx, list)).To(Succeed()) + owned := make([]nvidiacomv1beta1.DynamoComponentDeployment, 0, len(list.Items)) + for i := range list.Items { + if metav1.IsControlledBy(&list.Items[i], dgd) { + owned = append(owned, list.Items[i]) + } + } + return owned +} + +func markCutoverDCDsReady(ctx context.Context, dgd *nvidiacomv1beta1.DynamoGraphDeployment) { + for _, dcd := range ownedCutoverDCDs(ctx, dgd) { + replicas := ptr.Deref(dcd.Spec.Replicas, int32(1)) + dcd.Status.ObservedGeneration = dcd.Generation + dcd.Status.Conditions = []metav1.Condition{{ + Type: nvidiacomv1beta1.DynamoComponentDeploymentConditionTypeAvailable, + Status: metav1.ConditionTrue, + Reason: "CutoverTestReady", + LastTransitionTime: metav1.Now(), + }} + dcd.Status.Component = &nvidiacomv1beta1.ComponentReplicaStatus{ + ComponentKind: nvidiacomv1beta1.ComponentKindLeaderWorkerSet, + ComponentNames: []string{dcd.Name + "-0"}, + Replicas: replicas, + UpdatedReplicas: replicas, + ReadyReplicas: ptr.To(replicas), + AvailableReplicas: ptr.To(replicas), + } + Expect(k8sClient.Status().Update(ctx, &dcd)).To(Succeed()) + } +} + +func cutoverServiceUIDs(ctx context.Context, namespace string, expected map[string]types.UID) map[string]types.UID { + uids := make(map[string]types.UID, len(expected)) + for name := range expected { + service := &corev1.Service{} + Expect(k8sClient.Get(ctx, types.NamespacedName{Name: name, Namespace: namespace}, service)).To(Succeed()) + uids[name] = service.UID + } + return uids +} + +func newDSHappyPathDGD() *nvidiacomv1beta1.DynamoGraphDeployment { + return &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-ds", + Namespace: "default", + UID: "demo-ds-uid", + Annotations: map[string]string{ + consts.KubeAnnotationEnableDisaggregatedSet: consts.KubeLabelValueTrue, + }, + }, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + { + ComponentName: "prefill", + ComponentType: nvidiacomv1beta1.ComponentTypePrefill, + RuntimeVersionOverride: "1.0.0", + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + PodTemplate: dsTestPodTemplate(), + }, + { + ComponentName: "decode", + ComponentType: nvidiacomv1beta1.ComponentTypeDecode, + RuntimeVersionOverride: "1.0.0", + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + PodTemplate: dsTestPodTemplate(), + }, + }, + }, + } +} + +func dsTestPodTemplate() *corev1.PodTemplateSpec { + return &corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{ + Name: "main", + Image: "busybox:1.36", + Command: []string{"sh"}, + Args: []string{"-c", "sleep 3600"}, + }}, + }, + } +} diff --git a/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_real_crd_test.go b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_real_crd_test.go new file mode 100644 index 000000000000..e43b186895b0 --- /dev/null +++ b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_real_crd_test.go @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "os/exec" + "path/filepath" + "strings" + "testing" + + configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/testing/operatorenv" + "github.com/stretchr/testify/require" + "k8s.io/client-go/tools/record" +) + +func TestDisaggregatedSetRealLWSCRDValidationAndConvergence(t *testing.T) { + t.Log("start envtest with Dynamo and the pinned real LWS CRDs") + operatorRoot, err := filepath.Abs(filepath.Join("..", "..")) + require.NoError(t, err) + cmd := exec.Command("go", "list", "-m", "-f", "{{.Dir}}", "sigs.k8s.io/lws") + cmd.Dir = operatorRoot + output, err := cmd.Output() + require.NoError(t, err) + lwsModuleRoot := strings.TrimSpace(string(output)) + + env := operatorenv.New(operatorenv.Options{ + SetupWebhooks: setupProductionWebhooks, + CRDDirectoryPaths: []string{ + filepath.Join(operatorRoot, "config", "crd", "bases"), + filepath.Join(lwsModuleRoot, "config", "crd", "bases"), + }, + }) + testEnv := env.RunT(t) + + t.Log("generate and create a two-role DisaggregatedSet through the real API schema") + dgd := newDSHappyPathDGD() + dgd.Namespace = testEnv.Namespace() + dcds, err := dynamo.GenerateDynamoComponentsDeployments(dgd, nil, nil, dynamo.RollingUpdateContext{}) + require.NoError(t, err) + selection, reason := selectDisaggregatedSetComponents(dgd) + require.Empty(t, reason) + + reconciler := &DynamoGraphDeploymentReconciler{ + Client: testEnv.Client(), + Recorder: record.NewFakeRecorder(10), + Config: &configv1alpha1.OperatorConfiguration{}, + RuntimeConfig: &commoncontroller.RuntimeConfig{ + Gate: features.Gates{LWS: true, DisaggregatedSet: true}, + }, + } + desired, err := reconciler.generateDisaggregatedSet(t.Context(), dgd, dcds, selection) + require.NoError(t, err) + modified, _, err := reconciler.syncDisaggregatedSet(t.Context(), dgd, desired) + require.NoError(t, err) + require.True(t, modified) + + t.Log("reconcile identical desired state after API defaulting and verify convergence") + desired, err = reconciler.generateDisaggregatedSet(t.Context(), dgd, dcds, selection) + require.NoError(t, err) + modified, _, err = reconciler.syncDisaggregatedSet(t.Context(), dgd, desired) + require.NoError(t, err) + require.False(t, modified, "an identical desired DisaggregatedSet must converge after API defaulting") +} diff --git a/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_test.go b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_test.go new file mode 100644 index 000000000000..940004a7d088 --- /dev/null +++ b/deploy/operator/internal/controller/dynamographdeployment_disaggregatedset_test.go @@ -0,0 +1,1041 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "context" + "fmt" + "strings" + "testing" + + configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/apimachinery/pkg/types" + "k8s.io/utils/ptr" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/event" + disaggregatedsetv1 "sigs.k8s.io/lws/api/disaggregatedset/v1" + leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1" + disaggregatedsetutils "sigs.k8s.io/lws/pkg/utils/disaggregatedset" + + nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1" + nvidiacomv1beta1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta1" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/checkpoint" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/discovery" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" + "github.com/stretchr/testify/require" +) + +func TestSelectDisaggregatedSetComponents(t *testing.T) { + t.Run("selects multinode worker roles", func(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + { + ComponentName: "prefill", + ComponentType: nvidiacomv1beta1.ComponentTypePrefill, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + Replicas: ptr.To(int32(2)), + }, + { + ComponentName: "decode", + ComponentType: nvidiacomv1beta1.ComponentTypeDecode, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + Replicas: ptr.To(int32(2)), + }, + { + ComponentName: "frontend", + ComponentType: nvidiacomv1beta1.ComponentTypeFrontend, + }, + }, + }, + } + + selection, reason := selectDisaggregatedSetComponents(dgd) + require.Empty(t, reason) + require.Equal(t, "prefill", selection.componentToRole["prefill"]) + require.Equal(t, "decode", selection.componentToRole["decode"]) + require.Len(t, selection.componentToRole, 2) + }) + + t.Run("rejects scaling adapter", func(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + { + ComponentName: "prefill", + ComponentType: nvidiacomv1beta1.ComponentTypePrefill, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + ScalingAdapter: &nvidiacomv1beta1.ScalingAdapter{}, + Replicas: ptr.To(int32(2)), + }, + { + ComponentName: "decode", + ComponentType: nvidiacomv1beta1.ComponentTypeDecode, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + Replicas: ptr.To(int32(2)), + }, + }, + }, + } + + _, reason := selectDisaggregatedSetComponents(dgd) + require.Contains(t, reason, "scalingAdapter") + }) + + t.Run("rejects more roles than the DS API supports", func(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{} + for i := 0; i < maxDisaggregatedSetRoles+1; i++ { + dgd.Spec.Components = append(dgd.Spec.Components, nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + ComponentName: fmt.Sprintf("worker-%d", i), + ComponentType: nvidiacomv1beta1.ComponentTypeWorker, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }) + } + + _, reason := selectDisaggregatedSetComponents(dgd) + require.Contains(t, reason, "at most 10 roles") + }) +} + +func TestDisaggregatedSetChildNamesFitDNSLabelLimit(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: strings.Repeat("d", 63)}, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + { + ComponentName: strings.Repeat("p", 63), + ComponentType: nvidiacomv1beta1.ComponentTypeWorker, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }, + { + ComponentName: strings.Repeat("q", 63), + ComponentType: nvidiacomv1beta1.ComponentTypeWorker, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }, + }, + }, + } + + selection, reason := selectDisaggregatedSetComponents(dgd) + require.Empty(t, reason) + require.Len(t, selection.componentToRole, 2) + setName := disaggregatedSetName(dgd) + require.LessOrEqual(t, len(setName), maxDisaggregatedSetNameLength) + for _, roleName := range selection.componentToRole { + require.LessOrEqual(t, len(roleName), maxDisaggregatedSetRoleNameLength) + childName := disaggregatedsetutils.GenerateName(setName, roleName, strings.Repeat("a", disaggregatedSetRevisionLength)) + require.LessOrEqual(t, len(childName), 63) + } +} + +func TestCheckDisaggregatedSetReadiness(t *testing.T) { + ds := newDisaggregatedSetObject() + ds.SetName("demo") + ds.SetGeneration(3) + ds.Object["status"] = map[string]any{ + "observedGeneration": int64(2), + "roleStatuses": []any{ + map[string]any{"name": "prefill", "replicas": int64(2), "updatedReplicas": int64(2), "readyReplicas": int64(2)}, + map[string]any{"name": "decode", "replicas": int64(2), "updatedReplicas": int64(2), "readyReplicas": int64(1)}, + }, + } + selection := disaggregatedSetSelection{ + componentToRole: map[string]string{"prefill": "prefill", "decode": "decode"}, + desiredReplicas: map[string]int32{"prefill": 2, "decode": 2}, + } + + ready, reason, statuses := checkDisaggregatedSetReadiness(ds, selection) + require.False(t, ready) + require.Contains(t, reason, "observed generation") + require.Equal(t, int32(2), ptr.Deref(statuses["prefill"].ReadyReplicas, 0)) + + ds.Object["status"].(map[string]any)["observedGeneration"] = int64(3) + ready, reason, statuses = checkDisaggregatedSetReadiness(ds, selection) + require.False(t, ready) + require.Contains(t, reason, "decode") + require.Equal(t, int32(1), ptr.Deref(statuses["decode"].ReadyReplicas, 0)) + + ds.Object["status"].(map[string]any)["roleStatuses"] = []any{ + map[string]any{"name": "prefill", "replicas": int64(2), "updatedReplicas": int64(2), "readyReplicas": int64(2)}, + map[string]any{"name": "decode", "replicas": int64(2), "updatedReplicas": int64(2), "readyReplicas": int64(2)}, + } + ready, _, _ = checkDisaggregatedSetReadiness(ds, selection) + require.True(t, ready) + + t.Log("parent role status must not report readiness while extra replicas remain") + ds.Object["status"].(map[string]any)["roleStatuses"] = []any{ + map[string]any{"name": "prefill", "replicas": int64(3), "updatedReplicas": int64(2), "readyReplicas": int64(3)}, + map[string]any{"name": "decode", "replicas": int64(2), "updatedReplicas": int64(2), "readyReplicas": int64(2)}, + } + ready, reason, _ = checkDisaggregatedSetReadiness(ds, selection) + require.False(t, ready) + require.Contains(t, reason, "prefill") +} + +func TestGenerateDisaggregatedSetRolesUseDGDIdentityAndExplicitDefaults(t *testing.T) { + t.Log("build a two-role DisaggregatedSet with Kubernetes discovery enabled") + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + require.NoError(t, appsv1.AddToScheme(scheme)) + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, leaderworkersetv1.AddToScheme(scheme)) + + dgd := newDSHappyPathDGD() + dcds, err := dynamo.GenerateDynamoComponentsDeployments(dgd, nil, nil, dynamo.RollingUpdateContext{}) + require.NoError(t, err) + selection, reason := selectDisaggregatedSetComponents(dgd) + require.Empty(t, reason) + + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).Build(), + Config: &configv1alpha1.OperatorConfiguration{ + Discovery: configv1alpha1.DiscoveryConfiguration{Backend: configv1alpha1.DiscoveryBackendKubernetes}, + }, + RuntimeConfig: &commoncontroller.RuntimeConfig{Gate: features.Gates{LWS: true, DisaggregatedSet: true}}, + } + ds, err := reconciler.generateDisaggregatedSet(t.Context(), dgd, dcds, selection) + require.NoError(t, err) + + t.Log("verify every role carries schema-valid defaults and the stable DGD service account") + roles, found, err := unstructured.NestedSlice(ds.Object, "spec", "roles") + require.NoError(t, err) + require.True(t, found) + require.Len(t, roles, 2) + for _, rawRole := range roles { + role := rawRole.(map[string]any) + require.Equal(t, string(leaderworkersetv1.RollingUpdateStrategyType), nestedString(t, role, "spec", "rolloutStrategy", "type")) + require.Equal(t, string(leaderworkersetv1.RecreateGroupOnPodRestart), nestedString(t, role, "spec", "leaderWorkerTemplate", "restartPolicy")) + require.Equal(t, discovery.GetK8sDiscoveryServiceAccountName(dgd.Name), nestedString(t, role, "spec", "leaderWorkerTemplate", "leaderTemplate", "spec", "serviceAccountName")) + require.Equal(t, discovery.GetK8sDiscoveryServiceAccountName(dgd.Name), nestedString(t, role, "spec", "leaderWorkerTemplate", "workerTemplate", "spec", "serviceAccountName")) + } +} + +func nestedString(t *testing.T, obj map[string]any, fields ...string) string { + t.Helper() + value, found, err := unstructured.NestedString(obj, fields...) + require.NoError(t, err) + require.True(t, found, "missing field %s", strings.Join(fields, ".")) + return value +} + +func TestDisaggregatedSetPredicatesObserveRoutingMetadata(t *testing.T) { + t.Log("DisaggregatedSet updates enqueue on labels and controller ownership") + baseDS := newDisaggregatedSetObject() + baseDS.SetLabels(map[string]string{consts.KubeLabelDynamoGraphDeploymentName: "demo"}) + relabeledDS := baseDS.DeepCopy() + relabeledDS.SetLabels(map[string]string{consts.KubeLabelDynamoGraphDeploymentName: "other"}) + reownedDS := baseDS.DeepCopy() + reownedDS.SetOwnerReferences([]metav1.OwnerReference{{APIVersion: nvidiacomv1beta1.GroupVersion.String(), Kind: dynamoGraphDeploymentKind, Name: "demo", UID: "demo-uid", Controller: ptr.To(true)}}) + dsCases := []struct { + name string + updated *unstructured.Unstructured + changed bool + }{ + {name: "unchanged", updated: baseDS.DeepCopy(), changed: false}, + {name: "label changed", updated: relabeledDS, changed: true}, + {name: "owner changed", updated: reownedDS, changed: true}, + } + for _, testCase := range dsCases { + t.Run("DisaggregatedSet/"+testCase.name, func(t *testing.T) { + require.Equal(t, testCase.changed, disaggregatedSetStatusChanged(baseDS, testCase.updated)) + }) + } + + t.Log("LeaderWorkerSet updates enqueue on routing labels and controller ownership") + baseLWS := &leaderworkersetv1.LeaderWorkerSet{ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + disaggregatedsetv1.SetNameLabelKey: "demo", + disaggregatedsetv1.RoleLabelKey: "prefill", + disaggregatedsetv1.RevisionLabelKey: "revision-a", + }, + }} + relabeledLWS := baseLWS.DeepCopy() + relabeledLWS.Labels[disaggregatedsetv1.RevisionLabelKey] = "revision-b" + reownedLWS := baseLWS.DeepCopy() + reownedLWS.SetOwnerReferences([]metav1.OwnerReference{{APIVersion: disaggregatedSetGVK.GroupVersion().String(), Kind: "DisaggregatedSet", Name: "demo", UID: "demo-ds-uid", Controller: ptr.To(true)}}) + lwsCases := []struct { + name string + updated *leaderworkersetv1.LeaderWorkerSet + changed bool + }{ + {name: "unchanged", updated: baseLWS.DeepCopy(), changed: false}, + {name: "revision label changed", updated: relabeledLWS, changed: true}, + {name: "owner changed", updated: reownedLWS, changed: true}, + } + for _, testCase := range lwsCases { + t.Run("LeaderWorkerSet/"+testCase.name, func(t *testing.T) { + require.Equal(t, testCase.changed, leaderWorkerSetStatusChanged(baseLWS, testCase.updated)) + }) + } +} + +func TestEnsureControlledByDCDTransfersOnlyDGDOwnedResources(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + require.NoError(t, corev1.AddToScheme(scheme)) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "dgd-uid"}, + } + dcd := &nvidiacomv1beta1.DynamoComponentDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo-prefill", Namespace: "default", UID: "dcd-uid"}, + } + + t.Run("preserves non-controller owners during DGD to DCD handoff", func(t *testing.T) { + service := &corev1.Service{ObjectMeta: metav1.ObjectMeta{ + Name: "demo-prefill", + Namespace: "default", + OwnerReferences: []metav1.OwnerReference{ + *dgdControllerOwnerReference(dgd), + {APIVersion: "v1", Kind: "ConfigMap", Name: "keep", UID: "keep-uid"}, + }, + }} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(service).Build(), + } + + require.NoError(t, reconciler.ensureControlledByDCD(t.Context(), dgd, dcd, service)) + persisted := &corev1.Service{} + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(service), persisted)) + require.True(t, metav1.IsControlledBy(persisted, dcd)) + require.Contains(t, persisted.OwnerReferences, metav1.OwnerReference{ + APIVersion: "v1", + Kind: "ConfigMap", + Name: "keep", + UID: "keep-uid", + }) + }) + + t.Run("refuses to steal a resource from an unrelated controller", func(t *testing.T) { + foreignOwner := metav1.OwnerReference{ + APIVersion: "apps/v1", + Kind: "StatefulSet", + Name: "foreign", + UID: "foreign-uid", + Controller: ptr.To(true), + } + service := &corev1.Service{ObjectMeta: metav1.ObjectMeta{ + Name: "foreign-owned", + Namespace: "default", + OwnerReferences: []metav1.OwnerReference{foreignOwner}, + }} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(service).Build(), + } + + err := reconciler.ensureControlledByDCD(t.Context(), dgd, dcd, service) + require.ErrorContains(t, err, "resource is controlled by apps/v1/StatefulSet") + persisted := &corev1.Service{} + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(service), persisted)) + require.Equal(t, []metav1.OwnerReference{foreignOwner}, persisted.OwnerReferences) + }) + + t.Run("accepts a sibling DCD controlled by the same DGD", func(t *testing.T) { + sibling := &nvidiacomv1beta1.DynamoComponentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-decode", + Namespace: "default", + UID: "sibling-uid", + OwnerReferences: []metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}, + }, + } + service := &corev1.Service{ObjectMeta: metav1.ObjectMeta{ + Name: "shared-model", + Namespace: "default", + OwnerReferences: []metav1.OwnerReference{*dcdControllerOwnerReference(sibling)}, + }} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(sibling, service).Build(), + } + + require.NoError(t, reconciler.ensureControlledByDCD(t.Context(), dgd, dcd, service)) + persisted := &corev1.Service{} + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(service), persisted)) + require.True(t, metav1.IsControlledBy(persisted, sibling)) + }) +} + +func TestApplyDisaggregatedSetCheckpointStartupPoliciesCoordinatesSelectedRoles(t *testing.T) { + dcds := map[string]*nvidiacomv1beta1.DynamoComponentDeployment{ + "prefill": { + Spec: nvidiacomv1beta1.DynamoComponentDeploymentSpec{ + DynamoComponentDeploymentSharedSpec: nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{Replicas: ptr.To(int32(2))}, + }, + }, + "decode": { + Spec: nvidiacomv1beta1.DynamoComponentDeploymentSpec{ + DynamoComponentDeploymentSharedSpec: nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{Replicas: ptr.To(int32(2))}, + }, + }, + } + checkpointInfos := map[string]*checkpoint.CheckpointInfo{ + "prefill": { + Enabled: true, + Ready: false, + StartupPolicy: nvidiacomv1alpha1.CheckpointStartupPolicyWaitForCheckpoint, + }, + } + selection := disaggregatedSetSelection{ + componentToRole: map[string]string{"prefill": "prefill", "decode": "decode"}, + desiredReplicas: map[string]int32{"prefill": 2, "decode": 2}, + } + + gated, err := applyDisaggregatedSetCheckpointStartupPolicies(dcds, checkpointInfos, selection) + require.NoError(t, err) + require.True(t, gated) + require.Equal(t, int32(0), ptr.Deref(dcds["prefill"].Spec.Replicas, -1)) + require.Equal(t, int32(0), ptr.Deref(dcds["decode"].Spec.Replicas, -1)) + require.Equal(t, map[string]int32{"prefill": 0, "decode": 0}, selection.desiredReplicas) +} + +func TestDeleteStaleDisaggregatedSetComponentServices(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + require.NoError(t, corev1.AddToScheme(scheme)) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + } + labels := func(componentName string) map[string]string { + return map[string]string{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + consts.KubeLabelDynamoComponent: componentName, + } + } + service := func(name, componentName string, owned bool) *corev1.Service { + svc := &corev1.Service{ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: dgd.Namespace, Labels: labels(componentName)}} + if owned { + svc.OwnerReferences = []metav1.OwnerReference{*dgdControllerOwnerReference(dgd)} + } + return svc + } + current := service("demo-prefill-new", "prefill", true) + stale := service("demo-prefill-old", "prefill", true) + removedComponent := service("demo-removed-old", "removed", true) + userManaged := service("demo-prefill-user", "prefill", false) + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(dgd, current, stale, removedComponent, userManaged).Build(), + } + dcds := map[string]*nvidiacomv1beta1.DynamoComponentDeployment{ + "prefill": {ObjectMeta: metav1.ObjectMeta{Name: current.Name}}, + } + + require.NoError(t, reconciler.deleteStaleDisaggregatedSetComponentServices(t.Context(), dgd, dcds)) + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(current), &corev1.Service{})) + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(userManaged), &corev1.Service{})) + require.True(t, apierrors.IsNotFound(reconciler.Get(t.Context(), client.ObjectKeyFromObject(stale), &corev1.Service{}))) + require.True(t, apierrors.IsNotFound(reconciler.Get(t.Context(), client.ObjectKeyFromObject(removedComponent), &corev1.Service{}))) +} + +func TestAdoptSelectedModelServicesLeavesSharedForeignServiceOwner(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + require.NoError(t, corev1.AddToScheme(scheme)) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + {ComponentName: "prefill", ModelRef: &nvidiacomv1beta1.ModelReference{Name: "shared-model"}}, + }}, + } + foreignDGD := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "other", Namespace: "default", UID: "other-uid"}, + } + foreignDCD := &nvidiacomv1beta1.DynamoComponentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "other-prefill", + Namespace: "default", + UID: "other-prefill-uid", + OwnerReferences: []metav1.OwnerReference{*dgdControllerOwnerReference(foreignDGD)}, + }, + } + modelService := &corev1.Service{ObjectMeta: metav1.ObjectMeta{ + Name: dynamo.GenerateServiceName("shared-model"), + Namespace: "default", + OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef( + foreignDCD, + nvidiacomv1beta1.GroupVersion.WithKind(dynamoComponentDeploymentKind), + )}, + }} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(dgd, foreignDGD, foreignDCD, modelService).Build(), + } + selection := disaggregatedSetSelection{componentToRole: map[string]string{"prefill": "prefill"}} + + require.NoError(t, reconciler.adoptSelectedModelServices(t.Context(), dgd, selection)) + updated := &corev1.Service{} + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(modelService), updated)) + require.Equal(t, foreignDCD.UID, metav1.GetControllerOf(updated).UID) +} + +func TestCheckDisaggregatedSetReadinessFallsBackToTargetRevisionChildLWS(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, leaderworkersetv1.AddToScheme(scheme)) + + ds := newDisaggregatedSetObject() + ds.SetName("demo") + ds.SetNamespace("default") + ds.SetUID("ds-uid") + ds.SetGeneration(2) + typedDS := &disaggregatedsetv1.DisaggregatedSet{ + Spec: disaggregatedsetv1.DisaggregatedSetSpec{Roles: []disaggregatedsetv1.DisaggregatedRoleSpec{ + {Name: "prefill"}, + {Name: "decode"}, + }}, + } + typedObject, err := runtime.DefaultUnstructuredConverter.ToUnstructured(typedDS) + require.NoError(t, err) + ds.Object["spec"] = typedObject["spec"] + ds.Object["status"] = map[string]any{ + // These ready-looking statuses may belong to generation 1. Without an + // observation marker, readiness must still come from the target children. + "roleStatuses": []any{ + map[string]any{"name": "prefill", "replicas": int64(1), "updatedReplicas": int64(1), "readyReplicas": int64(1)}, + map[string]any{"name": "decode", "replicas": int64(1), "updatedReplicas": int64(1), "readyReplicas": int64(1)}, + }, + } + targetRevision := disaggregatedsetutils.ComputeRevision(typedDS.Spec.Roles) + selection := disaggregatedSetSelection{ + componentToRole: map[string]string{"prefill": "prefill", "decode": "decode"}, + desiredReplicas: map[string]int32{"prefill": 1, "decode": 1}, + } + owner := metav1.OwnerReference{ + APIVersion: disaggregatedSetGVK.GroupVersion().String(), + Kind: disaggregatedSetGVK.Kind, + Name: ds.GetName(), + UID: ds.GetUID(), + Controller: ptr.To(true), + } + child := func(name, role, revision string, ready int32) *leaderworkersetv1.LeaderWorkerSet { + return &leaderworkersetv1.LeaderWorkerSet{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: ds.GetNamespace(), + Generation: 1, + Labels: map[string]string{ + disaggregatedsetv1.SetNameLabelKey: ds.GetName(), + disaggregatedsetv1.RoleLabelKey: role, + disaggregatedsetv1.RevisionLabelKey: revision, + }, + OwnerReferences: []metav1.OwnerReference{owner}, + }, + Status: leaderworkersetv1.LeaderWorkerSetStatus{ + ObservedGeneration: 1, + Replicas: 1, + UpdatedReplicas: ready, + ReadyReplicas: ready, + }, + } + } + objects := []client.Object{ + child("demo-old-prefill", "prefill", "old", 1), + child("demo-new-prefill", "prefill", targetRevision, 0), + child("demo-decode", "decode", targetRevision, 1), + } + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(objects...).Build(), + } + + ready, reason, statuses, err := reconciler.checkDisaggregatedSetReadiness(t.Context(), ds, selection) + require.NoError(t, err) + require.False(t, ready) + require.Contains(t, reason, "demo-new-prefill") + require.Equal(t, []string{"demo-new-prefill", "demo-old-prefill"}, statuses["prefill"].ComponentNames) + require.Equal(t, int32(2), statuses["prefill"].Replicas) + + newPrefill := &leaderworkersetv1.LeaderWorkerSet{} + require.NoError(t, reconciler.Get(t.Context(), types.NamespacedName{Name: "demo-new-prefill", Namespace: ds.GetNamespace()}, newPrefill)) + newPrefill.Status.UpdatedReplicas = 1 + newPrefill.Status.ReadyReplicas = 1 + require.NoError(t, reconciler.Update(t.Context(), newPrefill)) + + ready, reason, _, err = reconciler.checkDisaggregatedSetReadiness(t.Context(), ds, selection) + require.NoError(t, err) + require.False(t, ready) + require.Contains(t, reason, "old child") + + oldPrefill := &leaderworkersetv1.LeaderWorkerSet{} + require.NoError(t, reconciler.Get(t.Context(), types.NamespacedName{Name: "demo-old-prefill", Namespace: ds.GetNamespace()}, oldPrefill)) + oldPrefill.Status.Replicas = 0 + oldPrefill.Status.UpdatedReplicas = 0 + oldPrefill.Status.ReadyReplicas = 0 + require.NoError(t, reconciler.Update(t.Context(), oldPrefill)) + + ready, _, _, err = reconciler.checkDisaggregatedSetReadiness(t.Context(), ds, selection) + require.NoError(t, err) + require.True(t, ready) +} + +func TestListOwnedSelectedDCDsSkipsUserManagedDCDs(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + require.NoError(t, corev1.AddToScheme(scheme)) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + UID: "demo-uid", + }, + } + owned := &nvidiacomv1beta1.DynamoComponentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-prefill", + Namespace: "default", + Labels: map[string]string{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + consts.KubeLabelDynamoComponent: "prefill", + }, + OwnerReferences: []metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}, + }, + Spec: nvidiacomv1beta1.DynamoComponentDeploymentSpec{ + DynamoComponentDeploymentSharedSpec: nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + ComponentName: "prefill", + ComponentType: nvidiacomv1beta1.ComponentTypePrefill, + }, + }, + } + userManaged := &nvidiacomv1beta1.DynamoComponentDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo-decode", + Namespace: "default", + Labels: map[string]string{ + consts.KubeLabelDynamoGraphDeploymentName: dgd.Name, + consts.KubeLabelDynamoComponent: "decode", + }, + }, + Spec: nvidiacomv1beta1.DynamoComponentDeploymentSpec{ + DynamoComponentDeploymentSharedSpec: nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + ComponentName: "decode", + ComponentType: nvidiacomv1beta1.ComponentTypeDecode, + }, + }, + } + + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(dgd, owned, userManaged).Build(), + } + + selection := disaggregatedSetSelection{ + componentToRole: map[string]string{ + dynamo.GetDCDComponentName(owned): "prefill", + dynamo.GetDCDComponentName(userManaged): "decode", + }, + desiredReplicas: map[string]int32{"prefill": 1, "decode": 1}, + } + + got, err := reconciler.listOwnedSelectedDCDs(t.Context(), dgd, selection) + require.NoError(t, err) + require.Len(t, got, 1) + require.Equal(t, owned.Name, got[0].Name) +} + +func TestReconcileReplacementBeforeDisaggregatedSetCleanup(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + } + ds := newDisaggregatedSetObject() + ds.SetName(disaggregatedSetName(dgd)) + ds.SetNamespace(dgd.Namespace) + ds.SetOwnerReferences([]metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}) + + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(ds).Build(), + RuntimeConfig: newTestRuntimeConfig(true), + } + key := types.NamespacedName{Name: ds.GetName(), Namespace: ds.GetNamespace()} + + result, err := reconciler.reconcileReplacementBeforeDisaggregatedSetCleanup(t.Context(), dgd, false, func() (ReconcileResult, error) { + existing := newDisaggregatedSetObject() + require.NoError(t, reconciler.Get(t.Context(), key, existing), "DisaggregatedSet must remain while its replacement is pending") + return ReconcileResult{State: nvidiacomv1beta1.DGDStatePending}, nil + }) + require.NoError(t, err) + require.Equal(t, nvidiacomv1beta1.DGDStatePending, result.State) + require.NoError(t, reconciler.Get(t.Context(), key, newDisaggregatedSetObject())) + + result, err = reconciler.reconcileReplacementBeforeDisaggregatedSetCleanup(t.Context(), dgd, false, func() (ReconcileResult, error) { + return ReconcileResult{State: nvidiacomv1beta1.DGDStateSuccessful}, nil + }) + require.NoError(t, err) + require.Equal(t, nvidiacomv1beta1.DGDStateSuccessful, result.State) + require.True(t, apierrors.IsNotFound(reconciler.Get(t.Context(), key, newDisaggregatedSetObject()))) +} + +type notFoundOnDeleteClient struct { + client.Client +} + +func (c notFoundOnDeleteClient) Delete(_ context.Context, obj client.Object, _ ...client.DeleteOption) error { + return apierrors.NewNotFound(schema.GroupResource{Group: disaggregatedSetGVK.Group, Resource: "disaggregatedsets"}, obj.GetName()) +} + +func TestDeleteDisaggregatedSetIgnoresNotFoundRace(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + } + ds := newDisaggregatedSetObject() + ds.SetName(disaggregatedSetName(dgd)) + ds.SetNamespace(dgd.Namespace) + ds.SetOwnerReferences([]metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}) + baseClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(ds).Build() + reconciler := &DynamoGraphDeploymentReconciler{Client: notFoundOnDeleteClient{Client: baseClient}} + + require.NoError(t, reconciler.deleteDisaggregatedSetIfExists(t.Context(), dgd)) +} + +func TestSyncDisaggregatedSetPreservesUnmanagedMetadata(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + scheme.AddKnownTypeWithName(disaggregatedSetGVK, &unstructured.Unstructured{}) + scheme.AddKnownTypeWithName(disaggregatedSetGVK.GroupVersion().WithKind("DisaggregatedSetList"), &unstructured.UnstructuredList{}) + + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + } + current := newDisaggregatedSetObject() + current.SetName("demo") + current.SetNamespace("default") + current.SetLabels(map[string]string{"example.com/keep": "label"}) + current.SetAnnotations(map[string]string{"example.com/keep": "annotation"}) + current.SetOwnerReferences([]metav1.OwnerReference{ + *dgdControllerOwnerReference(dgd), + {APIVersion: "v1", Kind: "ConfigMap", Name: "keep", UID: "keep-uid"}, + }) + current.Object["spec"] = map[string]any{"roles": []any{}} + desired := current.DeepCopy() + desired.SetLabels(map[string]string{consts.KubeLabelDynamoGraphDeploymentName: dgd.Name}) + desired.SetAnnotations(map[string]string{"example.com/desired": "annotation"}) + desired.SetOwnerReferences([]metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}) + desired.Object["spec"] = map[string]any{"roles": []any{map[string]any{"name": "prefill"}, map[string]any{"name": "decode"}}} + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(dgd, current).Build(), + } + + modified, synced, err := reconciler.syncDisaggregatedSet(t.Context(), dgd, desired) + require.NoError(t, err) + require.True(t, modified) + require.Equal(t, "label", synced.GetLabels()["example.com/keep"]) + require.Equal(t, dgd.Name, synced.GetLabels()[consts.KubeLabelDynamoGraphDeploymentName]) + require.Equal(t, "annotation", synced.GetAnnotations()["example.com/keep"]) + require.Equal(t, "annotation", synced.GetAnnotations()["example.com/desired"]) + require.Len(t, synced.GetOwnerReferences(), 2) + persisted := newDisaggregatedSetObject() + require.NoError(t, reconciler.Get(t.Context(), client.ObjectKeyFromObject(current), persisted)) + require.Equal(t, dgd.Name, persisted.GetLabels()[consts.KubeLabelDynamoGraphDeploymentName]) + require.Equal(t, "annotation", persisted.GetAnnotations()["example.com/desired"]) +} + +func TestMapDisaggregatedSetChildLWSToDGD(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiacomv1beta1.AddToScheme(scheme)) + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "demo", Namespace: "default", UID: "demo-uid"}, + } + ds := newDisaggregatedSetObject() + ds.SetName(disaggregatedSetName(dgd)) + ds.SetNamespace(dgd.Namespace) + ds.SetUID("ds-uid") + ds.SetOwnerReferences([]metav1.OwnerReference{*dgdControllerOwnerReference(dgd)}) + reconciler := &DynamoGraphDeploymentReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(ds).Build(), + } + child := &leaderworkersetv1.LeaderWorkerSet{ObjectMeta: metav1.ObjectMeta{ + Name: "demo-revision-prefill", + Namespace: dgd.Namespace, + Labels: map[string]string{disaggregatedsetv1.SetNameLabelKey: ds.GetName()}, + }} + + requests := reconciler.mapDisaggregatedSetChildLWSToDGD(t.Context(), child) + require.Len(t, requests, 1) + require.Equal(t, types.NamespacedName{Name: dgd.Name, Namespace: dgd.Namespace}, requests[0].NamespacedName) +} + +func TestShouldUseDisaggregatedSet(t *testing.T) { + twoEligibleDGD := func() *nvidiacomv1beta1.DynamoGraphDeployment { + return &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: "demo", + Namespace: "default", + Annotations: map[string]string{ + consts.KubeAnnotationEnableDisaggregatedSet: consts.KubeLabelValueTrue, + }, + }, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + { + ComponentName: "prefill", + ComponentType: nvidiacomv1beta1.ComponentTypePrefill, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }, + { + ComponentName: "decode", + ComponentType: nvidiacomv1beta1.ComponentTypeDecode, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }, + }, + }, + } + } + + t.Run("annotation missing falls back", func(t *testing.T) { + dgd := twoEligibleDGD() + delete(dgd.Annotations, consts.KubeAnnotationEnableDisaggregatedSet) + r := &DynamoGraphDeploymentReconciler{ + RuntimeConfig: newTestRuntimeConfig(true), + } + use, reason := r.shouldUseDisaggregatedSet(dgd) + require.False(t, use) + require.Empty(t, reason) + }) + + t.Run("API unavailable falls back with reason", func(t *testing.T) { + dgd := twoEligibleDGD() + r := &DynamoGraphDeploymentReconciler{ + RuntimeConfig: newTestRuntimeConfig(false), + } + use, reason := r.shouldUseDisaggregatedSet(dgd) + require.False(t, use) + require.Contains(t, reason, "DisaggregatedSet API is not available") + }) + + t.Run("only one eligible role falls back with reason", func(t *testing.T) { + dgd := twoEligibleDGD() + dgd.Spec.Components = dgd.Spec.Components[:1] + r := &DynamoGraphDeploymentReconciler{ + RuntimeConfig: newTestRuntimeConfig(true), + } + use, reason := r.shouldUseDisaggregatedSet(dgd) + require.False(t, use) + require.Contains(t, reason, "two eligible multinode worker roles") + }) + + t.Run("eligible DGD opts in", func(t *testing.T) { + dgd := twoEligibleDGD() + r := &DynamoGraphDeploymentReconciler{ + RuntimeConfig: newTestRuntimeConfig(true), + } + use, reason := r.shouldUseDisaggregatedSet(dgd) + require.True(t, use) + require.Empty(t, reason) + }) + + t.Run("non-selected multinode component requires LWS", func(t *testing.T) { + dgd := twoEligibleDGD() + dgd.Spec.Components = append(dgd.Spec.Components, nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + ComponentName: "frontend", + ComponentType: nvidiacomv1beta1.ComponentTypeFrontend, + Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}, + }) + runtimeConfig := newTestRuntimeConfig(true) + r := &DynamoGraphDeploymentReconciler{RuntimeConfig: runtimeConfig} + + use, reason := r.shouldUseDisaggregatedSet(dgd) + require.False(t, use) + require.Contains(t, reason, "requires LeaderWorkerSet support") + + runtimeConfig.Gate.LWS = true + use, reason = r.shouldUseDisaggregatedSet(dgd) + require.True(t, use) + require.Empty(t, reason) + }) +} + +func TestCoalesceDisaggregatedSetRestartState(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + {ComponentName: "frontend", ComponentType: nvidiacomv1beta1.ComponentTypeFrontend}, + {ComponentName: "prefill", ComponentType: nvidiacomv1beta1.ComponentTypePrefill, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + {ComponentName: "decode", ComponentType: nvidiacomv1beta1.ComponentTypeDecode, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + }, + }, + } + + state := &dynamo.RestartState{Timestamp: "restart-1", ComponentsToAnnotate: map[string]bool{"frontend": true}} + require.Same(t, state, coalesceDisaggregatedSetRestartState(dgd, state)) + require.Equal(t, map[string]bool{"frontend": true}, state.ComponentsToAnnotate) + + state.ComponentsToAnnotate["prefill"] = true + require.Same(t, state, coalesceDisaggregatedSetRestartState(dgd, state)) + require.True(t, state.ShouldAnnotateComponent("prefill")) + require.True(t, state.ShouldAnnotateComponent("decode"), "all DS roles must share one restart revision") +} + +func TestDisaggregatedSetRestartCoalescingRespectsGrovePrecedence(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + consts.KubeAnnotationEnableDisaggregatedSet: consts.KubeLabelValueTrue, + }}, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + {ComponentName: "prefill", ComponentType: nvidiacomv1beta1.ComponentTypePrefill, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + {ComponentName: "decode", ComponentType: nvidiacomv1beta1.ComponentTypeDecode, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + }, + }, + } + reconciler := &DynamoGraphDeploymentReconciler{RuntimeConfig: &commoncontroller.RuntimeConfig{ + Gate: features.Gates{Grove: true, DisaggregatedSet: true}, + }} + + groveState := &dynamo.RestartState{Timestamp: "restart-1", ComponentsToAnnotate: map[string]bool{"prefill": true}} + require.Same(t, groveState, reconciler.coalesceDisaggregatedSetRestartStateForPath(dgd, true, groveState)) + require.False(t, groveState.ShouldAnnotateComponent("decode"), "Grove precedence must retain sequential role restarts") + + dgd.Annotations[consts.KubeAnnotationEnableGrove] = consts.KubeLabelValueFalse + dsState := &dynamo.RestartState{Timestamp: "restart-1", ComponentsToAnnotate: map[string]bool{"prefill": true}} + require.Same(t, dsState, reconciler.coalesceDisaggregatedSetRestartStateForPath(dgd, true, dsState)) + require.True(t, dsState.ShouldAnnotateComponent("decode"), "the active DS path must restart selected roles as one unit") +} + +func TestSequentialRestartSkipsAlreadyRestartedDisaggregatedSetRoles(t *testing.T) { + dgd := &nvidiacomv1beta1.DynamoGraphDeployment{ + ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + consts.KubeAnnotationEnableDisaggregatedSet: consts.KubeLabelValueTrue, + consts.KubeAnnotationEnableGrove: consts.KubeLabelValueFalse, + }}, + Spec: nvidiacomv1beta1.DynamoGraphDeploymentSpec{ + Components: []nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{ + {ComponentName: "prefill", ComponentType: nvidiacomv1beta1.ComponentTypePrefill, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + {ComponentName: "decode", ComponentType: nvidiacomv1beta1.ComponentTypeDecode, Multinode: &nvidiacomv1beta1.MultinodeSpec{NodeCount: 2}}, + {ComponentName: "frontend", ComponentType: nvidiacomv1beta1.ComponentTypeFrontend}, + }, + }, + } + reconciler := &DynamoGraphDeploymentReconciler{RuntimeConfig: &commoncontroller.RuntimeConfig{ + Gate: features.Gates{LWS: true, DisaggregatedSet: true}, + }} + + next, found := reconciler.getNextSequentialRestartComponent( + dgd, + []string{"prefill", "decode", "frontend"}, + "prefill", + ) + require.True(t, found) + require.Equal(t, "frontend", next, "decode was restarted in the same DisaggregatedSet revision") + + next, found = reconciler.getNextSequentialRestartComponent( + dgd, + []string{"frontend", "prefill", "decode"}, + "frontend", + ) + require.True(t, found) + require.Equal(t, "prefill", next, "non-DisaggregatedSet components retain sequential ordering") + + next, found = reconciler.getNextSequentialRestartComponent( + dgd, + []string{"prefill", "decode"}, + "prefill", + ) + require.True(t, found) + require.Empty(t, next, "the restart completes after the shared DisaggregatedSet revision") +} + +func TestDisaggregatedSetRestartAnnotationsSurviveConsecutiveRequests(t *testing.T) { + dgd := newDSHappyPathDGD() + dgd.Spec.BackendFramework = string(dynamo.BackendFrameworkVLLM) + dgd.Spec.Components = append([]nvidiacomv1beta1.DynamoComponentDeploymentSharedSpec{{ + ComponentName: "frontend", + ComponentType: nvidiacomv1beta1.ComponentTypeFrontend, + PodTemplate: dsTestPodTemplate(), + }}, dgd.Spec.Components...) + selection, reason := selectDisaggregatedSetComponents(dgd) + require.Empty(t, reason) + + restartOneAnnotations := map[string]string{ + "prefill": "restart-1", + "decode": "restart-1", + } + ds := newDisaggregatedSetObject() + typedSpec := disaggregatedsetv1.DisaggregatedSetSpec{} + for componentName, roleName := range selection.componentToRole { + timestamp := restartOneAnnotations[componentName] + typedSpec.Roles = append(typedSpec.Roles, disaggregatedsetv1.DisaggregatedRoleSpec{ + Name: roleName, + LeaderWorkerSetTemplateSpec: leaderworkersetv1.LeaderWorkerSetTemplateSpec{ + Spec: leaderworkersetv1.LeaderWorkerSetSpec{ + LeaderWorkerTemplate: leaderworkersetv1.LeaderWorkerTemplate{ + LeaderTemplate: &corev1.PodTemplateSpec{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{consts.RestartAnnotation: timestamp}}}, + WorkerTemplate: corev1.PodTemplateSpec{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{consts.RestartAnnotation: timestamp}}}, + }, + }, + }, + }) + } + spec, err := runtime.DefaultUnstructuredConverter.ToUnstructured(&typedSpec) + require.NoError(t, err) + ds.Object["spec"] = spec + ds.SetName(disaggregatedSetName(dgd)) + ds.SetNamespace(dgd.Namespace) + + scheme := runtime.NewScheme() + scheme.AddKnownTypeWithName(disaggregatedSetGVK, &unstructured.Unstructured{}) + reconciler := &DynamoGraphDeploymentReconciler{Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(ds).Build()} + existingAnnotations, err := reconciler.getExistingRestartAnnotationsDisaggregatedSet(t.Context(), dgd, selection) + require.NoError(t, err) + require.Equal(t, restartOneAnnotations, existingAnnotations) + + // A second sequential request starts with the non-DS frontend. Until the + // restart reaches a selected role, the desired DS must retain restart-1. + restartTwoFrontend := &dynamo.RestartState{Timestamp: "restart-2", ComponentsToAnnotate: map[string]bool{"frontend": true}} + dcds, err := dynamo.GenerateDynamoComponentsDeployments(dgd, restartTwoFrontend, existingAnnotations, dynamo.RollingUpdateContext{}) + require.NoError(t, err) + require.Equal(t, "restart-2", dynamo.GetPodTemplateAnnotations(&dcds["frontend"].Spec.DynamoComponentDeploymentSharedSpec)[consts.RestartAnnotation]) + for componentName := range selection.componentToRole { + require.Equal(t, "restart-1", dynamo.GetPodTemplateAnnotations(&dcds[componentName].Spec.DynamoComponentDeploymentSharedSpec)[consts.RestartAnnotation]) + } + + // Once the first DS role is selected, coalescing produces exactly one new + // whole-set revision with restart-2 on every selected role. + restartTwoDS := &dynamo.RestartState{Timestamp: "restart-2", ComponentsToAnnotate: map[string]bool{"prefill": true}} + coalesceDisaggregatedSetRestartState(dgd, restartTwoDS) + dcds, err = dynamo.GenerateDynamoComponentsDeployments(dgd, restartTwoDS, existingAnnotations, dynamo.RollingUpdateContext{}) + require.NoError(t, err) + for componentName := range selection.componentToRole { + require.Equal(t, "restart-2", dynamo.GetPodTemplateAnnotations(&dcds[componentName].Spec.DynamoComponentDeploymentSharedSpec)[consts.RestartAnnotation]) + } +} + +func TestWorkloadRoutingAnnotationsChanged(t *testing.T) { + oldDGD := &nvidiacomv1beta1.DynamoGraphDeployment{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + consts.KubeAnnotationEnableGrove: consts.KubeLabelValueTrue, + }}} + newDGD := oldDGD.DeepCopy() + require.False(t, workloadRoutingAnnotationsChanged(event.UpdateEvent{ObjectOld: oldDGD, ObjectNew: newDGD})) + + newDGD.Annotations[consts.KubeAnnotationEnableGrove] = consts.KubeLabelValueFalse + newDGD.Annotations[consts.KubeAnnotationEnableDisaggregatedSet] = consts.KubeLabelValueTrue + require.True(t, workloadRoutingAnnotationsChanged(event.UpdateEvent{ObjectOld: oldDGD, ObjectNew: newDGD})) +} + +func TestDGDOwnedServiceEventPredicate(t *testing.T) { + p := dgdOwnedServiceEventPredicate() + service := &corev1.Service{} + require.False(t, p.Create(event.CreateEvent{Object: service})) + require.True(t, p.Update(event.UpdateEvent{ObjectOld: service, ObjectNew: service.DeepCopy()})) + require.True(t, p.Delete(event.DeleteEvent{Object: service})) + require.True(t, p.Generic(event.GenericEvent{Object: service})) +} + +func newTestRuntimeConfig(enabled bool) *commoncontroller.RuntimeConfig { + return &commoncontroller.RuntimeConfig{Gate: features.Gates{DisaggregatedSet: enabled}} +} diff --git a/deploy/operator/internal/controller/testing/disaggregatedset/disaggregatedsets.yaml b/deploy/operator/internal/controller/testing/disaggregatedset/disaggregatedsets.yaml new file mode 100644 index 000000000000..18089d140a4f --- /dev/null +++ b/deploy/operator/internal/controller/testing/disaggregatedset/disaggregatedsets.yaml @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal DisaggregatedSet CRD used by envtest. The real schema is owned by +# the LWS upstream; this fixture only declares the surface that the operator +# touches (spec.roles, status.roleStatuses, status.observedGeneration). +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: disaggregatedsets.disaggregatedset.x-k8s.io +spec: + group: disaggregatedset.x-k8s.io + scope: Namespaced + names: + plural: disaggregatedsets + singular: disaggregatedset + kind: DisaggregatedSet + listKind: DisaggregatedSetList + versions: + - name: v1 + served: true + storage: true + subresources: + status: {} + additionalPrinterColumns: [] + schema: + openAPIV3Schema: + type: object + x-kubernetes-preserve-unknown-fields: true diff --git a/deploy/operator/internal/controller/testing/disaggregatedset/leaderworkersets.yaml b/deploy/operator/internal/controller/testing/disaggregatedset/leaderworkersets.yaml new file mode 100644 index 000000000000..333e685fe541 --- /dev/null +++ b/deploy/operator/internal/controller/testing/disaggregatedset/leaderworkersets.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: leaderworkersets.leaderworkerset.x-k8s.io +spec: + group: leaderworkerset.x-k8s.io + names: + kind: LeaderWorkerSet + listKind: LeaderWorkerSetList + plural: leaderworkersets + singular: leaderworkerset + scope: Namespaced + versions: + - name: v1 + served: true + storage: true + schema: + openAPIV3Schema: + type: object + x-kubernetes-preserve-unknown-fields: true + subresources: + status: {} diff --git a/deploy/operator/internal/features/gates.go b/deploy/operator/internal/features/gates.go index b1efa7f8f4d8..771c3090e656 100644 --- a/deploy/operator/internal/features/gates.go +++ b/deploy/operator/internal/features/gates.go @@ -76,6 +76,18 @@ const ( // Default: true when both API groups are detected; false otherwise LWS Name = "lws" + // DisaggregatedSet enables the opt-in DisaggregatedSet workload pathway. + // + // Owner: @kay-yan + // Experimental since: v1.4.0 + // Beta since: N/A + // GA since: N/A + // Configuration: nvidia.com/enable-disaggregatedset annotation + // Auto-detection: disaggregatedset.x-k8s.io/v1 API + // Requires: LWS serving disaggregatedset.x-k8s.io/v1 + // Default: true when the API is detected; false otherwise + DisaggregatedSet Name = "disaggregatedSet" + // KaiScheduler enables Kai Scheduler integration. // // Owner: @julienmancuso @@ -143,6 +155,7 @@ var allNames = [...]Name{ Checkpoint, Grove, LWS, + DisaggregatedSet, KaiScheduler, VolcanoScheduler, DRA, @@ -164,6 +177,7 @@ type Gates struct { Checkpoint bool `json:"checkpoint"` Grove bool `json:"grove"` LWS bool `json:"lws"` + DisaggregatedSet bool `json:"disaggregatedSet"` KaiScheduler bool `json:"kaiScheduler"` VolcanoScheduler bool `json:"volcanoScheduler"` DRA bool `json:"dra"` @@ -198,6 +212,10 @@ func New(ctx context.Context, mgr ctrl.Manager, config *configv1alpha1.OperatorC lwsAvailable := detectAPIGroup(ctx, mgr, "leaderworkerset.x-k8s.io", "") volcanoAvailable := detectAPIGroup(ctx, mgr, "scheduling.volcano.sh", "") + disaggregatedSetAvailable := detectAPIGroup(ctx, mgr, "disaggregatedset.x-k8s.io", "v1") + // The DS pathway lists and watches the LWS children created by the DS + // controller. Do not register those watches when the LWS API is absent. + gates.DisaggregatedSet = lwsAvailable && disaggregatedSetAvailable if ptr.Deref(config.Orchestrators.LWS.Enabled, lwsAvailable && volcanoAvailable) { if !lwsAvailable { return Gates{}, fmt.Errorf("LWS is explicitly enabled in config but the LWS API group was not detected in the cluster") diff --git a/deploy/operator/internal/features/gates_test.go b/deploy/operator/internal/features/gates_test.go index 86556d642b1e..e9e93b88f1c1 100644 --- a/deploy/operator/internal/features/gates_test.go +++ b/deploy/operator/internal/features/gates_test.go @@ -56,6 +56,7 @@ func allEnabledGates() Gates { Checkpoint: true, Grove: true, LWS: true, + DisaggregatedSet: true, KaiScheduler: true, VolcanoScheduler: true, DRA: true, diff --git a/deploy/operator/internal/testing/operatorenv/env.go b/deploy/operator/internal/testing/operatorenv/env.go index ffa44079ab03..f36d11f41fb6 100644 --- a/deploy/operator/internal/testing/operatorenv/env.go +++ b/deploy/operator/internal/testing/operatorenv/env.go @@ -497,6 +497,7 @@ func crdDirectoryPaths(opts Options) []string { filepath.Join(root, "internal", "controller", "testing", "volcano.sh"), filepath.Join(root, "internal", "controller", "testing", "run.ai"), filepath.Join(root, "internal", "controller", "testing", "nvidia"), + filepath.Join(root, "internal", "controller", "testing", "disaggregatedset"), filepath.Join(root, "internal", "controller", "testing", "inference.networking.k8s.io"), filepath.Join(root, "internal", "controller", "testing", "grove.io"), } diff --git a/docs/fern/pages/developer-guide/knowledge-base/kubernetes/multinode/lws.md b/docs/fern/pages/developer-guide/knowledge-base/kubernetes/multinode/lws.md index 206c2120d0a1..fc2a061a992d 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/kubernetes/multinode/lws.md +++ b/docs/fern/pages/developer-guide/knowledge-base/kubernetes/multinode/lws.md @@ -2,11 +2,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 title: LWS -subtitle: LeaderWorkerSet integration for multinode Dynamo deployments +subtitle: LeaderWorkerSet and DisaggregatedSet integration for multinode Dynamo deployments --- Dynamo can use [LeaderWorkerSet (LWS)](https://lws.sigs.k8s.io/docs/) as the Kubernetes orchestration layer for multinode workloads. LWS is the lightweight path for spanning one Dynamo worker service across multiple nodes; Dynamo pairs it with [Volcano](https://volcano.sh/) for gang scheduling. +Recent LWS releases also serve the `disaggregatedset.x-k8s.io/v1` API. When that API is installed, Dynamo can place eligible multinode worker roles into a single DisaggregatedSet (DS) instead of creating one LeaderWorkerSet per component. + Use LWS when you want a simpler multinode orchestrator than Grove, or when your cluster already standardizes on LWS and Volcano. Grove remains the default when both Grove and LWS are available. ## Prerequisites @@ -16,18 +18,26 @@ Use LWS when you want a simpler multinode orchestrator than Grove, or when your - Volcano installed for gang scheduling. - Dynamo Kubernetes Platform installed. +If you want the DisaggregatedSet pathway, install LWS `v0.9.0` or newer. Dynamo builds and tests against `v0.9.0` as the current compatibility baseline for the `disaggregatedset.x-k8s.io/v1` API. DS availability is detected separately from the LWS + Volcano pathway. + The installation guide includes the exact Helm commands for [LWS and Volcano](../../../../kubernetes/installation/install-dynamo.md#lws--volcano). ## Orchestrator Selection -For multinode deployments, the Dynamo operator selects an orchestrator based on what is installed: +For multinode deployments, the operator applies this routing precedence: + +1. Grove +2. Opt-in DisaggregatedSet +3. The standard DynamoComponentDeployment (DCD) pathway + +This is an ordered routing decision, not mutual exclusion between Grove and DS. Installing the DS API does not change the pathway of existing DGDs because DS requires an explicit opt-in annotation. When Grove is available and enabled, it remains the selected pathway even if the DGD requests DS. | Cluster state | Operator behavior | | --- | --- | -| Grove and LWS installed | Uses Grove by default. | -| Grove and LWS installed, DGD has `nvidia.com/enable-grove: "false"` | Uses LWS. | -| Only LWS installed | Uses LWS. | -| Neither Grove nor LWS installed | Rejects multinode deployments. | +| Grove is available and `nvidia.com/enable-grove` is not `"false"` | Uses Grove. | +| Grove is disabled or unavailable, and the DGD sets `nvidia.com/enable-disaggregatedset: "true"` | Uses DS when the DS API and requested role configuration are supported. | +| Grove is disabled or unavailable, and DS is not requested or cannot be used | Uses the standard DCD pathway. Multinode components require LWS + Volcano. | +| No selected pathway can support the multinode components | Rejects the deployment with `no multinode orchestrator available`. | To force the LWS path when Grove is also present: @@ -42,6 +52,48 @@ spec: # ... ``` +## DisaggregatedSet Path + +Use the DisaggregatedSet path when you want one DS object to own multiple multinode worker roles. DS is an opt-in path below Grove in the routing precedence; it does not replace or disable Grove. + +To request the DS path: + +1. Install an LWS release that serves `disaggregatedset.x-k8s.io/v1`. +2. Add `nvidia.com/enable-disaggregatedset: "true"` to the DGD. +3. If Grove is available and enabled, also set `nvidia.com/enable-grove: "false"` so routing can continue to the DS decision. + +```yaml +apiVersion: nvidia.com/v1alpha1 +kind: DynamoGraphDeployment +metadata: + name: qwen3-disaggset + annotations: + nvidia.com/enable-grove: "false" + nvidia.com/enable-disaggregatedset: "true" +spec: + # ... +``` + +Using both annotations makes the DS selection explicit on clusters that run Grove. If Grove is unavailable, `nvidia.com/enable-grove: "false"` is not required. + +Dynamo falls back to the standard DCD pathway when the DS request cannot be honored. Multinode components then require the existing LWS + Volcano pathway; if that pathway is unavailable, reconciliation reports that no multinode orchestrator is available. Common fallback cases are: + +- The `disaggregatedset.x-k8s.io/v1` API is not installed. +- Fewer than two eligible multinode worker roles are selected. +- A selected component uses `scalingAdapter`. +- Selected roles mix zero replicas with positive replicas. +- More than 10 eligible multinode worker roles are selected. + +DS does not depend on Volcano for API detection. The current non-DS LWS path still requires both LWS and Volcano. + +### Compatibility and readiness + +Dynamo treats `status.roleStatuses` as authoritative when the installed DS controller publishes it. LWS `v0.9.0` defines this field, but deployments where the parent DS status remains empty are also supported: Dynamo reads the target-revision child LeaderWorkerSets by the upstream DS name, role, and revision labels. This compatibility path prevents an empty parent status from blocking readiness indefinitely. + +If neither status source matches the requested roles and revision, Dynamo keeps the DS pending and retains the previous serving workload. It does not infer readiness or perform cutover. Validate newer LWS releases against this behavior before upgrading because the DS API and controller are owned by the LWS project. + +The role-count and replica-uniformity checks are enforced by both Dynamo and the LWS `v0.9.0` or newer CRD. Dynamo checks them before creating the DS so an unsupported request can produce a clear fallback reason instead of relying on an API-server rejection. + ## Multinode Spec Set `multinode.nodeCount` on the service that should span nodes. The total GPU count is `multinode.nodeCount` multiplied by the per-node GPU limit: @@ -70,6 +122,8 @@ spec: In this example, Dynamo asks LWS to place the backend across 2 nodes with 4 GPUs per node, for 8 GPUs total. Make sure your backend's tensor parallel or distributed execution flags match that total. +For the DS path, Dynamo uses the same multinode pod-template rendering logic that the LWS path uses. The difference is the owning resource shape: one DS can hold multiple selected worker roles, while the LWS path creates one LeaderWorkerSet per selected component. + ## Backend Behavior The operator injects backend-specific multinode settings into the generated LeaderWorkerSet: diff --git a/docs/fern/pages/kubernetes/installation/install-dynamo.md b/docs/fern/pages/kubernetes/installation/install-dynamo.md index 3a197fa8a887..44983c3294f8 100644 --- a/docs/fern/pages/kubernetes/installation/install-dynamo.md +++ b/docs/fern/pages/kubernetes/installation/install-dynamo.md @@ -234,6 +234,26 @@ helm install lws oci://registry.k8s.io/lws/charts/lws \ See the [LWS docs](https://lws.sigs.k8s.io/docs/) and [Volcano docs](https://github.com/volcano-sh/volcano#quick-start-guide) for configuration options, and the [Multinode Deployment Guide](../model-deployment/multinode-deployments.md) for orchestrator selection. +#### DisaggregatedSet on top of LWS + +If you want Dynamo to place multiple multinode worker roles into a single DisaggregatedSet, install LWS `v0.9.0` or newer. Dynamo builds and tests against `v0.9.0` as the current compatibility baseline for the `disaggregatedset.x-k8s.io/v1` API. Dynamo detects that API at runtime; there is no Helm value for DS in the `dynamo-platform` chart. Validate DS behavior when upgrading beyond the tested baseline. + +To request the DS path on a `DynamoGraphDeployment`, add: + +```yaml +metadata: + annotations: + nvidia.com/enable-grove: "false" + nvidia.com/enable-disaggregatedset: "true" +``` + +The operator routes workloads in this order: Grove, opt-in DS, then the standard DCD pathway. This is routing precedence, not mutual exclusion between Grove and DS. When Grove is available and enabled, selecting DS requires both annotations shown above. If Grove is unavailable, `nvidia.com/enable-grove: "false"` is not required. + +Installing the DS API does not automatically move existing DGDs to DS. The `nvidia.com/enable-disaggregatedset: "true"` annotation is always required. + +> [!NOTE] +> The current non-DS LWS pathway requires both the LWS and Volcano APIs. The DS pathway detects `disaggregatedset.x-k8s.io/v1` separately because DS itself does not rely on Volcano. + ### Network Operator / RDMA RDMA setup is cloud-provider-specific. See the [Disaggregated Communication Guide](../../developer-guide/knowledge-base/kubernetes/kubernetes-operator/disagg-communication.md) for transport options, UCX configuration, and performance expectations, and your cloud provider guide for setup instructions: diff --git a/docs/fern/pages/kubernetes/model-deployment/multinode-deployments.md b/docs/fern/pages/kubernetes/model-deployment/multinode-deployments.md index c08cf0a19df3..d2bca29fb5c6 100644 --- a/docs/fern/pages/kubernetes/model-deployment/multinode-deployments.md +++ b/docs/fern/pages/kubernetes/model-deployment/multinode-deployments.md @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 title: Multinode Deployments -subtitle: Scales Dynamo inference across multiple GPU nodes with Grove and KAI-Scheduler for large-model tensor parallelism. +subtitle: Scales Dynamo inference across multiple GPU nodes with Grove, LWS, and DisaggregatedSet. --- This guide explains how to deploy Dynamo workloads across multiple nodes. Multinode deployments enable you to scale compute-intensive LLM workloads across multiple physical machines, maximizing GPU utilization and supporting larger models. @@ -70,19 +70,24 @@ LWS is a simple multinode deployment mechanism that allows you to deploy a workl Volcano is a Kubernetes native scheduler optimized for AI workloads at scale. It is used in conjunction with LWS to provide gang scheduling support. +#### Using DisaggregatedSet + +DisaggregatedSet (DS) is an opt-in multinode path for clusters that install the `disaggregatedset.x-k8s.io/v1` API. Dynamo uses it to group multiple eligible multinode worker roles under one DS object instead of creating one LeaderWorkerSet per selected component. + ## Core Concepts ### Orchestrator Selection Algorithm -Dynamo automatically selects the best available orchestrator for multinode deployments using the following logic: +Dynamo uses an ordered routing decision for multinode deployments: -#### When Both Grove and LWS are Available: -- **Grove is selected by default** (recommended for advanced AI workloads) -- **LWS is selected** if you explicitly set `nvidia.com/enable-grove: "false"` annotation on your DGD resource +1. **Grove:** selected when Grove is available and `nvidia.com/enable-grove` is not `"false"`. +2. **DisaggregatedSet:** considered only when Grove was not selected and the DGD sets `nvidia.com/enable-disaggregatedset: "true"`. The DS API and requested role configuration must also be supported. +3. **Standard DCD pathway:** used when neither Grove nor DS is selected. Multinode components on this pathway require LWS + Volcano. -#### When Only One Orchestrator is Available: -- The installed orchestrator (Grove or LWS) is automatically selected +Grove and DS are not mutually exclusive features. Grove has higher routing priority, while DS requires explicit opt-in. Installing the DS API alone does not move existing DGDs from Grove or DCD to DS. + +To select DS when Grove is available, set both `nvidia.com/enable-grove: "false"` and `nvidia.com/enable-disaggregatedset: "true"`. If Grove is unavailable, only the DS opt-in annotation is required. #### Scheduler Integration: - **With Grove**: Dynamo uses Grove for multinode orchestration when the Grove API is available, unless you set `nvidia.com/enable-grove: "false"` on the DGD resource. Scheduler integration is configured separately: @@ -90,6 +95,7 @@ Dynamo automatically selects the best available orchestrator for multinode deplo - **EXPERIMENTAL:** Volcano: Dynamo does not install Volcano. Set `global.volcano-scheduler.enabled=true` only when Volcano is already installed and the Volcano API is available. Select queues with `nvidia.com/volcano-queue`. - KAI-Scheduler and Volcano scheduler integration are mutually exclusive for a single Dynamo operator configuration because both set pod `schedulerName`. Helm rejects configurations that enable both integrations. - **With LWS**: Uses Volcano scheduler for gang scheduling and resource coordination +- **With DisaggregatedSet**: Reuses the same multinode pod-template rendering as the LWS path, but groups selected worker roles into one DS. DS API detection is separate from Volcano detection. > **EXPERIMENTAL:** The Dynamo/Grove Volcano scheduler integration is newly introduced and opt-in. Volcano itself is a mature CNCF scheduler, but this integration is intended for clusters where Volcano is already installed and understood by the platform operator. > @@ -153,6 +159,27 @@ spec: # ... your deployment spec ``` +**Request DisaggregatedSet:** +```yaml +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: my-ds-deployment + annotations: + nvidia.com/enable-grove: "false" + nvidia.com/enable-disaggregatedset: "true" +spec: + # ... your deployment spec +``` + +DS requests fall back to the standard DCD pathway when any of these conditions apply. Multinode components require LWS + Volcano after fallback; without that pathway, reconciliation reports that no multinode orchestrator is available. + +- The `disaggregatedset.x-k8s.io/v1` API is not available. +- Fewer than two eligible multinode worker roles are selected. +- A selected component uses `scalingAdapter`. +- Selected roles mix zero replicas with positive replicas. +- More than 10 eligible multinode worker roles are selected. + ### The `multinode` Section