kempnerforge.resilience.elastic

Elastic training and SLURM integration helpers.

Provides utilities for training jobs that may be preempted, requeued, or restarted with a different number of nodes:

  • SLURM job info detection

  • Requeue detection

  • Auto-resume path resolution

Functions

get_slurm_info()

Read SLURM job information from environment variables.

is_slurm_job()

Check if we are running under SLURM.

is_slurm_requeue()

Check if this is a requeued SLURM job.

log_job_info()

Log SLURM job information (if running under SLURM).

resolve_resume_path(checkpoint_dir[, ...])

Find the latest checkpoint for auto-resume.

Classes

SLURMInfo

Information about the current SLURM job.
class kempnerforge.resilience.elastic.SLURMInfo[source]

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

kempnerforge.resilience.elastic.get_slurm_info()[source]

Read SLURM job information from environment variables.

Returns:

SLURMInfo if running under SLURM, None otherwise.

Return type:

SLURMInfo | None

kempnerforge.resilience.elastic.is_slurm_job()[source]

Check if we are running under SLURM.

Return type:

bool

kempnerforge.resilience.elastic.is_slurm_requeue()[source]

Check if this is a requeued SLURM job.

Uses SLURM_RESTART_COUNT (set by SLURM on requeue).

Return type:

bool

kempnerforge.resilience.elastic.resolve_resume_path(checkpoint_dir, num_stages=None, *, require_train_state=True)[source]

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

kempnerforge.resilience.elastic.log_job_info()[source]

Log SLURM job information (if running under SLURM).

Return type:

None