wiki:Workshops/JobCheckpointing/Examples

Context Navigation

Version 5 (modified by Carl Baribault, 7 months ago) ( diff )
Ongoing update with submitter/runner implementation

Job Checkpointing Examples

Job Checkpointing Examples

For a quick start see Working Examples below.

Here is an implementation for running sequential workflows on Cypress - where each component result of the workflow depends the previous component result in the workflow.

It can be considered an implementation of an HPC workflow design pattern (see for example Design Patterns) that can be described by any of the following expressions.

submitter/runner pattern (see also Workflow for HPC)
job chaining (see (NERSC) Best Practices for Jobs)
sequential workflows
restart workflows
checkpoint/restart job series
using job dependencies to restart jobs that checkpoint their state”

But wait - before you adopt this type of checkpointing workflow, consider whether your workflow can be considered embarrassingly parallel, where the data elements can be processed entirely independently of each other - rather than each one depending on the previous. (See also Embarrassingly Parallel Algorithms Explained.)

If so, namely that your case is the former (embarrassingly parallel) in the above, then please first refer to the workshop Workshops/JobParallelism.

Checkpoint Submitter

Here is a job script for a fully self restarting job that controls the workflow by submitting the checkpoint runner job script (see Checkpoint Runner) further below.

checkpoint_submitter.sh

#!/bin/bash
###############################################################################
# CHECKPOINT SUBMITTER: self-restarting to re-submit runner
# Controls repeated submissions of checkpoint_runner.sh
###############################################################################
# User settings with defaults
###############################################################################
# CKPT_PATH            : path to checkpoint file
# MAX_ITER             : stop application after this many iterations total
# RUNNER_SCRIPT        : job script with checkpoint before end of time limit
# SUBMITTER_MARGIN_SEC : time margin to restart before time limit
# SUBMITTER_TIME_LIMIT : wall time <= 24 hours (format as --time in man sbatch)
###############################################################################
# Edit settings below
# Or pass values via...
#   sbatch --export=ALL,Var1=...,...,VarN=... checkpoint_submitter.sh
# Or pass values via...
#   Var1=... VarN=... sbatch checkpoint_submitter.sh
###############################################################################
CKPT_PATH="${CKPT_PATH:-state_iter.txt}"
MAX_ITER="${MAX_ITER:-500}"
RUNNER_SCRIPT="${RUNNER_SCRIPT:-checkpoint_runner.sh}"
SUBMITTER_MARGIN_SEC="${SUBMITTER_MARGIN_SEC:-60}"
SUBMITTER_TIME_LIMIT="${SUBMITTER_TIME_LIMIT:-3:30}"
###############################################################################
# Slurm directives (keep SUBMITTER_TIME_LIMIT <= 24 hours)
###############################################################################
#SBATCH --job-name=ckpt_submitter
#SBATCH --output=log_submitter_%j.out
#SBATCH --error=log_submitter_%j.err
#SBATCH --time=24:00:00
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=256M
#SBATCH --partition=centos7
#SBATCH --qos=normal
###############################################################################
# Exit on error
###############################################################################
set -euo pipefail
###############################################################################
# Helpers
###############################################################################
shopt -s expand_aliases
alias dtstamp='date +%Y%m%d-%H:%M:%S'
info(){ echo "Info[$(dtstamp)]: $*"; }
get_ckpt_iter() {
  [[ -f "${CKPT_PATH}" ]] || { echo 0; return; }
  tr -cd '0-9' < "${CKPT_PATH}" 2>/dev/null || echo 0
}
submit_worker() {
  sbatch "${RUNNER_SCRIPT}" | awk '{print $4}'
}
to_seconds() {
        # Acceptable time formats include "minutes", "minutes:seconds", "hours:min- utes:seconds",
        # "days-hours", "days-hours:minutes" and "days-hours:minutes:seconds"
        local input="$1"
        # Handle days-hours format by replacing '-' with ':' to unify the delimiter
        # Then split components into an array
        IFS=':' read -r -a parts <<< "${input//-/:}"
        case ${#parts[@]} in
            1) # Format: minutes
                seconds=$((parts[0] * 60))
                ;;
            2) # Formats: minutes:seconds OR days-hours
                if [[ $input == *"-"* ]]; then
                    seconds=$((parts[0] * 86400 + parts[1] * 3600))
                else
                    seconds=$((parts[0] * 60 + parts[1]))
                fi
                ;;
            3) # Formats: hours:minutes:seconds OR days-hours:minutes
                if [[ $input == *"-"* ]]; then
                    seconds=$((parts[0] * 86400 + parts[1] * 3600 + parts[2] * 60))
                else
                    seconds=$((parts[0] * 3600 + parts[1] * 60 + parts[2]))
                fi
                ;;
            4) # Format: days-hours:minutes:seconds
                seconds=$((parts[0] * 86400 + parts[1] * 3600 + parts[2] * 60 + parts[3]))
                ;;
            *)
                echo "Unsupported format"
                exit 1
                ;;
        esac
        echo "$seconds"
}
should_restart() {
  now=$(date +%s)
  elapsed=$(( now - restart_start_sec ))
  remaining=$(( restart_limit_sec - elapsed ))
  echo -ne "Time before restart: $remaining (sec)\r" >&2
  (( remaining < SUBMITTER_MARGIN_SEC ))
}
###############################################################################
# Log settings
###############################################################################
info "Start on $(hostname); JOBID=${SLURM_JOB_ID:-unknown}"
info "Submitter settings:"
info "CKPT_PATH=${CKPT_PATH}"
info "MAX_ITER=${MAX_ITER}"
info "RUNNER_SCRIPT=${RUNNER_SCRIPT}"
info "SUBMITTER_MARGIN_SEC=${SUBMITTER_MARGIN_SEC}"
info "SUBMITTER_TIME_LIMIT=${SUBMITTER_TIME_LIMIT}"
###############################################################################
# Restore state if present
###############################################################################
state_file="submitter_state.txt"
submitter_cycle=0
if [[ -f "${state_file}" ]]; then
  source "${state_file}"
  info "Loaded state: submitter_cycle=${submitter_cycle}"
