Merge pull request #70 from zqzten/cronjob

HorizontalPortrait: Support service account defaulting for algorithm job
traas-stack · Oct 19, 2023 · 786848e · 786848e
2 parents 4e6365b + 7a25c88
commit 786848e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/main.go b/main.go
@@ -96,8 +96,9 @@ type promConfig struct {
 
 type algorithmJobConfig struct {
 	Namespace                        string
+	DefaultServiceAccount            string
 	DefaultMetricsServerAddr         string
-	DefaultImagePredictiveHorizontal string
+	DefaultImageHorizontalPredictive string
 }
 
 func main() {
@@ -150,10 +151,12 @@ func main() {
 	flag.DurationVar(&promConfig.MetricsMaxAge, "prometheus-metrics-max-age", 0,
 		"The period for which to query the set of available metrics from Prometheus. If not set, it defaults to prometheus-metrics-relist-interval.")
 	flag.StringVar(&algorithmJobConfig.Namespace, "algorithm-job-namespace", "", "The namespace where algorithm jobs should run in.")
+	flag.StringVar(&algorithmJobConfig.DefaultServiceAccount, "algorithm-job-default-service-account", "",
+		"The default service account for algorithm jobs.")
 	flag.StringVar(&algorithmJobConfig.DefaultMetricsServerAddr, "algorithm-job-default-metrics-server-addr", "",
 		"The default metrics server address to inject to env of algorithm jobs.")
-	flag.StringVar(&algorithmJobConfig.DefaultImagePredictiveHorizontal, "algorithm-job-default-image-predictive-horizontal", "",
-		"The default image for algorithm jobs of predictive horizontal portrait.")
+	flag.StringVar(&algorithmJobConfig.DefaultImageHorizontalPredictive, "algorithm-job-default-image-horizontal-predictive", "",
+		"The default image for algorithm jobs of horizontal predictive portrait.")
 	opts := zap.Options{
 		Development: true,
 		TimeEncoder: zapcore.RFC3339TimeEncoder,
@@ -445,8 +448,8 @@ func initExternalHorizontalPortraitAlgorithmJobControllers(client client.Client,
 		return nil, fmt.Errorf("shall set namespace for algorithm jobs")
 	}
 	controllers := make(map[autoscalingv1alpha1.PortraitAlgorithmJobType]jobcontroller.Horizontal)
-	controllers[autoscalingv1alpha1.CronJobPortraitAlgorithmJobType] = jobcontroller.NewCronJobHorizontal(client, config.Namespace, config.DefaultMetricsServerAddr, map[autoscalingv1alpha1.PortraitType]string{
-		autoscalingv1alpha1.PredictivePortraitType: config.DefaultImagePredictiveHorizontal,
+	controllers[autoscalingv1alpha1.CronJobPortraitAlgorithmJobType] = jobcontroller.NewCronJobHorizontal(client, config.Namespace, config.DefaultServiceAccount, config.DefaultMetricsServerAddr, map[autoscalingv1alpha1.PortraitType]string{
+		autoscalingv1alpha1.PredictivePortraitType: config.DefaultImageHorizontalPredictive,
 	})
 	return controllers, nil
 }

diff --git a/pkg/portrait/algorithm/externaljob/jobcontroller/cronjob.go b/pkg/portrait/algorithm/externaljob/jobcontroller/cronjob.go
@@ -43,14 +43,16 @@ const (
 type CronJobHorizontal struct {
 	client                   client.Client
 	namespace                string
+	defaultServiceAccount    string
 	defaultMetricsServerAddr string
 	defaultImages            map[autoscalingv1alpha1.PortraitType]string
 }
 
-func NewCronJobHorizontal(client client.Client, namespace, defaultMetricsServerAddr string, defaultImages map[autoscalingv1alpha1.PortraitType]string) Horizontal {
+func NewCronJobHorizontal(client client.Client, namespace, defaultServiceAccount, defaultMetricsServerAddr string, defaultImages map[autoscalingv1alpha1.PortraitType]string) Horizontal {
 	return &CronJobHorizontal{
 		client:                   client,
 		namespace:                namespace,
+		defaultServiceAccount:    defaultServiceAccount,
 		defaultMetricsServerAddr: defaultMetricsServerAddr,
 		defaultImages:            defaultImages,
 	}
@@ -135,6 +137,10 @@ func (h *CronJobHorizontal) buildCronJobNamespacedName(hp *autoscalingv1alpha1.H
 }
 
 func (h *CronJobHorizontal) defaultingTemplatePodSpec(spec *corev1.PodSpec, hp *autoscalingv1alpha1.HorizontalPortrait) {
+	if h.defaultServiceAccount != "" && spec.ServiceAccountName == "" && spec.DeprecatedServiceAccount == "" {
+		spec.ServiceAccountName = h.defaultServiceAccount
+		spec.DeprecatedServiceAccount = h.defaultServiceAccount
+	}
 	for i := range spec.Containers {
 		container := &spec.Containers[i]
 		if container.Name != defaultAlgorithmContainerName {

diff --git a/pkg/portrait/algorithm/externaljob/jobcontroller/cronjob_test.go b/pkg/portrait/algorithm/externaljob/jobcontroller/cronjob_test.go
@@ -75,7 +75,7 @@ func TestCronJobHorizontal_Create(t *testing.T) {
 		},
 	}
 
-	cronJobHorizontal := NewCronJobHorizontal(fakeClient, cronJobNamespace, "", nil)
+	cronJobHorizontal := NewCronJobHorizontal(fakeClient, cronJobNamespace, "", "", nil)
 	err := cronJobHorizontal.UpdateJob(ctx, horizontalPortrait, cfg)
 	assert.Nil(t, err)
 
@@ -126,7 +126,7 @@ func TestCronJobHorizontal_Update(t *testing.T) {
 	}
 	_ = fakeClient.Create(ctx, cronJob)
 
-	cronJobHorizontal := NewCronJobHorizontal(fakeClient, cronJobNamespace, "", nil)
+	cronJobHorizontal := NewCronJobHorizontal(fakeClient, cronJobNamespace, "", "", nil)
 	err := cronJobHorizontal.UpdateJob(ctx, horizontalPortrait, cfg)
 	assert.Nil(t, err)