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)[source]

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

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

Log SLURM job information (if running under SLURM).

Return type:

None