else
  info "No previous state found; starting fresh."
fi
###############################################################################
# Track time for restart
###############################################################################
restart_start_sec=$(date +%s)
restart_limit_sec=$(to_seconds "${SUBMITTER_TIME_LIMIT}")
###############################################################################
# Main Loop
###############################################################################
info "Submitter cycle: ${submitter_cycle}"
module load anaconda3/2023.07 # using python3
while true; do
  current_iter=$(get_ckpt_iter)
  info "Current checkpoint iteration: ${current_iter}"
  if (( current_iter >= MAX_ITER )); then
          info "Reached MAX_ITER=${MAX_ITER}; running final analytics."
    # List of analytics scripts to check
    analytics_scripts=(
      "aggregate_usage.py"
      "plot_progress.py"
    )
          # Run resource usage aggregation
          if command -v python3 >/dev/null 2>&1; then
      # Loop over each script and run only if present
      for script in "${analytics_scripts[@]}"; do
        if [[ -f "${script}" ]]; then
          info "Running analytics ${script} ..."
          python3 "${script}"
        else
          info "WARNING: ${script} not found; skipping."
        fi
      done
          else
            info "WARNING: python3 not found; skipping analytics."
          fi
          info "Workflow complete. Exiting submitter."
    exit 0

  fi
  JOBID=$(submit_worker)
  info "Submitted worker: ${JOBID}"
        # Log worker JobIDs for later aggregation
        id_log="worker_ids.txt"
        echo "${JOBID}" >> "${id_log}"
  # Wait for worker to finish
  info "Waiting on first of worker ${JOBID} to finish or self-restart."
  while squeue -j "${JOBID}" >/dev/null 2>&1; do
    sleep 10
    if should_restart; then
      info "Submitter approaching timeout -> restarting self"
      # Save state
      echo "submitter_cycle=$((submitter_cycle+1))" > "${state_file}"
      # Submit a new submitter instance
      sbatch checkpoint_submitter.sh
      exit 0

    fi
  done
  # Read worker exit code
  rc=$(sacct -j "${JOBID}" --format=ExitCode -P | awk -F: 'NR==2{print $1}')
  info "Worker finished with code ${rc}"
  case "${rc}" in
    0)
      info "Worker reports completed workflow."
      exit 0

      ;;
    99|124)
      info "Checkpoint advance -> continuing workflow."
      ;;
    *)
      info "Unexpected worker failure (code ${rc}). Stopping."
      exit "${rc}"

      ;;
  esac
  # Loop continues to submit next worker
