kubeflow / trainer / 40014429491

77%

Build Type

push

github

Committed by web-flow

Commit Message

feat(controller): Support Job Replicas for multi-slice TPU (#4250)

* feat(operator): support multi-slice TPU training via trainer replicas

For multi-slice TPU, JobSet models each TPU slice as a ReplicatedJob
replica, with parallelism = hosts per slice and replicas = slice count.
The operator previously blocked this with two hard constraints:

1. builder.go unconditionally set trainer Replicas = 1, destroying any
   value from the runtime template.
2. trainingruntime_webhook.go rejected replicas != 1 for all ancestors
   including trainer.

Changes:
- builder.go: nil-guard for trainer Replicas, preserving the value from
  the runtime template instead of unconditional overwrite.
- jobset.go: in Build(), compute perSlice = numNodes / replicas for the
  trainer ancestor so each slice runs the correct number of hosts.
- trainingruntime_webhook.go: allow trainer ancestor replicas > 1 to
  enable multi-slice configurations to pass admission.
- trainingruntime_webhook_test.go: update invalid_replicas test to
  reflect that trainer replicas > 1 is now valid.
- trainingruntime_test.go: add test case for 4-slice x 8 hosts
  (NumNodes=32), verifying Parallelism=8 per slice and MinMember=34.

Semantics: numNodes = total hosts across all slices.
Per-slice hosts = numNodes / replicas.

REF: https://github.com/kubeflow/trainer/issues/3407
Signed-off-by: krishdef7 <gargkrish06@gmail.com>

* fix(webhook): reject trainer replicas below one

The trainer ancestor now allows replicas greater than one for
multi-slice TPU, but the webhook still admitted zero or negative
values. Add a lower bound check for the trainer ancestor and unit
tests covering replicas set to four (valid) and zero (invalid).

Signed-off-by: krishdef7 <gargkrish06@gmail.com>

* feat(controller): Support trainer replicas in JobSet

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

* fix(runtime): address review comments for multi-slice TPU

- Skip JAX_COORDINATOR_ADDRESS for multiple ... (continued)

Coverage Stats

89 of 99 new or added lines in 10 files covered. (89.9%)

3 existing lines in 1 file now uncovered.

3295 of 4268 relevant lines covered (77.2%)

0.78 hits per line