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.

If consecutive NaN count exceeds 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)

Check a loss value 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.

Parameters:

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

• step (int) – Current training step.

Returns:

True if the loss is valid (finite) on ALL ranks, False if any rank has NaN/Inf.

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)

Find the latest checkpoint for auto-resume.

Checks:

1. {checkpoint_dir}/latest symlink

2. Most recent step_N directory whose DCP shards are durable

CheckpointManager only ever points latest at a durable checkpoint, so (1) needs no further test. 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”. Skipping incomplete directories turns an unrecoverable run into one that resumes from the last durable checkpoint, whatever left the partial directory behind.

Parameters:

checkpoint_dir (str) – Base checkpoint directory.

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.