feat: bench job support inherit kubebench toleration (#74)

apecloud · Aug 13, 2024 · b31cdb4 · b31cdb4
1 parent bc82721
commit b31cdb4
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 2 deletions.
diff --git a/deploy/helm/templates/deployment.yaml b/deploy/helm/templates/deployment.yaml
@@ -55,6 +55,10 @@ spec:
               value: "{{ .Values.kubebenchImages.exporter }}"
             - name: KUBEBENCH_TOOLS_IMAGE
               value: "{{ .Values.kubebenchImages.tools }}"
+            {{- with .Values.tolerations }}
+            - name: CM_TOLERATIONS
+              value: {{ toJson . | quote }}
+            {{- end }}
           command:
           - /manager
           securityContext:

diff --git a/internal/utils/job.go b/internal/utils/job.go
@@ -18,19 +18,23 @@ package utils
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
-	"github.com/apecloud/kubebench/api/v1alpha1"
-	"github.com/apecloud/kubebench/pkg/constants"
 	"strings"
 
+	"github.com/spf13/viper"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/rest"
+	"k8s.io/klog/v2"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	"github.com/apecloud/kubebench/api/v1alpha1"
+	"github.com/apecloud/kubebench/pkg/constants"
 )
 
 func IsJobExisted(cli client.Client, reqCtx context.Context, jobName string, namespace string) (bool, error) {
@@ -200,6 +204,12 @@ func JobTemplate(name, namespace string) *batchv1.Job {
 
 func AddTolerationToJobs(jobs []*batchv1.Job, tolerations []corev1.Toleration) {
 	for _, job := range jobs {
+		if cmTolerations := viper.GetString(constants.CfgKeyCtrlrMgrTolerations); cmTolerations != "" {
+			if err := json.Unmarshal([]byte(cmTolerations), &job.Spec.Template.Spec.Tolerations); err != nil {
+				klog.Errorf("Failed to marshal CfgKeyCtrlrMgrTolerations for job, tolerations is %s ", cmTolerations)
+			}
+		}
+
 		job.Spec.Template.Spec.Tolerations = append(job.Spec.Template.Spec.Tolerations, tolerations...)
 	}
 }

diff --git a/pkg/constants/env.go b/pkg/constants/env.go
@@ -19,6 +19,10 @@ const (
 	KubebenchTools         = "KUBEBENCH_TOOLS_IMAGE"
 )
 
+const (
+	CfgKeyCtrlrMgrTolerations = "CM_TOLERATIONS"
+)
+
 func init() {
 	viper.SetDefault(KubebenchEnvPgbench, "registry.cn-hangzhou.aliyuncs.com/apecloud/spilo:14.8.0")
 	viper.SetDefault(KubebenchEnvSysbench, "registry.cn-hangzhou.aliyuncs.com/apecloud/customsuites:latest")
@@ -30,6 +34,7 @@ func init() {
 	viper.SetDefault(KubebenchEnvRedisBench, "registry.cn-hangzhou.aliyuncs.com/apecloud/redis:7.0.5")
 	viper.SetDefault(KubebenchExporter, "apecloud-registry.cn-zhangjiakou.cr.aliyuncs.com/apecloud/kubebench:0.0.10")
 	viper.SetDefault(KubebenchTools, "apecloud-registry.cn-zhangjiakou.cr.aliyuncs.com/apecloud/kubebench:0.0.10")
+	viper.SetDefault(CfgKeyCtrlrMgrTolerations, os.Getenv(CfgKeyCtrlrMgrTolerations))
 }
 
 // GetBenchmarkImage get benchmark image