kempnerforge.resilience.signal_handler

Graceful shutdown via signal handling.

Handles SIGTERM (SLURM preemption / graceful shutdown) and SIGUSR1 (SLURM requeue) by setting a flag that the training loop checks after each step. On signal, the loop saves an emergency checkpoint and exits.

Timeout protection ensures the process exits even if graceful shutdown stalls (e.g., stuck in NCCL collective).

Functions

ignore_shutdown_signals_in_worker(worker_id)

worker_init_fn that makes a DataLoader worker ignore shutdown signals.

Classes

ShutdownHandler

Cooperative shutdown handler for long-running training jobs.

kempnerforge.resilience.signal_handler.ignore_shutdown_signals_in_worker(worker_id)

worker_init_fn that makes a DataLoader worker ignore shutdown signals.

SLURM preemption and torchrun teardown signal the whole process group, not a single PID, so SIGTERM reaches every DataLoader worker as well as the rank’s main process. Workers install no handler, take the default action and die at once. The main process — which handles the signal cooperatively, by setting a flag and finishing its step — then fails fetching its next micro-batch, and that failure escapes before the loop reaches its should_shutdown() check, so no emergency checkpoint is written.

Ignoring the shutdown signals in the worker keeps it serving long enough for the step to finish, which puts the emergency save back on the loop’s normal path rather than on an exception path.

The cost: multiprocessing.Process.terminate() is itself SIGTERM, so a shielded worker cannot be reaped that way and the interpreter’s exit hook would block joining it. Workers must therefore be shut down explicitly through the loader’s sentinel path — see BatchStream.close.

Parameters:

worker_id (int) – Worker index, supplied by DataLoader. Unused.

Return type:

None

class kempnerforge.resilience.signal_handler.ShutdownHandler

Bases: object

Cooperative shutdown handler for long-running training jobs.

Register this handler before the training loop. The training loop checks should_shutdown() after each step and takes appropriate action (save checkpoint, clean up, exit).

If the graceful shutdown exceeds timeout_sec, a forced exit is triggered via os._exit to avoid hanging on stuck collectives.

Usage:

handler = ShutdownHandler(timeout_sec=120)
handler.register()

for step in range(max_steps):
    train_step()
    if handler.should_shutdown():
        save_checkpoint()
        handler.finish()
        break

Parameters:

timeout_sec – Maximum seconds allowed for graceful shutdown before forced exit. Set to 0 to disable the timeout.

__init__(timeout_sec=600.0)

Parameters:

timeout_sec (float)

Return type:

None

property shutdown_requested: bool

Whether a shutdown signal has been received.

property signal_received: Signals | None

The signal that triggered shutdown, or None.

should_shutdown()

Check if the training loop should exit.

Call this after each training step.

Return type:

bool

register()

Register signal handlers for SIGTERM and SIGUSR1.

Must be called from the main thread.

Return type:

None

unregister()

Restore original signal handlers.

Return type:

None

finish()

Call after graceful shutdown is complete.

Cancels the forced-exit timer and restores signal handlers.

Return type:

None