[disruption]: pipelined disruption budgets

pkg/controllers/controllers.go

@@ -98,7 +98,7 @@ func NewControllers(
 	deviceAllocationController := deviceallocation.NewController(kubeClient)
 	p := provisioning.NewProvisioner(kubeClient, recorder, cloudProvider, cluster, clock, deviceAllocationController)
 	evictionQueue := terminator.NewQueue(kubeClient, recorder)
-	disruptionQueue := disruption.NewQueue(kubeClient, recorder, cluster, clock, p)
+	disruptionQueue := disruption.NewQueue(kubeClient, recorder, cluster, clock, p, cloudProvider)
 	npState := nodepoolhealth.NewState()
 	clusterCost := cost.NewClusterCost(ctx, cloudProvider, kubeClient)
 	controllers := []controller.Controller{

pkg/controllers/disruption/consolidation_test.go

@@ -4210,7 +4210,7 @@ var _ = Describe("Consolidation", func() {
 				ExpectReconcileSucceeded(ctx, nodeClaimStateController, client.ObjectKeyFromObject(nc))
 			}
 			// Reset the disruption controller so consolidation methods are not cached as consolidated
-			*queue = lo.FromPtr(disruption.NewQueue(env.Client, recorder, cluster, env.Clock, prov))
+			*queue = lo.FromPtr(disruption.NewQueue(env.Client, recorder, cluster, env.Clock, prov, cloudProvider))
 			disruptionController = disruption.NewController(env.Clock, env.Client, prov, cloudProvider, recorder, cluster, queue, clusterCost,
 				disruption.WithMethods(NewMethodsWithNopValidator()...))
 			ExpectSingletonReconciled(ctx, disruptionController)

pkg/controllers/disruption/helpers.go

@@ -337,14 +337,89 @@ func BuildNodePoolMap(ctx context.Context, kubeClient client.Client, cloudProvid
 	return nodePoolMap, nodePoolToInstanceTypesMap, nil
 }
 
-// BuildDisruptionBudgets prepares our disruption budget mapping. The disruption budget maps each disruption reason to the number of allowed disruptions.
+// BuildDisruptionBudgetMapping prepares our disruption budget mapping. The disruption budget maps each disruption reason to the number of allowed disruptions.
 // We calculate allowed disruptions by taking the max disruptions allowed by disruption reason and subtracting the number of nodes that are NotReady and already being deleted by that disruption reason.
 //
+// With pipelined disruption budgets, a candidate whose command is still waiting on its replacements to initialize
+// only holds a slot in this mapping, which gates the start of new commands. The budget for actually terminating
+// nodes is checked separately by the disruption queue (see BuildTerminationBudgetMapping) right before it deletes a
+// command's candidates, so the NodePool's budget bounds how many nodes are draining at once rather than how many
+// commands are in flight.
+//
 //nolint:gocyclo
 func BuildDisruptionBudgetMapping(ctx context.Context, cluster *state.Cluster, clk clock.Clock, kubeClient client.Client, cloudProvider cloudprovider.CloudProvider, recorder events.Recorder, reason v1.DisruptionReason) (map[string]int, error) {
+	counts := countBudgetConsumers(cluster)
 	disruptionBudgetMapping := map[string]int{}
-	numNodes := map[string]int{}   // map[nodepool] -> node count in nodepool
-	disrupting := map[string]int{} // map[nodepool] -> nodes undergoing disruption
+	nodePools, err := nodepoolutils.ListManaged(ctx, kubeClient, cloudProvider)
+	if err != nil {
+		return disruptionBudgetMapping, fmt.Errorf("listing node pools, %w", err)
+	}
+	pipelined := options.FromContext(ctx).PipelinedDisruptionBudgets
+	for _, nodePool := range nodePools {
+		c := counts[nodePool.Name]
+		allowedDisruptions := nodePool.MustGetAllowedDisruptions(clk, c.total, reason)
+		// Without pipelining, every node that is marked, draining, or unhealthy consumes the single budget. With it,
+		// only the nodes still waiting on a replacement hold a slot here; draining nodes are accounted for at the
+		// termination gate instead.
+		consuming := c.pending + c.terminating
+		if pipelined {
+			consuming = c.pending
+		}
+		disruptionBudgetMapping[nodePool.Name] = lo.Max([]int{allowedDisruptions - consuming, 0})
+		NodePoolAllowedDisruptions.Set(float64(allowedDisruptions), map[string]string{
+			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
+		})
+		NodePoolNodesConsumingBudgets.Set(float64(consuming), map[string]string{
+			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
+		})
+		NodePoolNodesPendingReplacement.Set(float64(c.pending), map[string]string{
+			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
+		})
+		NodePoolNodesTerminating.Set(float64(c.terminating), map[string]string{
+			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
+		})
+		if c.total != 0 && allowedDisruptions == 0 {
+			recorder.Publish(disruptionevents.NodePoolBlockedForDisruptionReason(nodePool, reason))
+		}
+	}
+	return disruptionBudgetMapping, nil
+}
+
+// BuildTerminationBudgetMapping returns, per NodePool, how many more nodes may start draining right now for the
+// given reason: the NodePool's allowed disruptions less the nodes that are already terminating or NotReady. It is
+// the second stage of pipelined disruption budgets and is consulted by the disruption queue before it deletes a
+// command's candidates.
+func BuildTerminationBudgetMapping(ctx context.Context, cluster *state.Cluster, clk clock.Clock, kubeClient client.Client, cloudProvider cloudprovider.CloudProvider, reason v1.DisruptionReason) (map[string]int, error) {
+	counts := countBudgetConsumers(cluster)
+	terminationBudgetMapping := map[string]int{}
+	nodePools, err := nodepoolutils.ListManaged(ctx, kubeClient, cloudProvider)
+	if err != nil {
+		return terminationBudgetMapping, fmt.Errorf("listing node pools, %w", err)
+	}
+	for _, nodePool := range nodePools {
+		c := counts[nodePool.Name]
+		allowedDisruptions := nodePool.MustGetAllowedDisruptions(clk, c.total, reason)
+		terminationBudgetMapping[nodePool.Name] = lo.Max([]int{allowedDisruptions - c.terminating, 0})
+	}
+	return terminationBudgetMapping, nil
+}
+
+// budgetConsumers splits a NodePool's initialized nodes by how they relate to its disruption budget.
+type budgetConsumers struct {
+	// total is the number of initialized, non-terminated nodes the budget percentage is computed from.
+	total int
+	// pending is the number of nodes a disruption command has claimed but not yet deleted: their replacements are
+	// still booting, the node is tainted against new pods, and every pod on it is still running.
+	pending int
+	// terminating is the number of nodes that are actually being drained or are NotReady, i.e. whose pods are (or
+	// may be) unavailable.
+	terminating int
+}
+
+// countBudgetConsumers tallies budgetConsumers per NodePool name. NodePools with no initialized nodes are absent from
+// the map; the zero value is the correct count for them.
+func countBudgetConsumers(cluster *state.Cluster) map[string]budgetConsumers {
+	counts := map[string]budgetConsumers{}
 	for _, node := range cluster.DeepCopyNodes() {
 		// We only consider nodes that we own and are initialized towards the total.
 		// If a node is launched/registered, but not initialized, pods aren't scheduled
@@ -356,42 +431,26 @@ func BuildDisruptionBudgetMapping(ctx context.Context, cluster *state.Cluster, c
 		if !node.Managed() || !node.Initialized() {
 			continue
 		}
-
 		// Additionally, don't consider nodeclaims that have the terminating condition. A nodeclaim should have
 		// the Terminating condition only when the node is drained and cloudprovider.Delete() was successful
 		// on the underlying cloud provider machine.
 		if node.NodeClaim.StatusConditions().Get(v1.ConditionTypeInstanceTerminating).IsTrue() {
 			continue
 		}
-
 		nodePool := node.Labels()[v1.NodePoolLabelKey]
-		numNodes[nodePool]++
-
-		// If the node satisfies one of the following, we subtract it from the allowed disruptions.
-		// 1. Has a NotReady conditiion
-		// 2. Is marked as disrupting
-		if cond := nodeutils.GetCondition(node.Node, corev1.NodeReady); cond.Status != corev1.ConditionTrue || node.MarkedForDeletion() {
-			disrupting[nodePool]++
-		}
-	}
-	nodePools, err := nodepoolutils.ListManaged(ctx, kubeClient, cloudProvider)
-	if err != nil {
-		return disruptionBudgetMapping, fmt.Errorf("listing node pools, %w", err)
-	}
-	for _, nodePool := range nodePools {
-		allowedDisruptions := nodePool.MustGetAllowedDisruptions(clk, numNodes[nodePool.Name], reason)
-		disruptionBudgetMapping[nodePool.Name] = lo.Max([]int{allowedDisruptions - disrupting[nodePool.Name], 0})
-		NodePoolAllowedDisruptions.Set(float64(allowedDisruptions), map[string]string{
-			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
-		})
-		NodePoolNodesConsumingBudgets.Set(float64(disrupting[nodePool.Name]), map[string]string{
-			metrics.NodePoolLabel: nodePool.Name, metrics.ReasonLabel: string(reason),
-		})
-		if numNodes[nodePool.Name] != 0 && allowedDisruptions == 0 {
-			recorder.Publish(disruptionevents.NodePoolBlockedForDisruptionReason(nodePool, reason))
+		c := counts[nodePool]
+		c.total++
+		// A node is terminating when it is NotReady or its NodeClaim is being deleted (drain in progress). A node the
+		// disruption queue has only marked for deletion is pending: nothing on it has been disrupted yet.
+		switch {
+		case nodeutils.GetCondition(node.Node, corev1.NodeReady).Status != corev1.ConditionTrue || node.Deleted():
+			c.terminating++
+		case node.MarkedForDeletion():
+			c.pending++
 		}
+		counts[nodePool] = c
 	}
-	return disruptionBudgetMapping, nil
+	return counts
 }
 
 // mapCandidates maps the list of proposed candidates with the current state

pkg/controllers/disruption/metrics.go

@@ -304,6 +304,43 @@ var (
 		},
 		[]string{metrics.NodePoolLabel, metrics.ReasonLabel},
 	)
+	// NodePoolNodesPendingReplacement and NodePoolNodesTerminating split NodePoolNodesConsumingBudgets into the two
+	// stages of a replace command: candidates whose replacement is still booting (pending) and candidates that are
+	// actually draining or NotReady (terminating). With pipelined budgets only the former counts against the budget
+	// for starting commands and only the latter against the budget for terminating; without it their sum is the spend.
+	NodePoolNodesPendingReplacement = opmetrics.NewPrometheusGauge(
+		crmetrics.Registry,
+		prometheus.GaugeOpts{
+			Namespace: metrics.Namespace,
+			Subsystem: metrics.NodePoolSubsystem,
+			Name:      "nodes_pending_replacement",
+			Help:      "The number of nodes a disruption command has claimed whose replacements are still initializing. Labeled by NodePool and reason.",
+		},
+		[]string{metrics.NodePoolLabel, metrics.ReasonLabel},
+	)
+	NodePoolNodesTerminating = opmetrics.NewPrometheusGauge(
+		crmetrics.Registry,
+		prometheus.GaugeOpts{
+			Namespace: metrics.Namespace,
+			Subsystem: metrics.NodePoolSubsystem,
+			Name:      "nodes_terminating",
+			Help:      "The number of initialized nodes that are draining or NotReady. Labeled by NodePool and reason.",
+		},
+		[]string{metrics.NodePoolLabel, metrics.ReasonLabel},
+	)
+	// DisruptionQueueTerminationWaitsTotal counts the times the disruption queue had a command whose replacements were
+	// initialized but held off deleting its candidates because the NodePool's termination budget was full. A
+	// sustained rate means replacements are booting faster than drains complete and the budget is the limiter.
+	DisruptionQueueTerminationWaitsTotal = opmetrics.NewPrometheusCounter(
+		crmetrics.Registry,
+		prometheus.CounterOpts{
+			Namespace: metrics.Namespace,
+			Subsystem: voluntaryDisruptionSubsystem,
+			Name:      "queue_termination_waits_total",
+			Help:      "The number of times a disruption command with initialized replacements waited for termination budget before deleting its candidates. Labeled by NodePool and reason.",
+		},
+		[]string{metrics.NodePoolLabel, metrics.ReasonLabel},
+	)
 	DisruptionQueueFailuresTotal = opmetrics.NewPrometheusCounter(
 		crmetrics.Registry,
 		prometheus.CounterOpts{

pkg/controllers/disruption/queue.go

@@ -49,12 +49,14 @@ import (
 	operatorlogging "sigs.k8s.io/karpenter/pkg/operator/logging"
 
 	v1 "sigs.k8s.io/karpenter/pkg/apis/v1"
+	"sigs.k8s.io/karpenter/pkg/cloudprovider"
 	disruptionevents "sigs.k8s.io/karpenter/pkg/controllers/disruption/events"
 	"sigs.k8s.io/karpenter/pkg/controllers/provisioning"
 	"sigs.k8s.io/karpenter/pkg/controllers/state"
 	"sigs.k8s.io/karpenter/pkg/events"
 	"sigs.k8s.io/karpenter/pkg/metrics"
 	"sigs.k8s.io/karpenter/pkg/operator/injection"
+	"sigs.k8s.io/karpenter/pkg/operator/options"
 	utilscontroller "sigs.k8s.io/karpenter/pkg/utils/controller"
 	"sigs.k8s.io/karpenter/pkg/utils/pretty"
 )
@@ -66,6 +68,9 @@ const (
 	maxRetryDuration        = 1 * time.Hour
 	maxConcurrentReconciles = 100
 	retryDurationScale      = 80 * time.Millisecond
+	// terminationBudgetPollInterval is how often a command whose replacements are initialized re-checks the
+	// NodePool's termination budget when pipelined disruption budgets are enabled.
+	terminationBudgetPollInterval = 10 * time.Second
 )
 
 type UnrecoverableError struct {
@@ -76,6 +81,26 @@ func NewUnrecoverableError(err error) *UnrecoverableError {
 	return &UnrecoverableError{error: err}
 }
 
+// TerminationBudgetError is returned by waitOrTerminate when a command's replacements are initialized but the
+// NodePool's termination budget has no room to start draining its candidates. It is recoverable and, unlike other
+// recoverable errors, does not count against the command's retry duration: the replacement is up and nothing is
+// broken, the command is just queued behind drains that are already in progress.
+type TerminationBudgetError struct {
+	error
+}
+
+func NewTerminationBudgetError(err error) *TerminationBudgetError {
+	return &TerminationBudgetError{error: err}
+}
+
+func IsTerminationBudgetError(err error) bool {
+	if err == nil {
+		return false
+	}
+	var terminationBudgetError *TerminationBudgetError
+	return stderrors.As(err, &terminationBudgetError)
+}
+
 func IsUnrecoverableError(err error) bool {
 	if err == nil {
 		return false
@@ -101,11 +126,12 @@ type Queue struct {
 	cluster             *state.Cluster
 	clock               clock.Clock
 	provisioner         *provisioning.Provisioner
+	cloudProvider       cloudprovider.CloudProvider
 }
 
 // NewQueue creates a queue that will asynchronously orchestrate disruption commands
 func NewQueue(kubeClient client.Client, recorder events.Recorder, cluster *state.Cluster, clock clock.Clock,
-	provisioner *provisioning.Provisioner,
+	provisioner *provisioning.Provisioner, cloudProvider cloudprovider.CloudProvider,
 ) *Queue {
 	queue := &Queue{
 		// nolint:staticcheck
@@ -117,6 +143,7 @@ func NewQueue(kubeClient client.Client, recorder events.Recorder, cluster *state
 		cluster:             cluster,
 		clock:               clock,
 		provisioner:         provisioner,
+		cloudProvider:       cloudProvider,
 	}
 	return queue
 }
@@ -151,6 +178,11 @@ func (q *Queue) Reconcile(ctx context.Context, nodeClaim *v1.NodeClaim) (reconci
 	ctx = log.IntoContext(ctx, log.FromContext(ctx).WithValues(cmd.LogValues()...))
 
 	if err := q.waitOrTerminate(ctx, cmd); err != nil {
+		// Waiting on termination budget is a function of other nodes' drains finishing, which takes minutes, so poll
+		// it less aggressively than the replacement readiness wait.
+		if IsTerminationBudgetError(err) {
+			return reconcile.Result{RequeueAfter: terminationBudgetPollInterval}, nil
+		}
 		// If recoverable, re-queue and try again.
 		if !IsUnrecoverableError(err) {
 			return reconcile.Result{RequeueAfter: queueBaseDelay}, nil
@@ -195,6 +227,11 @@ func (q *Queue) waitOrTerminate(ctx context.Context, cmd *Command) (err error) {
 	retryDuration := q.GetMaxRetryDuration()
 	// Wrap an error in an unrecoverable error if it timed out
 	defer func() {
+		// A command held only by the termination budget has working replacements; failing it would orphan them and
+		// un-taint candidates that are about to drain, so the budget wait is not bounded by the retry duration.
+		if IsTerminationBudgetError(err) {
+			return
+		}
 		if q.clock.Since(cmd.CreationTimestamp) > retryDuration {
 			err = NewUnrecoverableError(serrors.Wrap(fmt.Errorf("command reached timeout, %w", err), "duration", q.clock.Since(cmd.CreationTimestamp)))
 		}
@@ -233,7 +270,13 @@ func (q *Queue) waitOrTerminate(ctx context.Context, cmd *Command) (err error) {
 		return fmt.Errorf("waiting for replacement initialization, %w", err)
 	}
 
-	// All replacements have been provisioned.
+	// All replacements have been provisioned. With pipelined budgets, deleting the candidates is what actually
+	// disrupts their pods, so this is where the NodePool's termination budget is spent.
+	if options.FromContext(ctx).PipelinedDisruptionBudgets {
+		if err := q.waitForTerminationBudget(ctx, cmd); err != nil {
+			return err
+		}
+	}
 	// All we need to do now is get a successful delete call for each node claim,
 	// then the termination controller will handle the eventual deletion of the nodes.
 	errs := make([]error, len(cmd.Candidates))
@@ -256,6 +299,36 @@ func (q *Queue) waitOrTerminate(ctx context.Context, cmd *Command) (err error) {
 	return multierr.Combine(errs...)
 }
 
+// waitForTerminationBudget returns a TerminationBudgetError when any NodePool among the command's candidates cannot
+// start draining all of the candidates it owns right now. Candidates that are already deleting (a previous attempt
+// got part way through) are not counted again; they already hold a terminating slot.
+func (q *Queue) waitForTerminationBudget(ctx context.Context, cmd *Command) error {
+	needed := map[string]int{}
+	for _, c := range cmd.Candidates {
+		if c.Deleted() {
+			continue
+		}
+		needed[c.NodePool.Name]++
+	}
+	if len(needed) == 0 {
+		return nil
+	}
+	terminationBudget, err := BuildTerminationBudgetMapping(ctx, q.cluster, q.clock, q.kubeClient, q.cloudProvider, cmd.Reason())
+	if err != nil {
+		return fmt.Errorf("building termination budgets, %w", err)
+	}
+	for nodePool, n := range needed {
+		if available := terminationBudget[nodePool]; available < n {
+			DisruptionQueueTerminationWaitsTotal.Inc(map[string]string{
+				metrics.NodePoolLabel: nodePool,
+				metrics.ReasonLabel:   pretty.ToSnakeCase(string(cmd.Reason())),
+			})
+			return NewTerminationBudgetError(serrors.Wrap(fmt.Errorf("waiting for termination budget"), "NodePool", klog.KRef("", nodePool), "needed", n, "available", available))
+		}
+	}
+	return nil
+}
+
 // markDisrupted taints the node and adds the Disrupted condition to the NodeClaim for a candidate that is about to be disrupted
 // For static NodeClaims, we mark NodeClaims as pendingdisruption in statenodepool
 func (q *Queue) markDisrupted(ctx context.Context, cmd *Command) ([]*Candidate, error) {