#698 Pytorch elastic hanging in case of cuda oom or non-rank0 exceptions Pull request opened by kumare3 ## Summary Fix PyTorch elastic training jobs hanging for up to 90 minutes when a worker hits CUDA OOM by adding NCCL timeout configuration and fixing several bugs in the

execute()

lifecycle. ### Problem When a DDP worker hits CUDA OOM and skips a collective (e.g.

loss.backward()

), the surviving workers block in NCCL all-reduce with PyTorch's default timeouts: • Collective timeout: 600s (10 min) — how long workers wait in a stuck all-reduce • Heartbeat timeout: 1800s (30 min) — how long before the heartbeat monitor SIGABRTs the process • Restart cycles:

max_restarts

× (collective + heartbeat) — with defaults: 3 × 2400s = ~2 hours Additionally,

execute()

never called

pre()

or

post()

, dropped

*args

in local mode, and didn't set up the task context properly. ### Changes

plugins/pytorch/src/flyteplugins/pytorch/task.py

• Add three NCCL timeout knobs to

Elastic

config: •

nccl_heartbeat_timeout_sec

(default 300s, PyTorch default is 1800s) •

nccl_async_error_handling

(default False, opt-in for faster crash-exit on stuck collectives) •

nccl_collective_timeout_sec

(default None/PyTorch default, opt-in to reduce the 600s collective timeout) •

pre()

propagates these as env vars (

TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC

,

TORCH_NCCL_ASYNC_ERROR_HANDLING

,

FLYTE_NCCL_COLLECTIVE_TIMEOUT_SEC

) — respects existing env vars •

launcher_entrypoint()

patches

torch.distributed.constants

AND

torch.distributed.distributed_c10d

to override the collective timeout before user code calls

init_process_group()

(both modules must be patched due to Python's

from X import Y

binding semantics) • Fix `execute()`: call `pre()`/`post()`, set up task context, fix local mode

*args

• Catch

ChildFailedError

and re-raise with human-readable error messages (e.g. SIGABRT → "NCCL collective timed out — common symptom of CUDA OOM causing ranks to desync") • Improve docstrings: document restart cycle cost formula, when to use

max_restarts=0

vs higher values

plugins/pytorch/tests/test_task.py

• 17 new tests (19 total) covering all NCCL timeout knobs, env var propagation, and the

launcher_entrypoint

monkey-patch

examples/plugins/pytorch/cuda_oom_hang.py

(new) • Reproducer that triggers CUDA OOM on rank 1 to desync DDP collectives • Demonstrates aggressive timeout settings:

nccl_collective_timeout_sec=60

,

nccl_heartbeat_timeout_sec=60

,

nccl_async_error_handling=True

,

max_restarts=0

### Results Tested on 2x T4 GPUs with the reproducer: | Setting | Time to failure | | --------------------------- | --------------- | | Before (PyTorch defaults) | ~90 min hang | | After (plugin defaults) | ~15 min | | After (aggressive settings) | ~2 min | flyteorg/flyte-sdk