Fix(gcluster): enforce strict GKE NAP accelerator validation (#5885)

Author: SwarnaBharathiMantena

pkg/orchestrator/gke/nap.go

@@ -30,28 +30,54 @@ func (g *GKEOrchestrator) isNAPEnabledForMachineType(machineType, zone string) (
 	resolvedType := config.ResolveMachineType(machineType)
 
 	if config.IsTPU(resolvedType) {
-		key := strings.ToLower(g.GenerateGKENodeSelectorLabel(resolvedType))
-		return g.napLimits[key] > 0 || g.napLimits["google.com/tpu"] > 0, nil
+		return g.validateTPUNAPLimit(resolvedType)
 	}
 
 	cap, err := g.FetchMachineCapabilities(resolvedType, zone)
 	if err != nil {
 		return false, err
 	}
 	if len(cap.Accelerators) > 0 {
-		key := strings.ToLower(g.GenerateGKENodeSelectorLabel(resolvedType))
-		if strings.EqualFold(key, resolvedType) {
-			key = strings.ToLower(g.GenerateGKENodeSelectorLabel(cap.Accelerators[0].Type))
-		}
-		if !isKnownGKEAccelerator(key) {
-			return false, fmt.Errorf("unknown accelerator label: %q", cap.Accelerators[0].Type)
-		}
-		return g.napLimits[key] > 0 || g.napLimits["nvidia.com/gpu"] > 0, nil
+		return g.validateGPUNAPLimit(resolvedType, cap)
 	}
 
 	return g.napLimits["cpu"] > 0, nil
 }
 
+func (g *GKEOrchestrator) validateTPUNAPLimit(resolvedType string) (bool, error) {
+	key := strings.ToLower(g.GenerateGKENodeSelectorLabel(resolvedType))
+	if limit, exists := g.napLimits[key]; exists {
+		return limit > 0, nil
+	}
+	// Fallback to generic TPU limit ONLY if no specific TPU limits are configured.
+	for k := range g.napLimits {
+		if isSpecificTPUKey(k) {
+			return false, nil
+		}
+	}
+	return g.napLimits["google.com/tpu"] > 0, nil
+}
+
+func (g *GKEOrchestrator) validateGPUNAPLimit(resolvedType string, cap MachineTypeCap) (bool, error) {
+	key := strings.ToLower(g.GenerateGKENodeSelectorLabel(resolvedType))
+	if strings.EqualFold(key, resolvedType) {
+		key = strings.ToLower(g.GenerateGKENodeSelectorLabel(cap.Accelerators[0].Type))
+	}
+	if !isKnownGKEAccelerator(key) {
+		return false, fmt.Errorf("unknown accelerator label: %q", cap.Accelerators[0].Type)
+	}
+	if limit, exists := g.napLimits[key]; exists {
+		return limit > 0, nil
+	}
+	// Fallback to generic GPU limit ONLY if no specific GPU limits are configured.
+	for k := range g.napLimits {
+		if isSpecificGPUKey(k) {
+			return false, nil
+		}
+	}
+	return g.napLimits["nvidia.com/gpu"] > 0, nil
+}
+
 func (g *GKEOrchestrator) checkNAPFlagsSupported(hasNAPFlags bool, job *orchestrator.JobDefinition) error {
 	if !g.napEnabled && hasNAPFlags {
 		return fmt.Errorf("GKE NAP provisioning options (--gke-nap-provisioning %q, --gke-nap-reservation %q) are only supported on GKE clusters with Node Auto-Provisioning (NAP) enabled. The current cluster does not have NAP enabled.\nRemediation: Enable Node Auto-Provisioning on your cluster to use these options, or submit your job without them", job.GKENAPProvisioning, job.GKENAPReservation)
@@ -126,6 +152,14 @@ func extractShortReservationName(resName string) string {
 	return parts[len(parts)-1]
 }
 
+func isSpecificGPUKey(key string) bool {
+	return strings.HasPrefix(key, "nvidia-")
+}
+
+func isSpecificTPUKey(key string) bool {
+	return strings.HasPrefix(key, "tpu-")
+}
+
 func isKnownGKEAccelerator(key string) bool {
 	switch key {
 	case "nvidia-tesla-t4", "nvidia-tesla-v100":

pkg/orchestrator/gke/resource_resolver_test.go

@@ -784,6 +784,46 @@ func TestValidateConsumptionForStaticCluster(t *testing.T) {
 			wantErr:     true,
 			expectedErr: "unknown accelerator label: \"unknown-gpu\"",
 		},
+		{
+			name:       "NAP Cluster - TPU: Specific limit configured, requesting different TPU (Should Fail)",
+			napEnabled: true,
+			napLimits: map[string]int64{
+				"tpu-v6e-slice":  8,
+				"google.com/tpu": 8,
+			},
+			job: orchestrator.JobDefinition{
+				MachineType:        "ct5lp-hightpu-4t", // TPU v5e (tpu-v5-lite-podslice)
+				GKENAPProvisioning: "spot",
+			},
+			wantErr:     true,
+			expectedErr: "is not configured within your cluster's Node Auto-Provisioning (NAP) limits",
+		},
+		{
+			name:       "NAP Cluster - GPU: Specific limit configured, requesting different GPU (Should Fail)",
+			napEnabled: true,
+			napLimits: map[string]int64{
+				"nvidia-h100-mega-80gb": 8,
+				"nvidia.com/gpu":        8,
+			},
+			job: orchestrator.JobDefinition{
+				MachineType:        "g2-standard-12", // L4 GPU (nvidia-l4)
+				GKENAPProvisioning: "spot",
+			},
+			wantErr:     true,
+			expectedErr: "is not configured within your cluster's Node Auto-Provisioning (NAP) limits",
+		},
+		{
+			name:       "NAP Cluster - GPU: Generic limit only, requesting GPU (Should Pass)",
+			napEnabled: true,
+			napLimits: map[string]int64{
+				"nvidia.com/gpu": 8,
+			},
+			job: orchestrator.JobDefinition{
+				MachineType:        "g2-standard-12",
+				GKENAPProvisioning: "spot",
+			},
+			wantErr: false,
+		},
 	}
 
 	for _, tt := range tests {
@@ -810,6 +850,19 @@ func TestValidateConsumptionForStaticCluster(t *testing.T) {
 						},
 					},
 				},
+				"g2-standard-12:": {
+					GuestCpus: 12,
+					MemoryMb:  48000,
+					Accelerators: []struct {
+						Count int    `json:"guestAcceleratorCount"`
+						Type  string `json:"guestAcceleratorType"`
+					}{
+						{
+							Count: 1,
+							Type:  "nvidia-l4",
+						},
+					},
+				},
 			}
 
 			err := orc.validateConsumptionForStaticCluster(&tt.job)