diff --git a/Makefile b/Makefile index e5d6f5044..bed38af0a 100644 --- a/Makefile +++ b/Makefile @@ -59,7 +59,7 @@ NFS_IMAGE_TAG = $(shell $(YQ) '.image.tag' helm/nfs-server/values.yaml) OPERATOR_IMAGE_TAG = $(VERSION) ifeq ($(shell uname), Darwin) - SED_COMMAND = sed -i '' -e + SED_COMMAND = sed -i -e else SED_COMMAND = sed -i -e endif diff --git a/api/v1alpha1/nodeset_types.go b/api/v1alpha1/nodeset_types.go index 5596b286d..dc86430ec 100644 --- a/api/v1alpha1/nodeset_types.go +++ b/api/v1alpha1/nodeset_types.go @@ -7,6 +7,8 @@ import ( "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/util/intstr" + + "nebius.ai/slurm-operator/internal/consts" ) // +kubebuilder:object:root=true @@ -300,6 +302,12 @@ type NodeSetSpec struct { CustomInitContainers []corev1.Container `json:"customInitContainers,omitempty"` // endregion Scheduling + + // UpdateStrategy for the advanced StatefulSet + // + // +kubebuilder:validation:Enum=rollingUpdate;onDelete + // +kubebuilder:default="rollingUpdate" + UpdateStrategy consts.UpdateStrategy `json:"updateStrategy"` } // ContainerSlurmdSpec defines the Slurm worker daemon configuration 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/cmd/soperatorchecks/main.go b/cmd/soperatorchecks/main.go index bb524cf85..4332d10b7 100644 --- a/cmd/soperatorchecks/main.go +++ b/cmd/soperatorchecks/main.go @@ -282,7 +282,6 @@ func main() { mgr.GetScheme(), mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName), slurmAPIClients, - corev1.NodeConditionType(maintenanceConditionType), ).SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { cli.Fail(setupLog, err, "unable to create slurm api clients controller", "controller", soperatorchecks.SlurmAPIClientsControllerName) } diff --git a/config/crd/bases/slurm.nebius.ai_nodesets.yaml b/config/crd/bases/slurm.nebius.ai_nodesets.yaml index 8480b7a13..345fb06c4 100644 --- a/config/crd/bases/slurm.nebius.ai_nodesets.yaml +++ b/config/crd/bases/slurm.nebius.ai_nodesets.yaml @@ -12083,6 +12083,13 @@ spec: Defaults to "root", which preserves the single-fabric behavior. type: string type: object + updateStrategy: + default: rollingUpdate + description: UpdateStrategy for the advanced StatefulSet + enum: + - rollingUpdate + - onDelete + type: string workerAnnotations: additionalProperties: type: string @@ -12105,6 +12112,7 @@ spec: required: - munge - slurmd + - updateStrategy type: object status: description: NodeSetStatus defines the observed state of SlurmCluster diff --git a/helm/nodesets/templates/nodeset.yaml b/helm/nodesets/templates/nodeset.yaml index 0689ebfb3..70deea760 100644 --- a/helm/nodesets/templates/nodeset.yaml +++ b/helm/nodesets/templates/nodeset.yaml @@ -31,6 +31,8 @@ spec: maxUnavailable: {{ . }} {{- end }} + updateStrategy: {{ (.updateStrategy | default "rollingUpdate") | quote }} + {{- with .maxConcurrentStartup }} maxConcurrentStartup: {{ . }} {{- end }} @@ -263,7 +265,7 @@ spec: sssdLdapCAConfigMapRefName: {{ . | quote }} {{- end }} - {{- with (.configMapRefSupervisord | default "") }} + {{- with (.configMapRefSupervisord | default $.Values.configMapRefSupervisord | default "") }} configMapRefSupervisord: {{ . | quote }} {{- end }} diff --git a/helm/nodesets/tests/basic_test.yaml b/helm/nodesets/tests/basic_test.yaml index 7d6015ca0..226c7323b 100644 --- a/helm/nodesets/tests/basic_test.yaml +++ b/helm/nodesets/tests/basic_test.yaml @@ -97,6 +97,63 @@ tests: path: kind value: NodeSet + - it: should use chart-level supervisord ConfigMap ref when NodeSet does not override it + set: + configMapRefSupervisord: custom-supervisord-config-4.0.2 + nodesets: + - name: test-workers + slurmd: + image: + repository: "test/slurm" + resources: + cpu: "1" + memory: "1Gi" + volumes: + spool: + emptyDir: {} + jail: + emptyDir: {} + jailSubMounts: [] + munge: + image: + repository: "test/munge" + resources: + cpu: "100m" + memory: "128Mi" + asserts: + - equal: + path: spec.configMapRefSupervisord + value: custom-supervisord-config-4.0.2 + + - it: should prefer NodeSet supervisord ConfigMap ref over chart-level default + set: + configMapRefSupervisord: custom-supervisord-config-4.0.2 + nodesets: + - name: test-workers + configMapRefSupervisord: per-nodeset-supervisord-config + slurmd: + image: + repository: "test/slurm" + resources: + cpu: "1" + memory: "1Gi" + volumes: + spool: + emptyDir: {} + jail: + emptyDir: {} + jailSubMounts: [] + munge: + image: + repository: "test/munge" + resources: + cpu: "100m" + memory: "128Mi" + asserts: + - equal: + path: spec.configMapRefSupervisord + value: per-nodeset-supervisord-config + - it: should include common labels set: nodesets: diff --git a/helm/nodesets/tests/custom_values_test.yaml b/helm/nodesets/tests/custom_values_test.yaml index 62a71d3ea..84b4f18ee 100644 --- a/helm/nodesets/tests/custom_values_test.yaml +++ b/helm/nodesets/tests/custom_values_test.yaml @@ -8,6 +8,7 @@ tests: - name: custom-workers replicas: 10 priorityClass: "custom-priority" + updateStrategy: onDelete slurmd: image: repository: "custom/slurm" @@ -48,6 +49,9 @@ tests: - equal: path: spec.priorityClass value: "custom-priority" + - equal: + path: spec.updateStrategy + value: "onDelete" - equal: path: spec.slurmd.image.repository value: "custom/slurm" @@ -99,6 +103,9 @@ tests: - equal: path: spec.replicas value: 1 # default value + - equal: + path: spec.updateStrategy + value: "rollingUpdate" - notExists: path: spec.priorityClass - notExists: diff --git a/helm/nodesets/values.yaml b/helm/nodesets/values.yaml index e11b8ebee..1feef95c9 100644 --- a/helm/nodesets/values.yaml +++ b/helm/nodesets/values.yaml @@ -6,6 +6,9 @@ nameOverride: "" fullnameOverride: "" clusterName: "" +# Default custom SupervisorD config ConfigMap for NodeSets. +# Individual NodeSets can override this with nodesets[*].configMapRefSupervisord. +configMapRefSupervisord: "" # Priority Classes configuration # Define priority classes that can be used by NodeSets priorityClasses: @@ -50,6 +53,10 @@ nodesets: # Could be a count (number) or percent (string) # Optional, defaults to 20% maxUnavailable: 1 + # Update strategy for the worker StatefulSet. + # Valid values: rollingUpdate, onDelete. + # 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. diff --git a/helm/soperator-crds/templates/slurmcluster-crd.yaml b/helm/soperator-crds/templates/slurmcluster-crd.yaml index 39445b875..3ea7e8b9f 100644 --- a/helm/soperator-crds/templates/slurmcluster-crd.yaml +++ b/helm/soperator-crds/templates/slurmcluster-crd.yaml @@ -27177,6 +27177,13 @@ spec: Defaults to "root", which preserves the single-fabric behavior. type: string type: object + updateStrategy: + default: rollingUpdate + description: UpdateStrategy for the advanced StatefulSet + enum: + - rollingUpdate + - onDelete + type: string workerAnnotations: additionalProperties: type: string @@ -27199,6 +27206,7 @@ spec: required: - munge - slurmd + - updateStrategy type: object status: description: NodeSetStatus defines the observed state of SlurmCluster diff --git a/helm/soperator-custom-configmaps/templates/_helpers.tpl b/helm/soperator-custom-configmaps/templates/_helpers.tpl index 3f84dbe45..14f57422d 100644 --- a/helm/soperator-custom-configmaps/templates/_helpers.tpl +++ b/helm/soperator-custom-configmaps/templates/_helpers.tpl @@ -12,6 +12,15 @@ Create chart name and version as used by the chart label. {{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }} {{- end }} +{{/* +Create a ConfigMap name with a chart-version suffix. +*/}} +{{- define "soperator-custom-configmaps.versionedName" -}} +{{- $name := index . 0 -}} +{{- $version := index . 1 -}} +{{- printf "%s-%s" $name $version | lower | replace "+" "-" | replace "_" "-" | trunc 253 | trimSuffix "-" }} +{{- end }} + {{/* Common labels */}} @@ -24,6 +33,15 @@ app.kubernetes.io/version: {{ .Chart.AppVersion | quote }} app.kubernetes.io/managed-by: {{ .Release.Service }} {{- end }} +{{/* +Stable labels for legacy kept resources. These intentionally exclude chart/app +version labels so chart upgrades do not keep patching the old ConfigMap. +*/}} +{{- define "soperator-custom-configmaps.stableLabels" -}} +{{ include "soperator-custom-configmaps.selectorLabels" . }} +app.kubernetes.io/managed-by: {{ .Release.Service }} +{{- end }} + {{/* Selector labels */}} diff --git a/helm/soperator-custom-configmaps/templates/configmap-supervisord.yaml b/helm/soperator-custom-configmaps/templates/configmap-supervisord.yaml index 72f1a2b84..f732cd166 100644 --- a/helm/soperator-custom-configmaps/templates/configmap-supervisord.yaml +++ b/helm/soperator-custom-configmaps/templates/configmap-supervisord.yaml @@ -1,14 +1,60 @@ {{- if .Values.configMaps.supervisord.enabled }} +{{- $baseName := .Values.configMaps.supervisord.name }} +{{- $versionedName := include "soperator-custom-configmaps.versionedName" (list $baseName .Chart.Version) }} +{{- $keepLegacy := true }} +{{- if hasKey .Values.configMaps.supervisord "keepLegacy" }} +{{- $keepLegacy = .Values.configMaps.supervisord.keepLegacy }} +{{- end }} apiVersion: v1 kind: ConfigMap metadata: - name: {{ .Values.configMaps.supervisord.name }} + name: {{ $versionedName }} namespace: {{ .Values.namespace }} + annotations: + helm.sh/resource-policy: keep + slurm.nebius.ai/dependency-version: name labels: {{- include "soperator-custom-configmaps.labels" . | nindent 4 }} +data: + {{- range .Values.configMaps.supervisord.files }} + {{ . }}: | +{{ tpl ($.Files.Get (printf "config-files/%s" .)) $ | indent 4 }} +{{- end }} +{{- if and $keepLegacy (ne $baseName $versionedName) }} +{{- $legacyConfigMap := lookup "v1" "ConfigMap" .Values.namespace $baseName }} +{{- $legacyLabels := include "soperator-custom-configmaps.stableLabels" . | fromYaml }} +{{- $legacyAnnotations := dict "helm.sh/resource-policy" "keep" }} +{{- if $legacyConfigMap }} +{{- with $legacyConfigMap.metadata.labels }} +{{- $legacyLabels = deepCopy . }} +{{- end }} +{{- with $legacyConfigMap.metadata.annotations }} +{{- $legacyAnnotations = mergeOverwrite (deepCopy .) $legacyAnnotations }} +{{- end }} +{{- end }} +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ $baseName }} + namespace: {{ .Values.namespace }} + annotations: + {{- toYaml $legacyAnnotations | nindent 4 }} + labels: + {{- toYaml $legacyLabels | nindent 4 }} +{{- if and $legacyConfigMap $legacyConfigMap.data }} +data: + {{- toYaml $legacyConfigMap.data | nindent 2 }} +{{- else }} data: {{- range .Values.configMaps.supervisord.files }} {{ . }}: | {{ tpl ($.Files.Get (printf "config-files/%s" .)) $ | indent 4 }} {{- end }} {{- end }} +{{- if and $legacyConfigMap $legacyConfigMap.binaryData }} +binaryData: + {{- toYaml $legacyConfigMap.binaryData | nindent 2 }} +{{- end }} +{{- end }} +{{- end }} diff --git a/helm/soperator-custom-configmaps/tests/supervisord_configmap_test.yaml b/helm/soperator-custom-configmaps/tests/supervisord_configmap_test.yaml new file mode 100644 index 000000000..693052d9b --- /dev/null +++ b/helm/soperator-custom-configmaps/tests/supervisord_configmap_test.yaml @@ -0,0 +1,38 @@ +suite: test supervisord configmap +templates: + - templates/configmap-supervisord.yaml +tests: + - it: should render versioned and legacy supervisord ConfigMaps + asserts: + - hasDocuments: + count: 2 + - equal: + path: metadata.name + value: custom-supervisord-config-4.0.2 + documentIndex: 0 + - equal: + path: metadata.annotations["helm.sh/resource-policy"] + value: keep + documentIndex: 0 + - equal: + path: metadata.annotations["slurm.nebius.ai/dependency-version"] + value: name + documentIndex: 0 + - equal: + path: metadata.name + value: custom-supervisord-config + documentIndex: 1 + - equal: + path: metadata.annotations["helm.sh/resource-policy"] + value: keep + documentIndex: 1 + + - it: should skip legacy supervisord ConfigMap when disabled + set: + configMaps.supervisord.keepLegacy: false + asserts: + - hasDocuments: + count: 1 + - equal: + path: metadata.name + value: custom-supervisord-config-4.0.2 diff --git a/helm/soperator-custom-configmaps/values.yaml b/helm/soperator-custom-configmaps/values.yaml index af0bace40..30707f9cd 100644 --- a/helm/soperator-custom-configmaps/values.yaml +++ b/helm/soperator-custom-configmaps/values.yaml @@ -10,6 +10,7 @@ configMaps: supervisord: enabled: true name: custom-supervisord-config + keepLegacy: true files: - supervisord.conf motd: diff --git a/helm/soperator-fluxcd/templates/_helpers.tpl b/helm/soperator-fluxcd/templates/_helpers.tpl index a7b93d2cf..921113732 100644 --- a/helm/soperator-fluxcd/templates/_helpers.tpl +++ b/helm/soperator-fluxcd/templates/_helpers.tpl @@ -21,6 +21,48 @@ Create chart name and version as used by the chart label. {{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }} {{- end }} +{{/* +Create a Kubernetes resource name with a version suffix. +*/}} +{{- define "soperator-fluxcd.versionedName" -}} +{{- $name := index . 0 -}} +{{- $version := index . 1 -}} +{{- printf "%s-%s" $name $version | lower | replace "+" "-" | replace "_" "-" | trunc 253 | trimSuffix "-" }} +{{- end }} + +{{/* +Base name for the custom supervisord ConfigMap created by the custom-configmaps chart. +*/}} +{{- define "soperator-fluxcd.customSupervisordConfigMapBaseName" -}} +{{- $name := .Values.customConfigmaps.supervisordConfigMapName | default "custom-supervisord-config" -}} +{{- with .Values.customConfigmaps.values -}} +{{- with .configMaps -}} +{{- with .supervisord -}} +{{- with .name -}} +{{- $name = . -}} +{{- end -}} +{{- end -}} +{{- end -}} +{{- end -}} +{{- with .Values.customConfigmaps.overrideValues -}} +{{- with .configMaps -}} +{{- with .supervisord -}} +{{- with .name -}} +{{- $name = . -}} +{{- end -}} +{{- end -}} +{{- end -}} +{{- end -}} +{{- $name -}} +{{- end }} + +{{/* +Versioned name for the custom supervisord ConfigMap. +*/}} +{{- define "soperator-fluxcd.customSupervisordConfigMapName" -}} +{{- include "soperator-fluxcd.versionedName" (list (include "soperator-fluxcd.customSupervisordConfigMapBaseName" .) .Values.customConfigmaps.version) -}} +{{- end }} + {{/* Common labels */}} diff --git a/helm/soperator-fluxcd/templates/nodesets.yaml b/helm/soperator-fluxcd/templates/nodesets.yaml index f2592a0f4..914d61fc7 100644 --- a/helm/soperator-fluxcd/templates/nodesets.yaml +++ b/helm/soperator-fluxcd/templates/nodesets.yaml @@ -24,6 +24,9 @@ spec: {{- if .Values.slurmCluster.enabled }} - name: {{ include "soperator-fluxcd.fullname" . }}-slurm-cluster {{- end }} +{{- if .Values.customConfigmaps.enabled }} + - name: {{ include "soperator-fluxcd.fullname" . }}-custom-configmaps +{{- end }} install: {{- toYaml .Values.nodesets.install | nindent 4 }} @@ -33,15 +36,21 @@ spec: targetNamespace: {{ .Values.nodesets.namespace }} upgrade: {{- toYaml .Values.nodesets.upgrade | nindent 4 }} - values: - clusterName: {{ .Values.nodesets.clusterName | quote }} + {{- $nodesetsValues := dict "clusterName" .Values.nodesets.clusterName }} + {{- if .Values.customConfigmaps.enabled }} + {{- $_ := set $nodesetsValues "configMapRefSupervisord" (include "soperator-fluxcd.customSupervisordConfigMapName" .) }} + {{- end }} {{- if .Values.nodesets.overrideValues }} - {{- toYaml .Values.nodesets.overrideValues | nindent 4 }} + {{- with .Values.nodesets.overrideValues }} + {{- $nodesetsValues = mergeOverwrite $nodesetsValues (deepCopy .) }} + {{- end }} {{- else }} - {{- if .Values.nodesets.values }} - {{- toYaml .Values.nodesets.values | nindent 4 }} - {{- end }} + {{- with .Values.nodesets.values }} + {{- $nodesetsValues = mergeOverwrite $nodesetsValues (deepCopy .) }} + {{- end }} {{- end }} + values: + {{- toYaml $nodesetsValues | nindent 4 }} valuesFrom: - kind: ConfigMap name: terraform-nodesets diff --git a/helm/soperator-fluxcd/templates/soperator-custom-configmaps.yaml b/helm/soperator-fluxcd/templates/soperator-custom-configmaps.yaml index 57d91d690..8307596cc 100644 --- a/helm/soperator-fluxcd/templates/soperator-custom-configmaps.yaml +++ b/helm/soperator-fluxcd/templates/soperator-custom-configmaps.yaml @@ -20,6 +20,9 @@ spec: version: {{ .Values.customConfigmaps.version }} dependsOn: - name: {{ include "soperator-fluxcd.fullname" . }}-ns +{{- if .Values.soperator.enabled }} + - name: {{ include "soperator-fluxcd.fullname" . }}-soperator +{{- end }} install: {{- toYaml .Values.customConfigmaps.install | nindent 4 }} upgrade: @@ -28,10 +31,18 @@ spec: timeout: {{ .Values.customConfigmaps.timeout }} releaseName: {{ .Values.customConfigmaps.releaseName }} targetNamespace: {{ .Values.customConfigmaps.namespace }} - values: + {{- $customConfigmapsValues := dict "configMaps" (dict "supervisord" (dict "name" (include "soperator-fluxcd.customSupervisordConfigMapBaseName" .))) }} {{- if .Values.customConfigmaps.overrideValues }} - {{- toYaml .Values.customConfigmaps.overrideValues | nindent 4 }} + {{- with .Values.customConfigmaps.overrideValues }} + {{- $customConfigmapsValues = mergeOverwrite $customConfigmapsValues (deepCopy .) }} + {{- end }} + {{- else }} + {{- with .Values.customConfigmaps.values }} + {{- $customConfigmapsValues = mergeOverwrite $customConfigmapsValues (deepCopy .) }} {{- end }} + {{- end }} + values: + {{- toYaml $customConfigmapsValues | nindent 4 }} valuesFrom: - kind: ConfigMap name: custom-configmaps diff --git a/helm/soperator-fluxcd/tests/custom_configmaps_test.yaml b/helm/soperator-fluxcd/tests/custom_configmaps_test.yaml index 5829fa3aa..182f33d80 100644 --- a/helm/soperator-fluxcd/tests/custom_configmaps_test.yaml +++ b/helm/soperator-fluxcd/tests/custom_configmaps_test.yaml @@ -89,6 +89,20 @@ tests: path: spec.dependsOn content: name: RELEASE-NAME-ns + - contains: + path: spec.dependsOn + content: + name: RELEASE-NAME-soperator + + - it: should pass the supervisord ConfigMap base name + set: + customConfigmaps: + enabled: true + supervisordConfigMapName: custom-supervisord-config + asserts: + - equal: + path: spec.values.configMaps.supervisord.name + value: custom-supervisord-config - it: should reference correct HelmRepository set: diff --git a/helm/soperator-fluxcd/tests/soperator_dependency_test.yaml b/helm/soperator-fluxcd/tests/soperator_dependency_test.yaml index b216e73db..3d49d9db6 100644 --- a/helm/soperator-fluxcd/tests/soperator_dependency_test.yaml +++ b/helm/soperator-fluxcd/tests/soperator_dependency_test.yaml @@ -119,6 +119,7 @@ tests: - it: should have dependencies on namespace and soperator set: nodesets.enabled: true + customConfigmaps.enabled: false asserts: - contains: path: spec.dependsOn @@ -129,6 +130,16 @@ tests: content: name: RELEASE-NAME-soperator + - it: should have dependency on custom configmaps when custom configmaps are enabled + set: + nodesets.enabled: true + customConfigmaps.enabled: true + asserts: + - contains: + path: spec.dependsOn + content: + name: RELEASE-NAME-custom-configmaps + - it: should have dependency on slurm-cluster when slurm cluster is enabled set: nodesets.enabled: true diff --git a/helm/soperator-fluxcd/tests/values_override_test.yaml b/helm/soperator-fluxcd/tests/values_override_test.yaml index 761274ee5..f84b678dc 100644 --- a/helm/soperator-fluxcd/tests/values_override_test.yaml +++ b/helm/soperator-fluxcd/tests/values_override_test.yaml @@ -448,6 +448,7 @@ tests: - it: should use custom values when provided set: nodesets.enabled: true + customConfigmaps.version: 4.0.2 nodesets.values: priorityClasses: - name: soperator-worker-high-priority @@ -468,10 +469,14 @@ tests: - equal: path: spec.values.nodesets.test-nodeset.size value: 3 + - equal: + path: spec.values.configMapRefSupervisord + value: custom-supervisord-config-4.0.2 - it: should use override values when provided set: nodesets.enabled: true + customConfigmaps.version: 4.0.2 nodesets.overrideValues: customConfig: true nodesets: @@ -487,3 +492,6 @@ tests: - equal: path: spec.values.nodesets.gpu-nodeset.nodesetName value: "gpu-workers" + - equal: + path: spec.values.configMapRefSupervisord + value: custom-supervisord-config-4.0.2 diff --git a/helm/soperator-fluxcd/values.yaml b/helm/soperator-fluxcd/values.yaml index 0987d2e66..7cd1e4b0f 100644 --- a/helm/soperator-fluxcd/values.yaml +++ b/helm/soperator-fluxcd/values.yaml @@ -505,7 +505,7 @@ slurmCluster: driftDetection: mode: warn nodesets: - enabled: false + enabled: true clusterName: "" driftDetection: mode: warn @@ -752,6 +752,7 @@ customConfigmaps: version: 4.0.2 namespace: soperator releaseName: soperator-custom-configmaps + supervisordConfigMapName: custom-supervisord-config overrideValues: null values: null install: diff --git a/helm/soperator/crds/slurmcluster-crd.yaml b/helm/soperator/crds/slurmcluster-crd.yaml index 39445b875..3ea7e8b9f 100644 --- a/helm/soperator/crds/slurmcluster-crd.yaml +++ b/helm/soperator/crds/slurmcluster-crd.yaml @@ -27177,6 +27177,13 @@ spec: Defaults to "root", which preserves the single-fabric behavior. type: string type: object + updateStrategy: + default: rollingUpdate + description: UpdateStrategy for the advanced StatefulSet + enum: + - rollingUpdate + - onDelete + type: string workerAnnotations: additionalProperties: type: string @@ -27199,6 +27206,7 @@ spec: required: - munge - slurmd + - updateStrategy type: object status: description: NodeSetStatus defines the observed state of SlurmCluster 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 ff100b400..d9a4f7f99 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/common/scripts/reboot.sh b/images/common/scripts/reboot.sh index d21dab328..a38140a52 100644 --- a/images/common/scripts/reboot.sh +++ b/images/common/scripts/reboot.sh @@ -17,9 +17,40 @@ do esac done +POD_NAME="$(hostname)" +NS="$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace)" +TOKEN="$(cat /var/run/secrets/kubernetes.io/serviceaccount/token)" +CACERT=/var/run/secrets/kubernetes.io/serviceaccount/ca.crt +POD_URL="https://kubernetes.default.svc/api/v1/namespaces/$NS/pods/$POD_NAME" +DELETE_CANDIDATE_LABEL="slurm.nebius.ai/delete-candidate" +DELETE_CANDIDATE_LABEL_VALUE_STOPPING="stopping" +DELETE_CANDIDATE_LABEL_VALUE_DELETING="deleting" + +if POD_JSON="$(curl -fsS --cacert "$CACERT" \ + -H "Authorization: Bearer $TOKEN" \ + "$POD_URL")"; then + DELETE_CANDIDATE_LABEL_VALUE="$(jq -r \ + --arg label "$DELETE_CANDIDATE_LABEL" \ + '.metadata.labels[$label] // ""' \ + <<<"$POD_JSON" 2>/dev/null || true)" + + if [[ "$DELETE_CANDIDATE_LABEL_VALUE" == "$DELETE_CANDIDATE_LABEL_VALUE_STOPPING" ]]; then + curl -fsS --cacert "$CACERT" \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/merge-patch+json" \ + -X PATCH \ + -d "{\"metadata\":{\"labels\":{\"$DELETE_CANDIDATE_LABEL\":\"$DELETE_CANDIDATE_LABEL_VALUE_DELETING\"}}}" \ + "$POD_URL" + exit 0 + fi + + if [[ "$DELETE_CANDIDATE_LABEL_VALUE" == "$DELETE_CANDIDATE_LABEL_VALUE_DELETING" ]]; then + exit 0 + fi +fi + if ! mountpoint -q /run/nvidia/driver; then echo "This command only works on GPU nodes" exit 1 fi - chroot /run/nvidia/driver nsenter -t 1 -m -u -i -n /usr/sbin/reboot diff --git a/images/controller/slurmctld.dockerfile b/images/controller/slurmctld.dockerfile index 43d0d077f..2df88ce5f 100644 --- a/images/controller/slurmctld.dockerfile +++ b/images/controller/slurmctld.dockerfile @@ -63,4 +63,3 @@ RUN chmod 755 /opt/soperator/bin/power_resume.sh /opt/soperator/bin/power_suspen COPY images/controller/slurmctld_entrypoint.sh /opt/bin/slurm/ RUN chmod +x /opt/bin/slurm/slurmctld_entrypoint.sh ENTRYPOINT ["/opt/bin/slurm/slurmctld_entrypoint.sh"] - diff --git a/images/worker/worker_init.py b/images/worker/worker_init.py index 7cd3331b7..f3e435a17 100644 --- a/images/worker/worker_init.py +++ b/images/worker/worker_init.py @@ -660,9 +660,6 @@ def apply_node_topology(hostname: str, topology: str) -> None: f"nodename={hostname}", f"{node_addr}", f"{topology}", - "state=UNDRAIN", - "reason=", - "comment=", ] logger.info("Running: %s", " ".join(cmd)) result: subprocess.CompletedProcess[str] = subprocess.run( diff --git a/internal/consts/annotation.go b/internal/consts/annotation.go index 12e58ba64..0c154ccda 100644 --- a/internal/consts/annotation.go +++ b/internal/consts/annotation.go @@ -6,6 +6,11 @@ const ( AnnotationDefaultContainerName = "kubectl.kubernetes.io/default-container" AnnotationClusterName = K8sGroupNameSoperator + "/cluster" AnnotationActiveCheckName = K8sGroupNameSoperator + "/activecheck" + AnnotationDependencyVersion = K8sGroupNameSoperator + "/dependency-version" + + AnnotationDependencyVersionName = "name" AnnotationParentalClusterRefName = K8sGroupNameSoperator + "/parental-cluster-ref" + + AnnotationSoperatorRollingUpdateMaxUnavailable = K8sGroupNameSoperator + "/rolling-update-max-unavailable" ) diff --git a/internal/consts/container.go b/internal/consts/container.go index d1ab46d25..bb5b21b09 100644 --- a/internal/consts/container.go +++ b/internal/consts/container.go @@ -27,4 +27,5 @@ const ( ContainerPortNameMonitoring = "monitoring" ContainerPortMonitoring = 8081 ContainerPathMonitoring = "/metrics" + ContainerPathHealthz = "/healthz" ) diff --git a/internal/consts/label.go b/internal/consts/label.go index dccd66667..d02afb84d 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -42,4 +42,10 @@ const ( LabelJailedAggregationCommonValue = "common" AnnotationConfigHash = K8sGroupNameSoperator + "/config-hash" + + LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" + LabelSoperatorRollingUpdateValue = "true" + LabelSoperatorDeleteCandidate = K8sGroupNameSoperator + "/delete-candidate" + LabelSoperatorDeleteCandidateValueStopping = "stopping" + LabelSoperatorDeleteCandidateValueDeleting = "deleting" ) diff --git a/internal/consts/statefulset.go b/internal/consts/statefulset.go index f1aefa117..b7a0dac44 100644 --- a/internal/consts/statefulset.go +++ b/internal/consts/statefulset.go @@ -3,3 +3,10 @@ package consts const ( PodManagementPolicy = "Parallel" ) + +type UpdateStrategy string + +const ( + UpdateStrategyRollingUpdate UpdateStrategy = "rollingUpdate" + UpdateStrategyOnDelete UpdateStrategy = "onDelete" +) diff --git a/internal/controller/clustercontroller/common.go b/internal/controller/clustercontroller/common.go index ccfe215cd..1172707bc 100644 --- a/internal/controller/clustercontroller/common.go +++ b/internal/controller/clustercontroller/common.go @@ -179,25 +179,6 @@ func (r SlurmClusterReconciler) ReconcileCommon( return nil }, }, - utils.MultiStepExecutionStep{ - Name: "Slurm Worker ServiceAccount", - Func: func(stepCtx context.Context) error { - stepLogger := log.FromContext(stepCtx) - stepLogger.V(1).Info("Reconciling") - - desired := worker.RenderServiceAccount(clusterValues.Namespace, clusterValues.Name) - stepLogger = stepLogger.WithValues(logfield.ResourceKV(&desired)...) - stepLogger.V(1).Info("Rendered") - - if err := r.ServiceAccount.Reconcile(stepCtx, cluster, desired); err != nil { - stepLogger.Error(err, "Failed to reconcile") - return fmt.Errorf("reconciling worker ServiceAccount: %w", err) - } - stepLogger.V(1).Info("Reconciled") - - return nil - }, - }, utils.MultiStepExecutionStep{ Name: "Slurm Worker sysctl ConfigMap", Func: func(stepCtx context.Context) error { diff --git a/internal/controller/nodesetcontroller/configmap_predicate_test.go b/internal/controller/nodesetcontroller/configmap_predicate_test.go new file mode 100644 index 000000000..375165d85 --- /dev/null +++ b/internal/controller/nodesetcontroller/configmap_predicate_test.go @@ -0,0 +1,42 @@ +package nodesetcontroller + +import ( + "testing" + + "github.com/stretchr/testify/assert" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/event" +) + +func TestConfigMapDataChangedPredicate(t *testing.T) { + predicate := configMapDataChangedPredicate() + + oldConfigMap := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: "custom-supervisord-config", + Namespace: "soperator", + Labels: map[string]string{ + "helm.sh/chart": "helm-soperator-custom-configmaps-4.0.1", + }, + }, + Data: map[string]string{ + "supervisord.conf": "same", + }, + } + + newMetadataOnly := oldConfigMap.DeepCopy() + newMetadataOnly.Labels["helm.sh/chart"] = "helm-soperator-custom-configmaps-4.0.2" + + newData := oldConfigMap.DeepCopy() + newData.Data["supervisord.conf"] = "changed" + + assert.False(t, predicate.Update(event.UpdateEvent{ + ObjectOld: oldConfigMap, + ObjectNew: newMetadataOnly, + })) + assert.True(t, predicate.Update(event.UpdateEvent{ + ObjectOld: oldConfigMap, + ObjectNew: newData, + })) +} diff --git a/internal/controller/nodesetcontroller/controller.go b/internal/controller/nodesetcontroller/controller.go index 0c6b706f5..bdf6f1658 100644 --- a/internal/controller/nodesetcontroller/controller.go +++ b/internal/controller/nodesetcontroller/controller.go @@ -24,6 +24,7 @@ import ( kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/equality" apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/runtime" "k8s.io/client-go/tools/record" @@ -31,6 +32,7 @@ import ( 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/handler" "sigs.k8s.io/controller-runtime/pkg/log" "sigs.k8s.io/controller-runtime/pkg/predicate" @@ -66,9 +68,12 @@ type NodeSetReconciler struct { AdvancedStatefulSet *reconciler.AdvancedStatefulSetReconciler Service *reconciler.ServiceReconciler + ServiceAccount *reconciler.ServiceAccountReconciler Secret *reconciler.SecretReconciler ConfigMap *reconciler.ConfigMapReconciler NodeSetPowerState *reconciler.NodeSetPowerStateReconciler + Role *reconciler.RoleReconciler + RoleBinding *reconciler.RoleBindingReconciler } func NewNodeSetReconciler(client client.Client, scheme *runtime.Scheme, recorder record.EventRecorder) *NodeSetReconciler { @@ -77,9 +82,12 @@ func NewNodeSetReconciler(client client.Client, scheme *runtime.Scheme, recorder Reconciler: r, AdvancedStatefulSet: reconciler.NewAdvancedStatefulSetReconciler(r), Service: reconciler.NewServiceReconciler(r), + ServiceAccount: reconciler.NewServiceAccountReconciler(r), Secret: reconciler.NewSecretReconciler(r), ConfigMap: reconciler.NewConfigMapReconciler(r), NodeSetPowerState: reconciler.NewNodeSetPowerStateReconciler(r), + Role: reconciler.NewRoleReconciler(r), + RoleBinding: reconciler.NewRoleBindingReconciler(r), } } @@ -102,7 +110,7 @@ func (r *NodeSetReconciler) SetupWithManager(mgr ctrl.Manager, name string, maxC controllerBuilder.Watches( &corev1.ConfigMap{}, handler.EnqueueRequestsFromMapFunc(r.findObjectsForConfigMap), - builder.WithPredicates(predicate.ResourceVersionChangedPredicate{}), + builder.WithPredicates(configMapDataChangedPredicate()), ) controllerBuilder.Watches( @@ -129,6 +137,23 @@ func (r *NodeSetReconciler) SetupWithManager(mgr ctrl.Manager, name string, maxC return controllerBuilder.Complete(r) } +func configMapDataChangedPredicate() predicate.Funcs { + return predicate.Funcs{ + CreateFunc: func(event.CreateEvent) bool { return true }, + DeleteFunc: func(event.DeleteEvent) bool { return true }, + UpdateFunc: func(e event.UpdateEvent) bool { + oldConfigMap, okOld := e.ObjectOld.(*corev1.ConfigMap) + newConfigMap, okNew := e.ObjectNew.(*corev1.ConfigMap) + if !okOld || !okNew { + return true + } + + return !equality.Semantic.DeepEqual(oldConfigMap.Data, newConfigMap.Data) || + !equality.Semantic.DeepEqual(oldConfigMap.BinaryData, newConfigMap.BinaryData) + }, + } +} + func (r *NodeSetReconciler) createResourceChecks(saPredicate predicate.Funcs) []controllercommon.ResourceCheck { return []controllercommon.ResourceCheck{ { diff --git a/internal/controller/nodesetcontroller/reconcile.go b/internal/controller/nodesetcontroller/reconcile.go index d7eb07fc4..a67040040 100644 --- a/internal/controller/nodesetcontroller/reconcile.go +++ b/internal/controller/nodesetcontroller/reconcile.go @@ -322,6 +322,63 @@ func (r NodeSetReconciler) executeReconciliation( return nil }, }, + { + Name: "Slurm NodeSet ServiceAccount", + Func: func(stepCtx context.Context) error { + stepLogger := log.FromContext(stepCtx) + stepLogger.V(1).Info("Reconciling") + + desired := worker.RenderServiceAccount(nodeSet.Namespace, cluster.Name, nodeSet.Name) + stepLogger = stepLogger.WithValues(logfield.ResourceKV(&desired)...) + stepLogger.V(1).Info("Rendered") + + if err := r.ServiceAccount.Reconcile(stepCtx, cluster, desired); err != nil { + stepLogger.Error(err, "Failed to reconcile") + return fmt.Errorf("reconciling nodeSet ServiceAccount: %w", err) + } + stepLogger.V(1).Info("Reconciled") + + return nil + }, + }, + { + Name: "Slurm NodeSet Role", + Func: func(stepCtx context.Context) error { + stepLogger := log.FromContext(stepCtx) + stepLogger.V(1).Info("Reconciling") + + desired := worker.RenderRole(nodeSet.Namespace, cluster.Name, nodeSetValues) + stepLogger = stepLogger.WithValues(logfield.ResourceKV(&desired)...) + stepLogger.V(1).Info("Rendered") + + if err := r.Role.Reconcile(stepCtx, cluster, desired); err != nil { + stepLogger.Error(err, "Failed to reconcile") + return fmt.Errorf("reconciling nodeSet Role: %w", err) + } + stepLogger.V(1).Info("Reconciled") + + return nil + }, + }, + { + Name: "Slurm NodeSet RoleBinding", + Func: func(stepCtx context.Context) error { + stepLogger := log.FromContext(stepCtx) + stepLogger.V(1).Info("Reconciling") + + desired := worker.RenderRoleBinding(nodeSet.Namespace, cluster.Name, nodeSet.Name) + stepLogger = stepLogger.WithValues(logfield.ResourceKV(&desired)...) + stepLogger.V(1).Info("Rendered") + + if err := r.RoleBinding.Reconcile(stepCtx, cluster, desired); err != nil { + stepLogger.Error(err, "Failed to reconcile") + return fmt.Errorf("reconciling nodeSet RoleBinding: %w", err) + } + stepLogger.V(1).Info("Reconciled") + + return nil + }, + }, { Name: "Worker sssd.conf Secret", Func: func(stepCtx context.Context) error { diff --git a/internal/controller/reconciler/versioning.go b/internal/controller/reconciler/versioning.go index 99758e79d..bd697b94a 100644 --- a/internal/controller/reconciler/versioning.go +++ b/internal/controller/reconciler/versioning.go @@ -34,12 +34,20 @@ func generateDependencyVersionMap(deps ...metav1.Object) map[string]string { res := map[string]string{} for _, d := range deps { - res[fmt.Sprintf("%s.%s", d.GetNamespace(), d.GetName())] = d.GetResourceVersion() + res[fmt.Sprintf("%s.%s", d.GetNamespace(), d.GetName())] = generateDependencyVersion(d) } return res } +func generateDependencyVersion(dep metav1.Object) string { + if dep.GetAnnotations()[consts.AnnotationDependencyVersion] == consts.AnnotationDependencyVersionName { + return dep.GetName() + } + + return dep.GetResourceVersion() +} + func getVersionsAnnotation(resource metav1.Object) (map[string]string, error) { getVersionsAnnotationImpl := func() (map[string]string, error) { if resource.GetAnnotations() == nil { diff --git a/internal/controller/reconciler/versioning_test.go b/internal/controller/reconciler/versioning_test.go new file mode 100644 index 000000000..a604d68ca --- /dev/null +++ b/internal/controller/reconciler/versioning_test.go @@ -0,0 +1,55 @@ +package reconciler + +import ( + "testing" + + "github.com/stretchr/testify/assert" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "nebius.ai/slurm-operator/internal/consts" +) + +func TestGenerateDependencyVersionMap(t *testing.T) { + tests := []struct { + name string + dependency metav1.Object + expected map[string]string + }{ + { + name: "uses resource version by default", + dependency: &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: "default-config", + Namespace: "soperator", + ResourceVersion: "12345", + }, + }, + expected: map[string]string{ + "soperator.default-config": "12345", + }, + }, + { + name: "uses name when dependency version annotation requests it", + dependency: &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: "custom-supervisord-config-4.0.2", + Namespace: "soperator", + ResourceVersion: "12345", + Annotations: map[string]string{ + consts.AnnotationDependencyVersion: consts.AnnotationDependencyVersionName, + }, + }, + }, + expected: map[string]string{ + "soperator.custom-supervisord-config-4.0.2": "custom-supervisord-config-4.0.2", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.expected, generateDependencyVersionMap(tt.dependency)) + }) + } +} diff --git a/internal/controller/soperatorchecks/maintenance_condition_test.go b/internal/controller/soperatorchecks/maintenance_condition_test.go index 2261e660d..d416019d4 100644 --- a/internal/controller/soperatorchecks/maintenance_condition_test.go +++ b/internal/controller/soperatorchecks/maintenance_condition_test.go @@ -55,17 +55,6 @@ func TestMaintenanceConditionTypeConfiguration(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - controller := NewSlurmAPIClientsController( - client, - scheme, - recorder, - slurmAPIClients, - corev1.NodeConditionType(tt.inputConditionType), - ) - - assert.Equal(t, tt.expectedConditionType, string(controller.MaintenanceConditionType), - "SlurmAPIClientsController: %s", tt.description) - k8sController := NewK8SNodesController( client, scheme, @@ -104,65 +93,13 @@ func TestDefaultMaintenanceConditionTypeConstant(t *testing.T) { recorder := record.NewFakeRecorder(10) slurmAPIClients := slurmapi.NewClientSet(context.Background()) - slurmAPIController := NewSlurmAPIClientsController(client, scheme, recorder, slurmAPIClients, "") k8sController := NewK8SNodesController(client, scheme, recorder, 15*time.Minute, true, "", "") slurmController := NewSlurmNodesController(client, scheme, recorder, slurmAPIClients, 30*time.Second, true, true, client, "") expectedDefault := string(consts.DefaultMaintenanceConditionType) - assert.Equal(t, expectedDefault, string(slurmAPIController.MaintenanceConditionType), - "SlurmAPIClientsController should use default maintenance condition type") assert.Equal(t, expectedDefault, string(k8sController.MaintenanceConditionType), "K8SNodesController should use default maintenance condition type") assert.Equal(t, expectedDefault, string(slurmController.MaintenanceConditionType), "SlurmNodesController should use default maintenance condition type") } - -func TestMaintenanceConditionTypeIntegration(t *testing.T) { - testCases := []struct { - name string - cmdLineArg string - expectedResult string - }{ - { - name: "command line flag with custom value", - cmdLineArg: "ProductionMaintenanceScheduled", - expectedResult: "ProductionMaintenanceScheduled", - }, - { - name: "command line flag with default value", - cmdLineArg: string(consts.DefaultMaintenanceConditionType), - expectedResult: "NebiusMaintenanceScheduled", - }, - { - name: "empty command line flag uses default", - cmdLineArg: "", - expectedResult: string(consts.DefaultMaintenanceConditionType), - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - scheme := runtime.NewScheme() - client := fake.NewClientBuilder().WithScheme(scheme).Build() - recorder := record.NewFakeRecorder(10) - slurmAPIClients := slurmapi.NewClientSet(context.Background()) - - maintenanceConditionType := tc.cmdLineArg - if maintenanceConditionType == "" { - maintenanceConditionType = string(consts.DefaultMaintenanceConditionType) - } - - controller := NewSlurmAPIClientsController( - client, - scheme, - recorder, - slurmAPIClients, - corev1.NodeConditionType(maintenanceConditionType), - ) - - assert.Equal(t, tc.expectedResult, string(controller.MaintenanceConditionType), - "Integration test failed for case: %s", tc.name) - }) - } -} diff --git a/internal/controller/soperatorchecks/slurm_api_clients_controller.go b/internal/controller/soperatorchecks/slurm_api_clients_controller.go index 82ff185f6..a7434aa16 100644 --- a/internal/controller/soperatorchecks/slurm_api_clients_controller.go +++ b/internal/controller/soperatorchecks/slurm_api_clients_controller.go @@ -5,15 +5,8 @@ import ( "fmt" "time" - corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/client-go/tools/record" - "nebius.ai/slurm-operator/internal/consts" - "nebius.ai/slurm-operator/internal/controller/reconciler" - "nebius.ai/slurm-operator/internal/controllerconfig" - "nebius.ai/slurm-operator/internal/jwt" - "nebius.ai/slurm-operator/internal/naming" - "nebius.ai/slurm-operator/internal/slurmapi" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" @@ -22,6 +15,12 @@ import ( "sigs.k8s.io/controller-runtime/pkg/predicate" slurmv1 "nebius.ai/slurm-operator/api/v1" + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/controller/reconciler" + "nebius.ai/slurm-operator/internal/controllerconfig" + "nebius.ai/slurm-operator/internal/jwt" + "nebius.ai/slurm-operator/internal/naming" + "nebius.ai/slurm-operator/internal/slurmapi" ) var ( @@ -31,8 +30,7 @@ var ( type SlurmAPIClientsController struct { *reconciler.Reconciler - slurmAPIClients *slurmapi.ClientSet - MaintenanceConditionType corev1.NodeConditionType + slurmAPIClients *slurmapi.ClientSet } func NewSlurmAPIClientsController( @@ -40,18 +38,12 @@ func NewSlurmAPIClientsController( scheme *runtime.Scheme, recorder record.EventRecorder, slurmAPIClients *slurmapi.ClientSet, - maintenanceConditionType corev1.NodeConditionType, ) *SlurmAPIClientsController { r := reconciler.NewReconciler(client, scheme, recorder) - if maintenanceConditionType == "" { - maintenanceConditionType = consts.DefaultMaintenanceConditionType - } - return &SlurmAPIClientsController{ - Reconciler: r, - slurmAPIClients: slurmAPIClients, - MaintenanceConditionType: maintenanceConditionType, + Reconciler: r, + slurmAPIClients: slurmAPIClients, } } 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/topologyconfcontroller/nodetopology_controller.go b/internal/controller/topologyconfcontroller/nodetopology_controller.go index 3d2407495..508d8c152 100644 --- a/internal/controller/topologyconfcontroller/nodetopology_controller.go +++ b/internal/controller/topologyconfcontroller/nodetopology_controller.go @@ -481,10 +481,12 @@ func (r *NodeTopologyReconciler) SetupWithManager(mgr ctrl.Manager, return false } - _, newHasLabel := newNode.Labels[r.tierOneLabel()] - _, oldHasLabel := oldNode.Labels[r.tierOneLabel()] + newLabel, newHasLabel := newNode.Labels[r.tierOneLabel()] + oldLabel, oldHasLabel := oldNode.Labels[r.tierOneLabel()] - return newHasLabel || (oldHasLabel && oldNode.Labels[r.tierOneLabel()] != newNode.Labels[r.tierOneLabel()]) + return (newHasLabel && !oldHasLabel) || + (!newHasLabel && oldHasLabel) || + (newHasLabel && oldHasLabel && newLabel != oldLabel) }, DeleteFunc: func(e event.DeleteEvent) bool { node, ok := e.Object.(*corev1.Node) diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go new file mode 100644 index 000000000..e7387b18f --- /dev/null +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -0,0 +1,355 @@ +/* +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" + 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 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 + } + + labels := sts.GetLabels() + if labels[consts.LabelSoperatorRollingUpdateEnabled] != consts.LabelSoperatorRollingUpdateValue { + logger.Info("rolling update is disabled", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + + 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 + } + + if sts.Status.UpdatedReplicas == replicas { + logger.Info("statefulset is up to date", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + + outdatedPodList, err := r.getOutdatedPodList(ctx, sts) + if err != nil { + return ctrl.Result{}, err + } + + if err := r.processRollingUpdate(ctx, clusterName, sts, outdatedPodList); err != nil { + return ctrl.Result{}, err + } + + return ctrl.Result{RequeueAfter: time.Minute}, nil +} + +func (r *RollingUpdateReconciler) getOutdatedPodList( + 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) + } + + var res []corev1.Pod + + for _, pod := range podList.Items { + // TODO: REVISION + podControllerRevisionHash := pod.Labels["controller-revision-hash"] + if podControllerRevisionHash == sts.Status.UpdateRevision { + continue + } + + res = append(res, pod) + } + + return res, nil +} + +func (r *RollingUpdateReconciler) processRollingUpdate( + ctx context.Context, + clusterName 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 pod.Labels[consts.LabelSoperatorDeleteCandidate] != consts.LabelSoperatorDeleteCandidateValueDeleting { + podsToStop = append(podsToStop, pod) + continue + } + + if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil { + return fmt.Errorf("delete pod %s/%s marked as delete candidate: %w", pod.Namespace, pod.Name, err) + } + deletedPods++ + } + if deletedPods > 0 { + logger.Info("deleted outdated pods marked as delete candidates", "count", deletedPods) + 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) + } + + type rebootCandidate struct { + pod corev1.Pod + slurmNode slurmapi.Node + } + + candidates := make([]rebootCandidate, 0, len(podsToStop)) + inFlightReady := 0 + for _, pod := range podsToStop { + slurmNode, err := slurmClient.GetNode(ctx, pod.Name) + if err != nil { + return err + } + + if slurmNode.IsRebootIssuedState() || slurmNode.IsRebootRequestedState() { + if podReady(&pod) { + inFlightReady++ + } + continue + } + + candidates = append(candidates, rebootCandidate{pod: pod, slurmNode: slurmNode}) + } + + budget := rebootBudget(sts) + unavailable := unavailableReplicas(sts) + availableSlots := budget - unavailable - inFlightReady + if availableSlots <= 0 { + logger.Info( + "rolling update budget is exhausted", + "budget", budget, + "unavailable", unavailable, + "inFlightReady", inFlightReady, + ) + return nil + } + + slurmNodesToReboot := make([]string, 0, availableSlots) + for _, candidate := range candidates { + if len(slurmNodesToReboot) >= availableSlots { + break + } + + pod := candidate.pod + if pod.Labels[consts.LabelSoperatorDeleteCandidate] != consts.LabelSoperatorDeleteCandidateValueStopping { + patchBase := pod.DeepCopy() + if pod.Labels == nil { + pod.Labels = map[string]string{} + } + pod.Labels[consts.LabelSoperatorDeleteCandidate] = consts.LabelSoperatorDeleteCandidateValueStopping + if err := r.Patch(ctx, &pod, client.MergeFrom(patchBase)); err != nil { + return fmt.Errorf("label pod %s/%s for reboot handoff: %w", 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, + }); 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 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 raw := sts.GetAnnotations()[consts.AnnotationSoperatorRollingUpdateMaxUnavailable]; raw != "" { + maxUnavailable = intstr.Parse(raw) + } + + 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 +} + +// 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 { + if obj == nil { + return false + } + return obj.GetLabels()[consts.LabelSoperatorRollingUpdateEnabled] == consts.LabelSoperatorRollingUpdateValue +} diff --git a/internal/naming/naming.go b/internal/naming/naming.go index ba040edeb..6103e1154 100644 --- a/internal/naming/naming.go +++ b/internal/naming/naming.go @@ -238,12 +238,16 @@ func BuildVolumeMountSpoolPath(directory string) string { return path.Join(consts.VolumeMountPathSpool, directory) } -func BuildServiceAccountWorkerName(clusterName string) string { - return clusterName + "-worker-sa" +func BuildServiceAccountNodeSetName(clusterName, nodeSetName string) string { + return fmt.Sprintf("%s-%s-nodeset-sa", clusterName, nodeSetName) } -func BuildRoleWorkerName(clusterName string) string { - return clusterName + "-worker-events-role" +func BuildRoleNodeSetName(clusterName, nodeSetName string) string { + return fmt.Sprintf("%s-%s-nodeset-role", clusterName, nodeSetName) +} + +func BuildRoleBindingNodeSetName(clusterName, nodeSetName string) string { + return fmt.Sprintf("%s-%s-nodeset-role-binding", clusterName, nodeSetName) } func BuildServiceAccountActiveCheckName(clusterName string) string { diff --git a/internal/render/controller/service.go b/internal/render/controller/service.go index 7e3c36ab2..79b2336fd 100644 --- a/internal/render/controller/service.go +++ b/internal/render/controller/service.go @@ -32,12 +32,14 @@ func RenderService(namespace, clusterName, svcName string, controller *values.Sl Spec: corev1.ServiceSpec{ Type: controller.Service.Type, Selector: selector, - Ports: []corev1.ServicePort{{ - Name: controller.ContainerSlurmctld.Name, - Protocol: controller.Service.Protocol, - Port: controller.ContainerSlurmctld.Port, - TargetPort: intstr.FromString(controller.ContainerSlurmctld.Name), - }}, + Ports: []corev1.ServicePort{ + { + Name: controller.ContainerSlurmctld.Name, + Protocol: controller.Service.Protocol, + Port: controller.ContainerSlurmctld.Port, + TargetPort: intstr.FromString(controller.ContainerSlurmctld.Name), + }, + }, }, } } diff --git a/internal/render/worker/role.go b/internal/render/worker/role.go index 15b3cc6d6..3a29ca3b4 100644 --- a/internal/render/worker/role.go +++ b/internal/render/worker/role.go @@ -1,28 +1,51 @@ package worker import ( + "fmt" + rbacv1 "k8s.io/api/rbac/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "nebius.ai/slurm-operator/internal/consts" "nebius.ai/slurm-operator/internal/naming" "nebius.ai/slurm-operator/internal/render/common" + "nebius.ai/slurm-operator/internal/values" ) -func RenderRole(namespace, clusterName string) rbacv1.Role { - labels := common.RenderLabels(consts.ComponentTypeWorker, clusterName) +func RenderRole(namespace, clusterName string, nodeSet *values.SlurmNodeSet) rbacv1.Role { + labels := common.RenderLabels(consts.ComponentTypeNodeSet, clusterName) + + ordinals := nodeSetOrdinals(nodeSet) + resourceNames := make([]string, 0, len(ordinals)) + for _, i := range ordinals { + resourceNames = append(resourceNames, fmt.Sprintf("%s-%d", naming.BuildNodeSetStatefulSetName(nodeSet.Name), i)) + } return rbacv1.Role{ ObjectMeta: metav1.ObjectMeta{ - Name: naming.BuildRoleWorkerName(clusterName), + Name: naming.BuildRoleNodeSetName(clusterName, nodeSet.Name), Namespace: namespace, Labels: labels, }, Rules: []rbacv1.PolicyRule{ { - APIGroups: []string{""}, - Resources: []string{"events"}, - Verbs: []string{"create"}, + APIGroups: []string{""}, + Resources: []string{"pods"}, + ResourceNames: resourceNames, + Verbs: []string{"get", "patch"}, }, }, } } + +func nodeSetOrdinals(nodeSet *values.SlurmNodeSet) []int32 { + if nodeSet.EphemeralNodes != nil && *nodeSet.EphemeralNodes { + return append([]int32(nil), nodeSet.ActiveNodes...) + } + + ordinals := make([]int32, 0, nodeSet.StatefulSet.Replicas) + for i := int32(0); i < nodeSet.StatefulSet.Replicas; i++ { + ordinals = append(ordinals, i) + } + return ordinals +} diff --git a/internal/render/worker/role_test.go b/internal/render/worker/role_test.go index c16ed167d..c8d26096c 100644 --- a/internal/render/worker/role_test.go +++ b/internal/render/worker/role_test.go @@ -1,29 +1,57 @@ -package worker +package worker_test import ( "testing" - "nebius.ai/slurm-operator/internal/naming" -) - -func Test_RenderRole(t *testing.T) { - namespace := "test-namespace" - clusterName := "test-cluster" + "github.com/stretchr/testify/assert" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/utils/ptr" - role := RenderRole(namespace, clusterName) + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/render/worker" + "nebius.ai/slurm-operator/internal/values" +) - // Check the name - if role.Name != naming.BuildRoleWorkerName(clusterName) { - t.Errorf("Unexpected name: got %v, want %v", role.Name, naming.BuildRoleWorkerName(clusterName)) +func TestRenderRole_ResourceNames(t *testing.T) { + tests := []struct { + name string + nodeSet *values.SlurmNodeSet + expectedPods []string + }{ + { + name: "regular nodeset ordinals", + nodeSet: &values.SlurmNodeSet{ + Name: "workers", + StatefulSet: values.StatefulSet{ + Replicas: 3, + }, + }, + expectedPods: []string{"workers-0", "workers-1", "workers-2"}, + }, + { + name: "ephemeral sparse active ordinals", + nodeSet: &values.SlurmNodeSet{ + Name: "workers", + EphemeralNodes: ptr.To(true), + ActiveNodes: []int32{0, 3, 12}, + StatefulSet: values.StatefulSet{ + Replicas: 3, + MaxUnavailable: intstr.FromInt32(1), + MaxConcurrentStartup: intstr.FromInt32(1), + }, + }, + expectedPods: []string{"workers-0", "workers-3", "workers-12"}, + }, } - // Check the namespace - if role.Namespace != namespace { - t.Errorf("Unexpected namespace: got %v, want %v", role.Namespace, namespace) - } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + role := worker.RenderRole("default", "soperator", tt.nodeSet) - // Check the rules - if len(role.Rules) != 1 || role.Rules[0].APIGroups[0] != "" || role.Rules[0].Resources[0] != "events" || role.Rules[0].Verbs[0] != "create" { - t.Errorf("Unexpected rules: got %v, want one rule with apiGroups=[\"\"], resources=[\"events\"], and verbs=[\"create\"]", role.Rules) + assert.Len(t, role.Rules, 1) + assert.Equal(t, []string{"get", "patch"}, role.Rules[0].Verbs) + assert.Equal(t, tt.expectedPods, role.Rules[0].ResourceNames) + assert.Equal(t, consts.ComponentTypeNodeSet.String(), role.Labels[consts.LabelComponentKey]) + }) } } diff --git a/internal/render/worker/rolebinding.go b/internal/render/worker/rolebinding.go new file mode 100644 index 000000000..b56d2d094 --- /dev/null +++ b/internal/render/worker/rolebinding.go @@ -0,0 +1,34 @@ +package worker + +import ( + rbacv1 "k8s.io/api/rbac/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/naming" + "nebius.ai/slurm-operator/internal/render/common" +) + +func RenderRoleBinding(namespace, clusterName, nodeSetName string) rbacv1.RoleBinding { + labels := common.RenderLabels(consts.ComponentTypeWorker, clusterName) + + return rbacv1.RoleBinding{ + ObjectMeta: metav1.ObjectMeta{ + Name: naming.BuildRoleBindingNodeSetName(clusterName, nodeSetName), + Namespace: namespace, + Labels: labels, + }, + Subjects: []rbacv1.Subject{ + { + Kind: rbacv1.ServiceAccountKind, + Name: naming.BuildServiceAccountNodeSetName(clusterName, nodeSetName), + Namespace: namespace, + }, + }, + RoleRef: rbacv1.RoleRef{ + Kind: "Role", + Name: naming.BuildRoleNodeSetName(clusterName, nodeSetName), + APIGroup: rbacv1.GroupName, + }} + +} diff --git a/internal/render/worker/serviceaccount.go b/internal/render/worker/serviceaccount.go index 102832ae7..657fc006b 100644 --- a/internal/render/worker/serviceaccount.go +++ b/internal/render/worker/serviceaccount.go @@ -9,12 +9,12 @@ import ( "nebius.ai/slurm-operator/internal/render/common" ) -func RenderServiceAccount(namespace, clusterName string) corev1.ServiceAccount { - labels := common.RenderLabels(consts.ComponentTypeWorker, clusterName) +func RenderServiceAccount(namespace, clusterName, nodeSetName string) corev1.ServiceAccount { + labels := common.RenderLabels(consts.ComponentTypeNodeSet, clusterName) return corev1.ServiceAccount{ ObjectMeta: metav1.ObjectMeta{ - Name: naming.BuildServiceAccountWorkerName(clusterName), + Name: naming.BuildServiceAccountNodeSetName(clusterName, nodeSetName), Namespace: namespace, Labels: labels, }, diff --git a/internal/render/worker/sssd_test.go b/internal/render/worker/sssd_test.go index 73f1ce6f5..f8b8b2a6f 100644 --- a/internal/render/worker/sssd_test.go +++ b/internal/render/worker/sssd_test.go @@ -67,6 +67,7 @@ func TestRenderNodeSetStatefulSet_SSSD(t *testing.T) { SSHDConfigMapName: "sshd-config", GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, CustomInitContainers: []corev1.Container{}, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } result, err := worker.RenderNodeSetStatefulSet( diff --git a/internal/render/worker/statefulset.go b/internal/render/worker/statefulset.go index 7386efdd5..efa7ee7d0 100644 --- a/internal/render/worker/statefulset.go +++ b/internal/render/worker/statefulset.go @@ -36,6 +36,7 @@ func RenderNodeSetStatefulSet( labels := common.RenderLabels(consts.ComponentTypeNodeSet, nodeSet.ParentalCluster.Name) labels[consts.LabelNodeSetKey] = nodeSet.Name labels[consts.LabelWorkerKey] = consts.LabelWorkerValue + labels[consts.LabelSoperatorRollingUpdateEnabled] = consts.LabelSoperatorRollingUpdateValue matchLabels := common.RenderMatchLabels(consts.ComponentTypeNodeSet, nodeSet.ParentalCluster.Name) matchLabels[consts.LabelNodeSetKey] = nodeSet.Name @@ -114,7 +115,7 @@ func RenderNodeSetStatefulSet( }, }, PriorityClassName: nodeSet.PriorityClass, - ServiceAccountName: naming.BuildServiceAccountWorkerName(nodeSet.ParentalCluster.Name), + ServiceAccountName: naming.BuildServiceAccountNodeSetName(nodeSet.ParentalCluster.Name, nodeSet.Name), Affinity: nodeSet.Affinity, NodeSelector: nodeSet.NodeSelector, Tolerations: nodeSet.Tolerations, @@ -143,8 +144,9 @@ func RenderNodeSetStatefulSet( } annotations := map[string]string{ - "kruise.io/auto-generate-persistent-pod-state": "true", - "kruise.io/preferred-persistent-topology": "kubernetes.io/hostname", + "kruise.io/auto-generate-persistent-pod-state": "true", + "kruise.io/preferred-persistent-topology": "kubernetes.io/hostname", + consts.AnnotationSoperatorRollingUpdateMaxUnavailable: nodeSet.StatefulSet.MaxUnavailable.String(), } pvcRetentionPolicy := nodeSet.PersistentVolumeClaimRetentionPolicy @@ -154,6 +156,10 @@ func RenderNodeSetStatefulSet( WhenScaled: kruisev1b1.DeletePersistentVolumeClaimRetentionPolicyType, } } + updateStrategy, volumeClaimUpdateStrategy, err := renderUpdateStrategies(nodeSet) + if err != nil { + return kruisev1b1.StatefulSet{}, fmt.Errorf("rendering update strategies: %w", err) + } return kruisev1b1.StatefulSet{ ObjectMeta: metav1.ObjectMeta{ @@ -170,15 +176,7 @@ func RenderNodeSetStatefulSet( ScaleStrategy: &kruisev1b1.StatefulSetScaleStrategy{ MaxUnavailable: &nodeSet.StatefulSet.MaxConcurrentStartup, }, - UpdateStrategy: kruisev1b1.StatefulSetUpdateStrategy{ - Type: appsv1.RollingUpdateStatefulSetStrategyType, - RollingUpdate: &kruisev1b1.RollingUpdateStatefulSetStrategy{ - MaxUnavailable: &nodeSet.StatefulSet.MaxUnavailable, - PodUpdatePolicy: kruisev1b1.InPlaceIfPossiblePodUpdateStrategyType, - Partition: ptr.To(int32(0)), - MinReadySeconds: ptr.To(int32(0)), - }, - }, + UpdateStrategy: updateStrategy, Selector: &metav1.LabelSelector{ MatchLabels: matchLabels, }, @@ -188,9 +186,7 @@ func RenderNodeSetStatefulSet( nodeSet.ParentalCluster.Name, pvcTemplateSpecs, ), - VolumeClaimUpdateStrategy: kruisev1b1.VolumeClaimUpdateStrategy{ - Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, - }, + VolumeClaimUpdateStrategy: volumeClaimUpdateStrategy, Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ Labels: labels, @@ -203,6 +199,37 @@ func RenderNodeSetStatefulSet( }, nil } +func renderUpdateStrategies(nodeSet *values.SlurmNodeSet) (kruisev1b1.StatefulSetUpdateStrategy, kruisev1b1.VolumeClaimUpdateStrategy, error) { + switch nodeSet.UpdateStrategy { + case consts.UpdateStrategyRollingUpdate: + return kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.RollingUpdateStatefulSetStrategyType, + RollingUpdate: &kruisev1b1.RollingUpdateStatefulSetStrategy{ + MaxUnavailable: &nodeSet.StatefulSet.MaxUnavailable, + PodUpdatePolicy: kruisev1b1.InPlaceIfPossiblePodUpdateStrategyType, + Partition: ptr.To(int32(0)), + MinReadySeconds: ptr.To(int32(0)), + }, + }, + kruisev1b1.VolumeClaimUpdateStrategy{ + Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, + }, + nil + case consts.UpdateStrategyOnDelete: + return kruisev1b1.StatefulSetUpdateStrategy{ + Type: appsv1.OnDeleteStatefulSetStrategyType, + }, + // This doesn't work in kruise + // TODO: fix not working on-flight pvc resize + kruisev1b1.VolumeClaimUpdateStrategy{ + Type: kruisev1b1.OnPodRollingUpdateVolumeClaimUpdateStrategyType, + }, + nil + default: + return kruisev1b1.StatefulSetUpdateStrategy{}, kruisev1b1.VolumeClaimUpdateStrategy{}, fmt.Errorf("kek") + } +} + func renderNodeSetAnnotations(nodeSet *values.SlurmNodeSet) map[string]string { annotations := common.RenderDefaultContainerAnnotation(consts.ContainerNameSlurmd) maps.Copy(annotations, nodeSet.Annotations) diff --git a/internal/render/worker/statefulset_test.go b/internal/render/worker/statefulset_test.go index 4dbb8952a..6e3c34b59 100644 --- a/internal/render/worker/statefulset_test.go +++ b/internal/render/worker/statefulset_test.go @@ -290,6 +290,7 @@ func TestRenderNodeSetStatefulSet_SlurmdGPUEnv(t *testing.T) { SupervisorDConfigMapName: "supervisord-config", SSHDConfigMapName: "sshd-config", GPU: &slurmv1alpha1.GPUSpec{Enabled: tt.nodeSetGPUEnabled}, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } result, err := worker.RenderNodeSetStatefulSet( @@ -347,6 +348,7 @@ func TestRenderNodeSetStatefulSet_TopologyPlugin(t *testing.T) { GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, EphemeralNodes: ephemeralNodes, EphemeralTopologyWaitTimeout: waitTimeout, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } } @@ -527,6 +529,7 @@ func TestRenderNodeSetStatefulSet_PersistentVolumeClaimRetentionPolicy(t *testin SSHDConfigMapName: "sshd-config", GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, EphemeralNodes: ephemeralNodes, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } } @@ -628,6 +631,7 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { SupervisorDConfigMapName: "supervisord-config", SSHDConfigMapName: "sshd-config", GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } } @@ -678,6 +682,13 @@ func TestRenderNodeSetStatefulSet_ScaleStrategy(t *testing.T) { 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.expectedMaxUnavailable.String(), + result.Annotations[consts.AnnotationSoperatorRollingUpdateMaxUnavailable], + "rolling update controller annotation should preserve NodeSet.MaxUnavailable", + ) }) } } @@ -722,6 +733,7 @@ func TestRenderNodeSetStatefulSet_EphemeralNodesReserveOrdinals(t *testing.T) { GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, EphemeralNodes: &ephemeralNodes, ActiveNodes: activeNodes, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } } diff --git a/internal/slurmapi/client.go b/internal/slurmapi/client.go index c48ff6bd7..2428c8932 100644 --- a/internal/slurmapi/client.go +++ b/internal/slurmapi/client.go @@ -1,6 +1,7 @@ package slurmapi import ( + "bytes" "cmp" "context" "encoding/json" @@ -8,6 +9,7 @@ import ( "io" "net/http" "reflect" + "strings" "time" api "github.com/SlinkyProject/slurm-client/api/v0041" @@ -80,6 +82,8 @@ const ( headerApplicationJson = "application/json" SlurmUserSoperatorchecks = "soperatorchecks" + + slurmV0044RebootNodesPath = "/slurm/v0.0.44/nodes/reboot" ) func DefaultHTTPClient() *http.Client { @@ -113,6 +117,9 @@ type client struct { client0043 api0043.ClientWithResponsesInterface + server string + httpClient *http.Client + tokenIssuer tokenIssuer } @@ -129,6 +136,8 @@ func NewClient(server string, tokenIssuer tokenIssuer, httpClient *http.Client) } apiClient := &client{ + server: strings.TrimRight(server, "/"), + httpClient: httpClient, tokenIssuer: normalizeIssuer(tokenIssuer), } @@ -157,6 +166,13 @@ func NewClient(server string, tokenIssuer tokenIssuer, httpClient *http.Client) return apiClient, nil } +type RebootNodesRequest struct { + NodeList string `json:"nodes"` + ASAP bool `json:"asap,omitempty"` + NextState string `json:"next_state,omitempty"` + Reason string `json:"reason,omitempty"` +} + // normalizeIssuer converts a typed-nil inside an interface to a real nil. func normalizeIssuer(t tokenIssuer) tokenIssuer { if t == nil { @@ -190,6 +206,52 @@ func (c *client) setHeaders(ctx context.Context, req *http.Request) error { return nil } +func (c *client) RebootNodes(ctx context.Context, request RebootNodesRequest) error { + if request.NodeList == "" { + return fmt.Errorf("reboot nodes: nodes field is required") + } + + body, err := json.Marshal(request) + if err != nil { + return fmt.Errorf("marshal reboot nodes request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.server+slurmV0044RebootNodesPath, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("create reboot nodes request: %w", err) + } + if err := c.setHeaders(ctx, httpReq); err != nil { + return fmt.Errorf("set reboot nodes request headers: %w", err) + } + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return fmt.Errorf("post reboot nodes request: %w", err) + } + defer resp.Body.Close() + + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("read reboot nodes response: %w", err) + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return fmt.Errorf("reboot nodes request failed: status=%d %s", resp.StatusCode, summarizeSlurmRESTBody(responseBody)) + } + if slurmRESTBodyHasErrors(responseBody) { + return fmt.Errorf("reboot nodes responded with errors: %s", summarizeSlurmRESTBody(responseBody)) + } + + return nil +} + +func slurmRESTBodyHasErrors(body []byte) bool { + var parsed struct { + Errors []struct{} `json:"errors"` + } + return json.Unmarshal(body, &parsed) == nil && len(parsed.Errors) > 0 +} + func (c *client) ListNodes(ctx context.Context) ([]Node, error) { getNodesResp, err := c.SlurmV0041GetNodesWithResponse(ctx, &api.SlurmV0041GetNodesParams{}) if err != nil { diff --git a/internal/slurmapi/fake/mock_client.go b/internal/slurmapi/fake/mock_client.go index ca7d63d94..f01c3fbea 100644 --- a/internal/slurmapi/fake/mock_client.go +++ b/internal/slurmapi/fake/mock_client.go @@ -480,6 +480,53 @@ func (_c *MockClient_PostMaintenanceReservation_Call) RunAndReturn(run func(cont return _c } +// RebootNodes provides a mock function with given fields: ctx, request +func (_m *MockClient) RebootNodes(ctx context.Context, request slurmapi.RebootNodesRequest) error { + ret := _m.Called(ctx, request) + + if len(ret) == 0 { + panic("no return value specified for RebootNodes") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, slurmapi.RebootNodesRequest) error); ok { + r0 = rf(ctx, request) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockClient_RebootNodes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RebootNodes' +type MockClient_RebootNodes_Call struct { + *mock.Call +} + +// RebootNodes is a helper method to define mock.On call +// - ctx context.Context +// - request slurmapi.RebootNodesRequest +func (_e *MockClient_Expecter) RebootNodes(ctx interface{}, request interface{}) *MockClient_RebootNodes_Call { + return &MockClient_RebootNodes_Call{Call: _e.mock.On("RebootNodes", ctx, request)} +} + +func (_c *MockClient_RebootNodes_Call) Run(run func(ctx context.Context, request slurmapi.RebootNodesRequest)) *MockClient_RebootNodes_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(slurmapi.RebootNodesRequest)) + }) + return _c +} + +func (_c *MockClient_RebootNodes_Call) Return(_a0 error) *MockClient_RebootNodes_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockClient_RebootNodes_Call) RunAndReturn(run func(context.Context, slurmapi.RebootNodesRequest) error) *MockClient_RebootNodes_Call { + _c.Call.Return(run) + return _c +} + // SlurmV0041DeleteJobWithResponse provides a mock function with given fields: ctx, jobId, params, reqEditors func (_m *MockClient) SlurmV0041DeleteJobWithResponse(ctx context.Context, jobId string, params *v0041.SlurmV0041DeleteJobParams, reqEditors ...v0041.RequestEditorFn) (*v0041.SlurmV0041DeleteJobResponse, error) { _va := make([]interface{}, len(reqEditors)) diff --git a/internal/slurmapi/interface.go b/internal/slurmapi/interface.go index c6b599bbb..8b2953b54 100644 --- a/internal/slurmapi/interface.go +++ b/internal/slurmapi/interface.go @@ -6,10 +6,12 @@ import ( api "github.com/SlinkyProject/slurm-client/api/v0041" ) +// nolint:interfacebloat type Client interface { api.ClientWithResponsesInterface ListNodes(ctx context.Context) ([]Node, error) GetNode(ctx context.Context, nodeName string) (Node, error) + RebootNodes(ctx context.Context, request RebootNodesRequest) error GetJobsByIDFromAccounting(ctx context.Context, jobID string) ([]Job, error) ListJobs(ctx context.Context) ([]Job, error) ListJobsWithParams(ctx context.Context, params ListJobsParams) ([]Job, error) diff --git a/internal/slurmapi/node.go b/internal/slurmapi/node.go index eb2140548..868078f14 100644 --- a/internal/slurmapi/node.go +++ b/internal/slurmapi/node.go @@ -165,6 +165,16 @@ func (n *Node) IsInvalidState() bool { return exists } +func (n *Node) IsRebootIssuedState() bool { + _, exists := n.States[api.V0041NodeStateREBOOTISSUED] + return exists +} + +func (n *Node) IsRebootRequestedState() bool { + _, exists := n.States[api.V0041NodeStateREBOOTREQUESTED] + 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/slurmapi/reboot_test.go b/internal/slurmapi/reboot_test.go new file mode 100644 index 000000000..b9d01e81d --- /dev/null +++ b/internal/slurmapi/reboot_test.go @@ -0,0 +1,100 @@ +package slurmapi + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type staticTokenIssuer string + +func (s staticTokenIssuer) Issue(context.Context) (string, error) { + return string(s), nil +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func jsonResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Body: io.NopCloser(strings.NewReader(body)), + Header: make(http.Header), + } +} + +func TestRebootNodesPostsV0044PayloadAndHeaders(t *testing.T) { + var gotBody map[string]any + httpClient := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, slurmV0044RebootNodesPath, r.URL.Path) + assert.Equal(t, headerApplicationJson, r.Header.Get(headerContentType)) + assert.Equal(t, "token-value", r.Header.Get(headerSlurmUserToken)) + + require.NoError(t, json.NewDecoder(r.Body).Decode(&gotBody)) + return jsonResponse(http.StatusOK, `{"errors":[]}`), nil + })} + + client, err := NewClient("http://slurmrestd", staticTokenIssuer("token-value"), httpClient) + require.NoError(t, err) + + err = client.RebootNodes(context.Background(), RebootNodesRequest{ + NodeList: "worker-[1-2]", + ASAP: true, + NextState: "RESUME", + Reason: "rolling update", + }) + require.NoError(t, err) + + assert.Equal(t, "worker-[1-2]", gotBody["nodes"]) + assert.Equal(t, true, gotBody["asap"]) + assert.Equal(t, "RESUME", gotBody["next_state"]) + assert.Equal(t, "rolling update", gotBody["reason"]) + assert.NotContains(t, gotBody, "force") + assert.NotContains(t, gotBody, "power_action") +} + +func TestRebootNodesRejectsEmptyNodeList(t *testing.T) { + client, err := NewClient("http://slurmrestd", nil, nil) + require.NoError(t, err) + + err = client.RebootNodes(context.Background(), RebootNodesRequest{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "nodes field is required") +} + +func TestRebootNodesStatusErrorSummarizesSlurmEnvelope(t *testing.T) { + httpClient := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return jsonResponse(http.StatusBadRequest, `{"errors":[{"description":"Invalid next_state"}]}`), nil + })} + + client, err := NewClient("http://slurmrestd", nil, httpClient) + require.NoError(t, err) + + err = client.RebootNodes(context.Background(), RebootNodesRequest{NodeList: "worker-1"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "status=400") + assert.Contains(t, err.Error(), "Invalid next_state") +} + +func TestRebootNodesDetectsSlurmErrorsOnOK(t *testing.T) { + httpClient := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return jsonResponse(http.StatusOK, `{"errors":[{"error":"ESLURM_INVALID_NODE_STATE"}]}`), nil + })} + + client, err := NewClient("http://slurmrestd", nil, httpClient) + require.NoError(t, err) + + err = client.RebootNodes(context.Background(), RebootNodesRequest{NodeList: "worker-1"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "ESLURM_INVALID_NODE_STATE") +} diff --git a/internal/values/slurm_nodeset.go b/internal/values/slurm_nodeset.go index ded9dfe69..1aca0841f 100644 --- a/internal/values/slurm_nodeset.go +++ b/internal/values/slurm_nodeset.go @@ -67,6 +67,8 @@ type SlurmNodeSet struct { TopologyFabric string ActiveNodes []int32 + + UpdateStrategy consts.UpdateStrategy } func BuildSlurmNodeSetFrom( @@ -145,6 +147,8 @@ func BuildSlurmNodeSetFrom( EphemeralNodes: nsSpec.EphemeralNodes, EphemeralTopologyWaitTimeout: nsSpec.EphemeralTopologyWaitTimeout, TopologyFabric: nsSpec.Topology.Fabric, + // + UpdateStrategy: nsSpec.UpdateStrategy, } // region Submounts