done

Checkpoint Runner

Here is the checkpoint runner job script for use with the following.

for jobs submitted by Checkpoint Submitter above
for running any of the following, minimal working examples of checkpointed applications in BASH, Python, or R.
- For BASH, see BASH Checkpointint Example
- For Python, see Python Checkpointint Example
- For R, see R Checkpointint Example

checkpoint_runner.sh

#!/bin/bash
###############################################################################
# CHECKPOINT RUNNER: run application w/ checkpoint + controlled stop
# Use with checkpoint_submitter.sh to re-submit as needed
###############################################################################
# User settings with defaults
###############################################################################
# APP_CMD          : command to run (string) - must handle signal SIGTERM
# CHECKPOINT_EVERY : checkpoint after every this many application iterations
# CKPT_PATH        : path to checkpoint file
# LAUNCH_MODE      : run app command directly or via srun
# MAX_ITER         : stop application after this many iterations total
# MODULE_LIST      : modules loaded before running app command
# RUNNER_MARGIN_SEC: seconds before wall time to checkpoint (for timeout)
# RUNNER_TIME_LIMIT: wall time <= 24 hours (format as --time in man sbatch)
# SRUN_ARGS        : arguments to srun, if needed
###############################################################################
# Edit settings below
# Or pass values via...
#   sbatch --export=ALL,Var1=...,...,VarN=... checkpoint_submitter.sh
# Or pass values via...
#   Var1=... VarN=... sbatch checkpoint_submitter.sh
###############################################################################
APP_CMD="${APP_CMD:-python3 checkpoint_signal_iter.py}"
CHECKPOINT_EVERY="${CHECKPOINT_EVERY:-20}"
CKPT_PATH="${CKPT_PATH:-state_iter.txt}"
LAUNCH_MODE="${LAUNCH_MODE:-direct}"
MAX_ITER="${MAX_ITER:-500}"
MODULE_LIST="${MODULE_LIST:-anaconda3/2023.07}"
RUNNER_MARGIN_SEC="${RUNNER_MARGIN_SEC:-60}"
RUNNER_TIME_LIMIT="${RUNNER_TIME_LIMIT:-00:03:00}"
SRUN_ARGS="${SRUN_ARGS:--n 1}"
###############################################################################
# Slurm directives (keep RUNNER_TIME_LIMIT <= 24 hours)
###############################################################################
#SBATCH --job-name=ckpt_runner_demo
#SBATCH --output=log_runner_%j.out
#SBATCH --error=log_runner_%j.err
#SBATCH --time=00:03:00
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=512M
#SBATCH --partition=centos7
#SBATCH --qos=normal
###############################################################################
# Exit on error
###############################################################################
set -euo pipefail
###############################################################################
# Helpers
###############################################################################
shopt -s expand_aliases
alias dtstamp="date +%Y%m%d-%H:%M:%S"
info(){ echo "Info[$(dtstamp)]: $*"; }
get_ckpt_iter() {
  [[ -f "${CKPT_PATH}" ]] || { echo 0; return; }
  tr -cd '0-9' < "${CKPT_PATH}" 2>/dev/null || echo 0
}
to_seconds() {
        # Acceptable time formats include "minutes", "minutes:seconds", "hours:min- utes:seconds",
        # "days-hours", "days-hours:minutes" and "days-hours:minutes:seconds"
        local input="$1"
        # Handle days-hours format by replacing '-' with ':' to unify the delimiter
        # Then split components into an array
        IFS=':' read -r -a parts <<< "${input//-/:}"
        case ${#parts[@]} in
            1) # Format: minutes
                seconds=$((parts[0] * 60))
                ;;
            2) # Formats: minutes:seconds OR days-hours
                if [[ $input == *"-"* ]]; then
                    seconds=$((parts[0] * 86400 + parts[1] * 3600))
                else
                    seconds=$((parts[0] * 60 + parts[1]))
                fi
                ;;
            3) # Formats: hours:minutes:seconds OR days-hours:minutes
                if [[ $input == *"-"* ]]; then
                    seconds=$((parts[0] * 86400 + parts[1] * 3600 + parts[2] * 60))
                else
                    seconds=$((parts[0] * 3600 + parts[1] * 60 + parts[2]))
                fi
                ;;
            4) # Format: days-hours:minutes:seconds
                seconds=$((parts[0] * 86400 + parts[1] * 3600 + parts[2] * 60 + parts[3]))
                ;;
            *)
                echo "Unsupported format"
                exit 1
                ;;
        esac
        echo "$seconds"
}
###############################################################################
# Log settings
###############################################################################
info "Start on $(hostname); JOBID=${SLURM_JOB_ID:-unknown}"
info "Runner settings:"
info "APP_CMD=${APP_CMD}"
info "CHECKPOINT_EVERY=${CHECKPOINT_EVERY}"
info "CKPT_PATH=${CKPT_PATH}"
info "LAUNCH_MODE=${LAUNCH_MODE}"
info "MAX_ITER=${MAX_ITER}"
info "MODULE_LIST=${MODULE_LIST}"
info "RUNNER_MARGIN_SEC=${RUNNER_MARGIN_SEC}"
info "RUNNER_TIME_LIMIT=${RUNNER_TIME_LIMIT}"
info "SRUN_ARGS=${SRUN_ARGS}"
###############################################################################
# Load requested modules
###############################################################################
info "Loading modules."
module load "${MODULE_LIST}"
###############################################################################
# Timeout window
###############################################################################
TOTAL_SEC=$(to_seconds "${RUNNER_TIME_LIMIT}")
RUN_WINDOW_SEC=$(( TOTAL_SEC - RUNNER_MARGIN_SEC ))
if (( RUN_WINDOW_SEC <= 0 )); then
  info "WARNING: RUN_WINDOW_SEC <= 0; using 1s."
  RUN_WINDOW_SEC=1
fi
before_iter=$(get_ckpt_iter)
set +e
if [[ "${LAUNCH_MODE}" == "srun" ]]; then
  # timeout -> srun -> bash -lc "APP_CMD"
  # On expiry, timeout sends TERM to srun; srun forwards signals to its job step tasks.
  timeout "${RUN_WINDOW_SEC}s" srun ${SRUN_ARGS} bash -lc "${APP_CMD}"
else
  # direct mode: run in the batch shell
  timeout "${RUN_WINDOW_SEC}s" bash -lc "${APP_CMD}"
fi
rc=$?
set -e
info "Program exit code (from timeout wrapper): ${rc}"
###############################################################################
# Interpret exit-code from timeout
#   0   = completed all iterations (no further submissions necessary)
#   99  = app voluntarily checkpointed early -> submitter should continue
#   124 = timeout before walltime expiry -> checkpoint likely written -> continue
###############################################################################
if [[ ${rc} -eq 0 ]]; then
  info "Application finished all iterations."
  exit 0
elif [[ ${rc} -eq 99 ]]; then
  info "Application exited after writing checkpoint."
  exit 99
elif [[ ${rc} -eq 124 ]]; then
  after_iter=$(get_ckpt_iter)
  if (( after_iter > before_iter )); then
    info "Timeout + checkpoint advance detected (${before_iter}->${after_iter})."
    exit 124
  else
    info "Timeout but checkpoint did not advance. Marking failure."
    exit 1
  fi
else
  info "Unexpected exit code: ${rc}"
  exit "${rc}"
fi

Working Examples

Note: See TracWiki for help on using the wiki.

Download in other formats:

Plain Text