Skip to content

Commit 8f30edb

Browse files
Neelabh94Neelabh94
authored
fix(job): improve slicing topology handling (#5963)
1 parent c891ffe commit 8f30edb

5 files changed

Lines changed: 77 additions & 21 deletions

File tree

‎pkg/orchestrator/gke/gke_job_orchestrator.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1111,6 +1111,14 @@ func (g *GKEOrchestrator) validateRequestedTopology(requested string, topologies
11111111
return fmt.Errorf("failed to check topology containment: %w", err)
11121112
}
11131113
if fit {
1114+
if !g.hasSlicingTopologies() {
1115+
var valid []string
1116+
for t := range topologies {
1117+
valid = append(valid, t)
1118+
}
1119+
slices.Sort(valid)
1120+
return fmt.Errorf("requested topology %s fits inside discovered limits but no slicing labels found. It must match discovered limits exactly: %v", requested, valid)
1121+
}
11141122
logging.Info("Validated provided Topology: %s", requested)
11151123
return nil
11161124
}

‎pkg/orchestrator/gke/gke_job_orchestrator_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1329,7 +1329,7 @@ func TestVerifyDynamicSlicingActive(t *testing.T) {
13291329
if tt.mockResponses != nil {
13301330
if _, ok := tt.mockResponses["kubectl get topologies.kueue.x-k8s.io -o json"]; !ok {
13311331
tt.mockResponses["kubectl get topologies.kueue.x-k8s.io -o json"] = []shell.CommandResult{
1332-
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`},
1332+
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x2-id"}]}}]}`},
13331333
}
13341334
}
13351335
}
@@ -1340,7 +1340,7 @@ func TestVerifyDynamicSlicingActive(t *testing.T) {
13401340
got, err := orc.verifyDynamicSlicingActive(tt.opts)
13411341

13421342
if (err != nil) != tt.wantErr {
1343-
t.Errorf("verifySuperSlicingActive() error = %v, wantErr %v", err, tt.wantErr)
1343+
t.Errorf("verifyDynamicSlicingActive() error = %v, wantErr %v", err, tt.wantErr)
13441344
return
13451345
}
13461346
if !tt.wantErr && got != tt.wantResult {
@@ -2035,7 +2035,7 @@ func TestGenerateGKEManifest_DynamicSlicingActive_TPU7x(t *testing.T) {
20352035

20362036
mockResponses := map[string][]shell.CommandResult{
20372037
"kubectl get resourceflavors": {{ExitCode: 0, Stdout: ""}},
2038-
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`}},
2038+
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-partition-4x4x4-id"}]}}]}`}},
20392039
"kubectl get admissioncheck": {{ExitCode: 0, Stdout: `{"items": [{"spec": {"controllerName": "accelerator.gke.io/slice"}}]}`}},
20402040
"kubectl get nodes -o jsonpath={range .items[*]}{.metadata.labels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end}": {{ExitCode: 0, Stdout: "8x8x8"}},
20412041
"gcloud compute machine-types describe tpu7x-standard-4t --zone=us-central1-a --format=json": {{ExitCode: 0, Stdout: `{"guestCpus": 8, "memoryMb": 32768, "accelerators": [{"guestAcceleratorCount": 4, "guestAcceleratorType": "tpu7x-standard-4t"}]}`}},
@@ -2116,7 +2116,7 @@ func TestGeneratePathwaysManifest_DynamicSlicing(t *testing.T) {
21162116

21172117
mockResponses := map[string][]shell.CommandResult{
21182118
"kubectl get resourceflavors": {{ExitCode: 0, Stdout: ""}},
2119-
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`}},
2119+
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-partition-4x4x4-id"}]}}]}`}},
21202120
"kubectl get admissioncheck": {{ExitCode: 0, Stdout: `{"items": [{"spec": {"controllerName": "accelerator.gke.io/slice"}}]}`}},
21212121
"kubectl get nodes -o jsonpath={range .items[*]}{.metadata.labels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end} -l cloud.google.com/gke-tpu-accelerator=tpu7x": {{ExitCode: 0, Stdout: "8x8x8"}},
21222122
"gcloud compute machine-types describe tpu7x-standard-4t --zone=us-central1-a --format=json": {{ExitCode: 0, Stdout: `{"guestCpus": 8, "memoryMb": 32768, "accelerators": [{"guestAcceleratorCount": 4, "guestAcceleratorType": "tpu7x-standard-4t"}]}`}},
@@ -2179,7 +2179,7 @@ func TestGenerateGKEManifest_StaticSlicingActive_v6e(t *testing.T) {
21792179

21802180
mockResponses := map[string][]shell.CommandResult{
21812181
"kubectl get resourceflavors": {{ExitCode: 0, Stdout: ""}, {ExitCode: 0, Stdout: ""}},
2182-
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`}, {ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`}},
2182+
"kubectl get topologies.kueue.x-k8s.io -o json": {{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x2-id"}]}}]}`}, {ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x2-id"}]}}]}`}},
21832183
"kubectl get nodes": {{ExitCode: 0, Stdout: "4x4"}, {ExitCode: 0, Stdout: "4x4"}},
21842184
"gcloud compute machine-types describe ct6e-standard-8t --zone=us-central1-a --format=json": {{ExitCode: 0, Stdout: `{"guestCpus": 8, "memoryMb": 32768, "accelerators": [{"guestAcceleratorCount": 4, "guestAcceleratorType": "tpu-v6e-slice"}]}`}},
21852185
}

‎pkg/orchestrator/gke/resource_resolver.go‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ func (g *GKEOrchestrator) verifyDynamicSlicingActive(opts ManifestOptions) (bool
170170
}
171171

172172
isTPU7x := strings.Contains(strings.ToLower(requestedMachineName), "tpu7x")
173-
if !isTPU7x || !g.hasKueueTopologies() || !g.hasSliceAdmissionCheck() {
173+
if !isTPU7x || !g.hasSlicingTopologies() || !g.hasSliceAdmissionCheck() {
174174
g.dynamicSlicingCache[cacheKey] = false
175175
return false, nil
176176
}
@@ -198,7 +198,7 @@ func (g *GKEOrchestrator) verifyStaticSlicingActive(job *orchestrator.JobDefinit
198198
return val, nil
199199
}
200200

201-
if !g.hasKueueTopologies() {
201+
if !g.hasSlicingTopologies() {
202202
g.staticSlicingCache[cacheKey] = false
203203
return false, nil
204204
}
@@ -287,15 +287,29 @@ func (g *GKEOrchestrator) hasSliceAdmissionCheck() bool {
287287
return false
288288
}
289289

290-
func (g *GKEOrchestrator) hasKueueTopologies() bool {
290+
func (g *GKEOrchestrator) hasSlicingTopologies() bool {
291+
if g.slicingTopologiesChecked {
292+
return g.slicingTopologiesDetected
293+
}
294+
295+
defer func() {
296+
g.slicingTopologiesChecked = true
297+
}()
298+
291299
tResult := g.executor.ExecuteCommand("kubectl", "get", "topologies.kueue.x-k8s.io", "-o", "json")
292300
if tResult.ExitCode != 0 {
293301
logging.Warn("Failed to query Kueue topologies. Assuming dynamic-slicing not active.")
294302
return false
295303
}
296304

297305
var tList struct {
298-
Items []interface{} `json:"items"`
306+
Items []struct {
307+
Spec struct {
308+
Levels []struct {
309+
NodeLabel string `json:"nodeLabel"`
310+
} `json:"levels"`
311+
} `json:"spec"`
312+
} `json:"items"`
299313
}
300314

301315
if err := json.Unmarshal([]byte(tResult.Stdout), &tList); err != nil {
@@ -308,7 +322,17 @@ func (g *GKEOrchestrator) hasKueueTopologies() bool {
308322
return false
309323
}
310324

311-
return true
325+
for _, t := range tList.Items {
326+
for _, l := range t.Spec.Levels {
327+
if (strings.HasPrefix(l.NodeLabel, "cloud.google.com/gke-tpu-slice-") || strings.HasPrefix(l.NodeLabel, "cloud.google.com/gke-tpu-partition-")) && strings.HasSuffix(l.NodeLabel, "-id") {
328+
g.slicingTopologiesDetected = true
329+
return true
330+
}
331+
}
332+
}
333+
334+
logging.Info("Kueue topologies found but they do not contain slice/partition labels. Assuming dynamic-slicing not active.")
335+
return false
312336
}
313337

314338
func (g *GKEOrchestrator) calculateResourceLimits(opts ManifestOptions, profile JobProfile) (cpu, mem, gpu, tpu string, err error) {

‎pkg/orchestrator/gke/resource_resolver_test.go‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,27 @@ func (m *MockMachineTypeClient) GetMachineType(project, zone, machineType string
225225
return nil, fmt.Errorf("mock not configured")
226226
}
227227

228+
func injectDefaultMocksForShorthand(mockResponses map[string][]shell.CommandResult) {
229+
if _, ok := mockResponses["kubectl get resourceflavors.kueue.x-k8s.io"]; !ok {
230+
mockResponses["kubectl get resourceflavors.kueue.x-k8s.io"] = []shell.CommandResult{
231+
{ExitCode: 0, Stdout: ""},
232+
{ExitCode: 0, Stdout: ""},
233+
}
234+
}
235+
if _, ok := mockResponses["kubectl get nodes -o jsonpath"]; !ok {
236+
mockResponses["kubectl get nodes -o jsonpath"] = []shell.CommandResult{
237+
{ExitCode: 0, Stdout: "16x16\n"},
238+
{ExitCode: 0, Stdout: "16x16\n"},
239+
}
240+
}
241+
if _, ok := mockResponses["kubectl get topologies.kueue.x-k8s.io -o json"]; !ok {
242+
mockResponses["kubectl get topologies.kueue.x-k8s.io -o json"] = []shell.CommandResult{
243+
{ExitCode: 0, Stdout: `{"items":[{"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-4x8-id"}]}}]}`},
244+
{ExitCode: 0, Stdout: `{"items":[{"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-4x8-id"}]}}]}`},
245+
}
246+
}
247+
}
248+
228249
func TestResolveAcceleratorShorthand(t *testing.T) {
229250
setupMockMachineConfig(t)
230251

@@ -338,12 +359,7 @@ func TestResolveAcceleratorShorthand(t *testing.T) {
338359
mockResponses = make(map[string][]shell.CommandResult)
339360
}
340361
// Add default mocks for topology discovery if not provided
341-
if _, ok := mockResponses["kubectl get resourceflavors"]; !ok {
342-
mockResponses["kubectl get resourceflavors"] = []shell.CommandResult{{ExitCode: 0, Stdout: ""}}
343-
}
344-
if _, ok := mockResponses["kubectl get nodes -o jsonpath"]; !ok {
345-
mockResponses["kubectl get nodes -o jsonpath"] = []shell.CommandResult{{ExitCode: 0, Stdout: "16x16\n"}}
346-
}
362+
injectDefaultMocksForShorthand(mockResponses)
347363

348364
mockExecutor := NewMockExecutor(mockResponses)
349365
orc := newTestGKEOrchestrator(mockExecutor)
@@ -439,7 +455,7 @@ func TestVerifyStaticSlicingActive(t *testing.T) {
439455
requestedTopo: "2x2",
440456
mockResponses: map[string][]shell.CommandResult{
441457
"kubectl get topologies.kueue.x-k8s.io -o json": {
442-
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`},
458+
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x2-id"}]}}]}`},
443459
},
444460
"kubectl get resourceflavors.kueue.x-k8s.io -o jsonpath={range .items[*]}{.spec.nodeLabels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end} -l cloud.google.com/gke-tpu-accelerator=tpu-v6e-slice": {
445461
{ExitCode: 0, Stdout: "4x4\n"},
@@ -454,7 +470,7 @@ func TestVerifyStaticSlicingActive(t *testing.T) {
454470
requestedTopo: "4x4",
455471
mockResponses: map[string][]shell.CommandResult{
456472
"kubectl get topologies.kueue.x-k8s.io -o json": {
457-
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`},
473+
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-4x4-id"}]}}]}`},
458474
},
459475
"kubectl get resourceflavors.kueue.x-k8s.io -o jsonpath={range .items[*]}{.spec.nodeLabels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end} -l cloud.google.com/gke-tpu-accelerator=tpu-v6e-slice": {
460476
{ExitCode: 0, Stdout: "4x4\n"},
@@ -468,7 +484,7 @@ func TestVerifyStaticSlicingActive(t *testing.T) {
468484
requestedTopo: "8x8",
469485
mockResponses: map[string][]shell.CommandResult{
470486
"kubectl get topologies.kueue.x-k8s.io -o json": {
471-
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`},
487+
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-4x4-id"}]}}]}`},
472488
},
473489
"kubectl get resourceflavors.kueue.x-k8s.io -o jsonpath={range .items[*]}{.spec.nodeLabels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end} -l cloud.google.com/gke-tpu-accelerator=tpu-v6e-slice": {
474490
{ExitCode: 0, Stdout: "4x4\n"},
@@ -482,7 +498,7 @@ func TestVerifyStaticSlicingActive(t *testing.T) {
482498
requestedTopo: "2x2",
483499
mockResponses: map[string][]shell.CommandResult{
484500
"kubectl get topologies.kueue.x-k8s.io -o json": {
485-
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"}}]}`},
501+
{ExitCode: 0, Stdout: `{"items":[{"metadata":{"name":"tpu-topology"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x2-id"}]}}]}`},
486502
},
487503
"kubectl get resourceflavors.kueue.x-k8s.io -o jsonpath={range .items[*]}{.spec.nodeLabels.cloud\\.google\\.com/gke-tpu-topology}{\"\\n\"}{end} -l cloud.google.com/gke-tpu-accelerator=tpu-v6e-slice": {
488504
{ExitCode: 0, Stdout: ""},
@@ -583,7 +599,13 @@ func TestResolveHardwareRequirements_NAPIncompatibilities(t *testing.T) {
583599
machineType: "v6e-standard-8t",
584600
topology: "2x4",
585601
mockResponses: map[string][]shell.CommandResult{
586-
"kubectl get topologies.kueue.x-k8s.io": {{ExitCode: 0, Stdout: `{"items": [{"metadata":{"name":"tpu-v6e-slice"},"spec":{"topologies":["4x8"]}}]}`}},
602+
"kubectl get topologies.kueue.x-k8s.io": {
603+
{ExitCode: 0, Stdout: `{"items": [{"metadata":{"name":"tpu-v6e-slice"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x4-id"}]}}]}`},
604+
{ExitCode: 0, Stdout: `{"items": [{"metadata":{"name":"tpu-v6e-slice"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x4-id"}]}}]}`},
605+
{ExitCode: 0, Stdout: `{"items": [{"metadata":{"name":"tpu-v6e-slice"},"spec":{"levels":[{"nodeLabel":"cloud.google.com/gke-tpu-slice-2x4-id"}]}}]}`},
606+
},
607+
"kubectl get resourceflavors": {{ExitCode: 0, Stdout: ""}},
608+
"kubectl get nodes": {{ExitCode: 0, Stdout: "4x8\n"}},
587609
},
588610
wantErr: true,
589611
expectedErrMatch: "TPU Static Sub-slicing is not supported on GKE Node Auto-Provisioning (NAP) workloads",

‎pkg/orchestrator/gke/types.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ type GKEOrchestrator struct {
8585
dynamicSlicingCache map[string]bool
8686
staticSlicingCache map[string]bool
8787
topologyCache map[string]string
88+
slicingTopologiesChecked bool
89+
slicingTopologiesDetected bool
8890
}
8991

9092
// Types for GetClusterInfo unmarshaling

0 commit comments

Comments
 (0)