feat(job): enable static sub-slicing for v6e and v5p(#5774)

Author: Neelabh94

pkg/config/hardware.go

@@ -436,11 +436,11 @@ func isValid3DDimension(x int) bool {
 }
 
 // Validate3DTopology validates the 3D topology dimensions for a given machine type.
-// If superSlicingCheck is true, it skips the common 3D shapes check and strictly validates
-// that each dimension is a multiple of 4 and >= 4.
-func Validate3DTopology(topology string, machineType string, superSlicingCheck bool) error {
+// If dynamicSlicingCheck is true, it validates that the topology is either a valid subslice
+// (2x2x1, 2x2x2, 2x2x4, 2x4x4) or a valid superslice (dimensions are multiples of 4 and >= 4).
+func Validate3DTopology(topology string, machineType string, dynamicSlicingCheck bool) error {
 
-	if !superSlicingCheck && valid3DShapes[topology] {
+	if !dynamicSlicingCheck && valid3DShapes[topology] {
 		return nil
 	}
 
@@ -452,23 +452,13 @@ func Validate3DTopology(topology string, machineType string, superSlicingCheck b
 	b, _ := strconv.Atoi(dims[1])
 	c, _ := strconv.Atoi(dims[2])
 
-	if !isValid3DDimension(a) || !isValid3DDimension(b) || !isValid3DDimension(c) {
-		return fmt.Errorf("topology %s dimensions must be >= 4 and multiples of 4 (unless it is a common base topology)", topology)
-	}
-
-	maxCubes := 0
-	found := false
-
-	for prefix, constraints := range tpu3DConstraintsMap {
-		if strings.HasPrefix(machineType, prefix) {
-			maxCubes = constraints.maxCubes
-			found = true
-			break
-		}
+	if err := validateDimensions(topology, a, b, c, dynamicSlicingCheck); err != nil {
+		return err
 	}
 
-	if !found {
-		return fmt.Errorf("unknown 3D TPU family for machine type %s", machineType)
+	maxCubes, err := getMaxCubes(machineType)
+	if err != nil {
+		return err
 	}
 
 	cubes := (a / 4) * (b / 4) * (c / 4)
@@ -477,12 +467,38 @@ func Validate3DTopology(topology string, machineType string, superSlicingCheck b
 	}
 
 	if !(a <= b && b <= c) {
-		return fmt.Errorf("topology %s dimensions must be in non-decreasing order (A <= B <= C)", topology)
+		if topology != "2x2x1" {
+			return fmt.Errorf("topology %s dimensions must be in non-decreasing order (A <= B <= C)", topology)
+		}
 	}
 
 	return nil
 }
 
+func validateDimensions(topology string, a, b, c int, dynamicSlicingCheck bool) error {
+	if dynamicSlicingCheck {
+		isSubSlice := (topology == "2x2x1" || topology == "2x2x2" || topology == "2x2x4" || topology == "2x4x4")
+		isSuperSlice := (isValid3DDimension(a) && isValid3DDimension(b) && isValid3DDimension(c))
+		if !isSubSlice && !isSuperSlice {
+			return fmt.Errorf("topology %s is not valid for dynamic slicing. It must be a valid subslice (2x2x1, 2x2x2, 2x2x4, 2x4x4) or superslice (dimensions >= 4 and multiples of 4)", topology)
+		}
+	} else {
+		if !isValid3DDimension(a) || !isValid3DDimension(b) || !isValid3DDimension(c) {
+			return fmt.Errorf("topology %s dimensions must be >= 4 and multiples of 4 (unless it is a common base topology)", topology)
+		}
+	}
+	return nil
+}
+
+func getMaxCubes(machineType string) (int, error) {
+	for prefix, constraints := range tpu3DConstraintsMap {
+		if strings.HasPrefix(machineType, prefix) {
+			return constraints.maxCubes, nil
+		}
+	}
+	return 0, fmt.Errorf("unknown 3D TPU family for machine type %s", machineType)
+}
+
 // CheckTopologyContainment returns true if the requested topology fits within the container topology.
 func CheckTopologyContainment(requested, container string, accelType string) (bool, error) {
 	reqDims := strings.Split(requested, "x")
@@ -506,3 +522,19 @@ func CheckTopologyContainment(requested, container string, accelType string) (bo
 
 	return true, nil
 }
+
+// Is3DTorusTPU returns true if the machine type belongs to a 3D Torus TPU family (v4 or v5p)
+// where static sub-slicing coordinate labeling is unsupported by GKE.
+func Is3DTorusTPU(machineType string) bool {
+	lower := strings.ToLower(machineType)
+	// Filter out tpu7x (v7) which uses its own dynamic coordinates/partition scheme.
+	if strings.Contains(lower, "tpu7") {
+		return false
+	}
+	for _, family := range []string{"v4", "v5p", "ct4p", "ct5p"} {
+		if strings.Contains(lower, family) {
+			return true
+		}
+	}
+	return false
+}

pkg/config/hardware_test.go

@@ -575,3 +575,61 @@ func TestCheckTopologyContainment(t *testing.T) {
 		})
 	}
 }
+
+func TestIs3DTorusTPU(t *testing.T) {
+	tests := []struct {
+		name        string
+		machineType string
+		wantResult  bool
+	}{
+		{
+			name:        "v4 chip name (Torus)",
+			machineType: "ct4p-hightpu-4t",
+			wantResult:  true,
+		},
+		{
+			name:        "v5p machine name (Torus)",
+			machineType: "ct5p-hightpu-4t",
+			wantResult:  true,
+		},
+		{
+			name:        "v5e machine name (2D Mesh)",
+			machineType: "ct5l-hightpu-4t",
+			wantResult:  false,
+		},
+		{
+			name:        "v6e machine name (2D Mesh)",
+			machineType: "ct6e-standard-8t",
+			wantResult:  false,
+		},
+		{
+			name:        "tpu7x machine name (Exclude)",
+			machineType: "tpu7x-standard-16t",
+			wantResult:  false,
+		},
+		{
+			name:        "Shorthand v4-8",
+			machineType: "v4-8",
+			wantResult:  true,
+		},
+		{
+			name:        "Shorthand v5p-8",
+			machineType: "v5p-8",
+			wantResult:  true,
+		},
+		{
+			name:        "Non-TPU machine type",
+			machineType: "a2-highgpu-1g",
+			wantResult:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := Is3DTorusTPU(tt.machineType)
+			if got != tt.wantResult {
+				t.Errorf("Is3DTorusTPU(%q) = %v, want %v", tt.machineType, got, tt.wantResult)
+			}
+		})
+	}
+}

pkg/orchestrator/gke/gke_job_orchestrator.go

@@ -60,6 +60,7 @@ func NewGKEOrchestrator() *GKEOrchestrator {
 		machineCapCache:          make(map[string]MachineTypeCap),
 		topologyCache:            make(map[string]string),
 		dynamicSlicingCache:      make(map[string]bool),
+		staticSlicingCache:       make(map[string]bool),
 	}
 }
 
@@ -102,7 +103,7 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
 		return err
 	}
 
-	profile, isDynamicSlicing, err := g.resolveHardwareRequirements(&job)
+	profile, isDynamicSlicing, isStaticSlicing, err := g.resolveHardwareRequirements(&job)
 	if err != nil {
 		return err
 	}
@@ -115,7 +116,7 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
 		return err
 	}
 
