fix batch slowdown and persistence

Author: t0mmylam

chart/templates/skyhook-crd.yaml

@@ -518,6 +518,12 @@ spec:
                     failedInBatch:
                       description: Number of failed nodes in current batch
                       type: integer
+                    lastBatchFailed:
+                      description: Whether the last batch failed (for slowdown logic)
+                      type: boolean
+                    lastBatchSize:
+                      description: Last successful batch size (for slowdown calculations)
+                      type: integer
                     processedNodes:
                       description: Total number of nodes processed so far
                       type: integer

operator/api/v1alpha1/deployment_policy_types.go

@@ -25,6 +25,7 @@ package v1alpha1
 
 import (
 	"fmt"
+	"math"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/utils/ptr"
@@ -274,8 +275,16 @@ func (s *DeploymentStrategy) EvaluateBatchResult(state *BatchProcessingState, ba
 	state.FailedInBatch = failureCount
 	state.ProcessedNodes += batchSize
 
+	// Avoid divide by zero
+	if batchSize == 0 {
+		return
+	}
+
 	successPercentage := (successCount * 100) / batchSize
-	progressPercent := (state.ProcessedNodes * 100) / totalNodes
+	var progressPercent int
+	if totalNodes > 0 {
+		progressPercent = (state.ProcessedNodes * 100) / totalNodes
+	}
 
 	if successPercentage >= s.getBatchThreshold() {
 		state.ConsecutiveFailures = 0
@@ -347,6 +356,11 @@ func (s *FixedStrategy) CalculateBatchSize(totalNodes int, state *BatchProcessin
 func (s *LinearStrategy) CalculateBatchSize(totalNodes int, state *BatchProcessingState) int {
 	var batchSize int
 
+	// Avoid divide by zero
+	if totalNodes == 0 {
+		return 0
+	}
+
 	// Check if we should slow down due to last batch failure
 	progressPercent := (state.ProcessedNodes * 100) / totalNodes
 	if state.LastBatchFailed && progressPercent < *s.SafetyLimit && state.LastBatchSize > 0 {
@@ -367,16 +381,27 @@ func (s *LinearStrategy) CalculateBatchSize(totalNodes int, state *BatchProcessi
 func (s *ExponentialStrategy) CalculateBatchSize(totalNodes int, state *BatchProcessingState) int {
 	var batchSize int
 
+	// Avoid divide by zero
+	if totalNodes == 0 {
+		return 0
+	}
+
 	// Check if we should slow down due to last batch failure
 	progressPercent := (state.ProcessedNodes * 100) / totalNodes
-	if state.LastBatchFailed && progressPercent < *s.SafetyLimit && state.LastBatchSize > 0 {
+	if state.LastBatchFailed && progressPercent < *s.SafetyLimit && state.LastBatchSize > 0 && *s.GrowthFactor > 0 {
 		// Slow down: divide last batch size by growth factor
 		batchSize = max(1, state.LastBatchSize / *s.GrowthFactor)
 	} else {
 		// Normal growth: initialBatch * (growthFactor ^ (currentBatch - 1))
-		batchSize = *s.InitialBatch
-		for i := 1; i < state.CurrentBatch; i++ {
-			batchSize *= *s.GrowthFactor
+		// Use math.Pow for efficiency and to avoid overflow issues with large batch numbers
+		exponent := state.CurrentBatch - 1
+		growthMultiplier := math.Pow(float64(*s.GrowthFactor), float64(exponent))
+		batchSize = int(float64(*s.InitialBatch) * growthMultiplier)
+
+		// Cap at remaining nodes to prevent unreasonably large batch sizes
+		// and potential integer overflow
+		if batchSize > totalNodes {
+			batchSize = totalNodes
 		}
 	}
 

operator/api/v1alpha1/zz_generated.deepcopy.go

@@ -519,6 +519,12 @@ spec:
                     failedInBatch:
                       description: Number of failed nodes in current batch
                       type: integer
+                    lastBatchFailed:
+                      description: Whether the last batch failed (for slowdown logic)
+                      type: boolean
+                    lastBatchSize:
+                      description: Last successful batch size (for slowdown calculations)
+                      type: integer
                     processedNodes:
                       description: Total number of nodes processed so far
                       type: integer

operator/config/crd/bases/skyhook.nvidia.com_skyhooks.yaml

@@ -465,6 +465,34 @@ func (np *NodePicker) selectNodesWithCompartments(s SkyhookNodes, compartments m
 	return selectedNodes
 }
 
+// PersistCompartmentBatchStates saves the current batch state for all compartments to the Skyhook status
+func PersistCompartmentBatchStates(skyhook SkyhookNodes) bool {
+	compartments := skyhook.GetCompartments()
+	if len(compartments) == 0 {
+		return false // No compartments, nothing to persist
+	}
+
+	// Initialize the batch states map if needed
+	if skyhook.GetSkyhook().Status.CompartmentBatchStates == nil {
+		skyhook.GetSkyhook().Status.CompartmentBatchStates = make(map[string]v1alpha1.BatchProcessingState)
+	}
+
+	changed := false
+	for _, compartment := range compartments {
+		// Only persist if there are nodes in the current batch
+		if len(compartment.GetBatchState().CurrentBatchNodes) > 0 {
+			skyhook.GetSkyhook().Status.CompartmentBatchStates[compartment.GetName()] = compartment.GetBatchState()
+			changed = true
+		}
+	}
+
+	if changed {
+		skyhook.GetSkyhook().Updated = true
+	}
+
+	return changed
+}
+
 // selectNodesLegacy implements the original node selection logic for backward compatibility
 func (np *NodePicker) selectNodesLegacy(s SkyhookNodes, tolerations []corev1.Toleration) []wrapper.SkyhookNode {
 	nodes := make([]wrapper.SkyhookNode, 0)

operator/internal/controller/cluster_state_v2.go

@@ -560,6 +560,10 @@ func (r *SkyhookReconciler) RunSkyhookPackages(ctx context.Context, clusterState
 	}
 
 	selectedNode := nodePicker.SelectNodes(skyhook)
+
+	// Persist compartment batch states after node selection
+	PersistCompartmentBatchStates(skyhook)
+
 	for _, node := range selectedNode {
 
 		if node.IsComplete() && !node.Changed() {

operator/internal/controller/skyhook_controller.go

@@ -135,7 +135,7 @@ func (c *Compartment) createNewBatch() []SkyhookNode {
 	}
 
 	selectedNodes := make([]SkyhookNode, 0)
-	priority := []v1alpha1.Status{v1alpha1.StatusInProgress, v1alpha1.StatusUnknown, v1alpha1.StatusBlocked, v1alpha1.StatusErroring}
+	priority := []v1alpha1.Status{v1alpha1.StatusInProgress, v1alpha1.StatusUnknown, v1alpha1.StatusErroring}
 
 	for _, status := range priority {
 		for _, node := range c.Nodes {