Fix torchrun injection: add ancestor-step label, use torchrun command
- Add trainer.kubeflow.org/trainjob-ancestor-step label to replicatedJob template so Trainer controller injects PET_* env vars - Change TrainJob command from python to torchrun so PET_* env vars are used for distributed training setup - Remove command from Runtime (keep it in TrainJob only) Verified: 2-node 16-GPU distributed training running successfully Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
2b1ad10c2c
commit
7cd7c1af7e
|
|
@ -310,6 +310,9 @@ spec:
|
||||||
replicatedJobs:
|
replicatedJobs:
|
||||||
- name: node
|
- name: node
|
||||||
template:
|
template:
|
||||||
|
metadata:
|
||||||
|
labels:
|
||||||
|
trainer.kubeflow.org/trainjob-ancestor-step: trainer
|
||||||
spec:
|
spec:
|
||||||
template:
|
template:
|
||||||
metadata:
|
metadata:
|
||||||
|
|
@ -376,7 +379,7 @@ spec:
|
||||||
trainer:
|
trainer:
|
||||||
image: nvcr.io/nvidia/pytorch:24.10-py3
|
image: nvcr.io/nvidia/pytorch:24.10-py3
|
||||||
command:
|
command:
|
||||||
- python
|
- torchrun
|
||||||
- /workspace/scripts/train_nccl.py
|
- /workspace/scripts/train_nccl.py
|
||||||
numNodes: 2
|
numNodes: 2
|
||||||
resourcesPerNode:
|
resourcesPerNode:
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue