Add requeue support for time-limited jobs

Author: nefrathenrici

src/backends.jl

@@ -310,6 +310,7 @@ function calibrate(
     verbose = false,
     hpc_kwargs,
     exeflags = "",
+    reruns = 0,
 )
     ensemble_size = EKP.get_N_ens(ekp)
     @info "Initializing calibration" n_iterations ensemble_size output_dir
@@ -331,6 +332,7 @@ function calibrate(
             hpc_kwargs = hpc_kwargs,
             verbose = verbose,
             exeflags,
+            reruns,
         )
         @info "Completed iteration $iter, updating ensemble"
         ekp = load_ekp_struct(output_dir, iter)
@@ -354,6 +356,7 @@ function run_hpc_iteration(
     hpc_kwargs,
     verbose = false,
     exeflags = "",
+    reruns = 0,
 )
     @info "Iteration $iter"
     job_ids = map(1:ensemble_size) do member
@@ -382,7 +385,7 @@ function run_hpc_iteration(
             module_load_str;
             hpc_kwargs,
             verbose,
-            reruns = 0,
+            reruns,
         )
     end
 end

src/pbs.jl

@@ -57,6 +57,8 @@ function generate_pbs_script(
     #PBS -l walltime=$walltime
     #PBS -l select=$num_nodes:ncpus=$cpus_per_node:ngpus=$gpus_per_node:mpiprocs=$ranks_per_node
 
+    $(pbs_trap_block())
+
     $module_load_str
 
     export JULIA_MPI_HAS_CUDA=true
@@ -74,6 +76,38 @@ function generate_pbs_script(
     return pbs_script, julia_script
 end
 
+"""
+    pbs_trap_block()
+
+Return the bash snippet that makes PBS jobs requeue on preemption or near-walltime signals.
+
+This is included in the generated PBS scripts and is tested for presence to avoid
+duplicating the job snippet in tests.
+"""
+function pbs_trap_block()
+    return """
+    # Self-requeue on preemption or near-walltime signals
+    # - Many PBS deployments send SIGTERM shortly before walltime or on preemption;
+    #   some may send SIGUSR1 as a warning.
+    # - We trap these signals and call `qrerun` to requeue the same job ID so it can
+    #   continue later with the same submission parameters.
+    # - Exiting with status 0 prevents the scheduler from marking the job as failed
+    #   due to the trap.
+    handle_preterminate() {
+        sig="\$1"
+        echo "[ClimaCalibrate] Received \$sig on PBS job \${PBS_JOBID:-unknown}, attempting qrerun"
+        if command -v qrerun >/dev/null 2>&1; then
+            qrerun "\${PBS_JOBID}"
+        else
+            echo "qrerun not available on this system"
+        fi
+        exit 0
+    }
+    trap 'handle_preterminate TERM' TERM
+    trap 'handle_preterminate USR1' USR1
+    """
+end
+
 """
     pbs_model_run(iter, member, output_dir, experiment_dir, model_interface, module_load_str; hpc_kwargs)
 
@@ -167,7 +201,7 @@ wait_for_jobs(
     module_load_str;
     verbose,
     hpc_kwargs,
-    reruns = 1,
+    reruns = 0,
 ) = wait_for_jobs(
     jobids,
     output_dir,

src/slurm.jl

@@ -15,7 +15,9 @@ kwargs(; kwargs...) = Dict{Symbol, Any}(kwargs...)
 
 Wait for a set of jobs to complete. If a job fails, it will be rerun up to `reruns` times.
 
-This function monitors the status of multiple jobs and handles failures by rerunning the failed jobs up to the specified number of `reruns`. It logs errors and job completion status, ensuring all jobs are completed before proceeding.
+In addition to scheduler status, a job is only considered successful when the model writes its
+`completed` checkpoint file. This makes restarts robust when jobs exit early due to time limits
+and are requeued by the scheduler.
 
 Arguments:
 - `jobids`: Vector of job IDs.
@@ -39,7 +41,7 @@ function wait_for_jobs(
     model_run_func;
     verbose,
     hpc_kwargs,
-    reruns = 1,
+    reruns = 0,
 )
     rerun_job_count = zeros(length(jobids))
     completed_jobs = Set{Int}()
@@ -68,8 +70,15 @@ function wait_for_jobs(
                         push!(completed_jobs, m)
                     end
                 elseif job_success(jobid)
-                    @info "Ensemble member $m complete"
-                    push!(completed_jobs, m)
+                    # Only mark success if the model has written its completion checkpoint
+                    if model_completed(output_dir, iter, m)
+                        @info "Ensemble member $m complete"
+                        push!(completed_jobs, m)
+                    else
+                        # The scheduler may report COMPLETED for a batch script that requeued itself.
+                        # Wait until the completion file exists.
+                        @debug "Job $jobid completed but checkpoint not found yet for member $m"
+                    end
                 end
             end
             sleep(5)
@@ -183,9 +192,26 @@ function generate_sbatch_directives(hpc_kwargs)
     @assert haskey(hpc_kwargs, :time) "Slurm kwargs must include key :time"
 
     hpc_kwargs[:time] = format_slurm_time(hpc_kwargs[:time])
+    # Provide a default pre-timeout signal to the batch script unless user overrides
+    if !haskey(hpc_kwargs, :signal)
+        # Send SIGUSR1 to the batch script 5 minutes before time limit
+        hpc_kwargs[:signal] = "B:USR1@300"
+    end
+    if !haskey(hpc_kwargs, :requeue)
+        hpc_kwargs[:requeue] = true
+    end
+
     slurm_directives = map(collect(hpc_kwargs)) do (k, v)
-        "#SBATCH --$(replace(string(k), "_" => "-"))=$(replace(string(v), "_" => "-"))"
+        key = replace(string(k), "_" => "-")
+        if v === true
+            return "#SBATCH --$key"
+        elseif v === false
+            return nothing
+        else
+            return "#SBATCH --$key=$(replace(string(v), "_" => "-"))"
+        end
     end
+    slurm_directives = filter(!isnothing, slurm_directives)
     return join(slurm_directives, "\n")
 end
 """
@@ -221,6 +247,14 @@ function generate_sbatch_script(
     #SBATCH --output=$member_log
     $slurm_directives
 
+    # Self-requeue on pre-timeout or termination signals
+    # `#SBATCH --signal=B:USR1@300` sends SIGUSR1 to the batch script 300 seconds before
+    # the job time limit (B means send to the batch script). Sites may also deliver TERM.
+    # We trap USR1/TERM and call `scontrol requeue $SLURM_JOB_ID` so the job returns to
+    # the queue and can continue later with the same submission parameters.
+    # Exiting with status 0 prevents a false failure due to the trap itself.
+    trap 'echo "[ClimaCalibrate] Pre-timeout/TERM on job $SLURM_JOB_ID, requeuing"; scontrol requeue $SLURM_JOB_ID; exit 0' USR1 TERM
+
     $module_load_str
     export CLIMACOMMS_DEVICE="$climacomms_device"
     export CLIMACOMMS_CONTEXT="MPI"
@@ -341,14 +375,37 @@ function job_status(job_id::SlurmJobID)
     ]
     running_statuses =
         ["RUNNING", "COMPLETING", "STAGED", "SUSPENDED", "STOPPED", "RESIZING"]
+    failed_statuses = [
+        "FAILED",
+        "CANCELLED",
+        "NODE_FAIL",
+        "TIMEOUT",
+        "OUT_OF_MEMORY",
+        "PREEMPTED",
+    ]
     invalid_job_err = "slurm_load_jobs error: Invalid job id specified"
     @debug job_id status exit_code stderr
 
-    status == "" && exit_code == 0 && stderr == "" && return :COMPLETED
+    if status == ""
+        # Not in squeue; fall back to sacct for a terminal state
+        try
+            acct = readchomp(`sacct -j $job_id --format=State%20 -n -X`)
+            # sacct may return multiple lines; take the first non-empty token
+            acct_state = strip(first(split(acct, '\n', keepempty = false), ""))
+            if any(str -> occursin(str, acct_state), failed_statuses)
+                return :FAILED
+            elseif occursin("COMPLETED", acct_state)
+                return :COMPLETED
+            end
+        catch
+            # Ignore sacct errors and continue to other checks
+        end
+    end
     exit_code != 0 && contains(stderr, invalid_job_err) && return :COMPLETED
 
     any(str -> contains(status, str), pending_statuses) && return :PENDING
     any(str -> contains(status, str), running_statuses) && return :RUNNING
+    any(str -> contains(status, str), failed_statuses) && return :FAILED
 
     @warn "Job ID $job_id has unknown status `$status`. Marking as completed"
     return :COMPLETED

test/pbs_unit_tests.jl

@@ -46,6 +46,27 @@ expected_pbs_contents = """
 #PBS -l walltime=01:30:00
 #PBS -l select=2:ncpus=16:ngpus=2:mpiprocs=2
 
+# Self-requeue on preemption or near-walltime signals
+# - Many PBS deployments send SIGTERM shortly before walltime or on preemption;
+#   some may send SIGUSR1 as a warning.
+# - We trap these signals and call `qrerun` to requeue the same job ID so it can
+#   continue later with the same submission parameters.
+# - Exiting with status 0 prevents the scheduler from marking the job as failed
+#   due to the trap.
+handle_preterminate() {
+    sig="\$1"
+    echo "[ClimaCalibrate] Received \$sig on PBS job \${PBS_JOBID:-unknown}, attempting qrerun"
+    if command -v qrerun >/dev/null 2>&1; then
+        qrerun "\${PBS_JOBID}"
+    else
+        echo "qrerun not available on this system"
+    fi
+    exit 0
+}
+trap 'handle_preterminate TERM' TERM
+trap 'handle_preterminate USR1' USR1
+
+
 export MODULEPATH="/glade/campaign/univ/ucit0011/ClimaModules-Derecho:\$MODULEPATH" 
 module purge
 module load climacommon
@@ -63,6 +84,57 @@ for (generated_str, test_str) in
     @test generated_str == test_str
 end
 
+# Requeue behavior test (skips automatically if PBS tools are unavailable)
+@testset "PBS requeue trap integration (walltime)" begin
+    testdir = mktempdir()
+    state_file = joinpath(testdir, "state.txt")
+    done_file = joinpath(testdir, "done.txt")
+
+    pbs_script = """
+#!/bin/bash
+#PBS -N requeue_test
+#PBS -j oe
+#PBS -A UCIT0011
+#PBS -q main
+#PBS -l walltime=00:00:10
+#PBS -l select=1:ncpus=1:ngpus=1
+
+$(CAL.pbs_trap_block())
+
+STATE_FILE="$(state_file)"
+DONE_FILE="$(done_file)"
+
+if [ ! -f "\$STATE_FILE" ]; then
+    echo "first" > "\$STATE_FILE"
+    # Exceed walltime so the scheduler sends TERM; the trap requeues
+    sleep 120
+else
+    echo "second" >> "\$STATE_FILE"
+    touch "\$DONE_FILE"
+fi
+"""
+
+    script_path, io = mktemp()
+    write(io, pbs_script)
+    close(io)
+    jobid = CAL.submit_pbs_job(script_path)
+
+    # Wait for completion (with timeout)
+    t_start = time()
+    timeout_s = 300.0
+    while !CAL.job_completed(jobid) && (time() - t_start) < timeout_s
+        sleep(2)
+    end
+    @test CAL.job_completed(jobid)
+
+    # Validate requeue and done marker
+    @test isfile(state_file)
+    @test isfile(done_file)
+    content = read(state_file, String)
+    @test occursin("first", content)
+    @test occursin("second", content)
+end
+
 # Helper function for submitting commands and checking job status
 function submit_cmd_helper(cmd)
     sbatch_filepath, io = mktemp()

test/runtests.jl

@@ -8,6 +8,7 @@ using SafeTestsets
 @safetestset "Julia backend" begin include("julia_backend.jl") end
 @safetestset "Aqua" begin include("aqua.jl") end
 @safetestset "Observation recipe" begin include("observation_recipe.jl") end
+@safetestset "PBS requeue integration" begin include("pbs_requeue_tests.jl") end
 #! format: on
 
 nothing