kfpytorch: make Elastic task_type dynamic so nnodes overrides flip single-pod vs distributed

[codegen-sh]

Summary

• Implement Option A: make the task type for Elastic PyTorch dynamic so nnodes overrides are respected at serialization time

Details

• Override task_type property in PytorchElasticFunctionTask to return "python-task" when current task_config.nnodes denotes single-node (1 or "1" or "1:1"), else "pytorch"
• Avoid torch imports in property; simple and resilient parsing for ints and range strings (e.g., "1:1")
• Keeps task_type_version unchanged and preserves existing get_custom branching behavior
• Single-pod execute path now ignores PET_NNODES and forces nnodes=1
• Single-pod execute path pins rdzv_endpoint=127.0.0.1:0 (avoids any DNS/IPv6 localhost issues)
• Added detailed debug logs of mode, raw nnodes string, parsed min/max, nproc_per_node, rdzv_backend, rdzv_endpoint, start_method to aid troubleshooting

Why

• Previously task_type was decided once in init from the initial config. Later overrides via Node.with_overrides(task_config=Elastic(...)) update _task_config but do not update task_type. This caused tasks defined with nnodes>1 to remain distributed even when overrides set nnodes=1.
• In some environments, env overrides (like PET_NNODES, PET_RDZV_ENDPOINT) can cause single-pod paths to wait for multiple nodes or use cross-pod rdzv, leading to rendezvous timeouts. We now pin nnodes and rdzv in single-pod mode.

Testing

• Existing tests around Elastic behavior should continue to pass. The dynamic property only affects the serialized TaskTemplate type when nnodes flips between single and multi.
• Single-pod mode now logs effective settings and no longer honors PET_NNODES>1 or external rdzv endpoints.

Notes

• This change ensures runtime overrides (e.g., in workflows/launch plans) can flip between single-pod torchrun and multi-node PyTorchJob without redefining the task, and prevents misconfig from env vars in single-pod mode.

---

💻 View my work • About Codegen