refactor: optimize loading historical metrics (#394)

* refactor: make it more object-oriented

* refactor: optimize loading historical metrics

* fix: linter issue

* fix: linter issue

* fix: handle close properly

Author: knave

internal/autoscaler/autoscaler.go

@@ -134,19 +134,9 @@ func (s *Autoscaler) loadWorkloads(ctx context.Context) {
 }
 
 func (s *Autoscaler) loadHistoryMetrics(ctx context.Context) error {
-	workersMetrics, err := s.metricsProvider.GetHistoryMetrics(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to get history metrics: %v", err)
-	}
-	for _, sample := range workersMetrics {
+	return s.metricsProvider.LoadHistoryMetrics(ctx, func(sample *metrics.WorkerUsage) {
 		s.findOrCreateWorkloadState(sample.Namespace, sample.WorkloadName).AddSample(sample)
-	}
-
-	if metricsCount := len(workersMetrics); metricsCount > 0 {
-		log.FromContext(ctx).Info("historical metrics loaded", "from",
-			workersMetrics[0].Timestamp, "to", workersMetrics[metricsCount-1].Timestamp, "metricsCount", metricsCount)
-	}
-	return nil
+	})
 }
 
 func (s *Autoscaler) loadRealTimeMetrics(ctx context.Context) {

internal/autoscaler/autoscaler_test.go

@@ -487,6 +487,28 @@ func (f *FakeMetricsProvider) GetWorkersMetrics(ctx context.Context) ([]*metrics
 	return f.Metrics, nil
 }
 
+func (f *FakeMetricsProvider) LoadHistoryMetrics(ctx context.Context, processMetricsFunc func(*metrics.WorkerUsage)) error {
+	startTime := time.Now().Add(-7 * 24 * time.Hour)
+	for day := 0; day < 7; day++ {
+		for hour := 0; hour < 1; hour++ {
+			for minute := 0; minute < 60; minute++ {
+				// idx := day*24 + hour
+				sample := &metrics.WorkerUsage{
+					Namespace:    "default",
+					WorkloadName: "workload-0",
+					WorkerName:   fmt.Sprintf("worker-%d", 1),
+					TflopsUsage:  100.0,
+					VramUsage:    1 * 1000 * 1000 * 1000,
+					Timestamp:    startTime.Add(time.Duration(day*24+hour)*time.Hour + time.Duration(minute)*time.Minute),
+				}
+				processMetricsFunc(sample)
+			}
+		}
+	}
+
+	return nil
+}
+
 func (f *FakeMetricsProvider) GetHistoryMetrics(ctx context.Context) ([]*metrics.WorkerUsage, error) {
 	sample := []*metrics.WorkerUsage{}
 	startTime := time.Now().Add(-8 * 24 * time.Hour)

internal/autoscaler/metrics/metrics_provider.go

@@ -27,6 +27,7 @@ type WorkerUsage struct {
 type Provider interface {
 	GetWorkersMetrics(context.Context) ([]*WorkerUsage, error)
 	GetHistoryMetrics(context.Context) ([]*WorkerUsage, error)
+	LoadHistoryMetrics(context.Context, func(*WorkerUsage)) error
 }
 
 type greptimeDBProvider struct {
@@ -47,12 +48,6 @@ func NewProvider() (Provider, error) {
 func (g *greptimeDBProvider) GetWorkersMetrics(ctx context.Context) ([]*WorkerUsage, error) {
 	now := time.Now()
 
-	log := log.FromContext(ctx)
-	log.V(6).Info("Started querying workers metrics", "startTime", now)
-	defer func() {
-		log.V(6).Info("Finished querying workers metrics", "duration", time.Since(now))
-	}()
-
 	timeoutCtx, cancel := context.WithTimeout(ctx, defaultQueryTimeout)
 	defer cancel()
 
@@ -99,17 +94,9 @@ type hypervisorWorkerUsageMetrics struct {
 func (g *greptimeDBProvider) GetHistoryMetrics(ctx context.Context) ([]*WorkerUsage, error) {
 	now := time.Now()
 
-	log := log.FromContext(ctx)
-	log.V(6).Info("Started querying history metrics", "startTime", now)
-	defer func() {
-		log.V(6).Info("Finished querying history metrics", "duration", time.Since(now))
-	}()
-
 	timeoutCtx, cancel := context.WithTimeout(ctx, defaultHistoryQueryTimeout)
 	defer cancel()
 
-	// TODO: replace using iteration for handling large datasets efficiently
-	// TODO: supply history resolution to config time window
 	data := []*hypervisorWorkerUsageMetrics{}
 	err := g.db.WithContext(timeoutCtx).
 		Select("namespace, workload, worker, max(compute_tflops) as compute_tflops, max(memory_bytes) as memory_bytes, date_bin('1 minute'::INTERVAL, ts) as time_window").
@@ -140,6 +127,47 @@ func (g *greptimeDBProvider) GetHistoryMetrics(ctx context.Context) ([]*WorkerUs
 	return workersMetrics, nil
 }
 
+func (g *greptimeDBProvider) LoadHistoryMetrics(ctx context.Context, processMetricsFunc func(*WorkerUsage)) error {
+	now := time.Now()
+
+	timeoutCtx, cancel := context.WithTimeout(ctx, defaultHistoryQueryTimeout)
+	defer cancel()
+
+	rows, err := g.db.WithContext(timeoutCtx).
+		Model(&hypervisorWorkerUsageMetrics{}).
+		Select("namespace, workload, worker, max(compute_tflops) as compute_tflops, max(memory_bytes) as memory_bytes, date_bin('1 minute'::INTERVAL, ts) as time_window").
+		Where("ts > ? and ts <= ?", now.Add(-time.Hour*24*7).UnixNano(), now.UnixNano()).
+		Group("namespace, workload, worker, time_window").
+		Order("time_window asc").
+		Rows()
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := rows.Close(); err != nil {
+			log.FromContext(ctx).Error(err, "failed to close rows")
+		}
+	}()
+
+	for rows.Next() {
+		var usage hypervisorWorkerUsageMetrics
+		if err := g.db.ScanRows(rows, &usage); err != nil {
+			return err
+		}
+		processMetricsFunc(&WorkerUsage{
+			Namespace:    usage.Namespace,
+			WorkloadName: usage.WorkloadName,
+			WorkerName:   usage.WorkerName,
+			TflopsUsage:  usage.ComputeTflops,
+			VramUsage:    usage.VRAMBytes,
+			Timestamp:    usage.TimeWindow,
+		})
+	}
+
+	g.lastQueryTime = now
+	return nil
+}
+
 // Setup GreptimeDB connection
 func setupTimeSeriesDB() (*metrics.TimeSeriesDB, error) {
 	timeSeriesDB := &metrics.TimeSeriesDB{}

internal/autoscaler/recommender/recommendation_test.go

@@ -43,7 +43,7 @@ var _ = Describe("Recommender", func() {
 			},
 		}
 
-		final := getResourcesFromRecResults(recs)
+		final := recommenderToRecResult(recs).generateRecommendation()
 		Expect(final.Equal(&tfv1.Resources{
 			Requests: tfv1.Resource{
 				Tflops: resource.MustParse("10"),
@@ -88,7 +88,7 @@ var _ = Describe("Recommender", func() {
 			},
 		}
 
-		Expect(getResourcesFromRecResults(recs)).To(BeNil())
+		Expect(recommenderToRecResult(recs).generateRecommendation()).To(BeNil())
 	})
 
 	It("should return recommendation that replaced with the current maximum allowable GPU resource", func() {

internal/autoscaler/recommender/recommender.go

@@ -20,59 +20,31 @@ type RecResult struct {
 	ScaleDownLocking bool
 }
 
-func GetRecommendation(ctx context.Context, workload *workload.State, recommenders []Interface) (*tfv1.Resources, error) {
-	recResults := map[string]*RecResult{}
-	for _, recommender := range recommenders {
-		result, err := recommender.Recommend(ctx, workload)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get recommendation from %s: %v", recommender.Name(), err)
-		}
-		if result != nil {
-			recResults[recommender.Name()] = result
-		}
-	}
+type recommenderToRecResult map[string]*RecResult
 
-	if len(recResults) <= 0 {
-		return nil, nil
-	}
-
-	resources := getResourcesFromRecResults(recResults)
-	if resources != nil {
-		curRes := workload.GetCurrentResourcesSpec()
-		// If a resource value is zero, replace it with current value
-		if resources.Requests.Tflops.IsZero() || resources.Limits.Tflops.IsZero() {
-			resources.Requests.Tflops = curRes.Requests.Tflops
-			resources.Limits.Tflops = curRes.Limits.Tflops
-		}
-
-		if resources.Requests.Vram.IsZero() || resources.Limits.Vram.IsZero() {
-			resources.Requests.Vram = curRes.Requests.Vram
-			resources.Limits.Vram = curRes.Limits.Vram
-		}
+func (r recommenderToRecResult) generateRecommendation() *tfv1.Resources {
+	if len(r) == 0 {
+		return nil
 	}
 
-	return resources, nil
-}
-
-func getResourcesFromRecResults(recResults map[string]*RecResult) *tfv1.Resources {
-	targetRes := &tfv1.Resources{}
+	recommendation := &tfv1.Resources{}
 	minRes := &tfv1.Resources{}
-	for _, rec := range recResults {
-		if !rec.HasApplied {
-			mergeResourcesByLargerRequests(targetRes, &rec.Resources)
+	for _, result := range r {
+		if !result.HasApplied {
+			mergeResourcesByLargerRequests(recommendation, &result.Resources)
 		}
-		if rec.ScaleDownLocking {
-			mergeResourcesByLargerRequests(minRes, &rec.Resources)
+		if result.ScaleDownLocking {
+			mergeResourcesByLargerRequests(minRes, &result.Resources)
 		}
 	}
 
-	if targetRes.IsZero() ||
-		(targetRes.Requests.Tflops.Cmp(minRes.Requests.Tflops) < 0 &&
-			targetRes.Requests.Vram.Cmp(minRes.Requests.Vram) < 0) {
+	if recommendation.IsZero() ||
+		(recommendation.Requests.Tflops.Cmp(minRes.Requests.Tflops) < 0 &&
+			recommendation.Requests.Vram.Cmp(minRes.Requests.Vram) < 0) {
 		return nil
 	}
 
-	return targetRes
+	return recommendation
 }
 
 func mergeResourcesByLargerRequests(src *tfv1.Resources, target *tfv1.Resources) {
@@ -85,3 +57,33 @@ func mergeResourcesByLargerRequests(src *tfv1.Resources, target *tfv1.Resources)
 		src.Limits.Vram = target.Limits.Vram
 	}
 }
+
+func GetRecommendation(ctx context.Context, workload *workload.State, recommenders []Interface) (*tfv1.Resources, error) {
+	recResults := make(recommenderToRecResult)
+	for _, recommender := range recommenders {
+		result, err := recommender.Recommend(ctx, workload)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get recommendation from %s: %v", recommender.Name(), err)
+		}
+		if result != nil {
+			recResults[recommender.Name()] = result
+		}
+	}
+
+	recommendation := recResults.generateRecommendation()
+	if recommendation != nil {
+		curRes := workload.GetCurrentResourcesSpec()
+		// If a resource value is zero, replace it with current value
+		if recommendation.Requests.Tflops.IsZero() || recommendation.Limits.Tflops.IsZero() {
+			recommendation.Requests.Tflops = curRes.Requests.Tflops
+			recommendation.Limits.Tflops = curRes.Limits.Tflops
+		}
+
+		if recommendation.Requests.Vram.IsZero() || recommendation.Limits.Vram.IsZero() {
+			recommendation.Requests.Vram = curRes.Requests.Vram
+			recommendation.Limits.Vram = curRes.Limits.Vram
+		}
+	}
+
+	return recommendation, nil
+}