kempnerforge.resilience

Fault tolerance and resilience for KempnerForge.

class kempnerforge.resilience.NaNDetector

Bases: object

Detects and tracks NaN/Inf values in loss and gradients.

Supports three responses to NaN:

• "warn": Log a warning and continue.

• "skip": Skip the optimizer step (zero gradients).

• "raise": Raise a RuntimeError.

Once the consecutive NaN count reaches max_consecutive, the detector signals that a checkpoint rollback is recommended.

Parameters:

• action – What to do when NaN is detected.

• max_consecutive – Consecutive NaN steps before recommending rollback.

• max_history – Number of NaN step indices to retain.

__init__(action='warn', max_consecutive=5, max_history=100)

Parameters:

• action (str)

• max_consecutive (int)

• max_history (int)

Return type:

None

check_loss(loss, step, grad_norm=None)

Check a step’s loss, and its gradient norm when given, for NaN/Inf.

When running distributed, all-reduces a NaN flag so ALL ranks agree on whether to skip. Prevents rank desync where one rank sees NaN and skips its optimizer step while others proceed normally.

A finite loss does not guarantee finite gradients: a micro-batch whose labels are all masked has an exact zero cross-entropy loss whatever its activations, so the caller passes the gradient norm as well. Each rank’s norm goes into the same reduced flag as the loss, so every rank skips the step together even when a non-finite gradient reaches only some ranks’ norms.

Parameters:

• loss (float) – The scalar loss value to check.

• step (int) – Current training step.

• grad_norm (float | None) – The step’s gradient norm, or None to check the loss alone.

Returns:

True if the loss and the gradient norm are finite on ALL ranks, False if any rank has NaN/Inf in either.

Raises:

RuntimeError – If action is “raise” and NaN is detected.

Return type:

bool

check_gradients(model, step)

Check model gradients for NaN/Inf before optimizer step.

Parameters:

• model (torch.nn.Module) – The model to check.

• step (int) – Current training step.

Returns:

True if all gradients are finite.

Return type:

bool

property should_rollback: bool

Whether consecutive NaN count suggests a checkpoint rollback.

reset()

Reset NaN tracking state (e.g., after a rollback).

Return type:

None

class kempnerforge.resilience.NaNState

Bases: object

Tracks NaN/Inf occurrences across training steps.

consecutive_nans: int = 0

total_nans: int = 0

last_good_loss: float = inf

last_good_step: int = 0

nan_steps: list[int]

__init__(consecutive_nans=0, total_nans=0, last_good_loss=inf, last_good_step=0, nan_steps=<factory>)

Parameters:

• consecutive_nans (int)

• total_nans (int)

• last_good_loss (float)

• last_good_step (int)

• nan_steps (list[int])

Return type:

None

class kempnerforge.resilience.SLURMInfo

Bases: object

Information about the current SLURM job.

job_id: str

job_name: str

node_list: str

num_nodes: int

ntasks_per_node: int

restart_count: int

partition: str

array_task_id: str | None

property is_requeued: bool

Whether this job has been requeued (restart_count > 0).

__init__(job_id, job_name, node_list, num_nodes, ntasks_per_node, restart_count, partition, array_task_id)

Parameters:

• job_id (str)

• job_name (str)

• node_list (str)

• num_nodes (int)

• ntasks_per_node (int)

• restart_count (int)

• partition (str)

• array_task_id (str | None)

Return type:

None

class kempnerforge.resilience.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

kempnerforge.resilience.check_gpu_health(device=0)

Run basic GPU health checks.

Performs:

1. CUDA availability check

2. Small test computation on the device

3. Memory allocation test

Returns:

Dict with health check results.

Parameters:

device (int)

Return type:

dict[str, bool | str]

kempnerforge.resilience.check_nccl_health(timeout_sec=10.0)

Check NCCL communication health via a lightweight all-reduce.

The all-reduce runs with async_op=True so work.wait(timeout=...) enforces the caller’s bound rather than falling back to the process-group default timeout (nccl_timeout_sec, 1800s). Without that, this function would sit for 30 minutes on a single stuck peer regardless of the timeout_sec argument.

Parameters:

timeout_sec (float) – Per-operation timeout for the collective. Returns False if the all-reduce does not complete within this budget.

Returns:

True on success, False on timeout, error, or world-size mismatch.

Return type:

bool

kempnerforge.resilience.get_slurm_info()

Read SLURM job information from environment variables.

Returns:

SLURMInfo if running under SLURM, None otherwise.

Return type:

SLURMInfo | None

kempnerforge.resilience.is_slurm_job()

Check if we are running under SLURM.

Return type:

bool

kempnerforge.resilience.is_slurm_requeue()

Check if this is a requeued SLURM job.

Uses SLURM_RESTART_COUNT (set by SLURM on requeue).

Return type:

bool

kempnerforge.resilience.log_job_info()

Log SLURM job information (if running under SLURM).

Return type:

None

kempnerforge.resilience.resolve_resume_path(checkpoint_dir, num_stages=None, *, require_train_state=True)

Find the latest checkpoint for auto-resume.

Checks:

1. {checkpoint_dir}/latest symlink

2. Most recent step_N directory whose save completed

CheckpointManager only ever points latest at a complete checkpoint, so (1) needs no further test; one that lost files since is returned as is, and the load reports it. The step_N fallback has no such guarantee: an interrupted save leaves a directory that exists but holds no .metadata, and resuming into it fails in dcp.load with “metadata is None”. A save that stopped after its DCP shards and before train_state.pt leaves optimizer state without the step, LR schedule and data position that go with it, which CheckpointManager.load refuses. Skipping both kinds turns an unrecoverable run into one that resumes from the last complete checkpoint, whatever left the partial directory behind. When step_N directories exist but none is complete, a run that meant to resume starts over, so that is logged at ERROR.

Parameters:

• checkpoint_dir (str) – Base checkpoint directory.

• num_stages (int | None) – The run’s pipeline degree. With it, a step_N in the per-stage layout counts only when pp0 through pp{num_stages - 1} all hold .metadata, including stages that never created their directory; see _dcp_durable. When no step has every stage, the newest one that is otherwise complete is returned as before, after an error log, rather than the run starting over: a checkpoint saved with another pipeline degree looks the same, and resume does not support changing the degree.

• require_train_state (bool) – Skip a step_N whose shards hold optimizer state but that has no train_state.pt. A caller that loads the weights alone, which such a checkpoint still holds in full, passes False to keep the newest one.

Returns:

Path to the latest usable checkpoint, or None if none found.

Return type:

Path | None

Modules

elastic	Elastic training and SLURM integration helpers.
health	GPU health monitoring and NaN detection.
signal_handler	Graceful shutdown via signal handling.