SCHED-2333: add rolling update controller

api/v1alpha1/nodeset_types.go

@@ -115,6 +115,7 @@ const (
 )
 
 // NodeSetSpec defines the desired state of NodeSet
+// +kubebuilder:validation:XValidation:rule="!has(self.maxUnavailable) || (type(self.maxUnavailable) == int ? (self.maxUnavailable >= 1 && self.maxUnavailable <= 500) : (self.maxUnavailable.matches('^[1-9][0-9]*%$') && int(self.maxUnavailable.find('^[0-9]+')) * self.replicas / 100 <= 500))",message="maxUnavailable must resolve to no more than 500 workers"
 type NodeSetSpec struct {
 	// ClusterName is the name of the SlurmCluster this NodeSet belongs to.
 	// Must be in the same namespace as the NodeSet.
@@ -130,27 +131,17 @@ type NodeSetSpec struct {
 	// +kubebuilder:default=1
 	Replicas int32 `json:"replicas,omitempty"`
 
-	// MaxUnavailable represents the maximum number of worker pods that can be unavailable during the update.
+	// MaxUnavailable represents the maximum number of worker pods that can be unavailable during scaling and updates.
 	// Value can be an absolute number (ex: 5) or a percentage of desired pods (ex: 10%).
 	// Absolute number is calculated from percentage by rounding down.
-	// Also, MaxUnavailable can just be allowed to work with [k8s.io/api/apps/v1.ParallelPodManagement].
-	// Defaults to 20%.
-	//
-	// +kubebuilder:validation:Optional
-	// +kubebuilder:default="20%"
-	MaxUnavailable *intstr.IntOrString `json:"maxUnavailable,omitempty"`
-
-	// MaxConcurrentStartup caps the number of worker pods created in parallel
-	// during initial NodeSet scale-out (i.e. cluster creation or NodeSet growth).
-	// Value can be an absolute number (ex: 500) or a percentage of desired pods (ex: 10%).
-	// Maps to the underlying kruise AdvancedStatefulSet's scaleStrategy.maxUnavailable.
-	// Prevents overloading the Slurm controller with simultaneous slurmd registrations
-	// on large clusters.
+	// It limits concurrent worker startup during scale-out and concurrent worker replacement during rolling updates,
+	// preventing simultaneous slurmd registrations from overloading the Slurm controller.
+	// MaxUnavailable can only be used with [k8s.io/api/apps/v1.ParallelPodManagement].
 	// Defaults to 500.
 	//
 	// +kubebuilder:validation:Optional
 	// +kubebuilder:default=500
-	MaxConcurrentStartup *intstr.IntOrString `json:"maxConcurrentStartup,omitempty"`
+	MaxUnavailable *intstr.IntOrString `json:"maxUnavailable,omitempty"`
 
 	// EphemeralNodes enables ephemeral node behavior for this NodeSet.
 	// When true, nodes will use dynamic topology injection instead of legacy topology.conf.
@@ -313,9 +304,11 @@ type NodeSetSpec struct {
 
 	// endregion Scheduling
 
-	// UpdateStrategy controls how the advanced StatefulSet updates worker pods.
+	// UpdateStrategy controls how worker pods are updated. The rollingUpdate strategy delegates
+	// updates to the advanced StatefulSet, while slurmAwareRollingUpdate coordinates each update
+	// with Slurm before replacing the worker pod.
 	//
-	// +kubebuilder:validation:Enum=rollingUpdate;onDelete
+	// +kubebuilder:validation:Enum=rollingUpdate;slurmAwareRollingUpdate
 	// +kubebuilder:default="rollingUpdate"
 	UpdateStrategy consts.UpdateStrategy `json:"updateStrategy"`
 }

api/v1alpha1/validation_test.go

@@ -0,0 +1,55 @@
+package v1alpha1_test
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	apiextensionstest "k8s.io/apiextensions-apiserver/pkg/test"
+)
+
+func TestNodeSetMaxUnavailableCELValidation(t *testing.T) {
+	validator, err := apiextensionstest.VersionValidatorFromFile(
+		t,
+		filepath.Join("..", "..", "config", "crd", "bases", "slurm.nebius.ai_nodesets.yaml"),
+		"v1alpha1",
+	)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name           string
+		replicas       int64
+		maxUnavailable any
+		wantValid      bool
+	}{
+		{name: "absolute value at limit", replicas: 10_000, maxUnavailable: int64(500), wantValid: true},
+		{name: "absolute value above limit", replicas: 10_000, maxUnavailable: int64(501), wantValid: false},
+		{name: "percentage resolving to limit", replicas: 10_000, maxUnavailable: "5%", wantValid: true},
+		{name: "percentage resolving above limit", replicas: 10_000, maxUnavailable: "10%", wantValid: false},
+		{name: "percentage rounded down to limit", replicas: 10_019, maxUnavailable: "5%", wantValid: true},
+		{name: "percentage rounded down above limit", replicas: 10_020, maxUnavailable: "5%", wantValid: false},
+		{name: "zero value", replicas: 10_000, maxUnavailable: int64(0), wantValid: false},
+		{name: "invalid percentage", replicas: 10_000, maxUnavailable: "5.5%", wantValid: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			errs := validator(map[string]any{
+				"spec": map[string]any{
+					"replicas":       tt.replicas,
+					"maxUnavailable": tt.maxUnavailable,
+				},
+			}, nil)
+
+			if tt.wantValid {
+				assert.Empty(t, errs)
+				return
+			}
+
+			if assert.NotEmpty(t, errs) {
+				assert.Contains(t, errs.ToAggregate().Error(), "maxUnavailable must resolve to no more than 500 workers")
+			}
+		})
+	}
+}

