From e3d86f5796b03b428d124f66948696abd33028ff Mon Sep 17 00:00:00 2001 From: Andrii Shestakov Date: Mon, 27 Jul 2026 10:31:54 +0200 Subject: [PATCH] Allow set extra labels/annotations --- api/v1/slurmcluster_types.go | 10 ++++ api/v1/zz_generated.deepcopy.go | 14 +++++ api/v1alpha1/nodeset_types.go | 5 ++ api/v1alpha1/zz_generated.deepcopy.go | 7 +++ .../crd/bases/slurm.nebius.ai_nodesets.yaml | 6 ++ .../bases/slurm.nebius.ai_slurmclusters.yaml | 12 ++++ helm/nodesets/templates/nodeset.yaml | 5 ++ helm/nodesets/values.yaml | 3 + helm/slurm-cluster/templates/pvc.yaml | 8 +++ .../templates/slurm-cluster-cr.yaml | 8 +++ helm/slurm-cluster/values.yaml | 3 + .../templates/slurmcluster-crd.yaml | 18 ++++++ helm/soperator/crds/slurmcluster-crd.yaml | 18 ++++++ .../controller/nodesetcontroller/reconcile.go | 2 + internal/render/accounting/mariadb.go | 4 ++ internal/render/accounting/mariadb_test.go | 31 ++++++++++ internal/render/accounting/pod.go | 11 +++- internal/render/accounting/pod_test.go | 15 +++++ internal/render/controller/statefulset.go | 13 +++-- .../render/controller/statefulset_test.go | 44 ++++++++++++++ internal/render/exporter/pod.go | 8 ++- internal/render/exporter/pod_test.go | 27 +++++++++ internal/render/login/statefulset.go | 7 ++- internal/render/login/statefulset_test.go | 58 +++++++++++++++++++ internal/render/populate_jail/job.go | 8 ++- internal/render/populate_jail/job_test.go | 37 ++++++++++++ internal/render/rest/pod.go | 12 +++- internal/render/rest/pod_test.go | 47 +++++++++++++++ internal/render/sconfigcontroller/pod.go | 16 +++-- internal/render/sconfigcontroller/pod_test.go | 38 ++++++++++++ internal/render/worker/statefulset.go | 1 + internal/render/worker/statefulset_test.go | 57 ++++++++++++++++++ internal/values/slurm_accounting.go | 4 ++ internal/values/slurm_cluster.go | 28 +++++++++ internal/values/slurm_cluster_test.go | 28 +++++++++ internal/values/slurm_controller.go | 2 + internal/values/slurm_exporter.go | 4 ++ internal/values/slurm_jail.go | 4 ++ internal/values/slurm_login.go | 4 ++ internal/values/slurm_nodeset.go | 6 +- internal/values/slurm_nodeset_test.go | 53 ++++++++++++++++- internal/values/slurm_rest.go | 4 ++ internal/values/slurm_sconfigcontroller.go | 4 ++ 43 files changed, 675 insertions(+), 19 deletions(-) create mode 100644 internal/render/rest/pod_test.go create mode 100644 internal/values/slurm_cluster_test.go diff --git a/api/v1/slurmcluster_types.go b/api/v1/slurmcluster_types.go index cb2af64bf..b254c411c 100644 --- a/api/v1/slurmcluster_types.go +++ b/api/v1/slurmcluster_types.go @@ -58,6 +58,16 @@ type SlurmClusterSpec struct { // +kubebuilder:validation:Required SlurmNodes SlurmNodes `json:"slurmNodes"` + // ExtraLabels are custom K8s labels added to every Pod and the spool PVC of this cluster. + // + // +kubebuilder:validation:Optional + ExtraLabels map[string]string `json:"extraLabels,omitempty"` + + // ExtraAnnotations are custom K8s annotations added to every Pod and the spool PVC of this cluster. + // + // +kubebuilder:validation:Optional + ExtraAnnotations map[string]string `json:"extraAnnotations,omitempty"` + // PartitionConfiguration define partition configuration of slurm worker nodes // https://slurm.schedmd.com/slurm.conf.html#SECTION_PARTITION-CONFIGURATION // +kubebuilder:validation:Optional diff --git a/api/v1/zz_generated.deepcopy.go b/api/v1/zz_generated.deepcopy.go index 23c3dc154..42d251fdb 100644 --- a/api/v1/zz_generated.deepcopy.go +++ b/api/v1/zz_generated.deepcopy.go @@ -829,6 +829,20 @@ func (in *SlurmClusterSpec) DeepCopyInto(out *SlurmClusterSpec) { } out.Secrets = in.Secrets in.SlurmNodes.DeepCopyInto(&out.SlurmNodes) + if in.ExtraLabels != nil { + in, out := &in.ExtraLabels, &out.ExtraLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } + if in.ExtraAnnotations != nil { + in, out := &in.ExtraAnnotations, &out.ExtraAnnotations + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } in.PartitionConfiguration.DeepCopyInto(&out.PartitionConfiguration) in.SlurmConfig.DeepCopyInto(&out.SlurmConfig) if in.Topology != nil { diff --git a/api/v1alpha1/nodeset_types.go b/api/v1alpha1/nodeset_types.go index e2a583fd7..19a88793c 100644 --- a/api/v1alpha1/nodeset_types.go +++ b/api/v1alpha1/nodeset_types.go @@ -299,6 +299,11 @@ type NodeSetSpec struct { // +kubebuilder:validation:Optional WorkerAnnotations map[string]string `json:"workerAnnotations,omitempty"` + // WorkerLabels represent K8S labels that should be added to the worker pods. + // + // +kubebuilder:validation:Optional + WorkerLabels map[string]string `json:"workerLabels,omitempty"` + // CustomInitContainers represent additional init containers which will be added to worker pods. // // +kubebuilder:validation:Optional diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 13939a4d1..75a82507b 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -1106,6 +1106,13 @@ func (in *NodeSetSpec) DeepCopyInto(out *NodeSetSpec) { (*out)[key] = val } } + if in.WorkerLabels != nil { + in, out := &in.WorkerLabels, &out.WorkerLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } if in.CustomInitContainers != nil { in, out := &in.CustomInitContainers, &out.CustomInitContainers *out = make([]v1.Container, len(*in)) diff --git a/config/crd/bases/slurm.nebius.ai_nodesets.yaml b/config/crd/bases/slurm.nebius.ai_nodesets.yaml index 4d8c3e4fd..d50e7b501 100644 --- a/config/crd/bases/slurm.nebius.ai_nodesets.yaml +++ b/config/crd/bases/slurm.nebius.ai_nodesets.yaml @@ -12134,6 +12134,12 @@ spec: format: int32 minimum: 0 type: integer + workerLabels: + additionalProperties: + type: string + description: WorkerLabels represent K8S labels that should be added + to the worker pods. + type: object required: - munge - slurmd diff --git a/config/crd/bases/slurm.nebius.ai_slurmclusters.yaml b/config/crd/bases/slurm.nebius.ai_slurmclusters.yaml index f73eb35c0..e8ad8123f 100644 --- a/config/crd/bases/slurm.nebius.ai_slurmclusters.yaml +++ b/config/crd/bases/slurm.nebius.ai_slurmclusters.yaml @@ -86,6 +86,18 @@ spec: Soperator does not guarantee the validity of the raw configuration. Raw config is merged with existing SlurmConfig values. type: string + extraAnnotations: + additionalProperties: + type: string + description: ExtraAnnotations are custom K8s annotations added to + every Pod and the spool PVC of this cluster. + type: object + extraLabels: + additionalProperties: + type: string + description: ExtraLabels are custom K8s labels added to every Pod + and the spool PVC of this cluster. + type: object healthCheckConfig: description: HealthCheckConfig defines Slurm health check configuration. properties: diff --git a/helm/nodesets/templates/nodeset.yaml b/helm/nodesets/templates/nodeset.yaml index 288694b70..60111cdae 100644 --- a/helm/nodesets/templates/nodeset.yaml +++ b/helm/nodesets/templates/nodeset.yaml @@ -307,6 +307,11 @@ spec: {{- toYaml . | nindent 4 }} {{- end }} + {{- with (.workerLabels | default dict) }} + workerLabels: + {{- toYaml . | nindent 4 }} + {{- end }} + {{- include "nodesets.customInitContainers" (dict "root" $ "customInitContainers" .customInitContainers "nodeExporter" .nodeExporter) | nindent 2 }} {{- end }} {{- end }} diff --git a/helm/nodesets/values.yaml b/helm/nodesets/values.yaml index afc81a1e7..42d1b1808 100644 --- a/helm/nodesets/values.yaml +++ b/helm/nodesets/values.yaml @@ -309,6 +309,9 @@ nodesets: workerAnnotations: prometheus.io/scrape: "true" prometheus.io/port: "9090" + # Additional labels to be added to the worker pods + # Optional, defaults to empty dict + workerLabels: {} # A list of custom init containers for worker pods # Each item must be a corev1.Container spec # Optional, defaults to empty list diff --git a/helm/slurm-cluster/templates/pvc.yaml b/helm/slurm-cluster/templates/pvc.yaml index f979a97e0..3d67f3fc7 100644 --- a/helm/slurm-cluster/templates/pvc.yaml +++ b/helm/slurm-cluster/templates/pvc.yaml @@ -5,6 +5,14 @@ kind: PersistentVolumeClaim metadata: namespace: {{ $.Release.Namespace }} name: {{ required "Claim name must be provided." $volume.persistentVolumeClaim.claimName | quote }} + {{- with ($.Values.extraLabels | default dict) }} + labels: + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with ($.Values.extraAnnotations | default dict) }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} spec: accessModes: - ReadWriteMany diff --git a/helm/slurm-cluster/templates/slurm-cluster-cr.yaml b/helm/slurm-cluster/templates/slurm-cluster-cr.yaml index 879a4c592..4155343d2 100644 --- a/helm/slurm-cluster/templates/slurm-cluster-cr.yaml +++ b/helm/slurm-cluster/templates/slurm-cluster-cr.yaml @@ -42,6 +42,14 @@ spec: maintenance: {{ default "none" .Values.maintenance | quote }} crVersion: {{ .Chart.Version }} useDefaultAppArmorProfile: {{ .Values.useDefaultAppArmorProfile }} + {{- with (.Values.extraLabels | default dict) }} + extraLabels: + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with (.Values.extraAnnotations | default dict) }} + extraAnnotations: + {{- toYaml . | nindent 4 }} + {{- end }} partitionConfiguration: configType: {{ (default "default" .Values.partitionConfiguration.configType) }} {{- if and .Values.partitionConfiguration.rawConfig (eq .Values.partitionConfiguration.configType "custom") }} diff --git a/helm/slurm-cluster/values.yaml b/helm/slurm-cluster/values.yaml index 47659c090..57a243d97 100644 --- a/helm/slurm-cluster/values.yaml +++ b/helm/slurm-cluster/values.yaml @@ -1,6 +1,9 @@ clusterName: "slurm1" # Additional annotations for the cluster annotations: {} +# Custom labels/annotations applied to every Pod and the spool PVC +extraLabels: {} +extraAnnotations: {} # Add appArmor profile to the cluster useDefaultAppArmorProfile: false # Maintenance defines the maintenance window for the cluster. diff --git a/helm/soperator-crds/templates/slurmcluster-crd.yaml b/helm/soperator-crds/templates/slurmcluster-crd.yaml index b61bb00d2..4a35b7612 100644 --- a/helm/soperator-crds/templates/slurmcluster-crd.yaml +++ b/helm/soperator-crds/templates/slurmcluster-crd.yaml @@ -27312,6 +27312,12 @@ spec: format: int32 minimum: 0 type: integer + workerLabels: + additionalProperties: + type: string + description: WorkerLabels represent K8S labels that should be added + to the worker pods. + type: object required: - munge - slurmd @@ -27502,6 +27508,18 @@ spec: Soperator does not guarantee the validity of the raw configuration. Raw config is merged with existing SlurmConfig values. type: string + extraAnnotations: + additionalProperties: + type: string + description: ExtraAnnotations are custom K8s annotations added to + every Pod and the spool PVC of this cluster. + type: object + extraLabels: + additionalProperties: + type: string + description: ExtraLabels are custom K8s labels added to every Pod + and the spool PVC of this cluster. + type: object healthCheckConfig: description: HealthCheckConfig defines Slurm health check configuration. properties: diff --git a/helm/soperator/crds/slurmcluster-crd.yaml b/helm/soperator/crds/slurmcluster-crd.yaml index b61bb00d2..4a35b7612 100644 --- a/helm/soperator/crds/slurmcluster-crd.yaml +++ b/helm/soperator/crds/slurmcluster-crd.yaml @@ -27312,6 +27312,12 @@ spec: format: int32 minimum: 0 type: integer + workerLabels: + additionalProperties: + type: string + description: WorkerLabels represent K8S labels that should be added + to the worker pods. + type: object required: - munge - slurmd @@ -27502,6 +27508,18 @@ spec: Soperator does not guarantee the validity of the raw configuration. Raw config is merged with existing SlurmConfig values. type: string + extraAnnotations: + additionalProperties: + type: string + description: ExtraAnnotations are custom K8s annotations added to + every Pod and the spool PVC of this cluster. + type: object + extraLabels: + additionalProperties: + type: string + description: ExtraLabels are custom K8s labels added to every Pod + and the spool PVC of this cluster. + type: object healthCheckConfig: description: HealthCheckConfig defines Slurm health check configuration. properties: diff --git a/internal/controller/nodesetcontroller/reconcile.go b/internal/controller/nodesetcontroller/reconcile.go index d7eb07fc4..3c7dfd99b 100644 --- a/internal/controller/nodesetcontroller/reconcile.go +++ b/internal/controller/nodesetcontroller/reconcile.go @@ -105,6 +105,8 @@ func (r *NodeSetReconciler) reconcile(ctx context.Context, nodeSet *slurmv1alpha cluster.Name, cluster.Spec.Maintenance, cluster.Spec.UseDefaultAppArmorProfile, + cluster.Spec.ExtraLabels, + cluster.Spec.ExtraAnnotations, ) nodeSets, err := resourcegetter.ListNodeSetsByClusterRef(ctx, r.Client, client.ObjectKeyFromObject(cluster)) diff --git a/internal/render/accounting/mariadb.go b/internal/render/accounting/mariadb.go index c3a3d2c53..6118fc300 100644 --- a/internal/render/accounting/mariadb.go +++ b/internal/render/accounting/mariadb.go @@ -102,6 +102,10 @@ func RenderMariaDb( SecurityContext: mariaDb.SecurityContext, }, PodTemplate: mariadbv1alpha1.PodTemplate{ + PodMetadata: &mariadbv1alpha1.Metadata{ + Labels: accounting.Labels, + Annotations: accounting.Annotations, + }, NodeSelector: nodeFilter.NodeSelector, Affinity: affinityConfig, Tolerations: nodeFilter.Tolerations, diff --git a/internal/render/accounting/mariadb_test.go b/internal/render/accounting/mariadb_test.go index 3fa33bcf4..2909f4652 100644 --- a/internal/render/accounting/mariadb_test.go +++ b/internal/render/accounting/mariadb_test.go @@ -3,8 +3,11 @@ package accounting import ( "testing" + mariadbv1alpha1 "github.com/mariadb-operator/mariadb-operator/v25/api/v1alpha1" "github.com/stretchr/testify/assert" corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + "k8s.io/utils/ptr" slurmv1 "nebius.ai/slurm-operator/api/v1" "nebius.ai/slurm-operator/internal/consts" @@ -116,3 +119,31 @@ func Test_RenderMariaDb(t *testing.T) { assert.Equal(t, consts.MariaDbSecretRootName, mariaDb.Spec.RootPasswordSecretKeyRef.SecretKeySelector.Name) assert.Equal(t, consts.MariaDbPasswordKey, mariaDb.Spec.RootPasswordSecretKeyRef.SecretKeySelector.Key) } + +func Test_RenderMariaDb_CustomLabelsAndAnnotations(t *testing.T) { + acc := &values.SlurmAccounting{ + SlurmNode: slurmv1.SlurmNode{ + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + MariaDb: slurmv1.MariaDbOperator{ + Enabled: true, + NodeContainer: slurmv1.NodeContainer{ + Image: "mariadb:10.5", + }, + Storage: mariadbv1alpha1.Storage{ + Size: ptr.To(resource.MustParse("1Gi")), + }, + }, + } + + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + + result, err := RenderMariaDb("test-namespace", "test-cluster", acc, nodeFilters) + assert.NoError(t, err) + + assert.NotNil(t, result.Spec.PodTemplate.PodMetadata) + assert.Equal(t, "123", result.Spec.PodTemplate.PodMetadata.Labels["gcore.com/project-id"]) + assert.Equal(t, "abc", result.Spec.PodTemplate.PodMetadata.Annotations["gcore.com/note"]) +} diff --git a/internal/render/accounting/pod.go b/internal/render/accounting/pod.go index 693d31f0f..10f25f5a5 100644 --- a/internal/render/accounting/pod.go +++ b/internal/render/accounting/pod.go @@ -1,6 +1,7 @@ package accounting import ( + "maps" "slices" corev1 "k8s.io/api/core/v1" @@ -65,10 +66,16 @@ func BasePodTemplateSpec( common.RenderContainerMunge(&accounting.ContainerMunge), ) + labels := maps.Clone(matchLabels) + maps.Copy(labels, accounting.Labels) + + annotations := common.RenderDefaultContainerAnnotation(consts.ContainerNameAccounting) + maps.Copy(annotations, accounting.Annotations) + res := &corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: matchLabels, - Annotations: common.RenderDefaultContainerAnnotation(consts.ContainerNameAccounting), + Labels: labels, + Annotations: annotations, }, Spec: corev1.PodSpec{ HostUsers: accounting.HostUsers, diff --git a/internal/render/accounting/pod_test.go b/internal/render/accounting/pod_test.go index e202768c3..a6680566a 100644 --- a/internal/render/accounting/pod_test.go +++ b/internal/render/accounting/pod_test.go @@ -115,3 +115,18 @@ func Test_BasePodTemplateSpec_PriorityClass(t *testing.T) { }) } } + +func Test_BasePodTemplateSpec_CustomLabelsAndAnnotations(t *testing.T) { + testAcc := *acc + testAcc.Labels = map[string]string{"gcore.com/project-id": "123"} + testAcc.Annotations = map[string]string{"gcore.com/note": "abc"} + + result, err := accounting.BasePodTemplateSpec( + defaultNameCluster, &testAcc, defaultNodeFilter, defaultVolumeSources, matchLabels, + ) + assert.NoError(t, err) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "value", result.Labels["key"]) // matchLabels preserved + assert.Equal(t, "abc", result.Annotations["gcore.com/note"]) +} diff --git a/internal/render/controller/statefulset.go b/internal/render/controller/statefulset.go index 5ae081b91..07500a9fe 100644 --- a/internal/render/controller/statefulset.go +++ b/internal/render/controller/statefulset.go @@ -2,6 +2,7 @@ package controller import ( "fmt" + "maps" "slices" appspub "github.com/openkruise/kruise-api/apps/pub" @@ -34,6 +35,12 @@ func RenderStatefulSet( labels[consts.LabelControllerType] = consts.LabelControllerTypeMain matchLabels[consts.LabelControllerType] = consts.LabelControllerTypeMain + maps.Copy(labels, controller.Labels) + + annotations := map[string]string{ + consts.AnnotationDefaultContainerName: consts.ContainerNameSlurmctld, + } + maps.Copy(annotations, controller.Annotations) nodeFilter := sliceutils.MustGetBy( nodeFilters, @@ -100,10 +107,8 @@ func RenderStatefulSet( }, Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: labels, - Annotations: map[string]string{ - consts.AnnotationDefaultContainerName: consts.ContainerNameSlurmctld, - }, + Labels: labels, + Annotations: annotations, }, Spec: corev1.PodSpec{ ReadinessGates: []corev1.PodReadinessGate{ diff --git a/internal/render/controller/statefulset_test.go b/internal/render/controller/statefulset_test.go index 8d7608cde..fd3b72e55 100644 --- a/internal/render/controller/statefulset_test.go +++ b/internal/render/controller/statefulset_test.go @@ -537,3 +537,47 @@ func TestRenderStatefulSetHostUsers(t *testing.T) { }) } } + +func TestRenderStatefulSet_CustomLabelsAndAnnotations(t *testing.T) { + controller := &values.SlurmController{ + K8sNodeFilterName: "test-filter", + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + StatefulSet: values.StatefulSet{ + Name: "test-controller-sts", + Replicas: 1, + MaxUnavailable: intstr.FromInt32(1), + }, + Service: values.Service{ + Name: "test-controller-svc", + }, + ContainerSlurmctld: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-image:latest", + }, + Name: "slurmctld", + }, + ContainerMunge: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "munge-image:latest", + }, + }, + VolumeSpool: slurmv1.NodeVolume{VolumeSourceName: ptr.To("test-volume")}, + VolumeJail: slurmv1.NodeVolume{VolumeSourceName: ptr.To("test-volume")}, + } + + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + volumeSources := []slurmv1.VolumeSource{ + { + Name: "test-volume", + VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}, + }, + } + + result, err := RenderStatefulSet("test-namespace", "test-cluster", nodeFilters, volumeSources, controller, false) + assert.NoError(t, err) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "123", result.Spec.Template.Labels["gcore.com/project-id"]) + assert.Equal(t, "abc", result.Spec.Template.Annotations["gcore.com/note"]) +} diff --git a/internal/render/exporter/pod.go b/internal/render/exporter/pod.go index 6e9d213b9..f573414fe 100644 --- a/internal/render/exporter/pod.go +++ b/internal/render/exporter/pod.go @@ -1,6 +1,8 @@ package exporter import ( + "maps" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -23,9 +25,13 @@ func renderPodTemplateSpec( _ = err // Ignore not found error, use "empty" node filter. nodeFilter = slurmv1.K8sNodeFilter{} } + labels := maps.Clone(matchLabels) + maps.Copy(labels, clusterValues.SlurmExporter.Labels) + result := corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: matchLabels, + Labels: labels, + Annotations: maps.Clone(clusterValues.SlurmExporter.Annotations), }, Spec: corev1.PodSpec{ HostUsers: clusterValues.SlurmExporter.HostUsers, diff --git a/internal/render/exporter/pod_test.go b/internal/render/exporter/pod_test.go index ac3bda5cb..8bfa13225 100644 --- a/internal/render/exporter/pod_test.go +++ b/internal/render/exporter/pod_test.go @@ -129,3 +129,30 @@ func Test_renderPodTemplateSpec_PriorityClass(t *testing.T) { }) } } + +func Test_renderPodTemplateSpec_CustomLabelsAndAnnotations(t *testing.T) { + clusterValues := &values.SlurmCluster{ + SlurmExporter: values.SlurmExporter{ + SlurmNode: slurmv1.SlurmNode{ + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + Container: slurmv1.NodeContainer{ + Image: "test-image", + }, + }, + NodeFilters: []slurmv1.K8sNodeFilter{ + {Name: "test-filter"}, + }, + } + + initContainers := []corev1.Container{} + matchLabels := map[string]string{"app": "test"} + + result := renderPodTemplateSpec(clusterValues, initContainers, matchLabels) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "test", result.Labels["app"]) // matchLabels preserved + assert.Equal(t, "abc", result.Annotations["gcore.com/note"]) +} diff --git a/internal/render/login/statefulset.go b/internal/render/login/statefulset.go index 2deb9ace9..c994d320d 100644 --- a/internal/render/login/statefulset.go +++ b/internal/render/login/statefulset.go @@ -2,6 +2,7 @@ package login import ( "fmt" + "maps" kruisev1b1 "github.com/openkruise/kruise-api/apps/v1beta1" appsv1 "k8s.io/api/apps/v1" @@ -30,6 +31,10 @@ func RenderStatefulSet( ) (kruisev1b1.StatefulSet, error) { labels := common.RenderLabels(consts.ComponentTypeLogin, clusterName) matchLabels := common.RenderMatchLabels(consts.ComponentTypeLogin, clusterName) + maps.Copy(labels, login.Labels) + + annotations := common.RenderDefaultContainerAnnotation(consts.ContainerNameSshd) + maps.Copy(annotations, login.Annotations) nodeFilter := utils.MustGetBy( nodeFilters, @@ -98,7 +103,7 @@ func RenderStatefulSet( Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ Labels: labels, - Annotations: common.RenderDefaultContainerAnnotation(consts.ContainerNameSshd), + Annotations: annotations, }, Spec: corev1.PodSpec{ HostUsers: login.HostUsers, diff --git a/internal/render/login/statefulset_test.go b/internal/render/login/statefulset_test.go index 24f1ce963..851cf7bcd 100644 --- a/internal/render/login/statefulset_test.go +++ b/internal/render/login/statefulset_test.go @@ -115,3 +115,61 @@ func TestRenderStatefulSet_PriorityClass(t *testing.T) { }) } } + +func TestRenderStatefulSet_CustomLabelsAndAnnotations(t *testing.T) { + namespace := "test-namespace" + clusterName := "test-cluster" + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + secrets := &slurmv1.Secrets{} + volumeSources := []slurmv1.VolumeSource{ + { + Name: "test-volume", + VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{}}, + }, + } + + login := &values.SlurmLogin{ + SlurmNode: slurmv1.SlurmNode{ + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + ContainerSshd: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-sshd-image", + Port: 22, + Resources: corev1.ResourceList{ + corev1.ResourceMemory: resource.MustParse("1Gi"), + corev1.ResourceCPU: resource.MustParse("100m"), + }, + }, + }, + ContainerMunge: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-munge-image", + Resources: corev1.ResourceList{ + corev1.ResourceMemory: resource.MustParse("1Gi"), + corev1.ResourceCPU: resource.MustParse("100m"), + }, + }, + }, + VolumeJail: slurmv1.NodeVolume{VolumeSourceName: &[]string{"test-volume"}[0]}, + StatefulSet: values.StatefulSet{Name: "test-login", Replicas: 1}, + HeadlessService: values.Service{Name: "test-headless"}, + } + + result, err := RenderStatefulSet(namespace, clusterName, true, nodeFilters, secrets, volumeSources, login) + if err != nil { + t.Fatalf("RenderStatefulSet() error = %v", err) + } + + if got := result.Labels["gcore.com/project-id"]; got != "123" { + t.Errorf("StatefulSet label = %v, want 123", got) + } + if got := result.Spec.Template.Labels["gcore.com/project-id"]; got != "123" { + t.Errorf("Pod template label = %v, want 123", got) + } + if got := result.Spec.Template.Annotations["gcore.com/note"]; got != "abc" { + t.Errorf("Pod template annotation = %v, want abc", got) + } +} diff --git a/internal/render/populate_jail/job.go b/internal/render/populate_jail/job.go index 4cd697ef5..ac892d2ae 100644 --- a/internal/render/populate_jail/job.go +++ b/internal/render/populate_jail/job.go @@ -1,6 +1,8 @@ package populate_jail import ( + "maps" + batchv1 "k8s.io/api/batch/v1" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -21,6 +23,10 @@ func RenderPopulateJailJob( populateJail *values.PopulateJail, ) batchv1.Job { labels := common.RenderLabels(consts.ComponentTypePopulateJail, clusterName) + maps.Copy(labels, populateJail.Labels) + + annotations := common.RenderDefaultContainerAnnotation(consts.ContainerNamePopulateJail) + maps.Copy(annotations, populateJail.Annotations) nodeFilter := utils.MustGetBy( nodeFilters, @@ -50,7 +56,7 @@ func RenderPopulateJailJob( Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ Labels: labels, - Annotations: common.RenderDefaultContainerAnnotation(consts.ContainerNamePopulateJail), + Annotations: annotations, }, Spec: corev1.PodSpec{ HostUsers: populateJail.HostUsers, diff --git a/internal/render/populate_jail/job_test.go b/internal/render/populate_jail/job_test.go index 67c351959..abf70c649 100644 --- a/internal/render/populate_jail/job_test.go +++ b/internal/render/populate_jail/job_test.go @@ -87,3 +87,40 @@ func Test_RenderPopulateJailJob_PriorityClass(t *testing.T) { }) } } + +func Test_RenderPopulateJailJob_CustomLabelsAndAnnotations(t *testing.T) { + namespace := "test-namespace" + clusterName := "test-cluster" + + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + volumeSources := []slurmv1.VolumeSource{ + { + Name: "test-volume-source", + VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{}}, + }, + } + + populateJail := &values.PopulateJail{ + PopulateJail: slurmv1.PopulateJail{ + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + Name: "test-populate-jail", + ContainerPopulateJail: values.Container{ + Name: "populate-jail", + NodeContainer: slurmv1.NodeContainer{ + Image: "nginx:latest", + }, + }, + VolumeJail: slurmv1.NodeVolume{ + VolumeSourceName: ptr.To("test-volume-source"), + }, + } + + result := populate_jail.RenderPopulateJailJob(namespace, clusterName, nodeFilters, volumeSources, populateJail) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "123", result.Spec.Template.Labels["gcore.com/project-id"]) + assert.Equal(t, "abc", result.Spec.Template.Annotations["gcore.com/note"]) +} diff --git a/internal/render/rest/pod.go b/internal/render/rest/pod.go index 0103771cd..a5ca2da28 100644 --- a/internal/render/rest/pod.go +++ b/internal/render/rest/pod.go @@ -1,6 +1,8 @@ package rest import ( + "maps" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -32,10 +34,16 @@ func BasePodTemplateSpec( return nil, err } + labels := maps.Clone(matchLabels) + maps.Copy(labels, valuesREST.Labels) + + annotations := common.RenderDefaultContainerAnnotation(consts.ContainerNameREST) + maps.Copy(annotations, valuesREST.Annotations) + res := &corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: matchLabels, - Annotations: common.RenderDefaultContainerAnnotation(consts.ContainerNameREST), + Labels: labels, + Annotations: annotations, }, Spec: corev1.PodSpec{ HostUsers: valuesREST.HostUsers, diff --git a/internal/render/rest/pod_test.go b/internal/render/rest/pod_test.go new file mode 100644 index 000000000..f1bc03182 --- /dev/null +++ b/internal/render/rest/pod_test.go @@ -0,0 +1,47 @@ +package rest + +import ( + "testing" + + "github.com/stretchr/testify/assert" + corev1 "k8s.io/api/core/v1" + + slurmv1 "nebius.ai/slurm-operator/api/v1" + "nebius.ai/slurm-operator/internal/values" +) + +func Test_BasePodTemplateSpec_CustomLabelsAndAnnotations(t *testing.T) { + valuesREST := &values.SlurmREST{ + SlurmNode: slurmv1.SlurmNode{ + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + ContainerREST: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-rest-image", + }, + }, + VolumeJail: slurmv1.NodeVolume{ + VolumeSourceName: &[]string{"test-volume-source"}[0], + }, + } + + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + volumeSources := []slurmv1.VolumeSource{ + { + Name: "test-volume-source", + VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{}}, + }, + } + matchLabels := map[string]string{"key": "value"} + + result, err := BasePodTemplateSpec( + "test-cluster", valuesREST, nodeFilters, volumeSources, matchLabels, + ) + assert.NoError(t, err) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "value", result.Labels["key"]) // matchLabels preserved + assert.Equal(t, "abc", result.Annotations["gcore.com/note"]) +} diff --git a/internal/render/sconfigcontroller/pod.go b/internal/render/sconfigcontroller/pod.go index 39fed3b1a..8fe63b867 100644 --- a/internal/render/sconfigcontroller/pod.go +++ b/internal/render/sconfigcontroller/pod.go @@ -1,6 +1,8 @@ package sconfigcontroller import ( + "maps" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -41,12 +43,18 @@ func BasePodTemplateSpec( } } + labels := maps.Clone(matchLabels) + maps.Copy(labels, sConfigController.Labels) + + annotations := map[string]string{ + consts.AnnotationDefaultContainerName: consts.ContainerNameSConfigController, + } + maps.Copy(annotations, sConfigController.Annotations) + res := &corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: matchLabels, - Annotations: map[string]string{ - consts.AnnotationDefaultContainerName: consts.ContainerNameSConfigController, - }, + Labels: labels, + Annotations: annotations, }, Spec: corev1.PodSpec{ HostUsers: sConfigController.HostUsers, diff --git a/internal/render/sconfigcontroller/pod_test.go b/internal/render/sconfigcontroller/pod_test.go index dabb80fe0..382816f05 100644 --- a/internal/render/sconfigcontroller/pod_test.go +++ b/internal/render/sconfigcontroller/pod_test.go @@ -94,3 +94,41 @@ func Test_BasePodTemplateSpec_PriorityClass(t *testing.T) { }) } } + +func Test_BasePodTemplateSpec_CustomLabelsAndAnnotations(t *testing.T) { + sConfigController := &values.SConfigController{ + SlurmNode: slurmv1.SlurmNode{ + Size: 1, + K8sNodeFilterName: "test-filter", + }, + Labels: map[string]string{"gcore.com/project-id": "123"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + Container: values.Container{ + Name: "test-container", + NodeContainer: slurmv1.NodeContainer{ + Image: "nginx:latest", + }, + }, + VolumeJail: slurmv1.NodeVolume{ + VolumeSourceName: ptr.To("test-volume-source"), + }, + } + + nodeFilters := []slurmv1.K8sNodeFilter{{Name: "test-filter"}} + volumeSources := []slurmv1.VolumeSource{ + { + Name: "test-volume-source", + VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{}}, + }, + } + matchLabels := map[string]string{"key": "value"} + + result, err := sconfigcontroller.BasePodTemplateSpec( + "test-namespace", "test-cluster", "http://slurm-api-server", sConfigController, nodeFilters, volumeSources, matchLabels, + ) + assert.NoError(t, err) + + assert.Equal(t, "123", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "value", result.Labels["key"]) // matchLabels preserved + assert.Equal(t, "abc", result.Annotations["gcore.com/note"]) +} diff --git a/internal/render/worker/statefulset.go b/internal/render/worker/statefulset.go index 0e0a88b3a..80f9a687a 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 + maps.Copy(labels, nodeSet.Labels) matchLabels := common.RenderMatchLabels(consts.ComponentTypeNodeSet, nodeSet.ParentalCluster.Name) matchLabels[consts.LabelNodeSetKey] = nodeSet.Name diff --git a/internal/render/worker/statefulset_test.go b/internal/render/worker/statefulset_test.go index 1fcc37a64..f08de070f 100644 --- a/internal/render/worker/statefulset_test.go +++ b/internal/render/worker/statefulset_test.go @@ -1293,3 +1293,60 @@ func TestRenderNodeSetStatefulSet_EphemeralNodesReserveOrdinals(t *testing.T) { }) } } + +func TestRenderNodeSetStatefulSet_CustomLabelsAndAnnotations(t *testing.T) { + nodeSet := &values.SlurmNodeSet{ + Name: "test-nodeset", + ParentalCluster: client.ObjectKey{ + Namespace: "test-namespace", + Name: "test-cluster", + }, + Labels: map[string]string{"gcore.com/flavor": "gpu-8x-h100"}, + Annotations: map[string]string{"gcore.com/note": "abc"}, + ContainerSlurmd: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "test-image", + ImagePullPolicy: corev1.PullIfNotPresent, + Resources: corev1.ResourceList{ + corev1.ResourceMemory: resource.MustParse("1Gi"), + corev1.ResourceCPU: resource.MustParse("100m"), + corev1.ResourceEphemeralStorage: resource.MustParse("1Gi"), + }, + }, + }, + ContainerMunge: values.Container{ + NodeContainer: slurmv1.NodeContainer{ + Image: "munge-image", + }, + }, + VolumeSpool: corev1.VolumeSource{ + HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/spool"}, + }, + VolumeJail: corev1.VolumeSource{ + HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/jail"}, + }, + StatefulSet: values.StatefulSet{ + Replicas: 1, + }, + ServiceUmbrella: values.Service{Name: "test-umbrella"}, + SupervisorDConfigMapName: "supervisord-config", + SSHDConfigMapName: "sshd-config", + GPU: &slurmv1alpha1.GPUSpec{Enabled: false}, + } + + result, err := worker.RenderNodeSetStatefulSet( + "test-cluster", + nodeSet, + &slurmv1.Secrets{}, + consts.CGroupV2, + false, + false, + "", + ) + assert.NoError(t, err) + + assert.Equal(t, "gpu-8x-h100", result.Labels["gcore.com/flavor"]) + assert.Equal(t, "gpu-8x-h100", result.Spec.Template.Labels["gcore.com/flavor"]) + assert.Equal(t, consts.LabelWorkerValue, result.Spec.Template.Labels[consts.LabelWorkerKey]) + assert.Equal(t, "abc", result.Spec.Template.Annotations["gcore.com/note"]) +} diff --git a/internal/values/slurm_accounting.go b/internal/values/slurm_accounting.go index 83af062d3..d168c8ef9 100644 --- a/internal/values/slurm_accounting.go +++ b/internal/values/slurm_accounting.go @@ -13,6 +13,10 @@ import ( type SlurmAccounting struct { slurmv1.SlurmNode + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + Enabled bool ContainerAccounting Container diff --git a/internal/values/slurm_cluster.go b/internal/values/slurm_cluster.go index 3e6379725..b81a77978 100644 --- a/internal/values/slurm_cluster.go +++ b/internal/values/slurm_cluster.go @@ -3,6 +3,7 @@ package values import ( "context" "fmt" + "maps" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/log" @@ -86,6 +87,22 @@ func BuildSlurmClusterFrom(ctx context.Context, cluster *slurmv1.SlurmCluster, n UseDefaultAppArmorProfile: cluster.Spec.UseDefaultAppArmorProfile, } + // Apply cluster-wide ExtraLabels/ExtraAnnotations to every component. + for _, m := range []*map[string]string{ + &res.NodeController.Labels, &res.NodeAccounting.Labels, &res.NodeRest.Labels, + &res.NodeLogin.Labels, &res.SlurmExporter.Labels, &res.SConfigController.Labels, + &res.PopulateJail.Labels, + } { + *m = mergeClusterExtra(cluster.Spec.ExtraLabels, *m) + } + for _, m := range []*map[string]string{ + &res.NodeController.Annotations, &res.NodeAccounting.Annotations, &res.NodeRest.Annotations, + &res.NodeLogin.Annotations, &res.SlurmExporter.Annotations, &res.SConfigController.Annotations, + &res.PopulateJail.Annotations, + } { + *m = mergeClusterExtra(cluster.Spec.ExtraAnnotations, *m) + } + if err := res.Validate(ctx); err != nil { logger.Error(err, "SlurmCluster validation failed") return res, fmt.Errorf("failed to validate SlurmCluster: %w", err) @@ -94,6 +111,17 @@ func BuildSlurmClusterFrom(ctx context.Context, cluster *slurmv1.SlurmCluster, n return res, nil } +// mergeClusterExtra merges cluster-wide extra key/value pairs with component-specific ones, +// with the component-specific values taking precedence on conflicting keys. +func mergeClusterExtra(clusterWide, componentSpecific map[string]string) map[string]string { + if len(clusterWide) == 0 { + return componentSpecific + } + merged := maps.Clone(clusterWide) + maps.Copy(merged, componentSpecific) + return merged +} + func BuildClusterWithGPUFromNodeSets(nodeSets []slurmav1alpha1.NodeSet) bool { for _, nodeSet := range nodeSets { if nodeSet.Spec.GPU.Enabled { diff --git a/internal/values/slurm_cluster_test.go b/internal/values/slurm_cluster_test.go new file mode 100644 index 000000000..e059cdd1a --- /dev/null +++ b/internal/values/slurm_cluster_test.go @@ -0,0 +1,28 @@ +package values + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func Test_mergeClusterExtra(t *testing.T) { + t.Run("component-specific wins on conflicting keys", func(t *testing.T) { + got := mergeClusterExtra( + map[string]string{"a": "cluster", "b": "cluster"}, + map[string]string{"a": "component"}, + ) + assert.Equal(t, map[string]string{"a": "component", "b": "cluster"}, got) + }) + + t.Run("nil cluster-wide map returns component-specific as-is", func(t *testing.T) { + componentSpecific := map[string]string{"a": "component"} + got := mergeClusterExtra(nil, componentSpecific) + assert.Equal(t, componentSpecific, got) + }) + + t.Run("nil component-specific map returns a clone of cluster-wide", func(t *testing.T) { + got := mergeClusterExtra(map[string]string{"a": "cluster"}, nil) + assert.Equal(t, map[string]string{"a": "cluster"}, got) + }) +} diff --git a/internal/values/slurm_controller.go b/internal/values/slurm_controller.go index 1566686e9..e7dff9426 100644 --- a/internal/values/slurm_controller.go +++ b/internal/values/slurm_controller.go @@ -13,6 +13,8 @@ type SlurmController struct { K8sNodeFilterName string CustomInitContainers []corev1.Container HostUsers *bool + Labels map[string]string + Annotations map[string]string ContainerSlurmctld Container ContainerMunge Container diff --git a/internal/values/slurm_exporter.go b/internal/values/slurm_exporter.go index 7922bd891..439b32884 100644 --- a/internal/values/slurm_exporter.go +++ b/internal/values/slurm_exporter.go @@ -13,6 +13,10 @@ import ( type SlurmExporter struct { slurmv1.SlurmNode + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + Enabled bool PodMonitorConfig slurmv1.PodMonitorConfig diff --git a/internal/values/slurm_jail.go b/internal/values/slurm_jail.go index 0fae24249..70863c455 100644 --- a/internal/values/slurm_jail.go +++ b/internal/values/slurm_jail.go @@ -11,6 +11,10 @@ import ( type PopulateJail struct { slurmv1.PopulateJail + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + Name string ContainerPopulateJail Container diff --git a/internal/values/slurm_login.go b/internal/values/slurm_login.go index 45179a78f..abf5ec97c 100644 --- a/internal/values/slurm_login.go +++ b/internal/values/slurm_login.go @@ -12,6 +12,10 @@ import ( type SlurmLogin struct { slurmv1.SlurmNode + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + ContainerSshd Container ContainerMunge Container ContainerSSSD *Container diff --git a/internal/values/slurm_nodeset.go b/internal/values/slurm_nodeset.go index 6eb1d0530..7a68bb814 100644 --- a/internal/values/slurm_nodeset.go +++ b/internal/values/slurm_nodeset.go @@ -24,6 +24,7 @@ type SlurmNodeSet struct { Tolerations []corev1.Toleration PriorityClass string Annotations map[string]string + Labels map[string]string ImagePullSecrets []corev1.LocalObjectReference ContainerSlurmd Container @@ -79,6 +80,8 @@ func BuildSlurmNodeSetFrom( clusterName string, maintenance *consts.MaintenanceMode, useDefaultAppArmorProfile bool, + clusterExtraLabels map[string]string, + clusterExtraAnnotations map[string]string, ) SlurmNodeSet { nsSpec := &nodeSet.Spec res := SlurmNodeSet{ @@ -92,7 +95,8 @@ func BuildSlurmNodeSetFrom( Affinity: nsSpec.Affinity.DeepCopy(), Tolerations: slices.Clone(nsSpec.Tolerations), PriorityClass: nsSpec.PriorityClass, - Annotations: maps.Clone(nsSpec.WorkerAnnotations), + Annotations: mergeClusterExtra(clusterExtraAnnotations, maps.Clone(nsSpec.WorkerAnnotations)), + Labels: mergeClusterExtra(clusterExtraLabels, maps.Clone(nsSpec.WorkerLabels)), ImagePullSecrets: slices.Clone(nsSpec.ImagePullSecrets), // ContainerSlurmd: buildContainerFrom( diff --git a/internal/values/slurm_nodeset_test.go b/internal/values/slurm_nodeset_test.go index e0688d3e0..b7ac7652b 100644 --- a/internal/values/slurm_nodeset_test.go +++ b/internal/values/slurm_nodeset_test.go @@ -39,7 +39,7 @@ func TestBuildSlurmNodeSetFrom_SSSD(t *testing.T) { Image: slurmv1alpha1.Image{Repository: "sssd", Tag: "latest"}, } - result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false) + result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false, nil, nil) if assert.NotNil(t, result.ContainerSSSD) { assert.Equal(t, "sssd:latest", result.ContainerSSSD.Image) @@ -56,7 +56,7 @@ func TestBuildSlurmNodeSetFrom_SSSD(t *testing.T) { } nodeSet.Spec.SSSDConfSecretRefName = "custom-worker-sssd" - result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false) + result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false, nil, nil) assert.False(t, result.IsSSSDSecretDefault) assert.Equal(t, "custom-worker-sssd", result.SSSDConfSecretName) @@ -65,7 +65,7 @@ func TestBuildSlurmNodeSetFrom_SSSD(t *testing.T) { t.Run("keeps sssd disabled when container is not configured", func(t *testing.T) { nodeSet := makeNodeSet() - result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false) + result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false, nil, nil) assert.Nil(t, result.ContainerSSSD) assert.Empty(t, result.SSSDConfSecretName) @@ -73,6 +73,53 @@ func TestBuildSlurmNodeSetFrom_SSSD(t *testing.T) { }) } +func TestBuildSlurmNodeSetFrom_ExtraLabelsAndAnnotations(t *testing.T) { + makeNodeSet := func() *slurmv1alpha1.NodeSet { + return &slurmv1alpha1.NodeSet{ + ObjectMeta: metav1ObjectMeta("worker-a", "test-ns"), + Spec: slurmv1alpha1.NodeSetSpec{ + Slurmd: slurmv1alpha1.ContainerSlurmdSpec{ + Image: slurmv1alpha1.Image{Repository: "slurmd", Tag: "latest"}, + Volumes: slurmv1alpha1.WorkerVolumesSpec{ + Spool: corev1.VolumeSource{}, + Jail: corev1.VolumeSource{}, + }, + }, + Munge: slurmv1alpha1.ContainerMungeSpec{ + Image: slurmv1alpha1.Image{Repository: "munge", Tag: "latest"}, + }, + }, + } + } + + t.Run("cluster-wide labels/annotations flow through when NodeSet sets none", func(t *testing.T) { + nodeSet := makeNodeSet() + + result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false, + map[string]string{"gcore.com/project-id": "proj-1"}, + map[string]string{"gcore.com/note": "cluster-wide"}, + ) + + assert.Equal(t, "proj-1", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "cluster-wide", result.Annotations["gcore.com/note"]) + }) + + t.Run("NodeSet WorkerLabels/WorkerAnnotations take precedence on conflicting keys", func(t *testing.T) { + nodeSet := makeNodeSet() + nodeSet.Spec.WorkerLabels = map[string]string{"gcore.com/project-id": "override", "gcore.com/flavor": "gpu-8x-h100"} + nodeSet.Spec.WorkerAnnotations = map[string]string{"gcore.com/note": "override"} + + result := BuildSlurmNodeSetFrom(nodeSet, "test-cluster", nil, false, + map[string]string{"gcore.com/project-id": "proj-1"}, + map[string]string{"gcore.com/note": "cluster-wide"}, + ) + + assert.Equal(t, "override", result.Labels["gcore.com/project-id"]) + assert.Equal(t, "gpu-8x-h100", result.Labels["gcore.com/flavor"]) + assert.Equal(t, "override", result.Annotations["gcore.com/note"]) + }) +} + func TestDefaultPersistentVolumeClaimRetentionPolicy(t *testing.T) { t.Run("defaults to delete for both fields when unset", func(t *testing.T) { got := defaultPersistentVolumeClaimRetentionPolicy(nil) diff --git a/internal/values/slurm_rest.go b/internal/values/slurm_rest.go index 64ad6bbf9..301cb0bb9 100644 --- a/internal/values/slurm_rest.go +++ b/internal/values/slurm_rest.go @@ -13,6 +13,10 @@ import ( type SlurmREST struct { slurmv1.SlurmNode + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + Enabled bool ThreadCount *int32 MaxConnections *int32 diff --git a/internal/values/slurm_sconfigcontroller.go b/internal/values/slurm_sconfigcontroller.go index 65b035019..b42ba7921 100644 --- a/internal/values/slurm_sconfigcontroller.go +++ b/internal/values/slurm_sconfigcontroller.go @@ -12,6 +12,10 @@ import ( type SConfigController struct { slurmv1.SlurmNode + // Set from the cluster's ExtraLabels/ExtraAnnotations in BuildSlurmClusterFrom. + Labels map[string]string + Annotations map[string]string + Container Container VolumeJail slurmv1.NodeVolume