-	if err := g.generateAndSubmitManifests(job, fullImageName, profile, isDynamicSlicing); err != nil {
+	if err := g.generateAndSubmitManifests(job, fullImageName, profile, isDynamicSlicing, isStaticSlicing); err != nil {
 		return err
 	}
 
@@ -250,16 +251,16 @@ func (g *GKEOrchestrator) GetJobLogs(name string, opts orchestrator.LogsOptions)
 	return res.Stdout, nil
 }
 
-func (g *GKEOrchestrator) generateAndSubmitManifests(job orchestrator.JobDefinition, fullImageName string, profile JobProfile, isDynamicSlicing bool) error {
+func (g *GKEOrchestrator) generateAndSubmitManifests(job orchestrator.JobDefinition, fullImageName string, profile JobProfile, isDynamicSlicing bool, isStaticSlicing bool) error {
 	if job.IsPathwaysJob {
-		manifestContent, err := g.GeneratePathwaysManifest(job, fullImageName, profile, isDynamicSlicing)
+		manifestContent, err := g.GeneratePathwaysManifest(job, fullImageName, profile, isDynamicSlicing, isStaticSlicing)
 		if err != nil {
 			return err
 		}
 		return g.ApplyManifest(manifestContent, job.DryRunManifest, job.WorkloadName)
 	}
 
-	manifestOpts, err := g.PrepareManifestOptions(job, fullImageName, profile, isDynamicSlicing)
+	manifestOpts, err := g.PrepareManifestOptions(job, fullImageName, profile, isDynamicSlicing, isStaticSlicing)
 	if err != nil {
 		return err
 	}
@@ -302,7 +303,7 @@ func (g *GKEOrchestrator) validateJobConflicts(workloadName string, clusterName
 	return nil
 }
 
-func (g *GKEOrchestrator) GeneratePathwaysManifest(job orchestrator.JobDefinition, fullImageName string, profile JobProfile, isDynamicSlicing bool) (string, error) {
+func (g *GKEOrchestrator) GeneratePathwaysManifest(job orchestrator.JobDefinition, fullImageName string, profile JobProfile, isDynamicSlicing bool, isStaticSlicing bool) (string, error) {
 	// Set default values for Pathways-specific fields if not provided
 	if job.Pathways.ProxyServerImage == "" {
 		job.Pathways.ProxyServerImage = defaultPathwaysProxyImage
@@ -320,7 +321,7 @@ func (g *GKEOrchestrator) GeneratePathwaysManifest(job orchestrator.JobDefinitio
 		return "", fmt.Errorf("failed to parse pathways jobset template: %w", err)
 	}
 
-	opts, err := g.PrepareManifestOptions(job, fullImageName, profile, isDynamicSlicing)
+	opts, err := g.PrepareManifestOptions(job, fullImageName, profile, isDynamicSlicing, isStaticSlicing)
 	if err != nil {
 		return "", err
 	}
@@ -1797,12 +1798,12 @@ func (g *GKEOrchestrator) addAcceleratorLabel(nodeSelector map[string]string, ac
 	}
 }
 
-func (g *GKEOrchestrator) addTopologyLabel(nodeSelector map[string]string, schedOpts SchedulingOptions, isGPU bool, isCPUMachine bool, isDynamicSlicing bool) error {
+func (g *GKEOrchestrator) addTopologyLabel(nodeSelector map[string]string, schedOpts SchedulingOptions, isGPU bool, isCPUMachine bool) error {
 	if schedOpts.Topology != "" {
 		if isGPU || isCPUMachine {
 			return fmt.Errorf("topology is not allowed for GPU and CPU jobs")
 		}
-		if !isDynamicSlicing {
+		if !schedOpts.IsDynamicSlicing && !schedOpts.IsStaticSlicing {
 			_, hasFallback := schedOpts.NodeAffinityLabels[tpuTopologyLabel]
 			if !hasFallback {
 				nodeSelector[tpuTopologyLabel] = schedOpts.Topology
@@ -1812,7 +1813,7 @@ func (g *GKEOrchestrator) addTopologyLabel(nodeSelector map[string]string, sched
 	return nil
 }
 
-func (g *GKEOrchestrator) buildNodeSelector(schedOpts SchedulingOptions, job orchestrator.JobDefinition, isDynamicSlicing bool, isCPUMachine bool) (string, error) {
+func (g *GKEOrchestrator) buildNodeSelector(schedOpts SchedulingOptions, job orchestrator.JobDefinition, isCPUMachine bool) (string, error) {
 	nodeSelector := GetNodeSelector(schedOpts)
 	if nodeSelector == nil {
 		nodeSelector = make(map[string]string)
@@ -1831,7 +1832,7 @@ func (g *GKEOrchestrator) buildNodeSelector(schedOpts SchedulingOptions, job orc
 
 	g.addAcceleratorLabel(nodeSelector, accelLabel, isCPUMachine, job.MachineType)
 
-	if err := g.addTopologyLabel(nodeSelector, schedOpts, isGPU, isCPUMachine, isDynamicSlicing); err != nil {
+	if err := g.addTopologyLabel(nodeSelector, schedOpts, isGPU, isCPUMachine); err != nil {
 		return "", err
 	}
 
@@ -1860,10 +1861,10 @@ func (g *GKEOrchestrator) buildAffinity(schedOpts SchedulingOptions) (string, er
 	return "", nil
 }
 
-func (g *GKEOrchestrator) buildTopologyAnnotation(topology string, numSlices int, isDynamicSlicing bool) string {
+func (g *GKEOrchestrator) buildTopologyAnnotation(topology string, machineType string, numSlices int, nodesPerSlice int, isSubSlicing bool) string {
 	var topologyAnnotation map[string]string
-	if isDynamicSlicing {
-		topologyAnnotation = GetTopologyAnnotation(topology, numSlices)
+	if isSubSlicing {
+		topologyAnnotation = GetTopologyAnnotation(topology, machineType, numSlices, nodesPerSlice)
 	} else if topology != "" {
 		topologyAnnotation = map[string]string{
 			"cloud.google.com/gke-tpu-slice-topology": topology,