🚀 The feature, motivation and pitch

I'm trying to implement graceful exit on SLURM pre-emption. SLURM gives 10 minutes upon pre-emption (i.e., SIGTERM; sleep 600; SIGKILL), but it seems that torchrun will send a SIGKILL from the launcher to all the child processes after 30 seconds (based on this PR). Since graceful termination (e.g., saving checkpoint) can take >30seconds when working with large models, the default 30s timeout should be a user-configurable parameter.

Minimal example/reproduction of this issue:

import time
from pathlib import Path
import signal

def ignore(s,f):
    print('!')

signal.signal(signal.SIGTERM, ignore)

p = Path("/tmp/finished")
p.touch()
p.unlink()
time.sleep(200)
p.touch()

1. Launch with python -u -m accelerate.commands.launch --num_processes 1 test.py
2. Send sigterm kill -TERM <pid>
3. Process receives SIGKILL from launcher (torch.distributed.elastic.multiprocessing.api: [WARNING] Unable to shutdown process 719448 via Signals.SIGTERM, forcefully exiting via Signals.SIGKILL) -- from this line

Tagging based on Slack discussion: @kurman @wconstab
Slack thread: https://pytorch.slack.com/archives/C3PDTEV8E/p1707866755665289

Alternatives

No response

Additional context

No response

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @penguinwu @XilunWu @tianyu-l @yf225