cmd/main.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"flag"
 	"os"
@@ -52,9 +53,12 @@ import (
 	"nebius.ai/slurm-operator/internal/controller/clustercontroller"
 	"nebius.ai/slurm-operator/internal/controller/nodeconfigurator"
 	"nebius.ai/slurm-operator/internal/controller/nodesetcontroller"
+	"nebius.ai/slurm-operator/internal/controller/soperatorchecks"
 	"nebius.ai/slurm-operator/internal/controller/topologyconfcontroller"
+	"nebius.ai/slurm-operator/internal/controller/updatecontroller"
 	"nebius.ai/slurm-operator/internal/controllersenabled"
 	metricsopts "nebius.ai/slurm-operator/internal/metrics"
+	"nebius.ai/slurm-operator/internal/slurmapi"
 	webhookv1 "nebius.ai/slurm-operator/internal/webhook/v1"
 	webhookv1alpha1 "nebius.ai/slurm-operator/internal/webhook/v1alpha1"
 	//+kubebuilder:scaffold:imports
@@ -180,7 +184,7 @@ func main() {
 		controllersSpec = controllersFlag
 		controllersSource = "flag"
 	}
-	availableControllers := []string{"cluster", "nodeconfigurator", "nodeset", "topology"}
+	availableControllers := []string{"cluster", "nodeconfigurator", "nodeset", "rollingupdate", "topology"}
 	controllersSet, err := controllersenabled.New(
 		controllersSpec,
 		availableControllers,
@@ -307,6 +311,29 @@ func main() {
 	}
 	// endregion Reconciler/NodeSet
 
+	slurmAPIClients := slurmapi.NewClientSet(context.Background())
+
+	if controllersSet.Enabled("rollingupdate") {
+		if err = soperatorchecks.NewSlurmAPIClientsController(
+			mgr.GetClient(),
+			mgr.GetScheme(),
+			mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName),
+			slurmAPIClients,
+		).SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil {
+			cli.Fail(setupLog, err, "unable to create slurm api clients controller", "controller", soperatorchecks.SlurmAPIClientsControllerName)
+		}
+
+		if err = updatecontroller.NewRollingUpdateReconciler(
+			mgr.GetClient(),
+			mgr.GetScheme(),
+			mgr.GetEventRecorderFor(updatecontroller.RollingUpdateControllerName),
+			slurmAPIClients,
+		).
+			SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil {
+			cli.Fail(setupLog, err, "unable to create controller", "controller", updatecontroller.RollingUpdateControllerName)
+		}
+	}
+
 	// region Reconciler/Topology
 	if controllersSet.Enabled("topology") {
 		if err = topologyconfcontroller.NewNodeTopologyReconciler(

config/crd/bases/slurm.nebius.ai_nodesets.yaml

@@ -2600,31 +2600,19 @@ spec:
                 format: int32
                 minimum: 0
                 type: integer
-              maxConcurrentStartup:
-                anyOf:
-                - type: integer
-                - type: string
-                default: 500
-                description: |-
-                  MaxConcurrentStartup caps the number of worker pods created in parallel
-                  during initial NodeSet scale-out (i.e. cluster creation or NodeSet growth).
-                  Value can be an absolute number (ex: 500) or a percentage of desired pods (ex: 10%).
-                  Maps to the underlying kruise AdvancedStatefulSet's scaleStrategy.maxUnavailable.
-                  Prevents overloading the Slurm controller with simultaneous slurmd registrations
-                  on large clusters.
-                  Defaults to 500.
-                x-kubernetes-int-or-string: true
               maxUnavailable:
                 anyOf:
                 - type: integer
                 - type: string
-                default: 20%
+                default: 500
                 description: |-
-                  MaxUnavailable represents the maximum number of worker pods that can be unavailable during the update.
+                  MaxUnavailable represents the maximum number of worker pods that can be unavailable during scaling and updates.
                   Value can be an absolute number (ex: 5) or a percentage of desired pods (ex: 10%).
                   Absolute number is calculated from percentage by rounding down.
-                  Also, MaxUnavailable can just be allowed to work with [k8s.io/api/apps/v1.ParallelPodManagement].
-                  Defaults to 20%.
+                  It limits concurrent worker startup during scale-out and concurrent worker replacement during rolling updates,
+                  preventing simultaneous slurmd registrations from overloading the Slurm controller.
+                  MaxUnavailable can only be used with [k8s.io/api/apps/v1.ParallelPodManagement].
+                  Defaults to 500.
                 x-kubernetes-int-or-string: true
               munge:
                 description: Munge defines the Slurm munge configuration.
@@ -12194,11 +12182,13 @@ spec:
                 type: object
               updateStrategy:
                 default: rollingUpdate
-                description: UpdateStrategy controls how the advanced StatefulSet
-                  updates worker pods.
+                description: |-
+                  UpdateStrategy controls how worker pods are updated. The rollingUpdate strategy delegates
+                  updates to the advanced StatefulSet, while slurmAwareRollingUpdate coordinates each update
+                  with Slurm before replacing the worker pod.
                 enum:
                 - rollingUpdate
-                - onDelete
+                - slurmAwareRollingUpdate
                 type: string
               workerAnnotations:
                 additionalProperties:
@@ -12224,6 +12214,12 @@ spec:
             - slurmd
             - updateStrategy
             type: object
+            x-kubernetes-validations:
+            - message: maxUnavailable must resolve to no more than 500 workers
+              rule: '!has(self.maxUnavailable) || (type(self.maxUnavailable) == int
+                ? (self.maxUnavailable >= 1 && self.maxUnavailable <= 500) : (self.maxUnavailable.matches(''^[1-9][0-9]*%$'')
+                && int(self.maxUnavailable.find(''^[0-9]+'')) * self.replicas / 100
+                <= 500))'
           status:
             description: NodeSetStatus defines the observed state of SlurmCluster
             properties:

go.mod

@@ -27,6 +27,7 @@ require (
 	google.golang.org/grpc v1.83.0
 	istio.io/pkg v0.0.0-20241216214326-d70796207df3
 	k8s.io/api v0.36.3
+	k8s.io/apiextensions-apiserver v0.36.3
 	k8s.io/apimachinery v0.36.3
 	k8s.io/client-go v0.36.3
 	k8s.io/utils v0.0.0-20260707023825-cf1189d6abe3
@@ -153,7 +154,6 @@ require (
 	google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-	k8s.io/apiextensions-apiserver v0.36.3 // indirect
 	k8s.io/apiserver v0.36.3 // indirect
 	k8s.io/klog/v2 v2.140.0
 	k8s.io/kube-openapi v0.0.0-20260603220949-865597e52e25 // indirect

helm/nodesets/templates/nodeset.yaml

@@ -32,16 +32,12 @@ spec:
   replicas: {{ . }}
   {{- end }}
 
-  {{- with (.maxUnavailable | default "20%") }}
+  {{- with (.maxUnavailable | default 500) }}
   maxUnavailable: {{ . }}
   {{- end }}
 
   updateStrategy: {{ (.updateStrategy | default "rollingUpdate") | quote }}
 
-  {{- with .maxConcurrentStartup }}
-  maxConcurrentStartup: {{ . }}
-  {{- end }}
-
   {{- if .ephemeralNodes }}
   ephemeralNodes: {{ .ephemeralNodes }}
   {{- end }}

helm/nodesets/tests/custom_values_test.yaml

@@ -8,7 +8,7 @@ tests:
         - name: custom-workers
           replicas: 10
           priorityClass: "custom-priority"
-          updateStrategy: onDelete
+          updateStrategy: slurmAwareRollingUpdate
           slurmd:
             image:
               repository: "custom/slurm"
@@ -52,7 +52,7 @@ tests:
           value: "custom-priority"
       - equal:
           path: spec.updateStrategy
-          value: "onDelete"
+          value: "slurmAwareRollingUpdate"
       - equal:
           path: spec.nodeConfig.autoResume
           value: false

helm/nodesets/tests/node_config_test.yaml

@@ -249,59 +249,6 @@ tests:
           value: 2
         documentIndex: 1
 
-  - it: should configure maxConcurrentStartup when set and omit it when unset
-    set:
-      nodesets:
-        - name: gpu-workers
-          replicas: 3
-          maxConcurrentStartup: 250
-          slurmd:
-            image:
-              repository: "test/slurm"
-            resources:
-              cpu: "4"
-              memory: "8Gi"
-            volumes:
-              spool:
-                emptyDir: {}
-              jail:
-                emptyDir: {}
-              jailSubMounts: []
-          munge:
-            image:
-              repository: "test/munge"
-            resources:
-              cpu: "100m"
-              memory: "128Mi"
-        - name: cpu-workers
-          replicas: 5
-          slurmd:
-            image:
-              repository: "test/slurm"
-            resources:
-              cpu: "2"
-              memory: "4Gi"
-            volumes:
-              spool:
-                emptyDir: {}
-              jail:
-                emptyDir: {}
-              jailSubMounts: []
-          munge:
-            image:
-              repository: "test/munge"
-            resources:
-              cpu: "50m"
-              memory: "64Mi"
-    asserts:
-      - equal:
-          path: spec.maxConcurrentStartup
-          value: 250
-        documentIndex: 0
-      - notExists:
-          path: spec.maxConcurrentStartup
-        documentIndex: 1
-
   - it: should configure worker annotations correctly
     set:
       nodesets:

helm/nodesets/values.yaml

@@ -62,20 +62,15 @@ nodesets:
     # A number of workers in the NodeSet
     # Optional, defaults to 1
     replicas: 3
-    # Maximum number of unavailable replicas during updates
+    # Maximum number of unavailable replicas during scaling and updates.
+    # This limits both concurrent worker startup and rolling replacement.
     # Could be a count (number) or percent (string)
-    # Optional, defaults to 20%
-    maxUnavailable: 1
+    # Optional, defaults to 500
+    maxUnavailable: 500
     # Update strategy for the worker StatefulSet.
-    # Valid values: rollingUpdate, onDelete.
+    # Valid values: rollingUpdate, slurmAwareRollingUpdate.
     # Optional, defaults to rollingUpdate
     updateStrategy: rollingUpdate
-    # Maximum number of worker pods that can be created in parallel during
-    # initial scale-out, to avoid overloading the Slurm controller with
-    # simultaneous slurmd registrations.
-    # Could be a count (number) or percent (string).
-    # Optional, defaults to 500
-    maxConcurrentStartup: 500
     # Enable ephemeral node behavior for this NodeSet.
     # When true, nodes will use dynamic topology injection instead of legacy topology.conf.
     # Topology data is read from the topology-node-labels ConfigMap at runtime.