Add data collection for measured pods to populate ci_operator_metrics table

Author: deepsm007

cmd/pod-scaler/admission.go

@@ -318,7 +318,7 @@ func preventUnschedulable(resources *corev1.ResourceRequirements, cpuCap int64,
 	}
 }
 
-func mutatePodResources(pod *corev1.Pod, server *resourceServer, mutateResourceLimits bool, cpuCap int64, memoryCap string, authoritativeCPU, authoritativeMemory bool, reporter results.PodScalerReporter, logger *logrus.Entry) {
+func mutatePodResources(pod *corev1.Pod, server *resourceServer, mutateResourceLimits bool, cpuCap int64, memoryCap string, _ /* authoritativeCPU */, _ /* authoritativeMemory */ bool, reporter results.PodScalerReporter, logger *logrus.Entry) {
 	mutateResources := func(containers []corev1.Container) {
 		for i := range containers {
 			meta := podscaler.MetadataFor(pod.ObjectMeta.Labels, pod.ObjectMeta.Name, containers[i].Name)

cmd/pod-scaler/main.go

@@ -52,9 +52,13 @@ type options struct {
 }
 
 type producerOptions struct {
-	kubernetesOptions prowflagutil.KubernetesOptions
-	once              bool
-	ignoreLatest      time.Duration
+	kubernetesOptions       prowflagutil.KubernetesOptions
+	once                    bool
+	ignoreLatest            time.Duration
+	enableMeasuredPods      bool
+	bigQueryProjectID       string
+	bigQueryDatasetID       string
+	bigQueryCredentialsFile string
 }
 
 type consumerOptions struct {
@@ -67,22 +71,19 @@ type consumerOptions struct {
 	cpuCap                int64
 	memoryCap             string
 	cpuPriorityScheduling int64
-
-	// Measured pods options - when enabled, pods are classified as "normal" or "measured"
-	// Measured pods run on isolated nodes to get accurate CPU/memory utilization data
-	enableMeasuredPods      bool
-	bigQueryProjectID       string
-	bigQueryDatasetID       string
-	bigQueryCredentialsFile string
 }
 
 func bindOptions(fs *flag.FlagSet) *options {
 	o := options{producerOptions: producerOptions{kubernetesOptions: prowflagutil.KubernetesOptions{NOInClusterConfigDefault: true}}}
 	o.instrumentationOptions.AddFlags(fs)
 	fs.StringVar(&o.mode, "mode", "", "Which mode to run in.")
 	o.producerOptions.kubernetesOptions.AddFlags(fs)
-	fs.DurationVar(&o.ignoreLatest, "ignore-latest", 0, "Duration of latest time series to ignore when querying Prometheus. For instance, 1h will ignore the latest hour of data.")
-	fs.BoolVar(&o.once, "produce-once", false, "Query Prometheus and refresh cached data only once before exiting.")
+	fs.DurationVar(&o.producerOptions.ignoreLatest, "ignore-latest", 0, "Duration of latest time series to ignore when querying Prometheus. For instance, 1h will ignore the latest hour of data.")
+	fs.BoolVar(&o.producerOptions.once, "produce-once", false, "Query Prometheus and refresh cached data only once before exiting.")
+	fs.BoolVar(&o.producerOptions.enableMeasuredPods, "enable-measured-pods", false, "Enable measured pods feature. When enabled, pods are classified as 'normal' or 'measured' and measured pods run on isolated nodes to get accurate CPU/memory utilization data.")
+	fs.StringVar(&o.producerOptions.bigQueryProjectID, "bigquery-project-id", "", "Google Cloud project ID for BigQuery (required if enable-measured-pods is true)")
+	fs.StringVar(&o.producerOptions.bigQueryDatasetID, "bigquery-dataset-id", "", "BigQuery dataset ID for pod metrics (required if enable-measured-pods is true)")
+	fs.StringVar(&o.producerOptions.bigQueryCredentialsFile, "bigquery-credentials-file", "", "Path to Google Cloud credentials file for BigQuery access")
 	fs.IntVar(&o.port, "port", 0, "Port to serve admission webhooks on.")
 	fs.IntVar(&o.uiPort, "ui-port", 0, "Port to serve frontend on.")
 	fs.StringVar(&o.certDir, "serving-cert-dir", "", "Path to directory with serving certificate and key for the admission webhook server.")
@@ -96,10 +97,6 @@ func bindOptions(fs *flag.FlagSet) *options {
 	fs.Int64Var(&o.cpuCap, "cpu-cap", 10, "The maximum CPU request value, ex: 10")
 	fs.StringVar(&o.memoryCap, "memory-cap", "20Gi", "The maximum memory request value, ex: '20Gi'")
 	fs.Int64Var(&o.cpuPriorityScheduling, "cpu-priority-scheduling", 8, "Pods with CPU requests at, or above, this value will be admitted with priority scheduling")
-	fs.BoolVar(&o.enableMeasuredPods, "enable-measured-pods", false, "Enable measured pods feature. When enabled, pods are classified as 'normal' or 'measured' and measured pods run on isolated nodes to get accurate CPU/memory utilization data.")
-	fs.StringVar(&o.bigQueryProjectID, "bigquery-project-id", "", "Google Cloud project ID for BigQuery queries (required if enable-measured-pods is true)")
-	fs.StringVar(&o.bigQueryDatasetID, "bigquery-dataset-id", "", "BigQuery dataset ID for pod metrics (required if enable-measured-pods is true)")
-	fs.StringVar(&o.bigQueryCredentialsFile, "bigquery-credentials-file", "", "Path to Google Cloud credentials file for BigQuery access")
 	o.resultsOptions.Bind(fs)
 	return &o
 }
@@ -255,7 +252,7 @@ func mainProduce(opts *options, cache Cache) {
 		logger.Debugf("Loaded Prometheus client.")
 	}
 
-	produce(clients, cache, opts.ignoreLatest, opts.once)
+	produce(clients, cache, opts.producerOptions.ignoreLatest, opts.producerOptions.once, opts.producerOptions.enableMeasuredPods, opts.producerOptions.bigQueryProjectID, opts.producerOptions.bigQueryDatasetID, opts.producerOptions.bigQueryCredentialsFile)
 
 }
 
@@ -279,7 +276,8 @@ func mainAdmission(opts *options, cache Cache) {
 		logrus.WithError(err).Fatal("Failed to create pod-scaler reporter.")
 	}
 
-	go admit(opts.port, opts.instrumentationOptions.HealthPort, opts.certDir, client, loaders(cache), opts.mutateResourceLimits, opts.cpuCap, opts.memoryCap, opts.cpuPriorityScheduling, false, false, opts.enableMeasuredPods, opts.bigQueryProjectID, opts.bigQueryDatasetID, opts.bigQueryCredentialsFile, reporter)
+	// Use producerOptions for measured pods config (shared between producer and consumer modes)
+	go admit(opts.port, opts.instrumentationOptions.HealthPort, opts.certDir, client, loaders(cache), opts.mutateResourceLimits, opts.cpuCap, opts.memoryCap, opts.cpuPriorityScheduling, false, false, opts.producerOptions.enableMeasuredPods, opts.producerOptions.bigQueryProjectID, opts.producerOptions.bigQueryDatasetID, opts.producerOptions.bigQueryCredentialsFile, reporter)
 }
 
 func loaders(cache Cache) map[string][]*cacheReloader {

cmd/pod-scaler/measured.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"sync"
 	"time"
@@ -140,27 +141,27 @@ func NewBigQueryClient(projectID, datasetID, credentialsFile string, cache *Meas
 func (bq *BigQueryClient) Refresh(ctx context.Context) error {
 	bq.logger.Info("Refreshing measured pod data from BigQuery")
 
-	// TODO: Replace with actual BigQuery query based on ci-metrics structure.
-	// This is a placeholder query - the actual query will depend on the BigQuery schema
-	// for ci-metrics pod CPU utilization data. We need to query the table that stores
-	// max CPU/memory utilization for pods that ran with the "measured" label.
+	// Query ci_operator_metrics table for measured pod data
+	// This queries the table that stores max CPU/memory utilization for pods that ran with the "measured" label
 	query := bq.client.Query(fmt.Sprintf(`
 		SELECT
 			org,
 			repo,
 			branch,
+			target,
 			container,
-			MAX(cpu_utilization) as max_cpu,
-			MAX(memory_utilization) as max_memory,
-			MAX(timestamp) as last_measured,
+			pod_name,
+			MAX(max_cpu) as max_cpu,
+			MAX(max_memory) as max_memory,
+			MAX(created) as last_measured,
 			ANY_VALUE(container_durations) as container_durations
 		FROM
-			`+"`%s.%s.pod_metrics`"+`
+			`+"`%s.%s.ci_operator_metrics`"+`
 		WHERE
 			pod_scaler_label = 'measured'
-			AND timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL %d DAY)
+			AND created >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL %d DAY)
 		GROUP BY
-			org, repo, branch, container
+			org, repo, branch, target, container, pod_name
 	`, bq.projectID, bq.datasetID, MeasuredPodDataRetentionDays))
 
 	query.QueryConfig.Labels = map[string]string{
@@ -179,7 +180,9 @@ func (bq *BigQueryClient) Refresh(ctx context.Context) error {
 			Org                string    `bigquery:"org"`
 			Repo               string    `bigquery:"repo"`
 			Branch             string    `bigquery:"branch"`
+			Target             string    `bigquery:"target"`
 			Container          string    `bigquery:"container"`
+			PodName            string    `bigquery:"pod_name"`
 			MaxCPU             float64   `bigquery:"max_cpu"`
 			MaxMemory          int64     `bigquery:"max_memory"`
 			LastMeasured       time.Time `bigquery:"last_measured"`
@@ -199,11 +202,33 @@ func (bq *BigQueryClient) Refresh(ctx context.Context) error {
 				Repo:   row.Repo,
 				Branch: row.Branch,
 			},
+			Target:    row.Target,
+			Pod:       row.PodName,
 			Container: row.Container,
 		}
 
-		// TODO: Parse container_durations JSON string into map[string]time.Duration
+		// Parse container_durations JSON string into map[string]time.Duration
+		// BigQuery stores this as a JSON string. time.Duration serializes as int64 (nanoseconds)
+		// so the format is: {"container1": 3600000000000, "container2": 245000000000}
 		containerDurations := make(map[string]time.Duration)
+		if row.ContainerDurations != "" {
+			var durationsMap map[string]int64
+			if err := json.Unmarshal([]byte(row.ContainerDurations), &durationsMap); err == nil {
+				for container, nanoseconds := range durationsMap {
+					containerDurations[container] = time.Duration(nanoseconds)
+				}
+			} else {
+				// Fallback: try parsing as string format (for backwards compatibility)
+				var durationsMapStr map[string]string
+				if err := json.Unmarshal([]byte(row.ContainerDurations), &durationsMapStr); err == nil {
+					for container, durationStr := range durationsMapStr {
+						if duration, err := time.ParseDuration(durationStr); err == nil {
+							containerDurations[container] = duration
+						}
+					}
+				}
+			}
+		}
 
 		data[meta] = &MeasuredPodData{
 			Metadata:             meta,
@@ -261,8 +286,9 @@ func ClassifyPod(pod *corev1.Pod, bqClient *BigQueryClient, logger *logrus.Entry
 			}
 		}
 	} else {
-		// If BigQuery isn't configured, default to measuring new pods.
-		shouldBeMeasured = true
+		// If BigQuery isn't configured, default to normal to avoid overwhelming isolated nodes.
+		shouldBeMeasured = false
+		logger.Warn("BigQuery client not available, defaulting pod to normal classification")
 	}
 
 	if shouldBeMeasured {
@@ -327,21 +353,28 @@ func AddPodAntiAffinity(pod *corev1.Pod, logger *logrus.Entry) {
 	}
 
 	if len(requiredTerms) > 0 {
+		// Merge with existing anti-affinity terms instead of overwriting
+		existingTerms := pod.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution
+		if existingTerms != nil {
+			requiredTerms = append(existingTerms, requiredTerms...)
+		}
 		pod.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = requiredTerms
 	}
 }
 
 // ApplyMeasuredPodResources uses the real resource data we collected when this pod ran in isolation.
 // We only increase resources for the longest-running container (the main workload), not all containers.
 // This is based on actual measured usage, not Prometheus data that might be skewed by node contention.
+// This function applies resources to pods labeled "normal" (which have measured data), not "measured" (which need measurement).
 func ApplyMeasuredPodResources(pod *corev1.Pod, bqClient *BigQueryClient, logger *logrus.Entry) {
 	if bqClient == nil {
 		return
 	}
 
 	podScalerLabel, hasLabel := pod.Labels[PodScalerLabelKey]
-	if !hasLabel || podScalerLabel != PodScalerLabelValueMeasured {
-		// Only apply measured resources to pods that are actually being measured.
+	// Apply measured resources to pods labeled "normal" (which have fresh measured data)
+	// Pods labeled "measured" don't have data yet - they're being measured for the first time.
+	if !hasLabel || podScalerLabel != PodScalerLabelValueNormal {
 		return
 	}
 

cmd/pod-scaler/measured_test.go

@@ -130,7 +130,7 @@ func TestClassifyPod(t *testing.T) {
 		t.Errorf("Expected pod to be classified as normal, got %s", pod2.Labels[PodScalerLabelKey])
 	}
 
-	// Test case 3: Pod with nil BigQuery client - should default to measured
+	// Test case 3: Pod with nil BigQuery client - should default to normal to avoid overwhelming isolated nodes
 	pod3 := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name: "test-pod-3",
@@ -143,8 +143,8 @@ func TestClassifyPod(t *testing.T) {
 	}
 
 	ClassifyPod(pod3, nil, logger)
-	if pod3.Labels[PodScalerLabelKey] != PodScalerLabelValueMeasured {
-		t.Errorf("Expected pod to be classified as measured when BigQuery client is nil, got %s", pod3.Labels[PodScalerLabelKey])
+	if pod3.Labels[PodScalerLabelKey] != PodScalerLabelValueNormal {
+		t.Errorf("Expected pod to be classified as normal when BigQuery client is nil (to avoid overwhelming isolated nodes), got %s", pod3.Labels[PodScalerLabelKey])
 	}
 }