feat: add auto-scaling switch config parsing and apply, TargetResource support value all

Author: knave

internal/autoscaler/autoscaler.go

@@ -220,6 +220,12 @@ func (s *Autoscaler) ProcessWorkloads(ctx context.Context) {
 		rr := s.ResourceRecommender.GetRecommendedResources(workloadState)
 		log.Info("recommend resources", "workload", workloadState.Name, "resources", rr)
 
+		// TODO: update recommmendation status of workload
+
+		if !workloadState.IsAutoScalingEnabled() {
+			continue
+		}
+
 		for _, worker := range podList.Items {
 			if !worker.DeletionTimestamp.IsZero() {
 				continue

internal/autoscaler/autoscaler_test.go

@@ -184,7 +184,7 @@ var _ = Describe("Autoscaler", func() {
 			}).Should(Succeed())
 		})
 
-		It("should update specific resources based on TargetResource", func() {
+		It("should update resources based on auto scaling config", func() {
 			tfEnv := NewTensorFusionEnvBuilder().
 				AddPoolWithNodeCount(1).SetGpuCountPerNode(1).
 				Build()
@@ -199,9 +199,22 @@ var _ = Describe("Autoscaler", func() {
 			rr := scaler.ResourceRecommender.GetRecommendedResources(nil)
 
 			workloadState := scaler.WorkloadStates[workload.Name]
-			workloadState.AutoScalingConfig.AutoSetResources.TargetResource = "tflops"
-
 			oldRes := workloadState.Resources
+
+			// verify IsAutoScalingEnabled
+			workloadState.AutoScalingConfig.AutoSetResources.Enable = false
+			scaler.ProcessWorkloads(ctx)
+			Eventually(func(g Gomega) {
+				tflopsRequest, tflopsLimit, vramRequest, vramLimit := parseResourceAnnotations(getWorkers(workload)[0])
+				Expect(tflopsRequest.Equal(oldRes.Requests.Tflops)).To(BeTrue())
+				Expect(tflopsLimit.Equal(oldRes.Limits.Tflops)).To(BeTrue())
+				Expect(vramRequest.Equal(oldRes.Requests.Vram)).To(BeTrue())
+				Expect(vramLimit.Equal(oldRes.Limits.Vram)).To(BeTrue())
+			}).Should(Succeed())
+
+			// verify IsTargetResource
+			workloadState.AutoScalingConfig.AutoSetResources.Enable = true
+			workloadState.AutoScalingConfig.AutoSetResources.TargetResource = "tflops"
 			scaler.ProcessWorkloads(ctx)
 			Eventually(func(g Gomega) {
 				tflopsRequest, tflopsLimit, vramRequest, vramLimit := parseResourceAnnotations(getWorkers(workload)[0])

internal/autoscaler/workloadstate.go

@@ -2,6 +2,7 @@ package autoscaler
 
 import (
 	"strconv"
+	"strings"
 	"time"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
@@ -100,5 +101,9 @@ func (w *WorkloadState) GetResourceRecommenderConfig() *ResourceRecommenderConfi
 
 func (w *WorkloadState) IsTargetResource(resourceName string) bool {
 	target := w.AutoScalingConfig.AutoSetResources.TargetResource
-	return target == "" || resourceName == target
+	return target == "" || strings.EqualFold(target, "all") || strings.EqualFold(resourceName, target)
+}
+
+func (w *WorkloadState) IsAutoScalingEnabled() bool {
+	return w.AutoScalingConfig.AutoSetResources.Enable
 }

internal/autoscaler/workloadstate_test.go

@@ -82,6 +82,13 @@ var _ = Describe("Workload State", func() {
 		Expect(ws.IsTargetResource("tflops")).To(BeTrue())
 		Expect(ws.IsTargetResource("vram")).To(BeTrue())
 
+		ws.AutoScalingConfig = tfv1.AutoScalingConfig{
+			AutoSetResources: tfv1.AutoSetResources{TargetResource: "all"},
+		}
+
+		Expect(ws.IsTargetResource("tflops")).To(BeTrue())
+		Expect(ws.IsTargetResource("vram")).To(BeTrue())
+
 		ws.AutoScalingConfig = tfv1.AutoScalingConfig{
 			AutoSetResources: tfv1.AutoSetResources{TargetResource: "tflops"},
 		}
@@ -94,4 +101,17 @@ var _ = Describe("Workload State", func() {
 		Expect(ws.IsTargetResource("tflops")).To(BeFalse())
 		Expect(ws.IsTargetResource("vram")).To(BeTrue())
 	})
+
+	It("should correctly determine if auto scaling is enabled based on config", func() {
+		ws := NewWorkloadState("test")
+
+		ws.AutoScalingConfig = tfv1.AutoScalingConfig{
+			AutoSetResources: tfv1.AutoSetResources{Enable: true},
+		}
+		Expect(ws.IsAutoScalingEnabled()).To(BeTrue())
+		ws.AutoScalingConfig = tfv1.AutoScalingConfig{
+			AutoSetResources: tfv1.AutoSetResources{Enable: false},
+		}
+		Expect(ws.IsAutoScalingEnabled()).To(BeFalse())
+	})
 })