diff --git a/api/v1alpha1/nodeset_types.go b/api/v1alpha1/nodeset_types.go index 26061006b..3f674fda7 100644 --- a/api/v1alpha1/nodeset_types.go +++ b/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"` } diff --git a/api/v1alpha1/validation_test.go b/api/v1alpha1/validation_test.go new file mode 100644 index 000000000..07f8a56ae --- /dev/null +++ b/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") + } + }) + } +} diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index b86ca2837..bbad75101 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -1064,11 +1064,6 @@ func (in *NodeSetSpec) DeepCopyInto(out *NodeSetSpec) { *out = new(intstr.IntOrString) **out = **in } - if in.MaxConcurrentStartup != nil { - in, out := &in.MaxConcurrentStartup, &out.MaxConcurrentStartup - *out = new(intstr.IntOrString) - **out = **in - } if in.EphemeralNodes != nil { in, out := &in.EphemeralNodes, &out.EphemeralNodes *out = new(bool) diff --git a/cmd/main.go b/cmd/main.go index 9677594f9..67a864392 100644 --- a/cmd/main.go +++ b/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( diff --git a/config/crd/bases/slurm.nebius.ai_nodesets.yaml b/config/crd/bases/slurm.nebius.ai_nodesets.yaml index eef74a5f3..152e275aa 100644 --- a/config/crd/bases/slurm.nebius.ai_nodesets.yaml +++ b/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: diff --git a/go.mod b/go.mod index f4f01613a..501e8e239 100644 --- a/go.mod +++ b/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 diff --git a/helm/nodesets/templates/nodeset.yaml b/helm/nodesets/templates/nodeset.yaml index b3681edb6..4a23181eb 100644 --- a/helm/nodesets/templates/nodeset.yaml +++ b/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 }} diff --git a/helm/nodesets/tests/custom_values_test.yaml b/helm/nodesets/tests/custom_values_test.yaml index 8ea3da15a..1cd2d1b38 100644 --- a/helm/nodesets/tests/custom_values_test.yaml +++ b/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 diff --git a/helm/nodesets/tests/node_config_test.yaml b/helm/nodesets/tests/node_config_test.yaml index f77b32cfe..5ae69c883 100644 --- a/helm/nodesets/tests/node_config_test.yaml +++ b/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: diff --git a/helm/nodesets/values.yaml b/helm/nodesets/values.yaml index f88e4ce5c..8c4954e51 100644 --- a/helm/nodesets/values.yaml +++ b/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. diff --git a/helm/soperator-crds/templates/slurmcluster-crd.yaml b/helm/soperator-crds/templates/slurmcluster-crd.yaml index 5a270691f..cc7e52440 100644 --- a/helm/soperator-crds/templates/slurmcluster-crd.yaml +++ b/helm/soperator-crds/templates/slurmcluster-crd.yaml @@ -17850,31 +17850,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. @@ -27444,11 +27432,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: @@ -27474,6 +27464,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: diff --git a/helm/soperator/crds/slurmcluster-crd.yaml b/helm/soperator/crds/slurmcluster-crd.yaml index 5a270691f..cc7e52440 100644 --- a/helm/soperator/crds/slurmcluster-crd.yaml +++ b/helm/soperator/crds/slurmcluster-crd.yaml @@ -17850,31 +17850,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. @@ -27444,11 +27432,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: @@ -27474,6 +27464,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: diff --git a/helm/soperator/templates/_helpers.tpl b/helm/soperator/templates/_helpers.tpl index 0112151df..0533f0eda 100644 --- a/helm/soperator/templates/_helpers.tpl +++ b/helm/soperator/templates/_helpers.tpl @@ -53,7 +53,7 @@ Create the name of the service account to use {{- end }} {{- define "soperator.controllersAvailable" -}} -cluster,nodeconfigurator,nodeset,topology +cluster,nodeconfigurator,nodeset,rollingupdate,topology {{- end }} {{- define "soperator.controllersSpec" -}} diff --git a/helm/soperator/values.yaml b/helm/soperator/values.yaml index 9a66dd8d8..36a7369bc 100644 --- a/helm/soperator/values.yaml +++ b/helm/soperator/values.yaml @@ -9,6 +9,7 @@ controllerManager: cluster: true nodeconfigurator: true nodeset: true + rollingupdate: true topology: true containerSecurityContext: allowPrivilegeEscalation: false diff --git a/images/worker/worker_init.py b/images/worker/worker_init.py index 4ead8a0df..cceeef276 100644 --- a/images/worker/worker_init.py +++ b/images/worker/worker_init.py @@ -675,13 +675,6 @@ def apply_node_topology( topology_plugin = topology_plugin or get_topology_plugin() if topology_plugin != TOPOLOGY_PLUGIN_BLOCK: cmd.append(f"{topology}") - cmd.extend( - [ - "state=UNDRAIN", - "reason=", - "comment=", - ] - ) logger.info("Running: %s", " ".join(cmd)) result: subprocess.CompletedProcess[str] = subprocess.run( cmd, diff --git a/images/worker/worker_init_test.py b/images/worker/worker_init_test.py index 9a87c703f..4d12834b5 100644 --- a/images/worker/worker_init_test.py +++ b/images/worker/worker_init_test.py @@ -768,9 +768,6 @@ def test_scontrol_update_uses_topology_argument_for_tree_plugin( "nodename=worker-0", "nodeaddr=worker-0.svc", "topology=default:root:leaf01", - "state=UNDRAIN", - "reason=", - "comment=", ], capture_output=True, text=True, @@ -798,9 +795,6 @@ def test_scontrol_update_omits_topology_argument_for_block_plugin( "update", "nodename=worker-0", "nodeaddr=worker-0.svc", - "state=UNDRAIN", - "reason=", - "comment=", ], capture_output=True, text=True, diff --git a/internal/consts/label.go b/internal/consts/label.go index dccd66667..4e83fa2af 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -42,4 +42,12 @@ const ( LabelJailedAggregationCommonValue = "common" AnnotationConfigHash = K8sGroupNameSoperator + "/config-hash" + + LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" + LabelSoperatorRollingUpdateValue = "true" + + LabelSoperatorWorkerOperationID = K8sGroupNameSoperator + "/worker-operation-id" + LabelSoperatorWorkerOperationPhase = K8sGroupNameSoperator + "/worker-operation-phase" + LabelSoperatorWorkerOperationPhaseStopping = "stopping" + LabelSoperatorWorkerOperationPhaseReady = "ready" ) diff --git a/internal/consts/statefulset.go b/internal/consts/statefulset.go index b7a0dac44..38b3ea385 100644 --- a/internal/consts/statefulset.go +++ b/internal/consts/statefulset.go @@ -7,6 +7,6 @@ const ( type UpdateStrategy string const ( - UpdateStrategyRollingUpdate UpdateStrategy = "rollingUpdate" - UpdateStrategyOnDelete UpdateStrategy = "onDelete" + UpdateStrategyRollingUpdate UpdateStrategy = "rollingUpdate" + UpdateStrategySlurmAwareRollingUpdate UpdateStrategy = "slurmAwareRollingUpdate" ) diff --git a/internal/controller/reconciler/k8s_statefulset_advanced.go b/internal/controller/reconciler/k8s_statefulset_advanced.go index 59cad0cec..dc6ccfa28 100644 --- a/internal/controller/reconciler/k8s_statefulset_advanced.go +++ b/internal/controller/reconciler/k8s_statefulset_advanced.go @@ -3,6 +3,7 @@ package reconciler import ( "context" "fmt" + "maps" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -48,6 +49,18 @@ func (r *AdvancedStatefulSetReconciler) patch(existing, desired client.Object) ( original := dst.DeepCopy() res := client.MergeFrom(original) + if len(src.Labels) > 0 { + if dst.Labels == nil { + dst.Labels = make(map[string]string, len(src.Labels)) + } + maps.Copy(dst.Labels, src.Labels) + } + if len(src.Annotations) > 0 { + if dst.Annotations == nil { + dst.Annotations = make(map[string]string, len(src.Annotations)) + } + maps.Copy(dst.Annotations, src.Annotations) + } dst.Spec.Template.ObjectMeta.Labels = src.Spec.Template.ObjectMeta.Labels // Copy annotations from the desired StatefulSet to the existing StatefulSet @@ -57,6 +70,7 @@ func (r *AdvancedStatefulSetReconciler) patch(existing, desired client.Object) ( } dst.Spec.Replicas = src.Spec.Replicas dst.Spec.UpdateStrategy = src.Spec.UpdateStrategy + dst.Spec.VolumeClaimUpdateStrategy = src.Spec.VolumeClaimUpdateStrategy dst.Spec.ScaleStrategy = src.Spec.ScaleStrategy dst.Spec.Template.Spec = src.Spec.Template.Spec dst.Spec.ReserveOrdinals = src.Spec.ReserveOrdinals diff --git a/internal/controller/reconciler/k8s_statefulset_test.go b/internal/controller/reconciler/k8s_statefulset_test.go index abf7b5bfd..812e32bff 100644 --- a/internal/controller/reconciler/k8s_statefulset_test.go +++ b/internal/controller/reconciler/k8s_statefulset_test.go @@ -185,6 +185,36 @@ func TestAdvancedStatefulSetPatchCopiesPVCDeletionPolicy(t *testing.T) { } } +func TestAdvancedStatefulSetPatchCopiesVolumeClaimUpdateStrategy(t *testing.T) { + existing := &kruisev1b1.StatefulSet{ + Spec: kruisev1b1.StatefulSetSpec{ + VolumeClaimUpdateStrategy: kruisev1b1.VolumeClaimUpdateStrategy{ + Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, + }, + }, + } + desired := &kruisev1b1.StatefulSet{ + Spec: kruisev1b1.StatefulSetSpec{ + VolumeClaimUpdateStrategy: kruisev1b1.VolumeClaimUpdateStrategy{ + Type: kruisev1b1.OnPVCDeleteVolumeClaimUpdateStrategyType, + }, + }, + } + + r := &AdvancedStatefulSetReconciler{} + if _, err := r.patch(existing, desired); err != nil { + t.Fatalf("patch returned error: %v", err) + } + + if existing.Spec.VolumeClaimUpdateStrategy != desired.Spec.VolumeClaimUpdateStrategy { + t.Fatalf( + "expected VolumeClaimUpdateStrategy=%+v, got %+v", + desired.Spec.VolumeClaimUpdateStrategy, + existing.Spec.VolumeClaimUpdateStrategy, + ) + } +} + func TestAdvancedStatefulSetPatchCopiesScaleStrategy(t *testing.T) { tests := []struct { name string @@ -225,6 +255,75 @@ func TestAdvancedStatefulSetPatchCopiesScaleStrategy(t *testing.T) { } } +func TestAdvancedStatefulSetPatchCopiesUpdateStrategy(t *testing.T) { + existingMaxUnavailable := intstr.FromString("10%") + desiredMaxUnavailable := intstr.FromString("40%") + existing := &kruisev1b1.StatefulSet{Spec: kruisev1b1.StatefulSetSpec{ + UpdateStrategy: kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.RollingUpdateStatefulSetStrategyType, + RollingUpdate: &kruisev1b1.RollingUpdateStatefulSetStrategy{ + MaxUnavailable: &existingMaxUnavailable, + }, + }, + }} + desired := &kruisev1b1.StatefulSet{Spec: kruisev1b1.StatefulSetSpec{ + UpdateStrategy: kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.RollingUpdateStatefulSetStrategyType, + RollingUpdate: &kruisev1b1.RollingUpdateStatefulSetStrategy{ + MaxUnavailable: &desiredMaxUnavailable, + }, + }, + }} + + r := &AdvancedStatefulSetReconciler{} + if _, err := r.patch(existing, desired); err != nil { + t.Fatalf("patch returned error: %v", err) + } + + if !equality.Semantic.DeepEqual(existing.Spec.UpdateStrategy, desired.Spec.UpdateStrategy) { + t.Fatalf("expected UpdateStrategy=%+v, got %+v", desired.Spec.UpdateStrategy, existing.Spec.UpdateStrategy) + } +} + +func TestAdvancedStatefulSetPatchUpdatesTopLevelMetadata(t *testing.T) { + existing := &kruisev1b1.StatefulSet{ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "rolling-update-enabled": "false", + "external-label": "preserved", + }, + Annotations: map[string]string{ + "managed-annotation": "old", + "versions": "preserved", + }, + }} + desired := &kruisev1b1.StatefulSet{ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "rolling-update-enabled": "true", + }, + Annotations: map[string]string{ + "managed-annotation": "new", + }, + }} + + r := &AdvancedStatefulSetReconciler{} + if _, err := r.patch(existing, desired); err != nil { + t.Fatalf("patch returned error: %v", err) + } + + if got := existing.Labels["rolling-update-enabled"]; got != "true" { + t.Fatalf("expected rolling-update-enabled=true, got %q", got) + } + if got := existing.Annotations["managed-annotation"]; got != "new" { + t.Fatalf("expected managed-annotation=new, got %q", got) + } + if got := existing.Labels["external-label"]; got != "preserved" { + t.Fatalf("expected external label to be preserved, got %q", got) + } + if got := existing.Annotations["versions"]; got != "preserved" { + t.Fatalf("expected versions annotation to be preserved, got %q", got) + } +} + func ptrIntOrString(v intstr.IntOrString) *intstr.IntOrString { return &v } func TestStatefulSetPatchCopiesPVCDeletionPolicy(t *testing.T) { diff --git a/internal/controller/suite_test.go b/internal/controller/suite_test.go index bde46c7a9..31b9cc7ee 100644 --- a/internal/controller/suite_test.go +++ b/internal/controller/suite_test.go @@ -25,6 +25,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + appsv1beta1 "k8s.io/api/apps/v1beta1" "k8s.io/client-go/kubernetes/scheme" "k8s.io/client-go/rest" "sigs.k8s.io/controller-runtime/pkg/client" @@ -57,6 +58,9 @@ var _ = BeforeSuite(func() { ctx, cancel = context.WithCancel(context.TODO()) var err error + err = appsv1beta1.AddToScheme(scheme.Scheme) + Expect(err).NotTo(HaveOccurred()) + // +kubebuilder:scaffold:scheme By("bootstrapping test environment") diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go new file mode 100644 index 000000000..f2229e114 --- /dev/null +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -0,0 +1,614 @@ +/* +Copyright 2025 Nebius B.V. + +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 updatecontroller + +import ( + "context" + "fmt" + "sort" + "strings" + "time" + + kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/client-go/tools/record" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/predicate" + + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/controller/reconciler" + "nebius.ai/slurm-operator/internal/controllerconfig" + "nebius.ai/slurm-operator/internal/slurmapi" +) + +const ( + RollingUpdateControllerName = "rollingupdate" +) + +const ( + defaultSTSReplicasCount = int32(1) + defaultRebootReason = "soperator rolling update" +) + +type workerUpdateAction int + +const ( + workerUpdateActionScheduleReboot workerUpdateAction = iota + workerUpdateActionWait + workerUpdateActionTrackInFlight + workerUpdateActionUndrain + workerUpdateActionDelete +) + +type workerUpdateDecision struct { + action workerUpdateAction + operationPhase string + slurmdCrashLooping bool + rebootHandoffInProgress bool + managedRebootInProgress bool +} + +type RollingUpdateReconciler struct { + *reconciler.Reconciler + + slurmAPIClients *slurmapi.ClientSet +} + +func NewRollingUpdateReconciler( + client client.Client, scheme *runtime.Scheme, + recorder record.EventRecorder, + slurmAPIClients *slurmapi.ClientSet, +) *RollingUpdateReconciler { + r := reconciler.NewReconciler(client, scheme, recorder) + return &RollingUpdateReconciler{ + Reconciler: r, + slurmAPIClients: slurmAPIClients, + } +} + +// +kubebuilder:rbac:groups=apps.kruise.io,resources=statefulsets,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups=apps.kruise.io,resources=statefulsets/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=apps.kruise.io,resources=statefulsets/finalizers,verbs=update +// +kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;update;patch;delete + +// Reconcile is part of the main kubernetes reconciliation loop which aims to +// move the current state of the cluster closer to the desired state. +// +// For more details, check Reconcile and its Result here: +// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.20.0/pkg/reconcile +func (r *RollingUpdateReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + logger := log.FromContext(ctx).WithName("rolling-update-reconciler") + logger.Info("reconciling statefulset", "namespace", req.Namespace, "name", req.Name) + + sts := &kruisev1b1.StatefulSet{} + err := r.Get(ctx, req.NamespacedName, sts) + if err != nil { + if client.IgnoreNotFound(err) == nil { + logger.Info("statefulset not found, might be deleted", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + return ctrl.Result{}, err + } + + if !rollingUpdateEnabled(sts) { + logger.Info("rolling update is disabled", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + + labels := sts.GetLabels() + clusterName, ok := labels[consts.LabelInstanceKey] + if !ok || clusterName == "" { + return ctrl.Result{}, fmt.Errorf("missing cluster name label %s on statefulset %s/%s", consts.LabelInstanceKey, sts.Namespace, sts.Name) + } + + replicas := defaultSTSReplicasCount + if sts.Spec.Replicas != nil { + replicas = *sts.Spec.Replicas + } + + podList, err := r.getPodList(ctx, sts) + if err != nil { + return ctrl.Result{}, err + } + + if sts.Status.UpdatedReplicas == replicas { + undrainedNodes, err := r.cleanupStaleRollingUpdateDrains(ctx, clusterName, sts, podList) + if err != nil { + return ctrl.Result{}, err + } + if undrainedNodes > 0 { + return ctrl.Result{RequeueAfter: time.Second}, nil + } + + logger.Info("statefulset is up to date", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + + outdatedPodList := filterOutdatedPods(podList, sts.Status.UpdateRevision) + + operationID := sts.Status.UpdateRevision + if operationID == "" { + return ctrl.Result{}, fmt.Errorf("missing update revision on statefulset %s/%s", sts.Namespace, sts.Name) + } + + if err := r.processRollingUpdate(ctx, clusterName, operationID, sts, outdatedPodList); err != nil { + return ctrl.Result{}, err + } + + return ctrl.Result{RequeueAfter: time.Minute}, nil +} + +func filterOutdatedPods(podList []corev1.Pod, updateRevision string) []corev1.Pod { + var res []corev1.Pod + + for _, pod := range podList { + podControllerRevisionHash := pod.Labels["controller-revision-hash"] + if podControllerRevisionHash == updateRevision { + continue + } + + res = append(res, pod) + } + + return res +} + +func indexSlurmNodesForPods(slurmNodes []slurmapi.Node, pods []corev1.Pod) map[string]slurmapi.Node { + podNames := make(map[string]struct{}, len(pods)) + for _, pod := range pods { + podNames[pod.Name] = struct{}{} + } + + nodesByName := make(map[string]slurmapi.Node, len(pods)) + for _, node := range slurmNodes { + if _, found := podNames[node.Name]; found { + nodesByName[node.Name] = node + } + } + + return nodesByName +} + +func (r *RollingUpdateReconciler) getPodList( + ctx context.Context, + sts *kruisev1b1.StatefulSet, +) ([]corev1.Pod, error) { + selector, err := metav1.LabelSelectorAsSelector(sts.Spec.Selector) + if err != nil { + return nil, fmt.Errorf("failed to convert label selector: %w", err) + } + + podList := &corev1.PodList{} + if err := r.List(ctx, podList, + client.InNamespace(sts.Namespace), + client.MatchingLabelsSelector{Selector: selector}, + ); err != nil { + return nil, fmt.Errorf("failed to list pods: %w", err) + } + + return podList.Items, nil +} + +func (r *RollingUpdateReconciler) processRollingUpdate( + ctx context.Context, + clusterName string, + operationID string, + sts *kruisev1b1.StatefulSet, + outdatedPods []corev1.Pod, +) error { + logger := log.FromContext(ctx).WithName("rolling-update-reconciler") + + if len(outdatedPods) == 0 { + logger.Info("no outdated pods found", "namespace", sts.Namespace, "name", sts.Name) + return nil + } + + sort.Slice(outdatedPods, func(i, j int) bool { + return outdatedPods[i].Name < outdatedPods[j].Name + }) + + podsToStop := make([]corev1.Pod, 0, len(outdatedPods)) + deletedPods := 0 + for _, pod := range outdatedPods { + if workerOperationPhase(&pod, operationID) != consts.LabelSoperatorWorkerOperationPhaseReady { + podsToStop = append(podsToStop, pod) + continue + } + + if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil { + return fmt.Errorf("delete pod %s/%s with completed worker handoff: %w", pod.Namespace, pod.Name, err) + } + deletedPods++ + } + if deletedPods > 0 { + logger.Info("deleted outdated pods with completed worker handoffs", "count", deletedPods, "operationID", operationID) + return nil + } + + for _, pod := range podsToStop { + if !containerCrashLoopBackOff(pod.Status.InitContainerStatuses, consts.ContainerNameWorkerInit) { + continue + } + + if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil { + return fmt.Errorf("delete outdated pod %s/%s with crash-looping worker init: %w", pod.Namespace, pod.Name, err) + } + logger.Info( + "deleted outdated pod with crash-looping worker init", + "namespace", pod.Namespace, + "pod", pod.Name, + ) + return nil + } + + slurmClient, ok := r.slurmAPIClients.GetClient(types.NamespacedName{ + Namespace: sts.Namespace, + Name: clusterName, + }) + if !ok { + logger.Info("no slurm api client", "namespace", sts.Namespace, "clusterName", clusterName) + return fmt.Errorf("no slurm api client for %s/%s", sts.Namespace, clusterName) + } + slurmNodes, err := slurmClient.ListNodes(ctx) + if err != nil { + return err + } + slurmNodesByName := indexSlurmNodesForPods(slurmNodes, podsToStop) + + type rebootCandidate struct { + pod corev1.Pod + slurmNode slurmapi.Node + } + + candidates := make([]rebootCandidate, 0, len(podsToStop)) + var undrainedNodes []string + readyPodsConsumingBudget := 0 + for _, pod := range podsToStop { + slurmNode, found := slurmNodesByName[pod.Name] + if !found { + return fmt.Errorf("slurm node %s is missing from list nodes response", pod.Name) + } + + decision := decideWorkerUpdateAction(&pod, &slurmNode, operationID) + switch decision.action { + case workerUpdateActionUndrain: + if err := slurmClient.UndrainNode(ctx, slurmNode.Name); err != nil { + return fmt.Errorf("undrain stale rolling update node %s: %w", slurmNode.Name, err) + } + undrainedNodes = append(undrainedNodes, slurmNode.Name) + if podReady(&pod) { + readyPodsConsumingBudget++ + } + case workerUpdateActionDelete: + if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil { + return fmt.Errorf("delete safely offline outdated pod %s/%s: %w", pod.Namespace, pod.Name, err) + } + logger.Info( + "deleted safely offline outdated pod with no allocations", + "namespace", pod.Namespace, + "pod", pod.Name, + "slurmNode", slurmNode.Name, + "slurmdCrashLooping", decision.slurmdCrashLooping, + "rebootHandoffInProgress", decision.rebootHandoffInProgress, + "managedRebootInProgress", decision.managedRebootInProgress, + "operationID", operationID, + "operationPhase", decision.operationPhase, + ) + if podReady(&pod) { + readyPodsConsumingBudget++ + } + case workerUpdateActionWait: + logger.Info( + "waiting to replace outdated pod with crash-looping slurmd", + "namespace", pod.Namespace, + "pod", pod.Name, + "slurmNode", slurmNode.Name, + "reason", "node is not safely offline with zero known allocations", + ) + case workerUpdateActionTrackInFlight: + if podReady(&pod) { + readyPodsConsumingBudget++ + } + case workerUpdateActionScheduleReboot: + candidates = append(candidates, rebootCandidate{pod: pod, slurmNode: slurmNode}) + } + } + if len(undrainedNodes) > 0 { + logger.Info("undrained stale rolling update nodes before reboot", "nodes", undrainedNodes) + } + + budget := rebootBudget(sts) + unavailable := unavailableReplicas(sts) + availableSlots := budget - unavailable - readyPodsConsumingBudget + if availableSlots <= 0 { + logger.Info( + "rolling update budget is exhausted", + "budget", budget, + "unavailable", unavailable, + "readyPodsConsumingBudget", readyPodsConsumingBudget, + ) + return nil + } + + slurmNodesToReboot := make([]string, 0, availableSlots) + for _, candidate := range candidates { + if len(slurmNodesToReboot) >= availableSlots { + break + } + + pod := candidate.pod + if workerOperationPhase(&pod, operationID) != consts.LabelSoperatorWorkerOperationPhaseStopping { + patchBase := pod.DeepCopy() + if pod.Labels == nil { + pod.Labels = map[string]string{} + } + pod.Labels[consts.LabelSoperatorWorkerOperationID] = operationID + pod.Labels[consts.LabelSoperatorWorkerOperationPhase] = + consts.LabelSoperatorWorkerOperationPhaseStopping + if err := r.Patch( + ctx, + &pod, + client.StrategicMergeFrom(patchBase, client.MergeFromWithOptimisticLock{}), + ); err != nil { + return fmt.Errorf("start worker operation %s on pod %s/%s: %w", operationID, pod.Namespace, pod.Name, err) + } + } + + slurmNodesToReboot = append(slurmNodesToReboot, candidate.slurmNode.Name) + } + + if len(slurmNodesToReboot) == 0 { + logger.Info("all outdated pods already have reboot requested", "namespace", sts.Namespace, "name", sts.Name) + return nil + } + + if err := slurmClient.RebootNodes(ctx, slurmapi.RebootNodesRequest{ + NodeList: strings.Join(slurmNodesToReboot, ","), + ASAP: true, + Reason: defaultRebootReason, + PowerAction: consts.SlurmPowerActionWorkerHandoff, + }); err != nil { + return fmt.Errorf("schedule slurm reboot through rest api: %w", err) + } + + logger.Info("scheduled slurm reboot through rest api", "nodes", slurmNodesToReboot) + + return nil +} + +func (r *RollingUpdateReconciler) cleanupStaleRollingUpdateDrains( + ctx context.Context, + clusterName string, + sts *kruisev1b1.StatefulSet, + pods []corev1.Pod, +) (int, error) { + logger := log.FromContext(ctx).WithName("rolling-update-reconciler") + eligibleNodeNames := make(map[string]struct{}, len(pods)) + for _, pod := range pods { + if pod.Labels["controller-revision-hash"] == sts.Status.UpdateRevision && podReady(&pod) { + eligibleNodeNames[pod.Name] = struct{}{} + } + } + if len(eligibleNodeNames) == 0 { + return 0, nil + } + + slurmClient, ok := r.slurmAPIClients.GetClient(types.NamespacedName{ + Namespace: sts.Namespace, + Name: clusterName, + }) + if !ok { + return 0, fmt.Errorf("no slurm api client for %s/%s", sts.Namespace, clusterName) + } + slurmNodes, err := slurmClient.ListNodes(ctx) + if err != nil { + return 0, err + } + + var undrainedNodes []string + for _, slurmNode := range slurmNodes { + if _, ok := eligibleNodeNames[slurmNode.Name]; !ok { + continue + } + if !staleRollingUpdateDrain(&slurmNode) { + continue + } + if err := slurmClient.UndrainNode(ctx, slurmNode.Name); err != nil { + return 0, fmt.Errorf("undrain stale rolling update node %s: %w", slurmNode.Name, err) + } + undrainedNodes = append(undrainedNodes, slurmNode.Name) + } + + if len(undrainedNodes) > 0 { + logger.Info("undrained stale rolling update nodes after update", "nodes", undrainedNodes) + } + return len(undrainedNodes), nil +} + +func rebootBudget(sts *kruisev1b1.StatefulSet) int { + replicas := defaultSTSReplicasCount + if sts.Spec.Replicas != nil { + replicas = *sts.Spec.Replicas + } + if replicas <= 0 { + return 0 + } + + maxUnavailable := intstr.FromInt32(1) + if sts.Spec.ScaleStrategy != nil && sts.Spec.ScaleStrategy.MaxUnavailable != nil { + maxUnavailable = *sts.Spec.ScaleStrategy.MaxUnavailable + } + + budget, err := intstr.GetScaledValueFromIntOrPercent(&maxUnavailable, int(replicas), false) + if err != nil || budget < 1 { + return 1 + } + if budget > int(replicas) { + return int(replicas) + } + return budget +} + +func unavailableReplicas(sts *kruisev1b1.StatefulSet) int { + replicas := defaultSTSReplicasCount + if sts.Spec.Replicas != nil { + replicas = *sts.Spec.Replicas + } + unavailable := replicas - sts.Status.ReadyReplicas + if unavailable < 0 { + return 0 + } + return int(unavailable) +} + +func podReady(pod *corev1.Pod) bool { + for _, condition := range pod.Status.Conditions { + if condition.Type == corev1.PodReady { + return condition.Status == corev1.ConditionTrue + } + } + return false +} + +func containerCrashLoopBackOff(statuses []corev1.ContainerStatus, containerName string) bool { + for _, status := range statuses { + if status.Name == containerName && + status.State.Waiting != nil && + status.State.Waiting.Reason == "CrashLoopBackOff" { + return true + } + } + return false +} + +func workerOperationPhase(pod *corev1.Pod, operationID string) string { + if pod.Labels[consts.LabelSoperatorWorkerOperationID] != operationID { + return "" + } + return pod.Labels[consts.LabelSoperatorWorkerOperationPhase] +} + +func decideWorkerUpdateAction( + pod *corev1.Pod, + node *slurmapi.Node, + operationID string, +) workerUpdateDecision { + rebootInProgress := node.IsRebootIssuedState() || node.IsRebootRequestedState() + decision := workerUpdateDecision{ + operationPhase: workerOperationPhase(pod, operationID), + slurmdCrashLooping: containerCrashLoopBackOff(pod.Status.ContainerStatuses, consts.ContainerNameSlurmd), + managedRebootInProgress: rebootInProgress && hasRollingUpdateReason(node), + } + decision.rebootHandoffInProgress = rebootInProgress && + decision.operationPhase == consts.LabelSoperatorWorkerOperationPhaseStopping + + switch { + case staleRollingUpdateDrain(node): + decision.action = workerUpdateActionUndrain + case (decision.slurmdCrashLooping || + decision.rebootHandoffInProgress || + decision.managedRebootInProgress) && safeToDeleteOfflineSlurmNode(node): + // Supervisord can keep the Pod Ready while repeatedly restarting slurmd. + // Slurm state is the source of truth for safely completing an in-flight handoff. + decision.action = workerUpdateActionDelete + case decision.slurmdCrashLooping: + decision.action = workerUpdateActionWait + case rebootInProgress: + decision.action = workerUpdateActionTrackInFlight + default: + decision.action = workerUpdateActionScheduleReboot + } + + return decision +} + +func hasRollingUpdateReason(node *slurmapi.Node) bool { + if node.Reason == nil { + return false + } + reason := node.Reason.Reason + return reason == defaultRebootReason || strings.HasPrefix(reason, defaultRebootReason+" : ") +} + +func staleRollingUpdateDrain(node *slurmapi.Node) bool { + if !node.IsDrainState() || !node.IsIdleState() || node.IsNotRespondingState() || + node.IsInvalidState() || node.IsCompletingState() { + return false + } + if node.IsRebootIssuedState() || node.IsRebootRequestedState() { + return false + } + return hasRollingUpdateReason(node) +} + +// safeToDeleteOfflineSlurmNode requires both zero known allocations and +// an offline Slurm state, so deleting the Pod cannot race with new scheduling. +func safeToDeleteOfflineSlurmNode(node *slurmapi.Node) bool { + allocatedCPUs, cpusKnown := node.CPUAllocated() + if !cpusKnown || allocatedCPUs != 0 { + return false + } + if node.AllocMemoryMB == nil || *node.AllocMemoryMB != 0 { + return false + } + if node.IsCompletingState() { + return false + } + return node.IsDownState() || (node.IsIdleState() && node.IsNotRespondingState()) +} + +// SetupWithManager sets up the controller with the Manager. +func (r *RollingUpdateReconciler) SetupWithManager( + mgr ctrl.Manager, + maxConcurrency int, + cacheSyncTimeout time.Duration, +) error { + + controllerBuilder := ctrl.NewControllerManagedBy(mgr). + For(&kruisev1b1.StatefulSet{}, builder.WithPredicates(predicate.Funcs{ + CreateFunc: func(tce event.TypedCreateEvent[client.Object]) bool { + return rollingUpdateEnabled(tce.Object) + }, + UpdateFunc: func(tue event.TypedUpdateEvent[client.Object]) bool { + return rollingUpdateEnabled(tue.ObjectNew) + }, + DeleteFunc: func(tde event.TypedDeleteEvent[client.Object]) bool { return false }, + GenericFunc: func(tge event.TypedGenericEvent[client.Object]) bool { return false }, + })). + Named(RollingUpdateControllerName). + WithOptions(controllerconfig.ControllerOptions(maxConcurrency, cacheSyncTimeout)) + + return controllerBuilder.Complete(r) +} + +func rollingUpdateEnabled(obj client.Object) bool { + sts, ok := obj.(*kruisev1b1.StatefulSet) + if !ok || sts == nil { + return false + } + return sts.Spec.UpdateStrategy.Type == appsv1.OnDeleteStatefulSetStrategyType && + sts.GetLabels()[consts.LabelSoperatorRollingUpdateEnabled] == consts.LabelSoperatorRollingUpdateValue +} diff --git a/internal/controller/updatecontroller/statefulset_controller_test.go b/internal/controller/updatecontroller/statefulset_controller_test.go new file mode 100644 index 000000000..cd77cecb5 --- /dev/null +++ b/internal/controller/updatecontroller/statefulset_controller_test.go @@ -0,0 +1,776 @@ +package updatecontroller + +import ( + "context" + "testing" + "time" + + api "github.com/SlinkyProject/slurm-client/api/v0044" + kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + 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/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/client-go/tools/record" + "k8s.io/utils/ptr" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + clientfake "sigs.k8s.io/controller-runtime/pkg/client/fake" + + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/slurmapi" + slurmapifake "nebius.ai/slurm-operator/internal/slurmapi/fake" +) + +func TestContainerCrashLoopBackOff(t *testing.T) { + statuses := []corev1.ContainerStatus{ + { + Name: "slurmd", + State: corev1.ContainerState{ + Waiting: &corev1.ContainerStateWaiting{Reason: "CrashLoopBackOff"}, + }, + }, + { + Name: "sidecar", + State: corev1.ContainerState{ + Waiting: &corev1.ContainerStateWaiting{Reason: "ImagePullBackOff"}, + }, + }, + } + + assert.True(t, containerCrashLoopBackOff(statuses, "slurmd")) + assert.False(t, containerCrashLoopBackOff(statuses, "sidecar")) + assert.False(t, containerCrashLoopBackOff(statuses, "missing")) +} + +func TestRollingUpdateEnabledRequiresOnDeleteStrategy(t *testing.T) { + sts := testStatefulSet() + sts.Labels = map[string]string{ + consts.LabelSoperatorRollingUpdateEnabled: consts.LabelSoperatorRollingUpdateValue, + } + + assert.True(t, rollingUpdateEnabled(sts)) + + sts.Spec.UpdateStrategy = kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.RollingUpdateStatefulSetStrategyType, + } + assert.False(t, rollingUpdateEnabled(sts)) + + sts.Spec.UpdateStrategy = kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.OnDeleteStatefulSetStrategyType, + } + sts.Labels[consts.LabelSoperatorRollingUpdateEnabled] = "false" + assert.False(t, rollingUpdateEnabled(sts)) +} + +func TestRebootBudgetUsesStatefulSetScaleStrategy(t *testing.T) { + tests := []struct { + name string + replicas int32 + maxUnavailable string + want int + }{ + {name: "percentage", replicas: 10, maxUnavailable: "40%", want: 4}, + {name: "absolute", replicas: 10, maxUnavailable: "3", want: 3}, + {name: "clamped to replicas", replicas: 2, maxUnavailable: "5", want: 2}, + {name: "invalid falls back to one", replicas: 10, maxUnavailable: "invalid", want: 1}, + {name: "zero replicas", replicas: 0, maxUnavailable: "40%", want: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sts := testStatefulSet() + sts.Spec.Replicas = ptr.To(tt.replicas) + setMaxUnavailable(sts, intstr.Parse(tt.maxUnavailable)) + + assert.Equal(t, tt.want, rebootBudget(sts)) + }) + } +} + +func TestSafeToDeleteOfflineSlurmNode(t *testing.T) { + tests := []struct { + name string + node slurmapi.Node + want bool + }{ + { + name: "down with no allocations", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateDOWN), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + want: true, + }, + { + name: "not responding with no allocations", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateNOTRESPONDING), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + want: true, + }, + { + name: "online idle node", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + }, + { + name: "allocated base state with stale zero allocations", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateALLOCATED, api.V0044NodeStateNOTRESPONDING), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + }, + { + name: "allocated CPUs", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateNOTRESPONDING), + AllocCPUs: ptr.To(int32(1)), + AllocMemoryMB: ptr.To(int64(0)), + }, + }, + { + name: "allocated memory", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateNOTRESPONDING), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(1)), + }, + }, + { + name: "unknown allocations", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateDOWN), + }, + }, + { + name: "completing", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateDOWN, api.V0044NodeStateCOMPLETING), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, safeToDeleteOfflineSlurmNode(&tt.node)) + }) + } +} + +func TestStaleRollingUpdateDrain(t *testing.T) { + tests := []struct { + name string + node slurmapi.Node + isStale bool + }{ + { + name: "exact rolling update reason", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason}, + }, + isStale: true, + }, + { + name: "slurm reboot suffix", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason + " : reboot issued [root@timestamp]"}, + }, + isStale: true, + }, + { + name: "manual drain", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: "hardware maintenance"}, + }, + }, + { + name: "reboot still in progress", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN, api.V0044NodeStateREBOOTISSUED), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason}, + }, + }, + { + name: "node not responding", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN, api.V0044NodeStateNOTRESPONDING), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason}, + }, + }, + { + name: "jobs still completing", + node: slurmapi.Node{ + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN, api.V0044NodeStateCOMPLETING), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.isStale, staleRollingUpdateDrain(&tt.node)) + }) + } +} + +func TestProcessRollingUpdateDeletesCrashLoopingWorkerInit(t *testing.T) { + pod := testOutdatedPod() + pod.Status.InitContainerStatuses = []corev1.ContainerStatus{ + crashLoopingContainerStatus(consts.ContainerNameWorkerInit), + } + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, nil) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &pod) +} + +func TestProcessRollingUpdateDeletesPodWithCompletedWorkerHandoff(t *testing.T) { + pod := testOutdatedPod() + pod.Labels = map[string]string{ + consts.LabelSoperatorWorkerOperationID: "new-revision", + consts.LabelSoperatorWorkerOperationPhase: consts.LabelSoperatorWorkerOperationPhaseReady, + } + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, nil) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &pod) +} + +func TestProcessRollingUpdateDeletesSafelyOfflineCrashLoopingSlurmd(t *testing.T) { + pod := testOutdatedPod() + pod.Status.ContainerStatuses = []corev1.ContainerStatus{ + crashLoopingContainerStatus(consts.ContainerNameSlurmd), + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates( + api.V0044NodeStateIDLE, + api.V0044NodeStateNOTRESPONDING, + api.V0044NodeStateREBOOTREQUESTED, + ), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }}, nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &pod) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateDeletesSafelyOfflineRebootHandoff(t *testing.T) { + pod := testOutdatedPod() + pod.Labels = map[string]string{ + consts.LabelSoperatorWorkerOperationID: "new-revision", + consts.LabelSoperatorWorkerOperationPhase: consts.LabelSoperatorWorkerOperationPhaseStopping, + } + pod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates( + api.V0044NodeStateDOWN, + api.V0044NodeStateNOTRESPONDING, + api.V0044NodeStateREBOOTISSUED, + ), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }}, nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &pod) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateContinuesWithinBudgetAfterSafeDelete(t *testing.T) { + deletingPod := testOutdatedPod() + deletingPod.Labels = map[string]string{ + consts.LabelSoperatorWorkerOperationID: "new-revision", + consts.LabelSoperatorWorkerOperationPhase: consts.LabelSoperatorWorkerOperationPhaseStopping, + } + deletingPod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + candidatePod := testOutdatedPod() + candidatePod.Name = "worker-1" + waitingPod := testOutdatedPod() + waitingPod.Name = "worker-2" + + sts := testStatefulSet() + sts.Spec.Replicas = ptr.To(int32(3)) + sts.Status.ReadyReplicas = 3 + setMaxUnavailable(sts, intstr.FromInt32(2)) + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{ + { + Name: deletingPod.Name, + States: nodeStates( + api.V0044NodeStateDOWN, + api.V0044NodeStateNOTRESPONDING, + api.V0044NodeStateREBOOTISSUED, + ), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }, + { + Name: candidatePod.Name, + States: nodeStates(api.V0044NodeStateIDLE), + }, + { + Name: waitingPod.Name, + States: nodeStates(api.V0044NodeStateIDLE), + }, + }, nil).Once() + slurmClient.On("RebootNodes", mock.Anything, slurmapi.RebootNodesRequest{ + NodeList: candidatePod.Name, + ASAP: true, + Reason: defaultRebootReason, + PowerAction: consts.SlurmPowerActionWorkerHandoff, + }).Return(nil).Once() + + reconciler, kubeClient := testRollingUpdateReconcilerWithPods( + t, + slurmClient, + &deletingPod, + &candidatePod, + &waitingPod, + ) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + sts, + []corev1.Pod{deletingPod, candidatePod, waitingPod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &deletingPod) + + gotCandidate := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&candidatePod), gotCandidate)) + assert.Equal(t, + consts.LabelSoperatorWorkerOperationPhaseStopping, + gotCandidate.Labels[consts.LabelSoperatorWorkerOperationPhase], + ) + gotWaiting := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&waitingPod), gotWaiting)) + assert.Empty(t, gotWaiting.Labels[consts.LabelSoperatorWorkerOperationPhase]) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateKeepsSafelyOfflineUnmanagedRebootWithoutHandoff(t *testing.T) { + pod := testOutdatedPod() + pod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates( + api.V0044NodeStateDOWN, + api.V0044NodeStateNOTRESPONDING, + api.V0044NodeStateREBOOTISSUED, + ), + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }}, nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + + got := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&pod), got)) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateDeletesSafelyOfflineManagedRebootWithoutHandoff(t *testing.T) { + pod := testOutdatedPod() + pod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates( + api.V0044NodeStateDOWN, + api.V0044NodeStateNOTRESPONDING, + api.V0044NodeStateREBOOTISSUED, + ), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason + " : reboot issued [root@timestamp]"}, + AllocCPUs: ptr.To(int32(0)), + AllocMemoryMB: ptr.To(int64(0)), + }}, nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + assertPodDeleted(t, kubeClient, &pod) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateKeepsCrashLoopingSlurmdWithAllocations(t *testing.T) { + pod := testOutdatedPod() + pod.Status.ContainerStatuses = []corev1.ContainerStatus{ + crashLoopingContainerStatus(consts.ContainerNameSlurmd), + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates(api.V0044NodeStateNOTRESPONDING), + AllocCPUs: ptr.To(int32(1)), + AllocMemoryMB: ptr.To(int64(1024)), + }}, nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + + got := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&pod), got)) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateStartsRevisionScopedWorkerOperation(t *testing.T) { + pod := testOutdatedPod() + pod.Labels = map[string]string{ + consts.LabelSoperatorWorkerOperationID: "old-revision", + consts.LabelSoperatorWorkerOperationPhase: consts.LabelSoperatorWorkerOperationPhaseReady, + } + sts := testStatefulSet() + sts.Status.ReadyReplicas = 1 + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates(api.V0044NodeStateIDLE), + }}, nil).Once() + slurmClient.On("RebootNodes", mock.Anything, slurmapi.RebootNodesRequest{ + NodeList: pod.Name, + ASAP: true, + Reason: defaultRebootReason, + PowerAction: consts.SlurmPowerActionWorkerHandoff, + }).Return(nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + sts, + []corev1.Pod{pod}, + ) + require.NoError(t, err) + + got := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&pod), got)) + assert.Equal(t, "new-revision", got.Labels[consts.LabelSoperatorWorkerOperationID]) + assert.Equal(t, + consts.LabelSoperatorWorkerOperationPhaseStopping, + got.Labels[consts.LabelSoperatorWorkerOperationPhase], + ) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateFailsWhenPodIsMissingFromSlurmNodeList(t *testing.T) { + pod := testOutdatedPod() + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{}, nil).Once() + + reconciler, _ := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.EqualError(t, err, "slurm node worker-0 is missing from list nodes response") + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateUndrainsStaleDrainBeforeReboot(t *testing.T) { + pod := testOutdatedPod() + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason + " : reboot issued [root@timestamp]"}, + }}, nil).Once() + slurmClient.On("UndrainNode", mock.Anything, pod.Name).Return(nil).Once() + + reconciler, kubeClient := testRollingUpdateReconciler(t, &pod, slurmClient) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + testStatefulSet(), + []corev1.Pod{pod}, + ) + require.NoError(t, err) + + got := &corev1.Pod{} + require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(&pod), got)) + slurmClient.AssertNotCalled(t, "RebootNodes", mock.Anything, mock.Anything) + slurmClient.AssertExpectations(t) +} + +func TestProcessRollingUpdateContinuesWithinBudgetAfterUndrain(t *testing.T) { + undrainedPod := testOutdatedPod() + undrainedPod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + candidatePod := testOutdatedPod() + candidatePod.Name = "worker-1" + + sts := testStatefulSet() + sts.Spec.Replicas = ptr.To(int32(2)) + sts.Status.ReadyReplicas = 2 + setMaxUnavailable(sts, intstr.FromInt32(2)) + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{ + { + Name: undrainedPod.Name, + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason + " : reboot issued [root@timestamp]"}, + }, + { + Name: candidatePod.Name, + States: nodeStates(api.V0044NodeStateIDLE), + }, + }, nil).Once() + slurmClient.On("UndrainNode", mock.Anything, undrainedPod.Name).Return(nil).Once() + slurmClient.On("RebootNodes", mock.Anything, slurmapi.RebootNodesRequest{ + NodeList: candidatePod.Name, + ASAP: true, + Reason: defaultRebootReason, + PowerAction: consts.SlurmPowerActionWorkerHandoff, + }).Return(nil).Once() + + reconciler, _ := testRollingUpdateReconcilerWithPods( + t, + slurmClient, + &undrainedPod, + &candidatePod, + ) + err := reconciler.processRollingUpdate( + context.Background(), + "cluster", + "new-revision", + sts, + []corev1.Pod{undrainedPod, candidatePod}, + ) + require.NoError(t, err) + slurmClient.AssertExpectations(t) +} + +func TestReconcileUndrainsStaleDrainAfterUpdate(t *testing.T) { + sts := testStatefulSet() + sts.Labels = map[string]string{ + consts.LabelSoperatorRollingUpdateEnabled: consts.LabelSoperatorRollingUpdateValue, + consts.LabelInstanceKey: "cluster", + } + sts.Spec.Selector = &metav1.LabelSelector{MatchLabels: map[string]string{"app": "worker"}} + sts.Status.UpdateRevision = "new-revision" + sts.Status.UpdatedReplicas = 1 + + pod := testOutdatedPod() + pod.Labels = map[string]string{ + "app": "worker", + "controller-revision-hash": sts.Status.UpdateRevision, + } + pod.Status.Conditions = []corev1.PodCondition{ + {Type: corev1.PodReady, Status: corev1.ConditionTrue}, + } + + slurmClient := &slurmapifake.MockClient{} + slurmClient.On("ListNodes", mock.Anything).Return([]slurmapi.Node{{ + Name: pod.Name, + States: nodeStates(api.V0044NodeStateIDLE, api.V0044NodeStateDRAIN), + Reason: &slurmapi.NodeReason{Reason: defaultRebootReason + " : reboot issued [root@timestamp]"}, + }}, nil).Once() + slurmClient.On("UndrainNode", mock.Anything, pod.Name).Return(nil).Once() + + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, kruisev1b1.AddToScheme(scheme)) + kubeClient := clientfake.NewClientBuilder().WithScheme(scheme).WithObjects(sts, &pod).Build() + slurmClients := slurmapi.NewClientSet(context.Background()) + slurmClients.AddClient(types.NamespacedName{Namespace: "default", Name: "cluster"}, slurmClient) + reconciler := NewRollingUpdateReconciler( + kubeClient, + scheme, + record.NewFakeRecorder(1), + slurmClients, + ) + + result, err := reconciler.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: client.ObjectKeyFromObject(sts), + }) + require.NoError(t, err) + assert.Equal(t, time.Second, result.RequeueAfter) + slurmClient.AssertExpectations(t) +} + +func nodeStates(states ...api.V0044NodeState) map[api.V0044NodeState]struct{} { + result := make(map[api.V0044NodeState]struct{}, len(states)) + for _, state := range states { + result[state] = struct{}{} + } + return result +} + +func testOutdatedPod() corev1.Pod { + return corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Name: "worker-0", + Namespace: "default", + ResourceVersion: "1", + }} +} + +func crashLoopingContainerStatus(name string) corev1.ContainerStatus { + return corev1.ContainerStatus{ + Name: name, + State: corev1.ContainerState{ + Waiting: &corev1.ContainerStateWaiting{Reason: "CrashLoopBackOff"}, + }, + } +} + +func testStatefulSet() *kruisev1b1.StatefulSet { + sts := &kruisev1b1.StatefulSet{ + ObjectMeta: metav1.ObjectMeta{Name: "workers", Namespace: "default"}, + Spec: kruisev1b1.StatefulSetSpec{ + Replicas: ptr.To(int32(1)), + UpdateStrategy: kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.OnDeleteStatefulSetStrategyType, + }, + }, + } + setMaxUnavailable(sts, intstr.FromInt32(1)) + return sts +} + +func setMaxUnavailable(sts *kruisev1b1.StatefulSet, maxUnavailable intstr.IntOrString) { + sts.Spec.ScaleStrategy = &kruisev1b1.StatefulSetScaleStrategy{ + MaxUnavailable: ptr.To(maxUnavailable), + } +} + +func testRollingUpdateReconciler( + t *testing.T, + pod *corev1.Pod, + slurmClient slurmapi.Client, +) (*RollingUpdateReconciler, client.Client) { + t.Helper() + return testRollingUpdateReconcilerWithPods(t, slurmClient, pod) +} + +func testRollingUpdateReconcilerWithPods( + t *testing.T, + slurmClient slurmapi.Client, + pods ...*corev1.Pod, +) (*RollingUpdateReconciler, client.Client) { + t.Helper() + + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + objects := make([]client.Object, 0, len(pods)) + for _, pod := range pods { + objects = append(objects, pod.DeepCopy()) + } + kubeClient := clientfake.NewClientBuilder().WithScheme(scheme).WithObjects(objects...).Build() + slurmClients := slurmapi.NewClientSet(context.Background()) + if slurmClient != nil { + slurmClients.AddClient(types.NamespacedName{Namespace: "default", Name: "cluster"}, slurmClient) + } + + return NewRollingUpdateReconciler( + kubeClient, + scheme, + record.NewFakeRecorder(1), + slurmClients, + ), kubeClient +} + +func assertPodDeleted(t *testing.T, kubeClient client.Client, pod *corev1.Pod) { + t.Helper() + err := kubeClient.Get(context.Background(), client.ObjectKeyFromObject(pod), &corev1.Pod{}) + assert.True(t, apierrors.IsNotFound(err), "expected pod to be deleted, got: %v", err) +} diff --git a/internal/render/worker/statefulset.go b/internal/render/worker/statefulset.go index 1e5fa5389..35ad34a19 100644 --- a/internal/render/worker/statefulset.go +++ b/internal/render/worker/statefulset.go @@ -5,6 +5,7 @@ import ( "maps" "slices" "sort" + "strconv" appspub "github.com/openkruise/kruise-api/apps/pub" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" @@ -36,6 +37,9 @@ func RenderNodeSetStatefulSet( labels := common.RenderLabels(consts.ComponentTypeNodeSet, nodeSet.ParentalCluster.Name) labels[consts.LabelNodeSetKey] = nodeSet.Name labels[consts.LabelWorkerKey] = consts.LabelWorkerValue + labels[consts.LabelSoperatorRollingUpdateEnabled] = strconv.FormatBool( + nodeSet.UpdateStrategy == consts.UpdateStrategySlurmAwareRollingUpdate, + ) matchLabels := common.RenderMatchLabels(consts.ComponentTypeNodeSet, nodeSet.ParentalCluster.Name) matchLabels[consts.LabelNodeSetKey] = nodeSet.Name @@ -174,7 +178,7 @@ func RenderNodeSetStatefulSet( Replicas: replicas, ReserveOrdinals: reserveOrdinals, ScaleStrategy: &kruisev1b1.StatefulSetScaleStrategy{ - MaxUnavailable: &nodeSet.StatefulSet.MaxConcurrentStartup, + MaxUnavailable: &nodeSet.StatefulSet.MaxUnavailable, }, UpdateStrategy: updateStrategy, Selector: &metav1.LabelSelector{ @@ -217,12 +221,12 @@ func renderUpdateStrategies(nodeSet *values.SlurmNodeSet) (kruisev1b1.StatefulSe Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, }, nil - case consts.UpdateStrategyOnDelete: + case consts.UpdateStrategySlurmAwareRollingUpdate: return kruisev1b1.StatefulSetUpdateStrategy{ Type: appsv1.OnDeleteStatefulSetStrategyType, }, kruisev1b1.VolumeClaimUpdateStrategy{ - Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, + Type: kruisev1b1.OnPVCDeleteVolumeClaimUpdateStrategyType, }, nil default: diff --git a/internal/render/worker/statefulset_test.go b/internal/render/worker/statefulset_test.go index 9a0a58058..6e6b2af07 100644 --- a/internal/render/worker/statefulset_test.go +++ b/internal/render/worker/statefulset_test.go @@ -1054,7 +1054,7 @@ func TestRenderNodeSetStatefulSet_PersistentVolumeClaimRetentionPolicy(t *testin } func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { - makeNodeSet := func(maxConcurrentStartup, maxUnavailable intstr.IntOrString) *values.SlurmNodeSet { + makeNodeSet := func(maxUnavailable intstr.IntOrString) *values.SlurmNodeSet { return &values.SlurmNodeSet{ Name: "test-nodeset", ParentalCluster: client.ObjectKey{ @@ -1078,9 +1078,8 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { VolumeSpool: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/spool"}}, VolumeJail: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/jail"}}, StatefulSet: values.StatefulSet{ - Replicas: 1, - MaxUnavailable: maxUnavailable, - MaxConcurrentStartup: maxConcurrentStartup, + Replicas: 1, + MaxUnavailable: maxUnavailable, }, SupervisorDConfigMapName: "supervisord-config", SSHDConfigMapName: "sshd-config", @@ -1089,25 +1088,16 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { } tests := []struct { - name string - maxConcurrentStartup intstr.IntOrString - maxUnavailable intstr.IntOrString - expectedMaxConcurrentStart intstr.IntOrString - expectedMaxUnavailable intstr.IntOrString + name string + maxUnavailable intstr.IntOrString }{ { - name: "absolute MaxConcurrentStartup is propagated to ScaleStrategy", - maxConcurrentStartup: intstr.FromInt32(500), - maxUnavailable: intstr.FromString("20%"), - expectedMaxConcurrentStart: intstr.FromInt32(500), - expectedMaxUnavailable: intstr.FromString("20%"), + name: "absolute MaxUnavailable is propagated to scale and update strategies", + maxUnavailable: intstr.FromInt32(500), }, { - name: "percentage MaxConcurrentStartup is propagated to ScaleStrategy", - maxConcurrentStartup: intstr.FromString("10%"), - maxUnavailable: intstr.FromInt32(1), - expectedMaxConcurrentStart: intstr.FromString("10%"), - expectedMaxUnavailable: intstr.FromInt32(1), + name: "percentage MaxUnavailable is propagated to scale and update strategies", + maxUnavailable: intstr.FromString("10%"), }, } @@ -1115,7 +1105,7 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { t.Run(tt.name, func(t *testing.T) { result, err := worker.RenderNodeSetStatefulSet( "test-cluster", - makeNodeSet(tt.maxConcurrentStartup, tt.maxUnavailable), + makeNodeSet(tt.maxUnavailable), &slurmv1.Secrets{}, consts.CGroupV2, true, @@ -1126,21 +1116,23 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { if assert.NotNil(t, result.Spec.ScaleStrategy) && assert.NotNil(t, result.Spec.ScaleStrategy.MaxUnavailable) { - assert.Equal(t, tt.expectedMaxConcurrentStart, *result.Spec.ScaleStrategy.MaxUnavailable, - "ScaleStrategy.MaxUnavailable should mirror StatefulSet.MaxConcurrentStartup") + assert.Equal(t, tt.maxUnavailable, *result.Spec.ScaleStrategy.MaxUnavailable, + "ScaleStrategy.MaxUnavailable should mirror StatefulSet.MaxUnavailable") } if assert.NotNil(t, result.Spec.UpdateStrategy.RollingUpdate) && assert.NotNil(t, result.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable) { - assert.Equal(t, tt.expectedMaxUnavailable, *result.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable, - "UpdateStrategy.RollingUpdate.MaxUnavailable governs the update path and must not be affected") + assert.Equal(t, tt.maxUnavailable, *result.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable, + "UpdateStrategy.RollingUpdate.MaxUnavailable should mirror StatefulSet.MaxUnavailable") } + + assert.Equal(t, "false", result.Labels[consts.LabelSoperatorRollingUpdateEnabled]) }) } - t.Run("onDelete strategy is propagated", func(t *testing.T) { - nodeSet := makeNodeSet(intstr.FromInt32(500), intstr.FromString("20%")) - nodeSet.UpdateStrategy = consts.UpdateStrategyOnDelete + t.Run("slurm-aware rolling update uses onDelete internally", func(t *testing.T) { + nodeSet := makeNodeSet(intstr.FromString("20%")) + nodeSet.UpdateStrategy = consts.UpdateStrategySlurmAwareRollingUpdate result, err := worker.RenderNodeSetStatefulSet( "test-cluster", @@ -1154,11 +1146,13 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { assert.NoError(t, err) assert.Equal(t, appsv1.OnDeleteStatefulSetStrategyType, result.Spec.UpdateStrategy.Type) assert.Nil(t, result.Spec.UpdateStrategy.RollingUpdate) - assert.Equal(t, kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, result.Spec.VolumeClaimUpdateStrategy.Type) + assert.Equal(t, kruisev1b1.OnPVCDeleteVolumeClaimUpdateStrategyType, result.Spec.VolumeClaimUpdateStrategy.Type) + assert.Equal(t, consts.LabelSoperatorRollingUpdateValue, result.Labels[consts.LabelSoperatorRollingUpdateEnabled]) + assert.Equal(t, intstr.FromString("20%"), *result.Spec.ScaleStrategy.MaxUnavailable) }) t.Run("unsupported strategy returns an error", func(t *testing.T) { - nodeSet := makeNodeSet(intstr.FromInt32(500), intstr.FromString("20%")) + nodeSet := makeNodeSet(intstr.FromString("20%")) nodeSet.UpdateStrategy = consts.UpdateStrategy("unsupported") _, err := worker.RenderNodeSetStatefulSet( diff --git a/internal/slurmapi/node.go b/internal/slurmapi/node.go index 2b54da7ad..545c7fe14 100644 --- a/internal/slurmapi/node.go +++ b/internal/slurmapi/node.go @@ -145,6 +145,11 @@ func (n *Node) IsFailState() bool { return exists } +func (n *Node) IsIdleState() bool { + _, exists := n.States[api.V0044NodeStateIDLE] + return exists +} + func (n *Node) IsPlannedState() bool { _, exists := n.States[api.V0044NodeStatePLANNED] return exists @@ -165,6 +170,16 @@ func (n *Node) IsInvalidState() bool { return exists } +func (n *Node) IsRebootIssuedState() bool { + _, exists := n.States[api.V0044NodeStateREBOOTISSUED] + return exists +} + +func (n *Node) IsRebootRequestedState() bool { + _, exists := n.States[api.V0044NodeStateREBOOTREQUESTED] + return exists +} + // IsCloudState reports whether the node is a cloud node, i.e. one that is // provisioned and powered on or off dynamically via Slurm power saving. func (n *Node) IsCloudState() bool { diff --git a/internal/values/slurm_controller.go b/internal/values/slurm_controller.go index 1566686e9..196222d8b 100644 --- a/internal/values/slurm_controller.go +++ b/internal/values/slurm_controller.go @@ -42,7 +42,6 @@ func buildSlurmControllerFrom(clusterName, namePrefix string, maintenance *const naming.BuildStatefulSetName(consts.ComponentTypeController, namePrefix), consts.SingleReplicas, nil, - nil, ) daemonSet := buildDaemonSetFrom( diff --git a/internal/values/slurm_nodeset.go b/internal/values/slurm_nodeset.go index 37c990b10..775bb32b9 100644 --- a/internal/values/slurm_nodeset.go +++ b/internal/values/slurm_nodeset.go @@ -134,7 +134,6 @@ func BuildSlurmNodeSetFrom( naming.BuildNodeSetStatefulSetName(nodeSet.Name), nsSpec.Replicas, nsSpec.MaxUnavailable, - nsSpec.MaxConcurrentStartup, ), UpdateStrategy: nsSpec.UpdateStrategy, Service: buildServiceFrom(naming.BuildNodeSetServiceName(clusterName, nodeSet.Name)), diff --git a/internal/values/types.go b/internal/values/types.go index 861a5ff0a..8eeafb044 100644 --- a/internal/values/types.go +++ b/internal/values/types.go @@ -64,10 +64,9 @@ func buildServiceFrom( // region StatefulSet type StatefulSet struct { - Name string - Replicas int32 - MaxUnavailable intstr.IntOrString - MaxConcurrentStartup intstr.IntOrString + Name string + Replicas int32 + MaxUnavailable intstr.IntOrString } func buildStatefulSetFrom( @@ -85,7 +84,6 @@ func buildStatefulSetWithMaxUnavailableFrom( name string, size int32, maxUnavailable *intstr.IntOrString, - maxConcurrentStartup *intstr.IntOrString, ) StatefulSet { result := StatefulSet{ Name: name, @@ -96,10 +94,6 @@ func buildStatefulSetWithMaxUnavailableFrom( result.MaxUnavailable = *maxUnavailable } - if maxConcurrentStartup != nil { - result.MaxConcurrentStartup = *maxConcurrentStartup - } - return result }