feat: auto-detect super-slicing via gcloud and omit standard node selectors

Author: Neelabh94

pkg/orchestrator/gke/gke_orchestrator.go

@@ -148,6 +148,7 @@ type ManifestOptions struct {
 	VolumesYAML             string
 	VolumeMountsYAML        string
 	GCSFuseEnabled          bool
+	IsSuperSlicing          bool // True if the cluster supports and is configured for Super-slicing.
 	Pathways                orchestrator.PathwaysJobDefinition
 }
 
@@ -1211,6 +1212,46 @@ func (g *GKEOrchestrator) FetchMachineCapacity(machineType, zone string) (int, e
 	return cap.Accelerators[0].Count, nil
 }
 
+func (g *GKEOrchestrator) verifySuperSlicingActive(opts ManifestOptions) (bool, error) {
+	// 1. TPU Focus: Return false immediately if not using TPUs.
+	if opts.AcceleratorType == "" || !strings.Contains(strings.ToLower(opts.AcceleratorType), "tpu") {
+		return false, nil
+	}
+
+	// 2. Machine Type/Profile Guard: Describe node pool to see if it uses PROVISION_ONLY!
+	// We check for placeholder variables for node pool name. In real usage, this should be resolved from the cluster.
+	poolName := os.Getenv("GKE_NODE_POOL_NAME")
+	if poolName == "" {
+		logging.Warn("GKE_NODE_POOL_NAME is not set. Assuming Super-slicing is not active for this node pool.")
+		return false, nil
+	}
+
+	result := g.executor.ExecuteCommand("gcloud", "container", "node-pools", "describe", poolName, "--cluster="+opts.ClusterName, "--zone="+opts.ClusterLocation, "--format=json(placementPolicy)")
+	if result.ExitCode != 0 {
+		logging.Warn("gcloud container node-pools describe failed: %s. Proceeding assuming no Super-slicing.", result.Stderr)
+		return false, nil
+	}
+
+	var policy map[string]interface{}
+	if err := json.Unmarshal([]byte(result.Stdout), &policy); err == nil {
+		if placement, ok := policy["placementPolicy"].(map[string]interface{}); ok {
+			if mode, ok := placement["acceleratorTopologyMode"].(string); ok && mode == "PROVISION_ONLY" {
+				logging.Info("Super-slicing PROVISION_ONLY mode detected for node pool %s.", poolName)
+				return true, nil
+			}
+		}
+	}
+
+	// 3. Kueue CRD Checks: Check for topologies.kueue.x-k8s.io and AdmissionChecks (simulated via shell commands)
+	crdResult := g.executor.ExecuteCommand("kubectl", "get", "crd", "topologies.kueue.x-k8s.io")
+	if crdResult.ExitCode != 0 {
+		logging.Warn("Topology CRD not found. Kueue Super-slicing not active.")
+		return false, nil
+	}
+
+	return true, nil
+}
+
 func (g *GKEOrchestrator) calculateResourceLimits(opts ManifestOptions) (cpu, mem, gpu, tpu string) {
 	mapped := g.GenerateGKENodeSelectorLabel(opts.AcceleratorType)
 
@@ -1323,7 +1364,14 @@ func (g *GKEOrchestrator) prepareManifestOptions(job orchestrator.JobDefinition,
 	}
 	schedOpts.Topology = topology
 
-	nodeSelectorStr, err := g.buildNodeSelector(schedOpts, job)
+	// Call verifySuperSlicingActive to determine if super-slicing is active!
+	isSuperSlicing, _ := g.verifySuperSlicingActive(ManifestOptions{
+		ClusterName:     job.ClusterName,
+		ClusterLocation: job.ClusterLocation,
+		AcceleratorType: job.AcceleratorType,
+	})
+
+	nodeSelectorStr, err := g.buildNodeSelector(schedOpts, job, isSuperSlicing)
 	if err != nil {
 		return ManifestOptions{}, err
 	}
@@ -1694,7 +1742,7 @@ func (g *GKEOrchestrator) waitForJobCompletion(workloadName, clusterName, cluste
 	return nil
 }
 
-func (g *GKEOrchestrator) buildNodeSelector(schedOpts scheduling.SchedulingOptions, job orchestrator.JobDefinition) (string, error) {
+func (g *GKEOrchestrator) buildNodeSelector(schedOpts scheduling.SchedulingOptions, job orchestrator.JobDefinition, isSuperSlicing bool) (string, error) {
 	nodeSelector := scheduling.GetNodeSelector(schedOpts)
 	accelLabel := g.GenerateGKENodeSelectorLabel(job.AcceleratorType)
 
@@ -1718,7 +1766,9 @@ func (g *GKEOrchestrator) buildNodeSelector(schedOpts scheduling.SchedulingOptio
 		if nodeSelector == nil {
 			nodeSelector = make(map[string]string)
 		}
-		nodeSelector["cloud.google.com/gke-tpu-topology"] = schedOpts.Topology
+		if !isSuperSlicing {
+			nodeSelector["cloud.google.com/gke-tpu-topology"] = schedOpts.Topology
+		}
 	}
 
 	if len(nodeSelector) > 0 {

pkg/orchestrator/gke/gke_orchestrator_test.go

@@ -485,4 +485,87 @@ func TestFetchMachineCapacity(t *testing.T) {
 		})
 	}
 }
+func TestVerifySuperSlicingActive(t *testing.T) {
+	tests := []struct {
+		name          string
+		opts          ManifestOptions
+		mockResponses map[string][]shell.CommandResult
+		envVars       map[string]string
+		wantResult    bool
+	}{
+		{
+			name: "Success - Super-slicing active",
+			opts: ManifestOptions{
+				ClusterName:     "test-cluster",
+				ClusterLocation: "us-central1-a",
+				AcceleratorType: "tpu-v6e-slice",
+			},
+			envVars: map[string]string{"GKE_NODE_POOL_NAME": "test-pool"},
+			mockResponses: map[string][]shell.CommandResult{
+				"gcloud container node-pools describe test-pool --cluster=test-cluster --zone=us-central1-a --format=json(placementPolicy)": {
+					{ExitCode: 0, Stdout: `{"placementPolicy": {"acceleratorTopologyMode": "PROVISION_ONLY"}}`},
+				},
+				"kubectl get crd topologies.kueue.x-k8s.io": {
+					{ExitCode: 0},
+				},
+			},
+			wantResult: true,
+		},
+		{
+			name: "Failure - No TPU",
+			opts: ManifestOptions{
+				ClusterName:     "test-cluster",
+				ClusterLocation: "us-central1-a",
+				AcceleratorType: "nvidia-l4",
+			},
+			envVars:       nil,
+			mockResponses: nil,
+			wantResult:    false,
+		},
+		{
+			name: "Failure - No Node Pool set",
+			opts: ManifestOptions{
+				ClusterName:     "test-cluster",
+				ClusterLocation: "us-central1-a",
+				AcceleratorType: "tpu-v6e-slice",
+			},
+			envVars:       nil,
+			mockResponses: nil,
+			wantResult:    false,
+		},
+		{
+			name: "Failure - gcloud fails",
+			opts: ManifestOptions{
+				ClusterName:     "test-cluster",
+				ClusterLocation: "us-central1-a",
+				AcceleratorType: "tpu-v6e-slice",
+			},
+			envVars: map[string]string{"GKE_NODE_POOL_NAME": "test-pool"},
+			mockResponses: map[string][]shell.CommandResult{
+				"gcloud container node-pools describe test-pool --cluster=test-cluster --zone=us-central1-a --format=json(placementPolicy)": {
+					{ExitCode: 1, Stderr: "slow network"},
+				},
+			},
+			wantResult: false,
+		},
+	}
 
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for k, v := range tt.envVars {
+				t.Setenv(k, v)
+			}
+			mockExecutor := NewMockExecutor(tt.mockResponses)
+			orc := &GKEOrchestrator{executor: mockExecutor}
+
+			got, err := orc.verifySuperSlicingActive(tt.opts)
+
+			if err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+			if got != tt.wantResult {
+				t.Errorf("Expected %t, got %t", tt.wantResult, got)
+			}
+		})
+	}
+}