Merge pull request #4369 from jmguzik/config_cluster

Extend cluster config: capacity, capabilities

cmd/prow-job-dispatcher/main.go

@@ -355,8 +355,13 @@ func dispatchJobs(prowJobConfigDir string, config *dispatcher.Config, jobVolumes
 	}
 
 	// cv stores the volume for each cluster in the build farm
-	cv := &clusterVolume{clusterVolumeMap: map[string]map[string]float64{}, cloudProviders: sets.New[string](), pjs: map[string]string{}, blocked: blocked,
-		volumePerCluster: volumePerCluster, specialClusters: map[string]float64{}}
+	cv := &clusterVolume{
+		clusterVolumeMap: map[string]map[string]float64{},
+		cloudProviders:   sets.New[string](),
+		pjs:              map[string]string{},
+		blocked:          blocked,
+		volumePerCluster: volumePerCluster,
+		specialClusters:  map[string]float64{}}
 	for cloudProvider, v := range config.BuildFarm {
 		for cluster := range v {
 			cloudProviderString := string(cloudProvider)
@@ -735,11 +740,11 @@ func main() {
 
 			addEnabledClusters(config, enabled,
 				func(cluster string) (api.Cloud, error) {
-					provider, exists := configClusterMap[cluster]
+					info, exists := configClusterMap[cluster]
 					if !exists {
 						return "", fmt.Errorf("have not found provider for cluster %s", cluster)
 					}
-					return api.Cloud(provider), nil
+					return api.Cloud(info.Provider), nil
 				})
 			pjs, err := dispatchJobs(o.prowJobConfigDir, config, jobVolumes, blocked, promVolumes.getTotalVolume()/float64(len(configClusterMap)))
 			if err != nil {

pkg/sanitizer/helpers.go

@@ -8,38 +8,61 @@ import (
 	"sigs.k8s.io/yaml"
 )
 
-const blocked = "blocked"
+// ClusterInfo holds the provider, capacity, and capabilities.
+type ClusterInfo struct {
+	Provider     string
+	Capacity     int
+	Capabilities []string
+}
 
-type ClusterMap map[string]string
+// ClusterMap maps a cluster name to its corresponding ClusterInfo.
+type ClusterMap map[string]ClusterInfo
 
-func LoadClusterConfig(filePath string) (ClusterMap, sets.Set[string], error) {
-	data, err := os.ReadFile(filePath)
-	if err != nil {
-		return nil, nil, err
+func loadClusterConfigFromBytes(data []byte) (ClusterMap, sets.Set[string], error) {
+	var clusters map[string][]struct {
+		Name         string   `yaml:"name"`
+		Capacity     int      `yaml:"capacity"`
+		Capabilities []string `yaml:"capabilities"`
+		Blocked      bool     `yaml:"blocked"`
 	}
-
-	var clusters map[string][]string
-	err = yaml.Unmarshal(data, &clusters)
-	if err != nil {
+	if err := yaml.Unmarshal(data, &clusters); err != nil {
 		return nil, nil, err
 	}
-
 	blockedClusters := sets.New[string]()
 	clusterMap := make(ClusterMap)
+
 	for provider, clusterList := range clusters {
-		if provider != blocked {
-			for _, cluster := range clusterList {
-				clusterMap[cluster] = provider
+		for _, cluster := range clusterList {
+			if cluster.Capacity == 0 || cluster.Capacity > 100 {
+				cluster.Capacity = 100
+			} else if cluster.Capacity < 0 {
+				cluster.Blocked = true
+			}
+			if cluster.Blocked {
+				blockedClusters.Insert(cluster.Name)
+				continue
+			}
+			clusterMap[cluster.Name] = ClusterInfo{
+				Provider:     provider,
+				Capacity:     cluster.Capacity,
+				Capabilities: cluster.Capabilities,
 			}
-		}
-		if provider == blocked {
-			blockedClusters.Insert(clusterList...)
 		}
 	}
 
 	return clusterMap, blockedClusters, nil
 }
 
+// LoadClusterConfig loads cluster configuration from a YAML file, returning a ClusterMap and a set of blocked clusters.
+func LoadClusterConfig(filePath string) (ClusterMap, sets.Set[string], error) {
+	data, err := os.ReadFile(filePath)
+	if err != nil {
+		return nil, nil, err
+	}
+	return loadClusterConfigFromBytes(data)
+
+}
+
 func FindMostUsedCluster(jc *prowconfig.JobConfig) string {
 	clusters := make(map[string]int)
 	for k := range jc.PresubmitsStatic {

pkg/sanitizer/helpers_test.go

@@ -179,3 +179,124 @@ func TestDetermineTargetCluster(t *testing.T) {
 		})
 	}
 }
+
+func TestLoadClusterConfigFromBytes(t *testing.T) {
+	tests := []struct {
+		name            string
+		yamlData        string
+		expectedCluster ClusterMap
+		expectedBlocked sets.Set[string]
+	}{
+		{
+			name: "Valid config with AWS and GCP",
+			yamlData: `
+aws:
+  - name: build01
+    capacity: 80
+    capabilities:
+      - aarch64
+      - amd64
+      - vpn
+  - name: build03
+  - name: build09
+    blocked: true
+  - name: build99
+    capacity: -99 #will be blocked as well
+gcp:
+  - name: build02
+    capacity: 60
+    capabilities:
+      - vpn
+`,
+			expectedCluster: ClusterMap{
+				"build01": {
+					Provider:     "aws",
+					Capacity:     80,
+					Capabilities: []string{"aarch64", "amd64", "vpn"},
+				},
+				"build03": {
+					Provider:     "aws",
+					Capacity:     100,
+					Capabilities: nil,
+				},
+				"build02": {
+					Provider:     "gcp",
+					Capacity:     60,
+					Capabilities: []string{"vpn"},
+				},
+			},
+			expectedBlocked: sets.New[string]("build09", "build99"),
+		},
+		{
+			name: "Config with missing capacities and capabilities",
+			yamlData: `
+aws:
+  - name: build01
+    capacity: 101 #capacity to 100
+gcp:
+  - name: build02
+    capabilities:
+      - vpn
+  - name: build03
+    blocked: true
+`,
+			expectedCluster: ClusterMap{
+				"build01": {
+					Provider:     "aws",
+					Capacity:     100,
+					Capabilities: nil,
+				},
+				"build02": {
+					Provider:     "gcp",
+					Capacity:     100,
+					Capabilities: []string{"vpn"},
+				},
+			},
+			expectedBlocked: sets.New[string]("build03"),
+		},
+		{
+			name: "Empty config",
+			yamlData: `
+aws: []
+gcp: []
+`,
+			expectedCluster: ClusterMap{},
+			expectedBlocked: sets.New[string](),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			data := []byte(tt.yamlData)
+
+			clusterMap, blockedClusters, err := loadClusterConfigFromBytes(data)
+			if err != nil {
+				t.Fatalf("Failed to load cluster config: %v", err)
+			}
+
+			for clusterName, expectedInfo := range tt.expectedCluster {
+				if info, exists := clusterMap[clusterName]; !exists {
+					t.Errorf("Expected cluster %s to be in clusterMap", clusterName)
+				} else {
+					if info.Provider != expectedInfo.Provider {
+						t.Errorf("Expected provider for %s: %s, got: %s", clusterName, expectedInfo.Provider, info.Provider)
+					}
+					if info.Capacity != expectedInfo.Capacity {
+						t.Errorf("Expected capacity for %s: %d, got: %d", clusterName, expectedInfo.Capacity, info.Capacity)
+					}
+					if len(info.Capabilities) != len(expectedInfo.Capabilities) {
+						t.Errorf("Expected capabilities length for %s: %d, got: %d", clusterName, len(expectedInfo.Capabilities), len(info.Capabilities))
+					}
+					for i, capability := range info.Capabilities {
+						if capability != expectedInfo.Capabilities[i] {
+							t.Errorf("Expected capability %d for %s: %s, got: %s", i, clusterName, expectedInfo.Capabilities[i], capability)
+						}
+					}
+				}
+			}
+			if !blockedClusters.Equal(tt.expectedBlocked) {
+				t.Errorf("Expected blocked clusters: %v, got: %v", tt.expectedBlocked.UnsortedList(), blockedClusters.UnsortedList())
+			}
+		})
+	}
+}