From 6cf180f20ae3e2142b50ca82d38e019134fa2f51 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Thu, 28 May 2026 18:42:15 +0200 Subject: [PATCH 01/10] poc seamless worker update --- Makefile | 2 +- PROJECT | 14 + api/v1alpha1/nodeset_types.go | 8 + cmd/main.go | 35 ++ cmd/soperatorchecks/main.go | 1 - .../crd/bases/slurm.nebius.ai_nodesets.yaml | 8 + config/crd/kustomization.yaml | 1 + config/rbac/kustomization.yaml | 3 + config/samples/kustomization.yaml | 1 + .../templates/slurmcluster-crd.yaml | 8 + helm/soperator/crds/slurmcluster-crd.yaml | 8 + images/common/scripts/reboot.sh | 30 +- internal/consts/label.go | 3 + internal/consts/statefulset.go | 7 + .../controller/clustercontroller/common.go | 19 -- .../nodesetcontroller/controller.go | 6 + .../controller/nodesetcontroller/reconcile.go | 57 ++++ .../maintenance_condition_test.go | 63 ---- .../slurm_api_clients_controller.go | 26 +- internal/controller/suite_test.go | 4 + .../statefulset_controller.go | 308 ++++++++++++++++++ internal/naming/naming.go | 12 +- internal/render/worker/role.go | 22 +- internal/render/worker/role_test.go | 29 -- internal/render/worker/rolebinding.go | 34 ++ internal/render/worker/serviceaccount.go | 6 +- internal/render/worker/sssd_test.go | 1 + internal/render/worker/statefulset.go | 52 ++- internal/render/worker/statefulset_test.go | 5 + internal/slurmapi/node.go | 10 + internal/values/slurm_nodeset.go | 4 + 31 files changed, 626 insertions(+), 161 deletions(-) create mode 100644 internal/controller/updatecontroller/statefulset_controller.go delete mode 100644 internal/render/worker/role_test.go create mode 100644 internal/render/worker/rolebinding.go 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/PROJECT b/PROJECT index c04b6a3a1..7621116f7 100644 --- a/PROJECT +++ b/PROJECT @@ -68,4 +68,18 @@ resources: kind: JailedConfig path: nebius.ai/slurm-operator/api/v1alpha1 version: v1alpha1 +- controller: true + core: true + group: apps + kind: StatefulSet + path: k8s.io/api/apps/v1beta1 + version: v1beta1 +- api: + crdVersion: v1 + namespaced: true + domain: nebius.ai + group: slurm + kind: RollingUpdateState + path: nebius.ai/slurm-operator/api/v1alpha1 + version: v1alpha1 version: "3" 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..9254d2ed8 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -27,6 +27,7 @@ import ( "go.uber.org/zap/zapcore" // Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.) // to ensure that exec-entrypoint and run can make use of them. + batchv1 "k8s.io/api/batch/v1" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" clientgoscheme "k8s.io/client-go/kubernetes/scheme" @@ -34,6 +35,7 @@ import ( "k8s.io/klog/v2" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/cache" + "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/healthz" "sigs.k8s.io/controller-runtime/pkg/log/zap" "sigs.k8s.io/controller-runtime/pkg/webhook" @@ -52,9 +54,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 @@ -244,6 +249,11 @@ func main() { Cache: cache.Options{ DefaultNamespaces: watchNsCacheByName, }, + Client: client.Options{ + Cache: &client.CacheOptions{ + DisableFor: []client.Object{&batchv1.Job{}, &batchv1.CronJob{}}, + }, + }, }) if err != nil { cli.Fail(setupLog, err, "unable to start manager") @@ -307,6 +317,31 @@ func main() { } // endregion Reconciler/NodeSet + { + ctx := ctrl.SetupSignalHandler() + slurmAPIClients := slurmapi.NewClientSet(ctx) + + 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("kekus"), + slurmAPIClients, + ). + SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { + cli.Fail(setupLog, err, "unable to create controller", "controller", "kek") + } + + } + // 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/config/crd/kustomization.yaml b/config/crd/kustomization.yaml index 9a504cc6d..48f73693a 100644 --- a/config/crd/kustomization.yaml +++ b/config/crd/kustomization.yaml @@ -8,6 +8,7 @@ resources: - bases/slurm.nebius.ai_slurmclusters.yaml - bases/slurm.nebius.ai_jailedconfigs.yaml +- bases/slurm.nebius.ai_rollingupdatestates.yaml #+kubebuilder:scaffold:crdkustomizeresource patches: [] diff --git a/config/rbac/kustomization.yaml b/config/rbac/kustomization.yaml index 1b7f9e4dc..3efcaf26a 100644 --- a/config/rbac/kustomization.yaml +++ b/config/rbac/kustomization.yaml @@ -24,6 +24,9 @@ resources: # default, aiding admins in cluster management. Those roles are # not used by the {{ .ProjectName }} itself. You can comment the following lines # if you do not want those helpers be installed with your Project. +- rollingupdatestate_admin_role.yaml +- rollingupdatestate_editor_role.yaml +- rollingupdatestate_viewer_role.yaml - activecheck_admin_role.yaml - activecheck_editor_role.yaml - activecheck_viewer_role.yaml diff --git a/config/samples/kustomization.yaml b/config/samples/kustomization.yaml index cc76b920b..2259f1f17 100644 --- a/config/samples/kustomization.yaml +++ b/config/samples/kustomization.yaml @@ -4,4 +4,5 @@ resources: - slurm_v1alpha1_nodeconfigurator.yaml - slurm_v1alpha1_nodeset.yaml - slurm_v1alpha1_jailedconfig.yaml +- slurm_v1alpha1_rollingupdatestate.yaml #+kubebuilder:scaffold:manifestskustomizesamples 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/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/images/common/scripts/reboot.sh b/images/common/scripts/reboot.sh index d21dab328..295be696f 100644 --- a/images/common/scripts/reboot.sh +++ b/images/common/scripts/reboot.sh @@ -17,9 +17,29 @@ do esac done -if ! mountpoint -q /run/nvidia/driver; then - echo "This command only works on GPU nodes" - exit 1 -fi +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" -chroot /run/nvidia/driver nsenter -t 1 -m -u -i -n /usr/sbin/reboot +if POD_JSON="$(curl -fsS --cacert "$CACERT" \ + -H "Authorization: Bearer $TOKEN" \ + "$POD_URL")" && [[ "$POD_JSON" == *'"kekus/todelete":"true"'* ]]; then + curl -fsS --cacert "$CACERT" \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/merge-patch+json" \ + -X PATCH \ + -d '{"metadata":{"labels":{"kekus/todelete":"false"}}}' \ + "$POD_URL" + curl -fsS --cacert "$CACERT" \ + -H "Authorization: Bearer $TOKEN" \ + -X DELETE \ + "$POD_URL" +else + 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 +fi diff --git a/internal/consts/label.go b/internal/consts/label.go index dccd66667..93366d230 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -42,4 +42,7 @@ const ( LabelJailedAggregationCommonValue = "common" AnnotationConfigHash = K8sGroupNameSoperator + "/config-hash" + + LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" + LabelSoperatorRollingUpdateValue = "true" ) 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/controller.go b/internal/controller/nodesetcontroller/controller.go index 0c6b706f5..682fb7f8a 100644 --- a/internal/controller/nodesetcontroller/controller.go +++ b/internal/controller/nodesetcontroller/controller.go @@ -66,9 +66,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 +80,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), } } diff --git a/internal/controller/nodesetcontroller/reconcile.go b/internal/controller/nodesetcontroller/reconcile.go index d7eb07fc4..26e0a47ee 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, nodeSet) + 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/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/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go new file mode 100644 index 000000000..35cec8684 --- /dev/null +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -0,0 +1,308 @@ +/* +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" + "strings" + "time" + + kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" + batchv1 "k8s.io/api/batch/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/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 ( + labelToDelete = "kekus/todelete" +) + +const ( + defaultSTSReplicasCount = int32(1) +) + +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;update;delete +// +kubebuilder:rbac:groups=slurm.nebius.ai,resources=rollingupdatestate,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups=slurm.nebius.ai,resources=rollingupdatestate/status,verbs=get;update;patch + +// 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 { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + if client.IgnoreNotFound(err) == nil { + logger.Info("statefulset not found, might be deleted", "namespace", req.Namespace, "name", req.Name) + return ctrl.Result{}, nil + } + } + + 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") + + slurmNodesToReboot := []string{} + for _, pod := range outdatedPods { + if pod.Labels[labelToDelete] == "true" { + if err := r.Delete(ctx, &pod); err != nil { + return err + } + continue + } + + slurmClient, ok := r.slurmAPIClients.GetClient(types.NamespacedName{ + Namespace: sts.Namespace, + Name: clusterName, + }) + if !ok { + logger.Info("no slurm clients", "namespace", sts.Namespace, "clusterName", clusterName) + return fmt.Errorf("no slurm clients") + } + + slurmNode, err := slurmClient.GetNode(ctx, pod.Name) + if err != nil { + return err + } + + if slurmNode.IsRebootIssuedState() || slurmNode.IsRebootRequestedState() { + continue + } + + if pod.Labels[labelToDelete] != "True" { + if pod.Labels == nil { + pod.Labels = map[string]string{} + } + pod.Labels[labelToDelete] = "True" + if err := r.Update(ctx, &pod); err != nil { + return err + } + } + + slurmNodesToReboot = append(slurmNodesToReboot, slurmNode.Name) + } + + rebooterCronJob := &batchv1.CronJob{} + err := r.Get(ctx, types.NamespacedName{ + Namespace: sts.Namespace, + Name: "rebooter", + }, rebooterCronJob) + if err != nil { + return err + } + + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: rebooterCronJob.Namespace, + GenerateName: rebooterCronJob.Name + "-", + Labels: rebooterCronJob.Spec.JobTemplate.Labels, + Annotations: rebooterCronJob.Spec.JobTemplate.Annotations, + OwnerReferences: []metav1.OwnerReference{ + *metav1.NewControllerRef(rebooterCronJob, batchv1.SchemeGroupVersion.WithKind("CronJob")), + }, + }, + Spec: *rebooterCronJob.Spec.JobTemplate.Spec.DeepCopy(), + } + appendSlurmNodeEnv(job, slurmNodesToReboot) + + if err := r.Create(ctx, job); err != nil { + return err + } + if err := r.waitForJobCompletion(ctx, client.ObjectKeyFromObject(job)); err != nil { + return err + } + + logger.Info("oki poki") + + return nil +} + +// 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 true }, + // TODO: update only in case of label added + UpdateFunc: func(tue event.TypedUpdateEvent[client.Object]) bool { return false }, + DeleteFunc: func(tde event.TypedDeleteEvent[client.Object]) bool { return false }, + GenericFunc: func(tge event.TypedGenericEvent[client.Object]) bool { return false }, + })). + Named("kekus"). + WithOptions(controllerconfig.ControllerOptions(maxConcurrency, cacheSyncTimeout)) + + return controllerBuilder.Complete(r) +} + +func appendSlurmNodeEnv(job *batchv1.Job, slurmNodes []string) { + + envVar := corev1.EnvVar{Name: "SLURM_NODES", Value: strings.Join(slurmNodes, ",")} + + for i := range job.Spec.Template.Spec.InitContainers { + job.Spec.Template.Spec.InitContainers[i].Env = append(job.Spec.Template.Spec.InitContainers[i].Env, envVar) + } + for i := range job.Spec.Template.Spec.Containers { + job.Spec.Template.Spec.Containers[i].Env = append(job.Spec.Template.Spec.Containers[i].Env, envVar) + } +} + +func (r *RollingUpdateReconciler) waitForJobCompletion(ctx context.Context, key client.ObjectKey) error { + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + + for { + job := &batchv1.Job{} + if err := r.Get(ctx, key, job); err != nil { + return err + } + + for _, condition := range job.Status.Conditions { + if condition.Status != corev1.ConditionTrue { + continue + } + switch condition.Type { + case batchv1.JobComplete, batchv1.JobSuccessCriteriaMet: + return nil + case batchv1.JobFailed, batchv1.JobFailureTarget: + return fmt.Errorf("job %s/%s failed: %s", job.Namespace, job.Name, condition.Message) + } + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + } +} 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/worker/role.go b/internal/render/worker/role.go index 15b3cc6d6..65ea425fd 100644 --- a/internal/render/worker/role.go +++ b/internal/render/worker/role.go @@ -1,27 +1,37 @@ package worker import ( + "fmt" + rbacv1 "k8s.io/api/rbac/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + slurmv1alpha1 "nebius.ai/slurm-operator/api/v1alpha1" "nebius.ai/slurm-operator/internal/consts" "nebius.ai/slurm-operator/internal/naming" "nebius.ai/slurm-operator/internal/render/common" ) -func RenderRole(namespace, clusterName string) rbacv1.Role { - labels := common.RenderLabels(consts.ComponentTypeWorker, clusterName) +func RenderRole(namespace, clusterName string, nodeSet *slurmv1alpha1.NodeSet) rbacv1.Role { + labels := common.RenderLabels(consts.ComponentTypeNodeSet, clusterName) + + resourceNames := make([]string, 0, nodeSet.Spec.Replicas) + for i := range nodeSet.Spec.Replicas { + 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", "update"}, }, }, } diff --git a/internal/render/worker/role_test.go b/internal/render/worker/role_test.go deleted file mode 100644 index c16ed167d..000000000 --- a/internal/render/worker/role_test.go +++ /dev/null @@ -1,29 +0,0 @@ -package worker - -import ( - "testing" - - "nebius.ai/slurm-operator/internal/naming" -) - -func Test_RenderRole(t *testing.T) { - namespace := "test-namespace" - clusterName := "test-cluster" - - role := RenderRole(namespace, clusterName) - - // Check the name - if role.Name != naming.BuildRoleWorkerName(clusterName) { - t.Errorf("Unexpected name: got %v, want %v", role.Name, naming.BuildRoleWorkerName(clusterName)) - } - - // Check the namespace - if role.Namespace != namespace { - t.Errorf("Unexpected namespace: got %v, want %v", role.Namespace, namespace) - } - - // 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) - } -} 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..c1bad1590 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, @@ -154,6 +155,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 +175,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 +185,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 +198,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..4f7a04da2 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, } } @@ -722,6 +726,7 @@ func TestRenderNodeSetStatefulSet_EphemeralNodesReserveOrdinals(t *testing.T) { GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, EphemeralNodes: &ephemeralNodes, ActiveNodes: activeNodes, + UpdateStrategy: consts.UpdateStrategyRollingUpdate, } } 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/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 From 71fe789687da053402b4b898f6e5bd859d176b7f Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Tue, 16 Jun 2026 18:27:48 +0200 Subject: [PATCH 02/10] add controller proxy --- PROJECT | 14 - cmd/main.go | 24 +- cmd/slurmcontrollerproxy/main.go | 49 ++++ cmd/soperatorchecks/main.go | 3 + config/crd/kustomization.yaml | 1 - config/rbac/kustomization.yaml | 3 - config/samples/kustomization.yaml | 1 - images/common/scripts/reboot.sh | 5 +- images/controller/slurmctld.dockerfile | 18 +- internal/consts/annotation.go | 2 + internal/consts/container.go | 37 +-- internal/consts/label.go | 1 + .../controller/nodesetcontroller/reconcile.go | 2 +- .../slurm_api_clients_controller.go | 30 +- .../statefulset_controller.go | 245 +++++++++------- internal/render/controller/container.go | 64 +++++ internal/render/controller/service.go | 20 +- internal/render/controller/statefulset.go | 3 +- .../render/controller/statefulset_test.go | 15 +- internal/render/worker/role.go | 23 +- internal/render/worker/role_test.go | 57 ++++ internal/render/worker/statefulset.go | 5 +- internal/render/worker/statefulset_test.go | 7 + internal/slurmproxy/client.go | 87 ++++++ internal/slurmproxy/client_set.go | 34 +++ internal/slurmproxy/server.go | 264 ++++++++++++++++++ internal/slurmproxy/server_test.go | 110 ++++++++ internal/slurmproxy/types.go | 33 +++ 28 files changed, 983 insertions(+), 174 deletions(-) create mode 100644 cmd/slurmcontrollerproxy/main.go create mode 100644 internal/render/worker/role_test.go create mode 100644 internal/slurmproxy/client.go create mode 100644 internal/slurmproxy/client_set.go create mode 100644 internal/slurmproxy/server.go create mode 100644 internal/slurmproxy/server_test.go create mode 100644 internal/slurmproxy/types.go diff --git a/PROJECT b/PROJECT index 7621116f7..c04b6a3a1 100644 --- a/PROJECT +++ b/PROJECT @@ -68,18 +68,4 @@ resources: kind: JailedConfig path: nebius.ai/slurm-operator/api/v1alpha1 version: v1alpha1 -- controller: true - core: true - group: apps - kind: StatefulSet - path: k8s.io/api/apps/v1beta1 - version: v1beta1 -- api: - crdVersion: v1 - namespaced: true - domain: nebius.ai - group: slurm - kind: RollingUpdateState - path: nebius.ai/slurm-operator/api/v1alpha1 - version: v1alpha1 version: "3" diff --git a/cmd/main.go b/cmd/main.go index 9254d2ed8..83a48806d 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" @@ -27,7 +28,6 @@ import ( "go.uber.org/zap/zapcore" // Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.) // to ensure that exec-entrypoint and run can make use of them. - batchv1 "k8s.io/api/batch/v1" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" clientgoscheme "k8s.io/client-go/kubernetes/scheme" @@ -35,7 +35,6 @@ import ( "k8s.io/klog/v2" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/cache" - "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/healthz" "sigs.k8s.io/controller-runtime/pkg/log/zap" "sigs.k8s.io/controller-runtime/pkg/webhook" @@ -60,6 +59,7 @@ import ( "nebius.ai/slurm-operator/internal/controllersenabled" metricsopts "nebius.ai/slurm-operator/internal/metrics" "nebius.ai/slurm-operator/internal/slurmapi" + "nebius.ai/slurm-operator/internal/slurmproxy" webhookv1 "nebius.ai/slurm-operator/internal/webhook/v1" webhookv1alpha1 "nebius.ai/slurm-operator/internal/webhook/v1alpha1" //+kubebuilder:scaffold:imports @@ -185,7 +185,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, @@ -249,11 +249,6 @@ func main() { Cache: cache.Options{ DefaultNamespaces: watchNsCacheByName, }, - Client: client.Options{ - Cache: &client.CacheOptions{ - DisableFor: []client.Object{&batchv1.Job{}, &batchv1.CronJob{}}, - }, - }, }) if err != nil { cli.Fail(setupLog, err, "unable to start manager") @@ -317,15 +312,16 @@ func main() { } // endregion Reconciler/NodeSet - { - ctx := ctrl.SetupSignalHandler() - slurmAPIClients := slurmapi.NewClientSet(ctx) + slurmAPIClients := slurmapi.NewClientSet(context.Background()) + slurmProxyClients := slurmproxy.NewClientSet() + if controllersSet.Enabled("rollingupdate") { if err = soperatorchecks.NewSlurmAPIClientsController( mgr.GetClient(), mgr.GetScheme(), mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName), slurmAPIClients, + slurmProxyClients, ).SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { cli.Fail(setupLog, err, "unable to create slurm api clients controller", "controller", soperatorchecks.SlurmAPIClientsControllerName) } @@ -333,13 +329,13 @@ func main() { if err = updatecontroller.NewRollingUpdateReconciler( mgr.GetClient(), mgr.GetScheme(), - mgr.GetEventRecorderFor("kekus"), + mgr.GetEventRecorderFor(updatecontroller.RollingUpdateControllerName), slurmAPIClients, + slurmProxyClients, ). SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { - cli.Fail(setupLog, err, "unable to create controller", "controller", "kek") + cli.Fail(setupLog, err, "unable to create controller", "controller", updatecontroller.RollingUpdateControllerName) } - } // region Reconciler/Topology diff --git a/cmd/slurmcontrollerproxy/main.go b/cmd/slurmcontrollerproxy/main.go new file mode 100644 index 000000000..24c44e541 --- /dev/null +++ b/cmd/slurmcontrollerproxy/main.go @@ -0,0 +1,49 @@ +package main + +import ( + "flag" + "log" + "net/http" + "os" + "time" + + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/slurmproxy" +) + +func main() { + var ( + listenAddress string + jwtKeyFile string + allowedUsers string + scontrolPath string + commandTimeout time.Duration + ) + + flag.StringVar(&listenAddress, "listen", slurmproxy.DefaultListenAddress, "HTTP listen address") + flag.StringVar(&jwtKeyFile, "jwt-key-file", consts.VolumeMountPathRESTJWTKey+"/"+consts.SecretRESTJWTKeyFileName, "JWT signing key file") + flag.StringVar(&allowedUsers, "allowed-users", slurmproxy.DefaultAllowedUsers, "Comma-separated Slurm usernames allowed to call privileged RPCs") + flag.StringVar(&scontrolPath, "scontrol-path", "/usr/bin/scontrol", "Path to scontrol") + flag.DurationVar(&commandTimeout, "command-timeout", 30*time.Second, "Timeout for each privileged Slurm command") + flag.Parse() + + jwtKey, err := os.ReadFile(jwtKeyFile) + if err != nil { + log.Fatalf("read jwt key file: %v", err) + } + + server, err := slurmproxy.NewServer(slurmproxy.ServerOptions{ + JWTKey: jwtKey, + AllowedUsers: allowedUsers, + ScontrolPath: scontrolPath, + CommandTimeout: commandTimeout, + }) + if err != nil { + log.Fatalf("create server: %v", err) + } + + log.Printf("starting slurm controller proxy on %s", listenAddress) + if err := http.ListenAndServe(listenAddress, server.Handler()); err != nil { + log.Fatalf("serve: %v", err) + } +} diff --git a/cmd/soperatorchecks/main.go b/cmd/soperatorchecks/main.go index 4332d10b7..e5a832c7a 100644 --- a/cmd/soperatorchecks/main.go +++ b/cmd/soperatorchecks/main.go @@ -50,6 +50,7 @@ import ( "nebius.ai/slurm-operator/internal/controllersenabled" metricsopts "nebius.ai/slurm-operator/internal/metrics" "nebius.ai/slurm-operator/internal/slurmapi" + "nebius.ai/slurm-operator/internal/slurmproxy" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" //+kubebuilder:scaffold:imports @@ -275,6 +276,7 @@ func main() { } slurmAPIClients := slurmapi.NewClientSet(ctx) + slurmProxyClients := slurmproxy.NewClientSet() if controllersSet.Enabled("slurmapiclients") { if err = soperatorchecks.NewSlurmAPIClientsController( @@ -282,6 +284,7 @@ func main() { mgr.GetScheme(), mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName), slurmAPIClients, + slurmProxyClients, ).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/kustomization.yaml b/config/crd/kustomization.yaml index 48f73693a..9a504cc6d 100644 --- a/config/crd/kustomization.yaml +++ b/config/crd/kustomization.yaml @@ -8,7 +8,6 @@ resources: - bases/slurm.nebius.ai_slurmclusters.yaml - bases/slurm.nebius.ai_jailedconfigs.yaml -- bases/slurm.nebius.ai_rollingupdatestates.yaml #+kubebuilder:scaffold:crdkustomizeresource patches: [] diff --git a/config/rbac/kustomization.yaml b/config/rbac/kustomization.yaml index 3efcaf26a..1b7f9e4dc 100644 --- a/config/rbac/kustomization.yaml +++ b/config/rbac/kustomization.yaml @@ -24,9 +24,6 @@ resources: # default, aiding admins in cluster management. Those roles are # not used by the {{ .ProjectName }} itself. You can comment the following lines # if you do not want those helpers be installed with your Project. -- rollingupdatestate_admin_role.yaml -- rollingupdatestate_editor_role.yaml -- rollingupdatestate_viewer_role.yaml - activecheck_admin_role.yaml - activecheck_editor_role.yaml - activecheck_viewer_role.yaml diff --git a/config/samples/kustomization.yaml b/config/samples/kustomization.yaml index 2259f1f17..cc76b920b 100644 --- a/config/samples/kustomization.yaml +++ b/config/samples/kustomization.yaml @@ -4,5 +4,4 @@ resources: - slurm_v1alpha1_nodeconfigurator.yaml - slurm_v1alpha1_nodeset.yaml - slurm_v1alpha1_jailedconfig.yaml -- slurm_v1alpha1_rollingupdatestate.yaml #+kubebuilder:scaffold:manifestskustomizesamples diff --git a/images/common/scripts/reboot.sh b/images/common/scripts/reboot.sh index 295be696f..901efe20f 100644 --- a/images/common/scripts/reboot.sh +++ b/images/common/scripts/reboot.sh @@ -22,15 +22,16 @@ 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_AFTER_REBOOT_LABEL="slurm.nebius.ai/delete-after-reboot" if POD_JSON="$(curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ - "$POD_URL")" && [[ "$POD_JSON" == *'"kekus/todelete":"true"'* ]]; then + "$POD_URL")" && [[ "$POD_JSON" == *"\"$DELETE_AFTER_REBOOT_LABEL\":\"true\""* ]]; then curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ -H "Content-Type: application/merge-patch+json" \ -X PATCH \ - -d '{"metadata":{"labels":{"kekus/todelete":"false"}}}' \ + -d "{\"metadata\":{\"labels\":{\"$DELETE_AFTER_REBOOT_LABEL\":\"false\"}}}" \ "$POD_URL" curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ diff --git a/images/controller/slurmctld.dockerfile b/images/controller/slurmctld.dockerfile index 43d0d077f..c102c2783 100644 --- a/images/controller/slurmctld.dockerfile +++ b/images/controller/slurmctld.dockerfile @@ -30,6 +30,20 @@ RUN --mount=type=cache,target=/root/.cache/go-build \ GOOS=$GOOS CGO_ENABLED=$CGO_ENABLED GO_LDFLAGS=$GO_LDFLAGS \ go build -v -o power-manager ./cmd/powermanager +FROM go-base AS slurmcontrollerproxy_builder + +ARG GO_LDFLAGS="" +ARG CGO_ENABLED=0 +ARG GOOS=linux + +COPY go.mod go.sum ./ +RUN go mod download + +COPY cmd/slurmcontrollerproxy cmd/slurmcontrollerproxy +RUN --mount=type=cache,target=/root/.cache/go-build \ + GOOS=$GOOS CGO_ENABLED=$CGO_ENABLED GO_LDFLAGS=$GO_LDFLAGS \ + go build -v -o slurm-controller-proxy ./cmd/slurmcontrollerproxy + # https://github.com/nebius/ml-containers/pull/81 FROM cr.eu-north1.nebius.cloud/ml-containers/slurm:${SLURM_VERSION}-20260421092341 AS controller_slurmctld @@ -54,6 +68,9 @@ RUN mkdir -p /var/log/slurm/multilog && \ COPY --from=powermanager_builder /build/power-manager /opt/soperator/bin/power-manager RUN chmod 755 /opt/soperator/bin/power-manager +COPY --from=slurmcontrollerproxy_builder /build/slurm-controller-proxy /opt/soperator/bin/slurm-controller-proxy +RUN chmod 755 /opt/soperator/bin/slurm-controller-proxy + # Copy power management scripts for Slurm ResumeProgram/SuspendProgram COPY images/controller/power_resume.sh /opt/soperator/bin/power_resume.sh COPY images/controller/power_suspend.sh /opt/soperator/bin/power_suspend.sh @@ -63,4 +80,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/internal/consts/annotation.go b/internal/consts/annotation.go index 12e58ba64..3691dd6b2 100644 --- a/internal/consts/annotation.go +++ b/internal/consts/annotation.go @@ -8,4 +8,6 @@ const ( AnnotationActiveCheckName = K8sGroupNameSoperator + "/activecheck" 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..c51cd2378 100644 --- a/internal/consts/container.go +++ b/internal/consts/container.go @@ -1,22 +1,23 @@ package consts const ( - ContainerNameSlurmctld = SlurmctldName - ContainerNameAccounting = AccountingName - ContainerNameMunge = Munge - ContainerNameSSSD = "sssd" - ContainerNameSlurmd = SlurmdName - ContainerNameREST = Slurmrestd - ContainerNameSshd = SshdName - ContainerNameWorkerInit = "worker-init" - ContainerNameDockerProxy = "docker-proxy" - ContainerNameWaitForDatabase = "wait-for-database" - ContainerNameWaitForAccounting = "wait-for-accounting" - ContainerNamePopulateJail = populateJail - ContainerNameExporter = Exporter - ContainerNameRebooter = "rebooter" - ContainerNameCustom = "custom-container" - ContainerNameSConfigController = SConfigControllerName + ContainerNameSlurmctld = SlurmctldName + ContainerNameSlurmControllerProxy = "slurm-controller-proxy" + ContainerNameAccounting = AccountingName + ContainerNameMunge = Munge + ContainerNameSSSD = "sssd" + ContainerNameSlurmd = SlurmdName + ContainerNameREST = Slurmrestd + ContainerNameSshd = SshdName + ContainerNameWorkerInit = "worker-init" + ContainerNameDockerProxy = "docker-proxy" + ContainerNameWaitForDatabase = "wait-for-database" + ContainerNameWaitForAccounting = "wait-for-accounting" + ContainerNamePopulateJail = populateJail + ContainerNameExporter = Exporter + ContainerNameRebooter = "rebooter" + ContainerNameCustom = "custom-container" + ContainerNameSConfigController = SConfigControllerName ContainerSecurityContextCapabilitySysAdmin = "SYS_ADMIN" @@ -27,4 +28,8 @@ const ( ContainerPortNameMonitoring = "monitoring" ContainerPortMonitoring = 8081 ContainerPathMonitoring = "/metrics" + ContainerPathHealthz = "/healthz" + + ContainerPortNameSlurmControllerProxy = "slurmctl-proxy" + ContainerPortSlurmControllerProxy = 6821 ) diff --git a/internal/consts/label.go b/internal/consts/label.go index 93366d230..21176bba5 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -45,4 +45,5 @@ const ( LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" LabelSoperatorRollingUpdateValue = "true" + LabelSoperatorDeleteAfterReboot = K8sGroupNameSoperator + "/delete-after-reboot" ) diff --git a/internal/controller/nodesetcontroller/reconcile.go b/internal/controller/nodesetcontroller/reconcile.go index 26e0a47ee..a67040040 100644 --- a/internal/controller/nodesetcontroller/reconcile.go +++ b/internal/controller/nodesetcontroller/reconcile.go @@ -347,7 +347,7 @@ func (r NodeSetReconciler) executeReconciliation( stepLogger := log.FromContext(stepCtx) stepLogger.V(1).Info("Reconciling") - desired := worker.RenderRole(nodeSet.Namespace, cluster.Name, nodeSet) + desired := worker.RenderRole(nodeSet.Namespace, cluster.Name, nodeSetValues) stepLogger = stepLogger.WithValues(logfield.ResourceKV(&desired)...) stepLogger.V(1).Info("Rendered") diff --git a/internal/controller/soperatorchecks/slurm_api_clients_controller.go b/internal/controller/soperatorchecks/slurm_api_clients_controller.go index a7434aa16..15f1fc776 100644 --- a/internal/controller/soperatorchecks/slurm_api_clients_controller.go +++ b/internal/controller/soperatorchecks/slurm_api_clients_controller.go @@ -21,6 +21,7 @@ import ( "nebius.ai/slurm-operator/internal/jwt" "nebius.ai/slurm-operator/internal/naming" "nebius.ai/slurm-operator/internal/slurmapi" + "nebius.ai/slurm-operator/internal/slurmproxy" ) var ( @@ -30,7 +31,8 @@ var ( type SlurmAPIClientsController struct { *reconciler.Reconciler - slurmAPIClients *slurmapi.ClientSet + slurmAPIClients *slurmapi.ClientSet + slurmProxyClients *slurmproxy.ClientSet } func NewSlurmAPIClientsController( @@ -38,12 +40,14 @@ func NewSlurmAPIClientsController( scheme *runtime.Scheme, recorder record.EventRecorder, slurmAPIClients *slurmapi.ClientSet, + slurmProxyClients *slurmproxy.ClientSet, ) *SlurmAPIClientsController { r := reconciler.NewReconciler(client, scheme, recorder) return &SlurmAPIClientsController{ - Reconciler: r, - slurmAPIClients: slurmAPIClients, + Reconciler: r, + slurmAPIClients: slurmAPIClients, + slurmProxyClients: slurmProxyClients, } } @@ -75,7 +79,8 @@ func (c *SlurmAPIClientsController) Reconcile(ctx context.Context, req ctrl.Requ logger.Info("Adding new slurm api client") - jwtToken := jwt.NewToken(c.Client).For(req.NamespacedName, "root").WithRegistry(jwt.NewTokenRegistry().Build()) + tokenRegistry := jwt.NewTokenRegistry().Build() + jwtToken := jwt.NewToken(c.Client).For(req.NamespacedName, "root").WithRegistry(tokenRegistry) slurmAPIServer := fmt.Sprintf("http://%s.%s:6820", naming.BuildServiceName(consts.ComponentTypeREST, req.Name), req.Namespace) slurmAPIClient, err := slurmapi.NewClient(slurmAPIServer, jwtToken, slurmapi.DefaultHTTPClient()) if err != nil { @@ -85,5 +90,22 @@ func (c *SlurmAPIClientsController) Reconcile(ctx context.Context, req ctrl.Requ c.slurmAPIClients.AddClient(req.NamespacedName, slurmAPIClient) + if c.slurmProxyClients != nil { + proxyToken := jwt.NewToken(c.Client).For(req.NamespacedName, slurmproxy.DefaultClientUser).WithRegistry(tokenRegistry) + proxyServer := fmt.Sprintf( + "http://%s.%s:%d", + naming.BuildServiceName(consts.ComponentTypeController, req.Name), + req.Namespace, + consts.ContainerPortSlurmControllerProxy, + ) + slurmProxyClient, err := slurmproxy.NewClient(proxyServer, proxyToken, slurmapi.DefaultHTTPClient()) + if err != nil { + logger.Error(err, "failed to create slurm controller proxy client") + return ctrl.Result{}, err + } + + c.slurmProxyClients.AddClient(req.NamespacedName, slurmProxyClient) + } + return ctrl.Result{}, nil } diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go index 35cec8684..2278b3a23 100644 --- a/internal/controller/updatecontroller/statefulset_controller.go +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -19,15 +19,15 @@ package updatecontroller import ( "context" "fmt" - "strings" + "sort" "time" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" - batchv1 "k8s.io/api/batch/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" @@ -40,10 +40,11 @@ import ( "nebius.ai/slurm-operator/internal/controller/reconciler" "nebius.ai/slurm-operator/internal/controllerconfig" "nebius.ai/slurm-operator/internal/slurmapi" + "nebius.ai/slurm-operator/internal/slurmproxy" ) const ( - labelToDelete = "kekus/todelete" + RollingUpdateControllerName = "rollingupdate" ) const ( @@ -53,27 +54,28 @@ const ( type RollingUpdateReconciler struct { *reconciler.Reconciler - slurmAPIClients *slurmapi.ClientSet + slurmAPIClients *slurmapi.ClientSet + slurmProxyClients *slurmproxy.ClientSet } func NewRollingUpdateReconciler( client client.Client, scheme *runtime.Scheme, recorder record.EventRecorder, slurmAPIClients *slurmapi.ClientSet, + slurmProxyClients *slurmproxy.ClientSet, ) *RollingUpdateReconciler { r := reconciler.NewReconciler(client, scheme, recorder) return &RollingUpdateReconciler{ - Reconciler: r, - slurmAPIClients: slurmAPIClients, + Reconciler: r, + slurmAPIClients: slurmAPIClients, + slurmProxyClients: slurmProxyClients, } } // +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;update;delete -// +kubebuilder:rbac:groups=slurm.nebius.ai,resources=rollingupdatestate,verbs=get;list;watch;create;update;patch;delete -// +kubebuilder:rbac:groups=slurm.nebius.ai,resources=rollingupdatestate/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;update;patch // Reconcile is part of the main kubernetes reconciliation loop which aims to // move the current state of the cluster closer to the desired state. @@ -87,13 +89,11 @@ func (r *RollingUpdateReconciler) Reconcile(ctx context.Context, req ctrl.Reques sts := &kruisev1b1.StatefulSet{} err := r.Get(ctx, req.NamespacedName, sts) if err != nil { - if client.IgnoreNotFound(err) != nil { - return ctrl.Result{}, client.IgnoreNotFound(err) - } 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() @@ -169,79 +169,154 @@ func (r *RollingUpdateReconciler) processRollingUpdate( ) error { logger := log.FromContext(ctx).WithName("rolling-update-reconciler") - slurmNodesToReboot := []string{} - for _, pod := range outdatedPods { - if pod.Labels[labelToDelete] == "true" { - if err := r.Delete(ctx, &pod); err != nil { - return err - } - continue - } + if len(outdatedPods) == 0 { + logger.Info("no outdated pods found", "namespace", sts.Namespace, "name", sts.Name) + return nil + } - slurmClient, ok := r.slurmAPIClients.GetClient(types.NamespacedName{ - Namespace: sts.Namespace, - Name: clusterName, - }) - if !ok { - logger.Info("no slurm clients", "namespace", sts.Namespace, "clusterName", clusterName) - return fmt.Errorf("no slurm clients") - } + 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) + } + + if r.slurmProxyClients == nil { + return fmt.Errorf("slurm controller proxy client set is not configured") + } + slurmProxyClient, ok := r.slurmProxyClients.GetClient(types.NamespacedName{ + Namespace: sts.Namespace, + Name: clusterName, + }) + if !ok { + logger.Info("no slurm controller proxy client", "namespace", sts.Namespace, "clusterName", clusterName) + return fmt.Errorf("no slurm controller proxy client for %s/%s", sts.Namespace, clusterName) + } + + sort.Slice(outdatedPods, func(i, j int) bool { + return outdatedPods[i].Name < outdatedPods[j].Name + }) + + type rebootCandidate struct { + pod corev1.Pod + slurmNode slurmapi.Node + } + candidates := make([]rebootCandidate, 0, len(outdatedPods)) + inFlightReady := 0 + for _, pod := range outdatedPods { slurmNode, err := slurmClient.GetNode(ctx, pod.Name) if err != nil { return err } if slurmNode.IsRebootIssuedState() || slurmNode.IsRebootRequestedState() { + if podReady(&pod) { + inFlightReady++ + } continue } - if pod.Labels[labelToDelete] != "True" { + 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.LabelSoperatorDeleteAfterReboot] != consts.LabelSoperatorRollingUpdateValue { + patchBase := pod.DeepCopy() if pod.Labels == nil { pod.Labels = map[string]string{} } - pod.Labels[labelToDelete] = "True" - if err := r.Update(ctx, &pod); err != nil { - return err + pod.Labels[consts.LabelSoperatorDeleteAfterReboot] = consts.LabelSoperatorRollingUpdateValue + 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, slurmNode.Name) + slurmNodesToReboot = append(slurmNodesToReboot, candidate.slurmNode.Name) } - rebooterCronJob := &batchv1.CronJob{} - err := r.Get(ctx, types.NamespacedName{ - Namespace: sts.Namespace, - Name: "rebooter", - }, rebooterCronJob) - if err != nil { - return err + if len(slurmNodesToReboot) == 0 { + logger.Info("all outdated pods already have reboot requested", "namespace", sts.Namespace, "name", sts.Name) + return nil } - job := &batchv1.Job{ - ObjectMeta: metav1.ObjectMeta{ - Namespace: rebooterCronJob.Namespace, - GenerateName: rebooterCronJob.Name + "-", - Labels: rebooterCronJob.Spec.JobTemplate.Labels, - Annotations: rebooterCronJob.Spec.JobTemplate.Annotations, - OwnerReferences: []metav1.OwnerReference{ - *metav1.NewControllerRef(rebooterCronJob, batchv1.SchemeGroupVersion.WithKind("CronJob")), - }, - }, - Spec: *rebooterCronJob.Spec.JobTemplate.Spec.DeepCopy(), + if err := slurmProxyClient.RebootNodes(ctx, slurmproxy.RebootNodesRequest{ + Nodes: slurmNodesToReboot, + Reason: slurmproxy.DefaultReason, + NextState: slurmproxy.RebootNextStateResume, + }); err != nil { + return fmt.Errorf("schedule slurm reboot through controller proxy: %w", err) } - appendSlurmNodeEnv(job, slurmNodesToReboot) - if err := r.Create(ctx, job); err != nil { - return err + logger.Info("scheduled slurm reboot through controller proxy", "nodes", slurmNodesToReboot) + + return nil +} + +func rebootBudget(sts *kruisev1b1.StatefulSet) int { + replicas := defaultSTSReplicasCount + if sts.Spec.Replicas != nil { + replicas = *sts.Spec.Replicas } - if err := r.waitForJobCompletion(ctx, client.ObjectKeyFromObject(job)); err != nil { - return err + if replicas <= 0 { + return 0 } - logger.Info("oki poki") + maxUnavailable := intstr.FromInt32(1) + if raw := sts.GetAnnotations()[consts.AnnotationSoperatorRollingUpdateMaxUnavailable]; raw != "" { + maxUnavailable = intstr.Parse(raw) + } - return nil + 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. @@ -253,56 +328,24 @@ func (r *RollingUpdateReconciler) SetupWithManager( controllerBuilder := ctrl.NewControllerManagedBy(mgr). For(&kruisev1b1.StatefulSet{}, builder.WithPredicates(predicate.Funcs{ - CreateFunc: func(tce event.TypedCreateEvent[client.Object]) bool { return true }, - // TODO: update only in case of label added - UpdateFunc: func(tue event.TypedUpdateEvent[client.Object]) bool { return false }, + 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("kekus"). + Named(RollingUpdateControllerName). WithOptions(controllerconfig.ControllerOptions(maxConcurrency, cacheSyncTimeout)) return controllerBuilder.Complete(r) } -func appendSlurmNodeEnv(job *batchv1.Job, slurmNodes []string) { - - envVar := corev1.EnvVar{Name: "SLURM_NODES", Value: strings.Join(slurmNodes, ",")} - - for i := range job.Spec.Template.Spec.InitContainers { - job.Spec.Template.Spec.InitContainers[i].Env = append(job.Spec.Template.Spec.InitContainers[i].Env, envVar) - } - for i := range job.Spec.Template.Spec.Containers { - job.Spec.Template.Spec.Containers[i].Env = append(job.Spec.Template.Spec.Containers[i].Env, envVar) - } -} - -func (r *RollingUpdateReconciler) waitForJobCompletion(ctx context.Context, key client.ObjectKey) error { - ticker := time.NewTicker(2 * time.Second) - defer ticker.Stop() - - for { - job := &batchv1.Job{} - if err := r.Get(ctx, key, job); err != nil { - return err - } - - for _, condition := range job.Status.Conditions { - if condition.Status != corev1.ConditionTrue { - continue - } - switch condition.Type { - case batchv1.JobComplete, batchv1.JobSuccessCriteriaMet: - return nil - case batchv1.JobFailed, batchv1.JobFailureTarget: - return fmt.Errorf("job %s/%s failed: %s", job.Namespace, job.Name, condition.Message) - } - } - - select { - case <-ctx.Done(): - return ctx.Err() - case <-ticker.C: - } +func rollingUpdateEnabled(obj client.Object) bool { + if obj == nil { + return false } + return obj.GetLabels()[consts.LabelSoperatorRollingUpdateEnabled] == consts.LabelSoperatorRollingUpdateValue } diff --git a/internal/render/controller/container.go b/internal/render/controller/container.go index 497ac7f45..53055309c 100644 --- a/internal/render/controller/container.go +++ b/internal/render/controller/container.go @@ -2,6 +2,7 @@ package controller import ( corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/util/intstr" slurmv1 "nebius.ai/slurm-operator/api/v1" "nebius.ai/slurm-operator/internal/utils/stringutils" @@ -52,6 +53,69 @@ func renderContainerSlurmctld(container *values.Container, customMounts []slurmv } } +func renderContainerSlurmControllerProxy(container *values.Container) corev1.Container { + return corev1.Container{ + Name: consts.ContainerNameSlurmControllerProxy, + Image: container.Image, + ImagePullPolicy: container.ImagePullPolicy, + Command: []string{ + "/bin/bash", "-c", + }, + Args: []string{ + // language=bash + stringutils.Dedent(` + set -eu + + ln -sf /mnt/jail/etc/passwd /etc/passwd + ln -sf /mnt/jail/etc/group /etc/group + ln -sf /mnt/jail/etc/shadow /etc/shadow + ln -sf /mnt/jail/etc/gshadow /etc/gshadow + chown -h 0:42 /etc/{shadow,gshadow} + + rm -rf /etc/slurm && ln -s /mnt/jail/etc/slurm /etc/slurm + + while [ ! -S "/run/munge/munge.socket.2" ]; do sleep 2; done + + exec /opt/soperator/bin/slurm-controller-proxy --listen=:6821 + `), + }, + Ports: []corev1.ContainerPort{{ + Name: consts.ContainerPortNameSlurmControllerProxy, + ContainerPort: consts.ContainerPortSlurmControllerProxy, + Protocol: corev1.ProtocolTCP, + }}, + VolumeMounts: []corev1.VolumeMount{ + common.RenderVolumeMountJail(), + common.RenderVolumeMountMungeSocket(), + common.RenderVolumeMountRESTJWTKey(), + }, + ReadinessProbe: &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{ + HTTPGet: &corev1.HTTPGetAction{ + Path: consts.ContainerPathHealthz, + Port: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), + }, + }, + PeriodSeconds: 5, + }, + LivenessProbe: &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{ + HTTPGet: &corev1.HTTPGetAction{ + Path: consts.ContainerPathHealthz, + Port: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), + }, + }, + InitialDelaySeconds: 30, + PeriodSeconds: 10, + }, + SecurityContext: &corev1.SecurityContext{ + AppArmorProfile: common.ParseAppArmorProfile(container.AppArmorProfile), + }, + TerminationMessagePath: corev1.TerminationMessagePathDefault, + TerminationMessagePolicy: corev1.TerminationMessageReadFile, + } +} + func renderContainerSlurmctldWithSSSD( container *values.Container, customMounts []slurmv1.NodeVolumeMount, diff --git a/internal/render/controller/service.go b/internal/render/controller/service.go index 7e3c36ab2..0d017f1c3 100644 --- a/internal/render/controller/service.go +++ b/internal/render/controller/service.go @@ -32,12 +32,20 @@ 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), + }, + { + Name: consts.ContainerPortNameSlurmControllerProxy, + Protocol: corev1.ProtocolTCP, + Port: consts.ContainerPortSlurmControllerProxy, + TargetPort: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), + }, + }, }, } } diff --git a/internal/render/controller/statefulset.go b/internal/render/controller/statefulset.go index 8d2b6c2aa..9400c5d41 100644 --- a/internal/render/controller/statefulset.go +++ b/internal/render/controller/statefulset.go @@ -69,6 +69,7 @@ func RenderStatefulSet( if controller.ContainerSSSD != nil { slurmctldContainer = renderContainerSlurmctldWithSSSD(&controller.ContainerSlurmctld, controller.CustomVolumeMounts) } + slurmControllerProxyContainer := renderContainerSlurmControllerProxy(&controller.ContainerSlurmctld) return kruisev1b1.StatefulSet{ ObjectMeta: metav1.ObjectMeta{ @@ -116,7 +117,7 @@ func RenderStatefulSet( NodeSelector: nodeFilter.NodeSelector, Tolerations: nodeFilter.Tolerations, InitContainers: initContainers, - Containers: []corev1.Container{slurmctldContainer}, + Containers: []corev1.Container{slurmctldContainer, slurmControllerProxyContainer}, Volumes: volumes, RestartPolicy: corev1.RestartPolicyAlways, TerminationGracePeriodSeconds: ptr.To(common.DefaultPodTerminationGracePeriodSeconds), diff --git a/internal/render/controller/statefulset_test.go b/internal/render/controller/statefulset_test.go index 30acba3ac..997613738 100644 --- a/internal/render/controller/statefulset_test.go +++ b/internal/render/controller/statefulset_test.go @@ -176,8 +176,8 @@ func TestRenderStatefulSet(t *testing.T) { // Check containers containers := result.Spec.Template.Spec.Containers - if len(containers) != 1 { - t.Errorf("Expected 1 container, got %d", len(containers)) + if len(containers) != 2 { + t.Errorf("Expected 2 containers, got %d", len(containers)) } else { container := containers[0] if container.Name != consts.ContainerNameSlurmctld { @@ -186,6 +186,14 @@ func TestRenderStatefulSet(t *testing.T) { if container.Image != tt.controller.ContainerSlurmctld.NodeContainer.Image { t.Errorf("Container image = %v, want %v", container.Image, tt.controller.ContainerSlurmctld.NodeContainer.Image) } + + proxyContainer := containers[1] + if proxyContainer.Name != consts.ContainerNameSlurmControllerProxy { + t.Errorf("Container name = %v, want %v", proxyContainer.Name, consts.ContainerNameSlurmControllerProxy) + } + if proxyContainer.Image != tt.controller.ContainerSlurmctld.NodeContainer.Image { + t.Errorf("Container image = %v, want %v", proxyContainer.Image, tt.controller.ContainerSlurmctld.NodeContainer.Image) + } } // Check init containers @@ -290,6 +298,9 @@ func TestRenderStatefulSet_SSSD(t *testing.T) { } assert.Contains(t, mountNames, consts.VolumeNameSSSDConf) assert.Contains(t, mountNames, consts.VolumeNameSSSDSocket) + + proxy := result.Spec.Template.Spec.Containers[1] + assert.Equal(t, consts.ContainerNameSlurmControllerProxy, proxy.Name) } func TestRenderStatefulSetWithMaintenance(t *testing.T) { diff --git a/internal/render/worker/role.go b/internal/render/worker/role.go index 65ea425fd..08386e431 100644 --- a/internal/render/worker/role.go +++ b/internal/render/worker/role.go @@ -6,17 +6,18 @@ import ( rbacv1 "k8s.io/api/rbac/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - slurmv1alpha1 "nebius.ai/slurm-operator/api/v1alpha1" "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, nodeSet *slurmv1alpha1.NodeSet) rbacv1.Role { +func RenderRole(namespace, clusterName string, nodeSet *values.SlurmNodeSet) rbacv1.Role { labels := common.RenderLabels(consts.ComponentTypeNodeSet, clusterName) - resourceNames := make([]string, 0, nodeSet.Spec.Replicas) - for i := range nodeSet.Spec.Replicas { + 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)) } @@ -31,8 +32,20 @@ func RenderRole(namespace, clusterName string, nodeSet *slurmv1alpha1.NodeSet) r APIGroups: []string{""}, Resources: []string{"pods"}, ResourceNames: resourceNames, - Verbs: []string{"get", "update"}, + Verbs: []string{"get", "patch", "update", "delete"}, }, }, } } + +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 new file mode 100644 index 000000000..478735b55 --- /dev/null +++ b/internal/render/worker/role_test.go @@ -0,0 +1,57 @@ +package worker_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/utils/ptr" + + "nebius.ai/slurm-operator/internal/consts" + "nebius.ai/slurm-operator/internal/render/worker" + "nebius.ai/slurm-operator/internal/values" +) + +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"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + role := worker.RenderRole("default", "soperator", tt.nodeSet) + + assert.Len(t, role.Rules, 1) + assert.Equal(t, []string{"get", "patch", "update", "delete"}, 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/statefulset.go b/internal/render/worker/statefulset.go index c1bad1590..efa7ee7d0 100644 --- a/internal/render/worker/statefulset.go +++ b/internal/render/worker/statefulset.go @@ -144,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 diff --git a/internal/render/worker/statefulset_test.go b/internal/render/worker/statefulset_test.go index 4f7a04da2..6e3c34b59 100644 --- a/internal/render/worker/statefulset_test.go +++ b/internal/render/worker/statefulset_test.go @@ -682,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", + ) }) } } diff --git a/internal/slurmproxy/client.go b/internal/slurmproxy/client.go new file mode 100644 index 000000000..f45aecc98 --- /dev/null +++ b/internal/slurmproxy/client.go @@ -0,0 +1,87 @@ +package slurmproxy + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +type tokenIssuer interface { + Issue(ctx context.Context) (string, error) +} + +type Interface interface { + RebootNodes(ctx context.Context, req RebootNodesRequest) error +} + +type Client struct { + server string + tokenIssuer tokenIssuer + httpClient *http.Client +} + +func NewClient(server string, tokenIssuer tokenIssuer, httpClient *http.Client) (*Client, error) { + if server == "" { + return nil, fmt.Errorf("unable to create slurm controller proxy client: empty server URL") + } + if httpClient == nil { + httpClient = http.DefaultClient + } + + return &Client{ + server: strings.TrimRight(server, "/"), + tokenIssuer: tokenIssuer, + httpClient: httpClient, + }, nil +} + +func (c *Client) RebootNodes(ctx context.Context, req RebootNodesRequest) error { + body, err := json.Marshal(req) + if err != nil { + return fmt.Errorf("marshal reboot nodes request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.server+EndpointRebootNodes, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("create reboot nodes request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + if c.tokenIssuer != nil { + token, err := c.tokenIssuer.Issue(ctx) + if err != nil { + return fmt.Errorf("issue controller proxy token: %w", err) + } + if token != "" { + httpReq.Header.Set("Authorization", "Bearer "+token) + } + } + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return fmt.Errorf("post reboot nodes request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return fmt.Errorf("reboot nodes request failed: status=%d: %s", resp.StatusCode, responseError(resp.Body)) + } + + return nil +} + +func responseError(body io.Reader) string { + data, err := io.ReadAll(body) + if err != nil { + return fmt.Sprintf("read response body: %v", err) + } + var apiErr ErrorResponse + if err := json.Unmarshal(data, &apiErr); err == nil && apiErr.Error != "" { + return apiErr.Error + } + return string(data) +} diff --git a/internal/slurmproxy/client_set.go b/internal/slurmproxy/client_set.go new file mode 100644 index 000000000..df1d4fe8a --- /dev/null +++ b/internal/slurmproxy/client_set.go @@ -0,0 +1,34 @@ +package slurmproxy + +import ( + "sync" + + "k8s.io/apimachinery/pkg/types" +) + +type ClientSet struct { + clients map[types.NamespacedName]Interface + mux *sync.Mutex +} + +func NewClientSet() *ClientSet { + return &ClientSet{ + clients: make(map[types.NamespacedName]Interface), + mux: &sync.Mutex{}, + } +} + +func (cs *ClientSet) AddClient(name types.NamespacedName, client Interface) { + cs.mux.Lock() + defer cs.mux.Unlock() + + cs.clients[name] = client +} + +func (cs *ClientSet) GetClient(name types.NamespacedName) (Interface, bool) { + cs.mux.Lock() + defer cs.mux.Unlock() + + client, found := cs.clients[name] + return client, found +} diff --git a/internal/slurmproxy/server.go b/internal/slurmproxy/server.go new file mode 100644 index 000000000..6aa0f528a --- /dev/null +++ b/internal/slurmproxy/server.go @@ -0,0 +1,264 @@ +package slurmproxy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os/exec" + "regexp" + "strings" + "time" + + jwtlib "github.com/golang-jwt/jwt/v5" + + internaljwt "nebius.ai/slurm-operator/internal/jwt" +) + +const ( + defaultScontrolPath = "/usr/bin/scontrol" + defaultCommandTimeout = 30 * time.Second + maxReasonLength = 512 +) + +var nodeNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]*$`) + +type CommandRunner interface { + Run(ctx context.Context, name string, args ...string) ([]byte, error) +} + +type ExecCommandRunner struct{} + +func (ExecCommandRunner) Run(ctx context.Context, name string, args ...string) ([]byte, error) { + cmd := exec.CommandContext(ctx, name, args...) + return cmd.CombinedOutput() +} + +type ServerOptions struct { + JWTKey []byte + AllowedUsers string + ScontrolPath string + CommandRunner CommandRunner + CommandTimeout time.Duration +} + +type Server struct { + jwtKey []byte + allowedUsers map[string]struct{} + allowAnyUser bool + scontrolPath string + commandRunner CommandRunner + commandTimeout time.Duration +} + +func NewServer(opts ServerOptions) (*Server, error) { + if len(opts.JWTKey) == 0 { + return nil, errors.New("jwt key is required") + } + + allowedUsers, allowAnyUser := parseAllowedUsers(opts.AllowedUsers) + if len(allowedUsers) == 0 && !allowAnyUser { + return nil, errors.New("at least one allowed user is required") + } + + if opts.ScontrolPath == "" { + opts.ScontrolPath = defaultScontrolPath + } + if opts.CommandRunner == nil { + opts.CommandRunner = ExecCommandRunner{} + } + if opts.CommandTimeout == 0 { + opts.CommandTimeout = defaultCommandTimeout + } + + return &Server{ + jwtKey: opts.JWTKey, + allowedUsers: allowedUsers, + allowAnyUser: allowAnyUser, + scontrolPath: opts.ScontrolPath, + commandRunner: opts.CommandRunner, + commandTimeout: opts.CommandTimeout, + }, nil +} + +func (s *Server) Handler() http.Handler { + mux := http.NewServeMux() + mux.HandleFunc(EndpointHealthz, s.handleHealthz) + mux.HandleFunc(EndpointRebootNodes, s.handleRebootNodes) + return mux +} + +func (s *Server) handleHealthz(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + writeError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + w.WriteHeader(http.StatusNoContent) +} + +func (s *Server) handleRebootNodes(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + writeError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + + if _, err := s.authorize(r); err != nil { + writeError(w, http.StatusUnauthorized, err.Error()) + return + } + + var req RebootNodesRequest + decoder := json.NewDecoder(r.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&req); err != nil { + writeError(w, http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err)) + return + } + if err := req.normalizeAndValidate(); err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + + commandCtx, cancel := context.WithTimeout(r.Context(), s.commandTimeout) + defer cancel() + + args := []string{"reboot"} + if req.Reason != "" { + args = append(args, "reason="+req.Reason) + } + args = append(args, strings.Join(req.Nodes, ",")) + + output, err := s.commandRunner.Run(commandCtx, s.scontrolPath, args...) + if err != nil { + status := http.StatusBadGateway + if errors.Is(commandCtx.Err(), context.DeadlineExceeded) { + status = http.StatusGatewayTimeout + } + writeError(w, status, strings.TrimSpace(fmt.Sprintf("%v: %s", err, output))) + return + } + + writeJSON(w, http.StatusOK, RebootNodesResponse{ + Nodes: req.Nodes, + Output: strings.TrimSpace(string(output)), + }) +} + +func (s *Server) authorize(r *http.Request) (string, error) { + raw := r.Header.Get("Authorization") + if raw == "" { + return "", errors.New("authorization header is required") + } + tokenString, ok := strings.CutPrefix(raw, "Bearer ") + if !ok || tokenString == "" { + return "", errors.New("bearer token is required") + } + + claims := &internaljwt.TokenClaims{} + token, err := jwtlib.ParseWithClaims( + tokenString, + claims, + func(token *jwtlib.Token) (interface{}, error) { + return s.jwtKey, nil + }, + jwtlib.WithValidMethods([]string{jwtlib.SigningMethodHS256.Alg()}), + jwtlib.WithIssuer("soperator"), + ) + if err != nil { + return "", fmt.Errorf("invalid token: %w", err) + } + if !token.Valid { + return "", errors.New("invalid token") + } + + user := claims.SlurmUsername + if user == "" { + user = claims.Subject + } + if user == "" { + return "", errors.New("token subject is empty") + } + if s.allowAnyUser { + return user, nil + } + if _, ok := s.allowedUsers[user]; !ok { + return "", fmt.Errorf("user %q is not allowed", user) + } + + return user, nil +} + +func (r *RebootNodesRequest) normalizeAndValidate() error { + if len(r.Nodes) == 0 { + return errors.New("nodes must not be empty") + } + + seen := make(map[string]struct{}, len(r.Nodes)) + nodes := make([]string, 0, len(r.Nodes)) + for _, node := range r.Nodes { + node = strings.TrimSpace(node) + if node == "" { + return errors.New("nodes must not contain empty values") + } + if !nodeNamePattern.MatchString(node) { + return fmt.Errorf("invalid node name %q", node) + } + if _, ok := seen[node]; ok { + continue + } + seen[node] = struct{}{} + nodes = append(nodes, node) + } + r.Nodes = nodes + + r.Reason = strings.TrimSpace(r.Reason) + if r.Reason == "" { + r.Reason = DefaultReason + } + if strings.ContainsAny(r.Reason, "\x00\r\n") { + return errors.New("reason must not contain control line breaks") + } + if len(r.Reason) > maxReasonLength { + return fmt.Errorf("reason is too long: got %d bytes, max %d", len(r.Reason), maxReasonLength) + } + + if r.NextState == "" { + r.NextState = RebootNextStateResume + } + switch r.NextState { + case RebootNextStateResume, RebootNextStateDown: + default: + return fmt.Errorf("unsupported nextState %q", r.NextState) + } + + return nil +} + +func parseAllowedUsers(raw string) (map[string]struct{}, bool) { + if raw == "" { + raw = DefaultAllowedUsers + } + res := map[string]struct{}{} + for _, part := range strings.Split(raw, ",") { + user := strings.TrimSpace(part) + if user == "" { + continue + } + if user == "*" { + return nil, true + } + res[user] = struct{}{} + } + return res, false +} + +func writeJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(v) +} + +func writeError(w http.ResponseWriter, status int, message string) { + writeJSON(w, status, ErrorResponse{Error: message}) +} diff --git a/internal/slurmproxy/server_test.go b/internal/slurmproxy/server_test.go new file mode 100644 index 000000000..9b925a378 --- /dev/null +++ b/internal/slurmproxy/server_test.go @@ -0,0 +1,110 @@ +package slurmproxy + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + jwtlib "github.com/golang-jwt/jwt/v5" + "github.com/stretchr/testify/require" + + internaljwt "nebius.ai/slurm-operator/internal/jwt" +) + +type recordingRunner struct { + name string + args []string +} + +func (r *recordingRunner) Run(_ context.Context, name string, args ...string) ([]byte, error) { + r.name = name + r.args = append([]string(nil), args...) + return []byte("scheduled"), nil +} + +func TestServer_RebootNodes(t *testing.T) { + key := []byte("12345678901234567890123456789012") + runner := &recordingRunner{} + server, err := NewServer(ServerOptions{ + JWTKey: key, + AllowedUsers: "root", + ScontrolPath: "/bin/scontrol", + CommandRunner: runner, + }) + require.NoError(t, err) + + body := `{"nodes":["worker-1","worker-1","worker-2"],"reason":"rolling update"}` + req := httptest.NewRequest(http.MethodPost, EndpointRebootNodes, stringsReader(body)) + req.Header.Set("Authorization", "Bearer "+signedToken(t, key, "root")) + resp := httptest.NewRecorder() + + server.Handler().ServeHTTP(resp, req) + + require.Equal(t, http.StatusOK, resp.Code) + require.Equal(t, "/bin/scontrol", runner.name) + require.Equal(t, []string{ + "reboot", + "reason=rolling update", + "worker-1,worker-2", + }, runner.args) + + var response RebootNodesResponse + require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &response)) + require.Equal(t, []string{"worker-1", "worker-2"}, response.Nodes) + require.Equal(t, "scheduled", response.Output) +} + +func TestServer_RebootNodesRejectsUnauthorizedUser(t *testing.T) { + key := []byte("12345678901234567890123456789012") + server, err := NewServer(ServerOptions{ + JWTKey: key, + AllowedUsers: "root", + }) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, EndpointRebootNodes, stringsReader(`{"nodes":["worker-1"]}`)) + req.Header.Set("Authorization", "Bearer "+signedToken(t, key, "alice")) + resp := httptest.NewRecorder() + + server.Handler().ServeHTTP(resp, req) + + require.Equal(t, http.StatusUnauthorized, resp.Code) +} + +func TestRebootNodesRequestValidate(t *testing.T) { + req := RebootNodesRequest{ + Nodes: []string{"worker-0"}, + } + + require.NoError(t, req.normalizeAndValidate()) + require.Equal(t, DefaultReason, req.Reason) + require.Equal(t, RebootNextStateResume, req.NextState) +} + +func signedToken(t *testing.T, key []byte, user string) string { + t.Helper() + + now := time.Now() + claims := internaljwt.TokenClaims{ + RegisteredClaims: jwtlib.RegisteredClaims{ + Issuer: "soperator", + Subject: user, + IssuedAt: jwtlib.NewNumericDate(now), + NotBefore: jwtlib.NewNumericDate(now.Add(-time.Minute)), + ExpiresAt: jwtlib.NewNumericDate(now.Add(time.Hour)), + }, + SlurmUsername: user, + } + + token, err := jwtlib.NewWithClaims(jwtlib.SigningMethodHS256, claims).SignedString(key) + require.NoError(t, err) + return token +} + +func stringsReader(s string) *strings.Reader { + return strings.NewReader(s) +} diff --git a/internal/slurmproxy/types.go b/internal/slurmproxy/types.go new file mode 100644 index 000000000..86929f6b4 --- /dev/null +++ b/internal/slurmproxy/types.go @@ -0,0 +1,33 @@ +package slurmproxy + +const ( + DefaultListenAddress = ":6821" + DefaultAllowedUsers = "root,soperatorchecks" + DefaultClientUser = "root" + DefaultReason = "soperator rolling update" + + EndpointHealthz = "/healthz" + EndpointRebootNodes = "/v1/nodes/reboot" +) + +type RebootNextState string + +const ( + RebootNextStateResume RebootNextState = "RESUME" + RebootNextStateDown RebootNextState = "DOWN" +) + +type RebootNodesRequest struct { + Nodes []string `json:"nodes"` + Reason string `json:"reason,omitempty"` + NextState RebootNextState `json:"nextState,omitempty"` +} + +type RebootNodesResponse struct { + Nodes []string `json:"nodes"` + Output string `json:"output,omitempty"` +} + +type ErrorResponse struct { + Error string `json:"error"` +} From 145dc67bb91956d8ef1451b11f20bbd759e2313f Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Wed, 17 Jun 2026 15:51:51 +0200 Subject: [PATCH 03/10] fix reboot program --- helm/nodesets/templates/nodeset.yaml | 2 ++ helm/nodesets/tests/custom_values_test.yaml | 7 +++++++ helm/nodesets/values.yaml | 4 ++++ helm/soperator/templates/_helpers.tpl | 2 +- helm/soperator/values.yaml | 1 + images/common/scripts/reboot.sh | 11 ++++++++--- internal/consts/label.go | 3 ++- .../updatecontroller/statefulset_controller.go | 9 ++++----- internal/slurmproxy/server.go | 9 --------- internal/slurmproxy/server_test.go | 1 - internal/slurmproxy/types.go | 12 ++---------- 11 files changed, 31 insertions(+), 30 deletions(-) diff --git a/helm/nodesets/templates/nodeset.yaml b/helm/nodesets/templates/nodeset.yaml index 0689ebfb3..4734cb1c1 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 }} 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..3a2eab2d8 100644 --- a/helm/nodesets/values.yaml +++ b/helm/nodesets/values.yaml @@ -50,6 +50,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/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 901efe20f..cc3072e45 100644 --- a/images/common/scripts/reboot.sh +++ b/images/common/scripts/reboot.sh @@ -22,16 +22,21 @@ 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_AFTER_REBOOT_LABEL="slurm.nebius.ai/delete-after-reboot" +DELETE_CANDIDATE_LABEL="slurm.nebius.ai/delete-candidate" +DELETE_CANDIDATE_LABEL_VALUE="true" if POD_JSON="$(curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ - "$POD_URL")" && [[ "$POD_JSON" == *"\"$DELETE_AFTER_REBOOT_LABEL\":\"true\""* ]]; then + "$POD_URL")" && jq -e \ + --arg label "$DELETE_CANDIDATE_LABEL" \ + --arg value "$DELETE_CANDIDATE_LABEL_VALUE" \ + '.metadata.labels[$label] == $value' \ + <<<"$POD_JSON" >/dev/null; then curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ -H "Content-Type: application/merge-patch+json" \ -X PATCH \ - -d "{\"metadata\":{\"labels\":{\"$DELETE_AFTER_REBOOT_LABEL\":\"false\"}}}" \ + -d "{\"metadata\":{\"labels\":{\"$DELETE_CANDIDATE_LABEL\":\"false\"}}}" \ "$POD_URL" curl -fsS --cacert "$CACERT" \ -H "Authorization: Bearer $TOKEN" \ diff --git a/internal/consts/label.go b/internal/consts/label.go index 21176bba5..da6f7270b 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -45,5 +45,6 @@ const ( LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" LabelSoperatorRollingUpdateValue = "true" - LabelSoperatorDeleteAfterReboot = K8sGroupNameSoperator + "/delete-after-reboot" + LabelSoperatorDeleteCandidate = K8sGroupNameSoperator + "/delete-candidate" + LabelSoperatorDeleteCandidateValue = "true" ) diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go index 2278b3a23..b69584339 100644 --- a/internal/controller/updatecontroller/statefulset_controller.go +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -242,12 +242,12 @@ func (r *RollingUpdateReconciler) processRollingUpdate( } pod := candidate.pod - if pod.Labels[consts.LabelSoperatorDeleteAfterReboot] != consts.LabelSoperatorRollingUpdateValue { + if pod.Labels[consts.LabelSoperatorDeleteCandidate] != consts.LabelSoperatorDeleteCandidateValue { patchBase := pod.DeepCopy() if pod.Labels == nil { pod.Labels = map[string]string{} } - pod.Labels[consts.LabelSoperatorDeleteAfterReboot] = consts.LabelSoperatorRollingUpdateValue + pod.Labels[consts.LabelSoperatorDeleteCandidate] = consts.LabelSoperatorDeleteCandidateValue 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) } @@ -262,9 +262,8 @@ func (r *RollingUpdateReconciler) processRollingUpdate( } if err := slurmProxyClient.RebootNodes(ctx, slurmproxy.RebootNodesRequest{ - Nodes: slurmNodesToReboot, - Reason: slurmproxy.DefaultReason, - NextState: slurmproxy.RebootNextStateResume, + Nodes: slurmNodesToReboot, + Reason: slurmproxy.DefaultReason, }); err != nil { return fmt.Errorf("schedule slurm reboot through controller proxy: %w", err) } diff --git a/internal/slurmproxy/server.go b/internal/slurmproxy/server.go index 6aa0f528a..7d577dbc5 100644 --- a/internal/slurmproxy/server.go +++ b/internal/slurmproxy/server.go @@ -223,15 +223,6 @@ func (r *RebootNodesRequest) normalizeAndValidate() error { return fmt.Errorf("reason is too long: got %d bytes, max %d", len(r.Reason), maxReasonLength) } - if r.NextState == "" { - r.NextState = RebootNextStateResume - } - switch r.NextState { - case RebootNextStateResume, RebootNextStateDown: - default: - return fmt.Errorf("unsupported nextState %q", r.NextState) - } - return nil } diff --git a/internal/slurmproxy/server_test.go b/internal/slurmproxy/server_test.go index 9b925a378..a2bd44459 100644 --- a/internal/slurmproxy/server_test.go +++ b/internal/slurmproxy/server_test.go @@ -82,7 +82,6 @@ func TestRebootNodesRequestValidate(t *testing.T) { require.NoError(t, req.normalizeAndValidate()) require.Equal(t, DefaultReason, req.Reason) - require.Equal(t, RebootNextStateResume, req.NextState) } func signedToken(t *testing.T, key []byte, user string) string { diff --git a/internal/slurmproxy/types.go b/internal/slurmproxy/types.go index 86929f6b4..9210e1760 100644 --- a/internal/slurmproxy/types.go +++ b/internal/slurmproxy/types.go @@ -10,17 +10,9 @@ const ( EndpointRebootNodes = "/v1/nodes/reboot" ) -type RebootNextState string - -const ( - RebootNextStateResume RebootNextState = "RESUME" - RebootNextStateDown RebootNextState = "DOWN" -) - type RebootNodesRequest struct { - Nodes []string `json:"nodes"` - Reason string `json:"reason,omitempty"` - NextState RebootNextState `json:"nextState,omitempty"` + Nodes []string `json:"nodes"` + Reason string `json:"reason,omitempty"` } type RebootNodesResponse struct { From ce12507ff5b22a70eeb32864c3857d8419c85b11 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Wed, 17 Jun 2026 17:42:39 +0200 Subject: [PATCH 04/10] delete pods from controller --- images/common/scripts/reboot.sh | 45 ++++++++++--------- internal/consts/label.go | 9 ++-- .../statefulset_controller.go | 36 +++++++++++---- internal/render/worker/role.go | 2 +- internal/render/worker/role_test.go | 2 +- 5 files changed, 59 insertions(+), 35 deletions(-) diff --git a/images/common/scripts/reboot.sh b/images/common/scripts/reboot.sh index cc3072e45..a38140a52 100644 --- a/images/common/scripts/reboot.sh +++ b/images/common/scripts/reboot.sh @@ -23,29 +23,34 @@ 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="true" +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")" && jq -e \ + "$POD_URL")"; then + DELETE_CANDIDATE_LABEL_VALUE="$(jq -r \ --arg label "$DELETE_CANDIDATE_LABEL" \ - --arg value "$DELETE_CANDIDATE_LABEL_VALUE" \ - '.metadata.labels[$label] == $value' \ - <<<"$POD_JSON" >/dev/null; then - curl -fsS --cacert "$CACERT" \ - -H "Authorization: Bearer $TOKEN" \ - -H "Content-Type: application/merge-patch+json" \ - -X PATCH \ - -d "{\"metadata\":{\"labels\":{\"$DELETE_CANDIDATE_LABEL\":\"false\"}}}" \ - "$POD_URL" - curl -fsS --cacert "$CACERT" \ - -H "Authorization: Bearer $TOKEN" \ - -X DELETE \ - "$POD_URL" -else - if ! mountpoint -q /run/nvidia/driver; then - echo "This command only works on GPU nodes" - exit 1 + '.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 - chroot /run/nvidia/driver nsenter -t 1 -m -u -i -n /usr/sbin/reboot 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/internal/consts/label.go b/internal/consts/label.go index da6f7270b..d02afb84d 100644 --- a/internal/consts/label.go +++ b/internal/consts/label.go @@ -43,8 +43,9 @@ const ( AnnotationConfigHash = K8sGroupNameSoperator + "/config-hash" - LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" - LabelSoperatorRollingUpdateValue = "true" - LabelSoperatorDeleteCandidate = K8sGroupNameSoperator + "/delete-candidate" - LabelSoperatorDeleteCandidateValue = "true" + LabelSoperatorRollingUpdateEnabled = K8sGroupNameSoperator + "/rolling-update-enabled" + LabelSoperatorRollingUpdateValue = "true" + LabelSoperatorDeleteCandidate = K8sGroupNameSoperator + "/delete-candidate" + LabelSoperatorDeleteCandidateValueStopping = "stopping" + LabelSoperatorDeleteCandidateValueDeleting = "deleting" ) diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go index b69584339..55ce3f7b7 100644 --- a/internal/controller/updatecontroller/statefulset_controller.go +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -75,7 +75,7 @@ func NewRollingUpdateReconciler( // +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 +// +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. @@ -174,6 +174,28 @@ func (r *RollingUpdateReconciler) processRollingUpdate( 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, @@ -195,18 +217,14 @@ func (r *RollingUpdateReconciler) processRollingUpdate( return fmt.Errorf("no slurm controller proxy client for %s/%s", sts.Namespace, clusterName) } - sort.Slice(outdatedPods, func(i, j int) bool { - return outdatedPods[i].Name < outdatedPods[j].Name - }) - type rebootCandidate struct { pod corev1.Pod slurmNode slurmapi.Node } - candidates := make([]rebootCandidate, 0, len(outdatedPods)) + candidates := make([]rebootCandidate, 0, len(podsToStop)) inFlightReady := 0 - for _, pod := range outdatedPods { + for _, pod := range podsToStop { slurmNode, err := slurmClient.GetNode(ctx, pod.Name) if err != nil { return err @@ -242,12 +260,12 @@ func (r *RollingUpdateReconciler) processRollingUpdate( } pod := candidate.pod - if pod.Labels[consts.LabelSoperatorDeleteCandidate] != consts.LabelSoperatorDeleteCandidateValue { + if pod.Labels[consts.LabelSoperatorDeleteCandidate] != consts.LabelSoperatorDeleteCandidateValueStopping { patchBase := pod.DeepCopy() if pod.Labels == nil { pod.Labels = map[string]string{} } - pod.Labels[consts.LabelSoperatorDeleteCandidate] = consts.LabelSoperatorDeleteCandidateValue + 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) } diff --git a/internal/render/worker/role.go b/internal/render/worker/role.go index 08386e431..3a29ca3b4 100644 --- a/internal/render/worker/role.go +++ b/internal/render/worker/role.go @@ -32,7 +32,7 @@ func RenderRole(namespace, clusterName string, nodeSet *values.SlurmNodeSet) rba APIGroups: []string{""}, Resources: []string{"pods"}, ResourceNames: resourceNames, - Verbs: []string{"get", "patch", "update", "delete"}, + Verbs: []string{"get", "patch"}, }, }, } diff --git a/internal/render/worker/role_test.go b/internal/render/worker/role_test.go index 478735b55..c8d26096c 100644 --- a/internal/render/worker/role_test.go +++ b/internal/render/worker/role_test.go @@ -49,7 +49,7 @@ func TestRenderRole_ResourceNames(t *testing.T) { role := worker.RenderRole("default", "soperator", tt.nodeSet) assert.Len(t, role.Rules, 1) - assert.Equal(t, []string{"get", "patch", "update", "delete"}, role.Rules[0].Verbs) + 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]) }) From b6ab122183334dc04e45b761318017dee432efe4 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Fri, 19 Jun 2026 15:32:33 +0200 Subject: [PATCH 05/10] patched slurm with reboot via API --- cmd/main.go | 4 - cmd/slurmcontrollerproxy/main.go | 49 ---- cmd/soperatorchecks/main.go | 3 - images/controller/slurmctld.dockerfile | 17 -- internal/consts/container.go | 36 ++- .../slurm_api_clients_controller.go | 30 +-- .../statefulset_controller.go | 34 +-- internal/render/controller/container.go | 64 ----- internal/render/controller/service.go | 6 - internal/render/controller/statefulset.go | 3 +- .../render/controller/statefulset_test.go | 15 +- internal/slurmapi/client.go | 62 +++++ internal/slurmapi/fake/mock_client.go | 47 ++++ internal/slurmapi/interface.go | 1 + internal/slurmapi/reboot_test.go | 100 +++++++ internal/slurmproxy/client.go | 87 ------ internal/slurmproxy/client_set.go | 34 --- internal/slurmproxy/server.go | 255 ------------------ internal/slurmproxy/server_test.go | 109 -------- internal/slurmproxy/types.go | 25 -- 20 files changed, 243 insertions(+), 738 deletions(-) delete mode 100644 cmd/slurmcontrollerproxy/main.go create mode 100644 internal/slurmapi/reboot_test.go delete mode 100644 internal/slurmproxy/client.go delete mode 100644 internal/slurmproxy/client_set.go delete mode 100644 internal/slurmproxy/server.go delete mode 100644 internal/slurmproxy/server_test.go delete mode 100644 internal/slurmproxy/types.go diff --git a/cmd/main.go b/cmd/main.go index 83a48806d..67a864392 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -59,7 +59,6 @@ import ( "nebius.ai/slurm-operator/internal/controllersenabled" metricsopts "nebius.ai/slurm-operator/internal/metrics" "nebius.ai/slurm-operator/internal/slurmapi" - "nebius.ai/slurm-operator/internal/slurmproxy" webhookv1 "nebius.ai/slurm-operator/internal/webhook/v1" webhookv1alpha1 "nebius.ai/slurm-operator/internal/webhook/v1alpha1" //+kubebuilder:scaffold:imports @@ -313,7 +312,6 @@ func main() { // endregion Reconciler/NodeSet slurmAPIClients := slurmapi.NewClientSet(context.Background()) - slurmProxyClients := slurmproxy.NewClientSet() if controllersSet.Enabled("rollingupdate") { if err = soperatorchecks.NewSlurmAPIClientsController( @@ -321,7 +319,6 @@ func main() { mgr.GetScheme(), mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName), slurmAPIClients, - slurmProxyClients, ).SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { cli.Fail(setupLog, err, "unable to create slurm api clients controller", "controller", soperatorchecks.SlurmAPIClientsControllerName) } @@ -331,7 +328,6 @@ func main() { mgr.GetScheme(), mgr.GetEventRecorderFor(updatecontroller.RollingUpdateControllerName), slurmAPIClients, - slurmProxyClients, ). SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { cli.Fail(setupLog, err, "unable to create controller", "controller", updatecontroller.RollingUpdateControllerName) diff --git a/cmd/slurmcontrollerproxy/main.go b/cmd/slurmcontrollerproxy/main.go deleted file mode 100644 index 24c44e541..000000000 --- a/cmd/slurmcontrollerproxy/main.go +++ /dev/null @@ -1,49 +0,0 @@ -package main - -import ( - "flag" - "log" - "net/http" - "os" - "time" - - "nebius.ai/slurm-operator/internal/consts" - "nebius.ai/slurm-operator/internal/slurmproxy" -) - -func main() { - var ( - listenAddress string - jwtKeyFile string - allowedUsers string - scontrolPath string - commandTimeout time.Duration - ) - - flag.StringVar(&listenAddress, "listen", slurmproxy.DefaultListenAddress, "HTTP listen address") - flag.StringVar(&jwtKeyFile, "jwt-key-file", consts.VolumeMountPathRESTJWTKey+"/"+consts.SecretRESTJWTKeyFileName, "JWT signing key file") - flag.StringVar(&allowedUsers, "allowed-users", slurmproxy.DefaultAllowedUsers, "Comma-separated Slurm usernames allowed to call privileged RPCs") - flag.StringVar(&scontrolPath, "scontrol-path", "/usr/bin/scontrol", "Path to scontrol") - flag.DurationVar(&commandTimeout, "command-timeout", 30*time.Second, "Timeout for each privileged Slurm command") - flag.Parse() - - jwtKey, err := os.ReadFile(jwtKeyFile) - if err != nil { - log.Fatalf("read jwt key file: %v", err) - } - - server, err := slurmproxy.NewServer(slurmproxy.ServerOptions{ - JWTKey: jwtKey, - AllowedUsers: allowedUsers, - ScontrolPath: scontrolPath, - CommandTimeout: commandTimeout, - }) - if err != nil { - log.Fatalf("create server: %v", err) - } - - log.Printf("starting slurm controller proxy on %s", listenAddress) - if err := http.ListenAndServe(listenAddress, server.Handler()); err != nil { - log.Fatalf("serve: %v", err) - } -} diff --git a/cmd/soperatorchecks/main.go b/cmd/soperatorchecks/main.go index e5a832c7a..4332d10b7 100644 --- a/cmd/soperatorchecks/main.go +++ b/cmd/soperatorchecks/main.go @@ -50,7 +50,6 @@ import ( "nebius.ai/slurm-operator/internal/controllersenabled" metricsopts "nebius.ai/slurm-operator/internal/metrics" "nebius.ai/slurm-operator/internal/slurmapi" - "nebius.ai/slurm-operator/internal/slurmproxy" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" //+kubebuilder:scaffold:imports @@ -276,7 +275,6 @@ func main() { } slurmAPIClients := slurmapi.NewClientSet(ctx) - slurmProxyClients := slurmproxy.NewClientSet() if controllersSet.Enabled("slurmapiclients") { if err = soperatorchecks.NewSlurmAPIClientsController( @@ -284,7 +282,6 @@ func main() { mgr.GetScheme(), mgr.GetEventRecorderFor(soperatorchecks.SlurmAPIClientsControllerName), slurmAPIClients, - slurmProxyClients, ).SetupWithManager(mgr, maxConcurrency, cacheSyncTimeout); err != nil { cli.Fail(setupLog, err, "unable to create slurm api clients controller", "controller", soperatorchecks.SlurmAPIClientsControllerName) } diff --git a/images/controller/slurmctld.dockerfile b/images/controller/slurmctld.dockerfile index c102c2783..2df88ce5f 100644 --- a/images/controller/slurmctld.dockerfile +++ b/images/controller/slurmctld.dockerfile @@ -30,20 +30,6 @@ RUN --mount=type=cache,target=/root/.cache/go-build \ GOOS=$GOOS CGO_ENABLED=$CGO_ENABLED GO_LDFLAGS=$GO_LDFLAGS \ go build -v -o power-manager ./cmd/powermanager -FROM go-base AS slurmcontrollerproxy_builder - -ARG GO_LDFLAGS="" -ARG CGO_ENABLED=0 -ARG GOOS=linux - -COPY go.mod go.sum ./ -RUN go mod download - -COPY cmd/slurmcontrollerproxy cmd/slurmcontrollerproxy -RUN --mount=type=cache,target=/root/.cache/go-build \ - GOOS=$GOOS CGO_ENABLED=$CGO_ENABLED GO_LDFLAGS=$GO_LDFLAGS \ - go build -v -o slurm-controller-proxy ./cmd/slurmcontrollerproxy - # https://github.com/nebius/ml-containers/pull/81 FROM cr.eu-north1.nebius.cloud/ml-containers/slurm:${SLURM_VERSION}-20260421092341 AS controller_slurmctld @@ -68,9 +54,6 @@ RUN mkdir -p /var/log/slurm/multilog && \ COPY --from=powermanager_builder /build/power-manager /opt/soperator/bin/power-manager RUN chmod 755 /opt/soperator/bin/power-manager -COPY --from=slurmcontrollerproxy_builder /build/slurm-controller-proxy /opt/soperator/bin/slurm-controller-proxy -RUN chmod 755 /opt/soperator/bin/slurm-controller-proxy - # Copy power management scripts for Slurm ResumeProgram/SuspendProgram COPY images/controller/power_resume.sh /opt/soperator/bin/power_resume.sh COPY images/controller/power_suspend.sh /opt/soperator/bin/power_suspend.sh diff --git a/internal/consts/container.go b/internal/consts/container.go index c51cd2378..bb5b21b09 100644 --- a/internal/consts/container.go +++ b/internal/consts/container.go @@ -1,23 +1,22 @@ package consts const ( - ContainerNameSlurmctld = SlurmctldName - ContainerNameSlurmControllerProxy = "slurm-controller-proxy" - ContainerNameAccounting = AccountingName - ContainerNameMunge = Munge - ContainerNameSSSD = "sssd" - ContainerNameSlurmd = SlurmdName - ContainerNameREST = Slurmrestd - ContainerNameSshd = SshdName - ContainerNameWorkerInit = "worker-init" - ContainerNameDockerProxy = "docker-proxy" - ContainerNameWaitForDatabase = "wait-for-database" - ContainerNameWaitForAccounting = "wait-for-accounting" - ContainerNamePopulateJail = populateJail - ContainerNameExporter = Exporter - ContainerNameRebooter = "rebooter" - ContainerNameCustom = "custom-container" - ContainerNameSConfigController = SConfigControllerName + ContainerNameSlurmctld = SlurmctldName + ContainerNameAccounting = AccountingName + ContainerNameMunge = Munge + ContainerNameSSSD = "sssd" + ContainerNameSlurmd = SlurmdName + ContainerNameREST = Slurmrestd + ContainerNameSshd = SshdName + ContainerNameWorkerInit = "worker-init" + ContainerNameDockerProxy = "docker-proxy" + ContainerNameWaitForDatabase = "wait-for-database" + ContainerNameWaitForAccounting = "wait-for-accounting" + ContainerNamePopulateJail = populateJail + ContainerNameExporter = Exporter + ContainerNameRebooter = "rebooter" + ContainerNameCustom = "custom-container" + ContainerNameSConfigController = SConfigControllerName ContainerSecurityContextCapabilitySysAdmin = "SYS_ADMIN" @@ -29,7 +28,4 @@ const ( ContainerPortMonitoring = 8081 ContainerPathMonitoring = "/metrics" ContainerPathHealthz = "/healthz" - - ContainerPortNameSlurmControllerProxy = "slurmctl-proxy" - ContainerPortSlurmControllerProxy = 6821 ) diff --git a/internal/controller/soperatorchecks/slurm_api_clients_controller.go b/internal/controller/soperatorchecks/slurm_api_clients_controller.go index 15f1fc776..a7434aa16 100644 --- a/internal/controller/soperatorchecks/slurm_api_clients_controller.go +++ b/internal/controller/soperatorchecks/slurm_api_clients_controller.go @@ -21,7 +21,6 @@ import ( "nebius.ai/slurm-operator/internal/jwt" "nebius.ai/slurm-operator/internal/naming" "nebius.ai/slurm-operator/internal/slurmapi" - "nebius.ai/slurm-operator/internal/slurmproxy" ) var ( @@ -31,8 +30,7 @@ var ( type SlurmAPIClientsController struct { *reconciler.Reconciler - slurmAPIClients *slurmapi.ClientSet - slurmProxyClients *slurmproxy.ClientSet + slurmAPIClients *slurmapi.ClientSet } func NewSlurmAPIClientsController( @@ -40,14 +38,12 @@ func NewSlurmAPIClientsController( scheme *runtime.Scheme, recorder record.EventRecorder, slurmAPIClients *slurmapi.ClientSet, - slurmProxyClients *slurmproxy.ClientSet, ) *SlurmAPIClientsController { r := reconciler.NewReconciler(client, scheme, recorder) return &SlurmAPIClientsController{ - Reconciler: r, - slurmAPIClients: slurmAPIClients, - slurmProxyClients: slurmProxyClients, + Reconciler: r, + slurmAPIClients: slurmAPIClients, } } @@ -79,8 +75,7 @@ func (c *SlurmAPIClientsController) Reconcile(ctx context.Context, req ctrl.Requ logger.Info("Adding new slurm api client") - tokenRegistry := jwt.NewTokenRegistry().Build() - jwtToken := jwt.NewToken(c.Client).For(req.NamespacedName, "root").WithRegistry(tokenRegistry) + jwtToken := jwt.NewToken(c.Client).For(req.NamespacedName, "root").WithRegistry(jwt.NewTokenRegistry().Build()) slurmAPIServer := fmt.Sprintf("http://%s.%s:6820", naming.BuildServiceName(consts.ComponentTypeREST, req.Name), req.Namespace) slurmAPIClient, err := slurmapi.NewClient(slurmAPIServer, jwtToken, slurmapi.DefaultHTTPClient()) if err != nil { @@ -90,22 +85,5 @@ func (c *SlurmAPIClientsController) Reconcile(ctx context.Context, req ctrl.Requ c.slurmAPIClients.AddClient(req.NamespacedName, slurmAPIClient) - if c.slurmProxyClients != nil { - proxyToken := jwt.NewToken(c.Client).For(req.NamespacedName, slurmproxy.DefaultClientUser).WithRegistry(tokenRegistry) - proxyServer := fmt.Sprintf( - "http://%s.%s:%d", - naming.BuildServiceName(consts.ComponentTypeController, req.Name), - req.Namespace, - consts.ContainerPortSlurmControllerProxy, - ) - slurmProxyClient, err := slurmproxy.NewClient(proxyServer, proxyToken, slurmapi.DefaultHTTPClient()) - if err != nil { - logger.Error(err, "failed to create slurm controller proxy client") - return ctrl.Result{}, err - } - - c.slurmProxyClients.AddClient(req.NamespacedName, slurmProxyClient) - } - return ctrl.Result{}, nil } diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go index 55ce3f7b7..746e75f76 100644 --- a/internal/controller/updatecontroller/statefulset_controller.go +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -20,6 +20,7 @@ import ( "context" "fmt" "sort" + "strings" "time" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" @@ -40,7 +41,6 @@ import ( "nebius.ai/slurm-operator/internal/controller/reconciler" "nebius.ai/slurm-operator/internal/controllerconfig" "nebius.ai/slurm-operator/internal/slurmapi" - "nebius.ai/slurm-operator/internal/slurmproxy" ) const ( @@ -49,26 +49,24 @@ const ( const ( defaultSTSReplicasCount = int32(1) + defaultRebootReason = "soperator rolling update" ) type RollingUpdateReconciler struct { *reconciler.Reconciler - slurmAPIClients *slurmapi.ClientSet - slurmProxyClients *slurmproxy.ClientSet + slurmAPIClients *slurmapi.ClientSet } func NewRollingUpdateReconciler( client client.Client, scheme *runtime.Scheme, recorder record.EventRecorder, slurmAPIClients *slurmapi.ClientSet, - slurmProxyClients *slurmproxy.ClientSet, ) *RollingUpdateReconciler { r := reconciler.NewReconciler(client, scheme, recorder) return &RollingUpdateReconciler{ - Reconciler: r, - slurmAPIClients: slurmAPIClients, - slurmProxyClients: slurmProxyClients, + Reconciler: r, + slurmAPIClients: slurmAPIClients, } } @@ -205,18 +203,6 @@ func (r *RollingUpdateReconciler) processRollingUpdate( return fmt.Errorf("no slurm api client for %s/%s", sts.Namespace, clusterName) } - if r.slurmProxyClients == nil { - return fmt.Errorf("slurm controller proxy client set is not configured") - } - slurmProxyClient, ok := r.slurmProxyClients.GetClient(types.NamespacedName{ - Namespace: sts.Namespace, - Name: clusterName, - }) - if !ok { - logger.Info("no slurm controller proxy client", "namespace", sts.Namespace, "clusterName", clusterName) - return fmt.Errorf("no slurm controller proxy client for %s/%s", sts.Namespace, clusterName) - } - type rebootCandidate struct { pod corev1.Pod slurmNode slurmapi.Node @@ -279,14 +265,14 @@ func (r *RollingUpdateReconciler) processRollingUpdate( return nil } - if err := slurmProxyClient.RebootNodes(ctx, slurmproxy.RebootNodesRequest{ - Nodes: slurmNodesToReboot, - Reason: slurmproxy.DefaultReason, + if err := slurmClient.RebootNodes(ctx, slurmapi.RebootNodesRequest{ + NodeList: strings.Join(slurmNodesToReboot, ","), + Reason: defaultRebootReason, }); err != nil { - return fmt.Errorf("schedule slurm reboot through controller proxy: %w", err) + return fmt.Errorf("schedule slurm reboot through rest api: %w", err) } - logger.Info("scheduled slurm reboot through controller proxy", "nodes", slurmNodesToReboot) + logger.Info("scheduled slurm reboot through rest api", "nodes", slurmNodesToReboot) return nil } diff --git a/internal/render/controller/container.go b/internal/render/controller/container.go index 53055309c..497ac7f45 100644 --- a/internal/render/controller/container.go +++ b/internal/render/controller/container.go @@ -2,7 +2,6 @@ package controller import ( corev1 "k8s.io/api/core/v1" - "k8s.io/apimachinery/pkg/util/intstr" slurmv1 "nebius.ai/slurm-operator/api/v1" "nebius.ai/slurm-operator/internal/utils/stringutils" @@ -53,69 +52,6 @@ func renderContainerSlurmctld(container *values.Container, customMounts []slurmv } } -func renderContainerSlurmControllerProxy(container *values.Container) corev1.Container { - return corev1.Container{ - Name: consts.ContainerNameSlurmControllerProxy, - Image: container.Image, - ImagePullPolicy: container.ImagePullPolicy, - Command: []string{ - "/bin/bash", "-c", - }, - Args: []string{ - // language=bash - stringutils.Dedent(` - set -eu - - ln -sf /mnt/jail/etc/passwd /etc/passwd - ln -sf /mnt/jail/etc/group /etc/group - ln -sf /mnt/jail/etc/shadow /etc/shadow - ln -sf /mnt/jail/etc/gshadow /etc/gshadow - chown -h 0:42 /etc/{shadow,gshadow} - - rm -rf /etc/slurm && ln -s /mnt/jail/etc/slurm /etc/slurm - - while [ ! -S "/run/munge/munge.socket.2" ]; do sleep 2; done - - exec /opt/soperator/bin/slurm-controller-proxy --listen=:6821 - `), - }, - Ports: []corev1.ContainerPort{{ - Name: consts.ContainerPortNameSlurmControllerProxy, - ContainerPort: consts.ContainerPortSlurmControllerProxy, - Protocol: corev1.ProtocolTCP, - }}, - VolumeMounts: []corev1.VolumeMount{ - common.RenderVolumeMountJail(), - common.RenderVolumeMountMungeSocket(), - common.RenderVolumeMountRESTJWTKey(), - }, - ReadinessProbe: &corev1.Probe{ - ProbeHandler: corev1.ProbeHandler{ - HTTPGet: &corev1.HTTPGetAction{ - Path: consts.ContainerPathHealthz, - Port: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), - }, - }, - PeriodSeconds: 5, - }, - LivenessProbe: &corev1.Probe{ - ProbeHandler: corev1.ProbeHandler{ - HTTPGet: &corev1.HTTPGetAction{ - Path: consts.ContainerPathHealthz, - Port: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), - }, - }, - InitialDelaySeconds: 30, - PeriodSeconds: 10, - }, - SecurityContext: &corev1.SecurityContext{ - AppArmorProfile: common.ParseAppArmorProfile(container.AppArmorProfile), - }, - TerminationMessagePath: corev1.TerminationMessagePathDefault, - TerminationMessagePolicy: corev1.TerminationMessageReadFile, - } -} - func renderContainerSlurmctldWithSSSD( container *values.Container, customMounts []slurmv1.NodeVolumeMount, diff --git a/internal/render/controller/service.go b/internal/render/controller/service.go index 0d017f1c3..79b2336fd 100644 --- a/internal/render/controller/service.go +++ b/internal/render/controller/service.go @@ -39,12 +39,6 @@ func RenderService(namespace, clusterName, svcName string, controller *values.Sl Port: controller.ContainerSlurmctld.Port, TargetPort: intstr.FromString(controller.ContainerSlurmctld.Name), }, - { - Name: consts.ContainerPortNameSlurmControllerProxy, - Protocol: corev1.ProtocolTCP, - Port: consts.ContainerPortSlurmControllerProxy, - TargetPort: intstr.FromString(consts.ContainerPortNameSlurmControllerProxy), - }, }, }, } diff --git a/internal/render/controller/statefulset.go b/internal/render/controller/statefulset.go index 9400c5d41..8d2b6c2aa 100644 --- a/internal/render/controller/statefulset.go +++ b/internal/render/controller/statefulset.go @@ -69,7 +69,6 @@ func RenderStatefulSet( if controller.ContainerSSSD != nil { slurmctldContainer = renderContainerSlurmctldWithSSSD(&controller.ContainerSlurmctld, controller.CustomVolumeMounts) } - slurmControllerProxyContainer := renderContainerSlurmControllerProxy(&controller.ContainerSlurmctld) return kruisev1b1.StatefulSet{ ObjectMeta: metav1.ObjectMeta{ @@ -117,7 +116,7 @@ func RenderStatefulSet( NodeSelector: nodeFilter.NodeSelector, Tolerations: nodeFilter.Tolerations, InitContainers: initContainers, - Containers: []corev1.Container{slurmctldContainer, slurmControllerProxyContainer}, + Containers: []corev1.Container{slurmctldContainer}, Volumes: volumes, RestartPolicy: corev1.RestartPolicyAlways, TerminationGracePeriodSeconds: ptr.To(common.DefaultPodTerminationGracePeriodSeconds), diff --git a/internal/render/controller/statefulset_test.go b/internal/render/controller/statefulset_test.go index 997613738..30acba3ac 100644 --- a/internal/render/controller/statefulset_test.go +++ b/internal/render/controller/statefulset_test.go @@ -176,8 +176,8 @@ func TestRenderStatefulSet(t *testing.T) { // Check containers containers := result.Spec.Template.Spec.Containers - if len(containers) != 2 { - t.Errorf("Expected 2 containers, got %d", len(containers)) + if len(containers) != 1 { + t.Errorf("Expected 1 container, got %d", len(containers)) } else { container := containers[0] if container.Name != consts.ContainerNameSlurmctld { @@ -186,14 +186,6 @@ func TestRenderStatefulSet(t *testing.T) { if container.Image != tt.controller.ContainerSlurmctld.NodeContainer.Image { t.Errorf("Container image = %v, want %v", container.Image, tt.controller.ContainerSlurmctld.NodeContainer.Image) } - - proxyContainer := containers[1] - if proxyContainer.Name != consts.ContainerNameSlurmControllerProxy { - t.Errorf("Container name = %v, want %v", proxyContainer.Name, consts.ContainerNameSlurmControllerProxy) - } - if proxyContainer.Image != tt.controller.ContainerSlurmctld.NodeContainer.Image { - t.Errorf("Container image = %v, want %v", proxyContainer.Image, tt.controller.ContainerSlurmctld.NodeContainer.Image) - } } // Check init containers @@ -298,9 +290,6 @@ func TestRenderStatefulSet_SSSD(t *testing.T) { } assert.Contains(t, mountNames, consts.VolumeNameSSSDConf) assert.Contains(t, mountNames, consts.VolumeNameSSSDSocket) - - proxy := result.Spec.Template.Spec.Containers[1] - assert.Equal(t, consts.ContainerNameSlurmControllerProxy, proxy.Name) } func TestRenderStatefulSetWithMaintenance(t *testing.T) { 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..4a022fbe3 100644 --- a/internal/slurmapi/interface.go +++ b/internal/slurmapi/interface.go @@ -10,6 +10,7 @@ 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/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/slurmproxy/client.go b/internal/slurmproxy/client.go deleted file mode 100644 index f45aecc98..000000000 --- a/internal/slurmproxy/client.go +++ /dev/null @@ -1,87 +0,0 @@ -package slurmproxy - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "strings" -) - -type tokenIssuer interface { - Issue(ctx context.Context) (string, error) -} - -type Interface interface { - RebootNodes(ctx context.Context, req RebootNodesRequest) error -} - -type Client struct { - server string - tokenIssuer tokenIssuer - httpClient *http.Client -} - -func NewClient(server string, tokenIssuer tokenIssuer, httpClient *http.Client) (*Client, error) { - if server == "" { - return nil, fmt.Errorf("unable to create slurm controller proxy client: empty server URL") - } - if httpClient == nil { - httpClient = http.DefaultClient - } - - return &Client{ - server: strings.TrimRight(server, "/"), - tokenIssuer: tokenIssuer, - httpClient: httpClient, - }, nil -} - -func (c *Client) RebootNodes(ctx context.Context, req RebootNodesRequest) error { - body, err := json.Marshal(req) - if err != nil { - return fmt.Errorf("marshal reboot nodes request: %w", err) - } - - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.server+EndpointRebootNodes, bytes.NewReader(body)) - if err != nil { - return fmt.Errorf("create reboot nodes request: %w", err) - } - httpReq.Header.Set("Content-Type", "application/json") - - if c.tokenIssuer != nil { - token, err := c.tokenIssuer.Issue(ctx) - if err != nil { - return fmt.Errorf("issue controller proxy token: %w", err) - } - if token != "" { - httpReq.Header.Set("Authorization", "Bearer "+token) - } - } - - resp, err := c.httpClient.Do(httpReq) - if err != nil { - return fmt.Errorf("post reboot nodes request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { - return fmt.Errorf("reboot nodes request failed: status=%d: %s", resp.StatusCode, responseError(resp.Body)) - } - - return nil -} - -func responseError(body io.Reader) string { - data, err := io.ReadAll(body) - if err != nil { - return fmt.Sprintf("read response body: %v", err) - } - var apiErr ErrorResponse - if err := json.Unmarshal(data, &apiErr); err == nil && apiErr.Error != "" { - return apiErr.Error - } - return string(data) -} diff --git a/internal/slurmproxy/client_set.go b/internal/slurmproxy/client_set.go deleted file mode 100644 index df1d4fe8a..000000000 --- a/internal/slurmproxy/client_set.go +++ /dev/null @@ -1,34 +0,0 @@ -package slurmproxy - -import ( - "sync" - - "k8s.io/apimachinery/pkg/types" -) - -type ClientSet struct { - clients map[types.NamespacedName]Interface - mux *sync.Mutex -} - -func NewClientSet() *ClientSet { - return &ClientSet{ - clients: make(map[types.NamespacedName]Interface), - mux: &sync.Mutex{}, - } -} - -func (cs *ClientSet) AddClient(name types.NamespacedName, client Interface) { - cs.mux.Lock() - defer cs.mux.Unlock() - - cs.clients[name] = client -} - -func (cs *ClientSet) GetClient(name types.NamespacedName) (Interface, bool) { - cs.mux.Lock() - defer cs.mux.Unlock() - - client, found := cs.clients[name] - return client, found -} diff --git a/internal/slurmproxy/server.go b/internal/slurmproxy/server.go deleted file mode 100644 index 7d577dbc5..000000000 --- a/internal/slurmproxy/server.go +++ /dev/null @@ -1,255 +0,0 @@ -package slurmproxy - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "net/http" - "os/exec" - "regexp" - "strings" - "time" - - jwtlib "github.com/golang-jwt/jwt/v5" - - internaljwt "nebius.ai/slurm-operator/internal/jwt" -) - -const ( - defaultScontrolPath = "/usr/bin/scontrol" - defaultCommandTimeout = 30 * time.Second - maxReasonLength = 512 -) - -var nodeNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]*$`) - -type CommandRunner interface { - Run(ctx context.Context, name string, args ...string) ([]byte, error) -} - -type ExecCommandRunner struct{} - -func (ExecCommandRunner) Run(ctx context.Context, name string, args ...string) ([]byte, error) { - cmd := exec.CommandContext(ctx, name, args...) - return cmd.CombinedOutput() -} - -type ServerOptions struct { - JWTKey []byte - AllowedUsers string - ScontrolPath string - CommandRunner CommandRunner - CommandTimeout time.Duration -} - -type Server struct { - jwtKey []byte - allowedUsers map[string]struct{} - allowAnyUser bool - scontrolPath string - commandRunner CommandRunner - commandTimeout time.Duration -} - -func NewServer(opts ServerOptions) (*Server, error) { - if len(opts.JWTKey) == 0 { - return nil, errors.New("jwt key is required") - } - - allowedUsers, allowAnyUser := parseAllowedUsers(opts.AllowedUsers) - if len(allowedUsers) == 0 && !allowAnyUser { - return nil, errors.New("at least one allowed user is required") - } - - if opts.ScontrolPath == "" { - opts.ScontrolPath = defaultScontrolPath - } - if opts.CommandRunner == nil { - opts.CommandRunner = ExecCommandRunner{} - } - if opts.CommandTimeout == 0 { - opts.CommandTimeout = defaultCommandTimeout - } - - return &Server{ - jwtKey: opts.JWTKey, - allowedUsers: allowedUsers, - allowAnyUser: allowAnyUser, - scontrolPath: opts.ScontrolPath, - commandRunner: opts.CommandRunner, - commandTimeout: opts.CommandTimeout, - }, nil -} - -func (s *Server) Handler() http.Handler { - mux := http.NewServeMux() - mux.HandleFunc(EndpointHealthz, s.handleHealthz) - mux.HandleFunc(EndpointRebootNodes, s.handleRebootNodes) - return mux -} - -func (s *Server) handleHealthz(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodGet { - writeError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - w.WriteHeader(http.StatusNoContent) -} - -func (s *Server) handleRebootNodes(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodPost { - writeError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - - if _, err := s.authorize(r); err != nil { - writeError(w, http.StatusUnauthorized, err.Error()) - return - } - - var req RebootNodesRequest - decoder := json.NewDecoder(r.Body) - decoder.DisallowUnknownFields() - if err := decoder.Decode(&req); err != nil { - writeError(w, http.StatusBadRequest, fmt.Sprintf("invalid request body: %v", err)) - return - } - if err := req.normalizeAndValidate(); err != nil { - writeError(w, http.StatusBadRequest, err.Error()) - return - } - - commandCtx, cancel := context.WithTimeout(r.Context(), s.commandTimeout) - defer cancel() - - args := []string{"reboot"} - if req.Reason != "" { - args = append(args, "reason="+req.Reason) - } - args = append(args, strings.Join(req.Nodes, ",")) - - output, err := s.commandRunner.Run(commandCtx, s.scontrolPath, args...) - if err != nil { - status := http.StatusBadGateway - if errors.Is(commandCtx.Err(), context.DeadlineExceeded) { - status = http.StatusGatewayTimeout - } - writeError(w, status, strings.TrimSpace(fmt.Sprintf("%v: %s", err, output))) - return - } - - writeJSON(w, http.StatusOK, RebootNodesResponse{ - Nodes: req.Nodes, - Output: strings.TrimSpace(string(output)), - }) -} - -func (s *Server) authorize(r *http.Request) (string, error) { - raw := r.Header.Get("Authorization") - if raw == "" { - return "", errors.New("authorization header is required") - } - tokenString, ok := strings.CutPrefix(raw, "Bearer ") - if !ok || tokenString == "" { - return "", errors.New("bearer token is required") - } - - claims := &internaljwt.TokenClaims{} - token, err := jwtlib.ParseWithClaims( - tokenString, - claims, - func(token *jwtlib.Token) (interface{}, error) { - return s.jwtKey, nil - }, - jwtlib.WithValidMethods([]string{jwtlib.SigningMethodHS256.Alg()}), - jwtlib.WithIssuer("soperator"), - ) - if err != nil { - return "", fmt.Errorf("invalid token: %w", err) - } - if !token.Valid { - return "", errors.New("invalid token") - } - - user := claims.SlurmUsername - if user == "" { - user = claims.Subject - } - if user == "" { - return "", errors.New("token subject is empty") - } - if s.allowAnyUser { - return user, nil - } - if _, ok := s.allowedUsers[user]; !ok { - return "", fmt.Errorf("user %q is not allowed", user) - } - - return user, nil -} - -func (r *RebootNodesRequest) normalizeAndValidate() error { - if len(r.Nodes) == 0 { - return errors.New("nodes must not be empty") - } - - seen := make(map[string]struct{}, len(r.Nodes)) - nodes := make([]string, 0, len(r.Nodes)) - for _, node := range r.Nodes { - node = strings.TrimSpace(node) - if node == "" { - return errors.New("nodes must not contain empty values") - } - if !nodeNamePattern.MatchString(node) { - return fmt.Errorf("invalid node name %q", node) - } - if _, ok := seen[node]; ok { - continue - } - seen[node] = struct{}{} - nodes = append(nodes, node) - } - r.Nodes = nodes - - r.Reason = strings.TrimSpace(r.Reason) - if r.Reason == "" { - r.Reason = DefaultReason - } - if strings.ContainsAny(r.Reason, "\x00\r\n") { - return errors.New("reason must not contain control line breaks") - } - if len(r.Reason) > maxReasonLength { - return fmt.Errorf("reason is too long: got %d bytes, max %d", len(r.Reason), maxReasonLength) - } - - return nil -} - -func parseAllowedUsers(raw string) (map[string]struct{}, bool) { - if raw == "" { - raw = DefaultAllowedUsers - } - res := map[string]struct{}{} - for _, part := range strings.Split(raw, ",") { - user := strings.TrimSpace(part) - if user == "" { - continue - } - if user == "*" { - return nil, true - } - res[user] = struct{}{} - } - return res, false -} - -func writeJSON(w http.ResponseWriter, status int, v any) { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(status) - _ = json.NewEncoder(w).Encode(v) -} - -func writeError(w http.ResponseWriter, status int, message string) { - writeJSON(w, status, ErrorResponse{Error: message}) -} diff --git a/internal/slurmproxy/server_test.go b/internal/slurmproxy/server_test.go deleted file mode 100644 index a2bd44459..000000000 --- a/internal/slurmproxy/server_test.go +++ /dev/null @@ -1,109 +0,0 @@ -package slurmproxy - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" - - jwtlib "github.com/golang-jwt/jwt/v5" - "github.com/stretchr/testify/require" - - internaljwt "nebius.ai/slurm-operator/internal/jwt" -) - -type recordingRunner struct { - name string - args []string -} - -func (r *recordingRunner) Run(_ context.Context, name string, args ...string) ([]byte, error) { - r.name = name - r.args = append([]string(nil), args...) - return []byte("scheduled"), nil -} - -func TestServer_RebootNodes(t *testing.T) { - key := []byte("12345678901234567890123456789012") - runner := &recordingRunner{} - server, err := NewServer(ServerOptions{ - JWTKey: key, - AllowedUsers: "root", - ScontrolPath: "/bin/scontrol", - CommandRunner: runner, - }) - require.NoError(t, err) - - body := `{"nodes":["worker-1","worker-1","worker-2"],"reason":"rolling update"}` - req := httptest.NewRequest(http.MethodPost, EndpointRebootNodes, stringsReader(body)) - req.Header.Set("Authorization", "Bearer "+signedToken(t, key, "root")) - resp := httptest.NewRecorder() - - server.Handler().ServeHTTP(resp, req) - - require.Equal(t, http.StatusOK, resp.Code) - require.Equal(t, "/bin/scontrol", runner.name) - require.Equal(t, []string{ - "reboot", - "reason=rolling update", - "worker-1,worker-2", - }, runner.args) - - var response RebootNodesResponse - require.NoError(t, json.Unmarshal(resp.Body.Bytes(), &response)) - require.Equal(t, []string{"worker-1", "worker-2"}, response.Nodes) - require.Equal(t, "scheduled", response.Output) -} - -func TestServer_RebootNodesRejectsUnauthorizedUser(t *testing.T) { - key := []byte("12345678901234567890123456789012") - server, err := NewServer(ServerOptions{ - JWTKey: key, - AllowedUsers: "root", - }) - require.NoError(t, err) - - req := httptest.NewRequest(http.MethodPost, EndpointRebootNodes, stringsReader(`{"nodes":["worker-1"]}`)) - req.Header.Set("Authorization", "Bearer "+signedToken(t, key, "alice")) - resp := httptest.NewRecorder() - - server.Handler().ServeHTTP(resp, req) - - require.Equal(t, http.StatusUnauthorized, resp.Code) -} - -func TestRebootNodesRequestValidate(t *testing.T) { - req := RebootNodesRequest{ - Nodes: []string{"worker-0"}, - } - - require.NoError(t, req.normalizeAndValidate()) - require.Equal(t, DefaultReason, req.Reason) -} - -func signedToken(t *testing.T, key []byte, user string) string { - t.Helper() - - now := time.Now() - claims := internaljwt.TokenClaims{ - RegisteredClaims: jwtlib.RegisteredClaims{ - Issuer: "soperator", - Subject: user, - IssuedAt: jwtlib.NewNumericDate(now), - NotBefore: jwtlib.NewNumericDate(now.Add(-time.Minute)), - ExpiresAt: jwtlib.NewNumericDate(now.Add(time.Hour)), - }, - SlurmUsername: user, - } - - token, err := jwtlib.NewWithClaims(jwtlib.SigningMethodHS256, claims).SignedString(key) - require.NoError(t, err) - return token -} - -func stringsReader(s string) *strings.Reader { - return strings.NewReader(s) -} diff --git a/internal/slurmproxy/types.go b/internal/slurmproxy/types.go deleted file mode 100644 index 9210e1760..000000000 --- a/internal/slurmproxy/types.go +++ /dev/null @@ -1,25 +0,0 @@ -package slurmproxy - -const ( - DefaultListenAddress = ":6821" - DefaultAllowedUsers = "root,soperatorchecks" - DefaultClientUser = "root" - DefaultReason = "soperator rolling update" - - EndpointHealthz = "/healthz" - EndpointRebootNodes = "/v1/nodes/reboot" -) - -type RebootNodesRequest struct { - Nodes []string `json:"nodes"` - Reason string `json:"reason,omitempty"` -} - -type RebootNodesResponse struct { - Nodes []string `json:"nodes"` - Output string `json:"output,omitempty"` -} - -type ErrorResponse struct { - Error string `json:"error"` -} From c64aaa2281443c2f31169185013aa6472ca5e4a7 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Fri, 19 Jun 2026 17:57:10 +0200 Subject: [PATCH 06/10] asap --- internal/controller/updatecontroller/statefulset_controller.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/controller/updatecontroller/statefulset_controller.go b/internal/controller/updatecontroller/statefulset_controller.go index 746e75f76..e7387b18f 100644 --- a/internal/controller/updatecontroller/statefulset_controller.go +++ b/internal/controller/updatecontroller/statefulset_controller.go @@ -267,6 +267,7 @@ func (r *RollingUpdateReconciler) processRollingUpdate( 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) From c6659d5aadb873bc0ecbc9236b0c62ba83a93b16 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Mon, 22 Jun 2026 15:08:02 +0200 Subject: [PATCH 07/10] fix nodes watch in topology controller --- .../topologyconfcontroller/nodetopology_controller.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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) From 6bd647cc316a0e764e084ac7097de9b9f2cd28c1 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Mon, 22 Jun 2026 16:56:25 +0200 Subject: [PATCH 08/10] lint --- internal/slurmapi/interface.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/slurmapi/interface.go b/internal/slurmapi/interface.go index 4a022fbe3..8b2953b54 100644 --- a/internal/slurmapi/interface.go +++ b/internal/slurmapi/interface.go @@ -6,6 +6,7 @@ import ( api "github.com/SlinkyProject/slurm-client/api/v0041" ) +// nolint:interfacebloat type Client interface { api.ClientWithResponsesInterface ListNodes(ctx context.Context) ([]Node, error) From eb672ffe03944783994b737eba1b1f8d06085440 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Tue, 23 Jun 2026 16:57:53 +0200 Subject: [PATCH 09/10] remove undrains --- images/worker/worker_init.py | 3 --- 1 file changed, 3 deletions(-) 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( From 21b32f632450b5094892c9fd4b864e5813ec7aa8 Mon Sep 17 00:00:00 2001 From: Dzmitry Amialiusik Date: Wed, 1 Jul 2026 16:06:13 +0200 Subject: [PATCH 10/10] change dependency source for nodeset and supervisord cm --- helm/nodesets/templates/nodeset.yaml | 2 +- helm/nodesets/tests/basic_test.yaml | 57 +++++++++++++++++++ helm/nodesets/values.yaml | 3 + .../templates/_helpers.tpl | 18 ++++++ .../templates/configmap-supervisord.yaml | 48 +++++++++++++++- .../tests/supervisord_configmap_test.yaml | 38 +++++++++++++ helm/soperator-custom-configmaps/values.yaml | 1 + helm/soperator-fluxcd/templates/_helpers.tpl | 42 ++++++++++++++ helm/soperator-fluxcd/templates/nodesets.yaml | 21 +++++-- .../soperator-custom-configmaps.yaml | 15 ++++- .../tests/custom_configmaps_test.yaml | 14 +++++ .../tests/soperator_dependency_test.yaml | 11 ++++ .../tests/values_override_test.yaml | 8 +++ helm/soperator-fluxcd/values.yaml | 3 +- internal/consts/annotation.go | 3 + .../configmap_predicate_test.go | 42 ++++++++++++++ .../nodesetcontroller/controller.go | 21 ++++++- internal/controller/reconciler/versioning.go | 10 +++- .../controller/reconciler/versioning_test.go | 55 ++++++++++++++++++ 19 files changed, 399 insertions(+), 13 deletions(-) create mode 100644 helm/soperator-custom-configmaps/tests/supervisord_configmap_test.yaml create mode 100644 internal/controller/nodesetcontroller/configmap_predicate_test.go create mode 100644 internal/controller/reconciler/versioning_test.go diff --git a/helm/nodesets/templates/nodeset.yaml b/helm/nodesets/templates/nodeset.yaml index 4734cb1c1..70deea760 100644 --- a/helm/nodesets/templates/nodeset.yaml +++ b/helm/nodesets/templates/nodeset.yaml @@ -265,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/values.yaml b/helm/nodesets/values.yaml index 3a2eab2d8..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: 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/internal/consts/annotation.go b/internal/consts/annotation.go index 3691dd6b2..0c154ccda 100644 --- a/internal/consts/annotation.go +++ b/internal/consts/annotation.go @@ -6,6 +6,9 @@ const ( AnnotationDefaultContainerName = "kubectl.kubernetes.io/default-container" AnnotationClusterName = K8sGroupNameSoperator + "/cluster" AnnotationActiveCheckName = K8sGroupNameSoperator + "/activecheck" + AnnotationDependencyVersion = K8sGroupNameSoperator + "/dependency-version" + + AnnotationDependencyVersionName = "name" AnnotationParentalClusterRefName = K8sGroupNameSoperator + "/parental-cluster-ref" 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 682fb7f8a..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" @@ -108,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( @@ -135,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/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)) + }) + } +}