Refactor/consolidate neq engine

README.md

@@ -158,7 +158,7 @@ We are working on the documentation and tutorials for QligFEP. In the meantime,
 
 Alongside the standard windowed (equilibrium) protocol, QligFEP supports a **non-equilibrium** alchemical workflow, referred to as **NEQ²**. Rather than sampling many fixed-λ windows, NEQ² drives λ continuously from one end state to the other over many short, independent switching trajectories, and recovers ΔΔG from the Bennett Acceptance Ratio (BAR) over the resulting forward and reverse work distributions. Because the switching trajectories are independent, they parallelize trivially across a cluster.
 
-Set up a non-equilibrium calculation by passing `--neq` to `setupFEP`, and analyze the accumulated switching work with the `qligfep_neq_analyze` CLI. The non-equilibrium engine (`qdyn_neq`) is built together with the other Q binaries by `make all` in `src/q6`. See the [Tyk2 NEQ² tutorial](/tutorials/Tyk2/neq2/README.md) for an end-to-end walkthrough.
+Set up a non-equilibrium calculation by passing `--neq` to `setupFEP`, and analyze the accumulated switching work with the `qligfep_neq_analyze` CLI. Non-equilibrium switching runs as a mode of the serial `qdyn` engine (selected by a `[lambda_scaling]` section in the input), which is built together with the other Q binaries by `make all` in `src/q6`. See the [Tyk2 NEQ² tutorial](/tutorials/Tyk2/neq2/README.md) for an end-to-end walkthrough.
 
 # 📊 Benchmarking
 

src/QligFEP/CLI/parser_base.py

@@ -234,7 +234,8 @@ def parse_arguments(program: str) -> argparse.Namespace:
         action="store_true",
         help=(
             "Set up a non-equilibrium (NEQ) FEP instead of the windowed equilibrium approach. "
-            "Instead of many fixed-lambda windows, NEQ runs the `qdyn_neq` engine to drive lambda "
+            "Instead of many fixed-lambda windows, NEQ runs the serial `qdyn` engine (fast-switching "
+            "mode) to drive lambda "
             "from one endpoint to the other over a single simulation, accumulating the switching "
             "work. Free energies are obtained from BAR over the forward/reverse work distributions "
             "(see `qligfep_neq_analyze`). When set, the windowed parameters `--windows` and "

src/QligFEP/INPUTS/run.sh

@@ -69,9 +69,31 @@ MODULES
 ## define qdynp location
 QDYN
 
+# Fail on the first error, so a failed mpirun can't exit 0 and be recorded COMPLETED. Set after
+# `module`, which returns non-zero on some stacks just for unload warnings.
+set -eo pipefail
+
 starttime=$(date +%s)
 starttime_readable=$(date)
 
+# Emitted from an EXIT trap so the footer is written even when set -e aborts the job midway.
+# Keep in sync with run_neq.sh; parsed by QligFEP.IO.read_slurm_diagnostics.
+write_footer() {
+    rc=$?
+    endtime=$(date +%s)
+    runtime=$((endtime - starttime))
+    echo "#    EXPRESS LOG for jobid: $SLURM_JOB_ID"
+    echo "#    Slurm tasks: $SLURM_NTASKS"
+    echo "#    Starttime: $starttime_readable"
+    echo "#    Endtime: $(date)"
+    echo "#    Runtime: $((runtime / 3600))h:$((runtime % 3600 / 60))m:$((runtime % 60))s"
+    echo "#    Random seed: ${seed:-unknown}"
+    echo "#    Replicate Number: ${run_num:-unknown}"
+    echo "#    Working Directory: ${workdir:-unknown}"
+    echo "#    Exit status: $rc"
+}
+trap write_footer EXIT
+
 length=${#fepfiles[@]}
 length=$((length-1))
 for index in $(seq 0 $length); do
@@ -97,11 +119,12 @@ echo "slurm cpus per node: $SLURM_JOB_CPUS_PER_NODE"
 echo
 
 echo -e "\n=== CPU Model Information ==="
-lscpu | grep -E "Model name|Architecture|CPU op|Thread|Core|Socket|NUMA|CPU(s)"
+lscpu | grep -E "Model name|Architecture|CPU op|Thread|Core|Socket|NUMA|CPU(s)" || true
 echo
 
+# Diagnostics only -- must never be fatal under set -e, hence the `|| true`.
 echo -e "\n=== Available CPU List ==="
-cpu_list=$(cat /proc/self/status | grep Cpus_allowed_list | awk '{print $2}')
+cpu_list=$(grep Cpus_allowed_list /proc/self/status | awk '{print $2}') || true
 echo "CPU list: $cpu_list"
 echo
 
@@ -124,24 +147,21 @@ if [ $index -lt 1 ]; then
 fi
 #RUN_FILES
 timeout 3m QFEP < qfep.inp > qfep.out || [ $? -eq 124 ]
+
+# qdyn's serial abort still exits 0, so a completed run is confirmed by "terminated normally" in
+# its log, not by the exit code. Failing here also stops FEP N+1 from starting on a bad eq5.re.
+incomplete=""
+logs_to_check=(md*.log)  # only the first FEP equilibrates; later ones inherit eq5.re
+if [ $index -lt 1 ]; then
+    logs_to_check+=(eq*.log)
+fi
+for log in "${logs_to_check[@]}"; do
+    grep -q "terminated normally" "$log" 2>/dev/null || incomplete="$incomplete $log"
+done
+if [ -n "$incomplete" ]; then
+    echo "ERROR: FEP$((index+1)) replicate $run_num (T=$temperature, seed=$seed) is incomplete."
+    echo "ERROR: logs without a normal qdyn termination:$incomplete"
+    exit 1
+fi
 done
 #CLEANUP
-
-endtime=$(date +%s)
-endtime_readable=$(date)
-# Calculate runtime
-runtime=$((endtime - starttime))
-
-# Convert runtime to hours:minutes:seconds
-hours=$(($runtime / 3600))
-minutes=$(($runtime % 3600 / 60))
-seconds=$(($runtime % 60))
-
-echo "#    EXPRESS LOG for jobid: $SLURM_JOB_ID"
-echo "#    Slurm tasks: $SLURM_NTASKS"
-echo "#    Starttime: $starttime_readable"
-echo "#    Endtime: $endtime_readable"
-echo "#    Runtime: ${hours}h:${minutes}m:${seconds}s"
-echo "#    Random seed: $seed"
-echo "#    Replicate Number: $run_num"
-echo "#    Working Directory: $workdir"

src/QligFEP/INPUTS/run_neq.sh

@@ -10,6 +10,27 @@
 #SBATCH --array=1-TOTAL_JOBS
 #SBATCH -o slurm.run%a.%N.%j.out
 
+starttime=$(date +%s)
+starttime_readable=$(date)
+
+# Emitted from an EXIT trap so the footer is written even when set -e aborts the job midway.
+# Keep in sync with run.sh; parsed by QligFEP.IO.read_slurm_diagnostics.
+write_footer() {
+    rc=$?
+    endtime=$(date +%s)
+    runtime=$((endtime - starttime))
+    echo "#    EXPRESS LOG for jobid: $SLURM_JOB_ID"
+    echo "#    Slurm tasks: $SLURM_NTASKS"
+    echo "#    Starttime: $starttime_readable"
+    echo "#    Endtime: $(date)"
+    echo "#    Runtime: $((runtime / 3600))h:$((runtime % 3600 / 60))m:$((runtime % 60))s"
+    echo "#    Random seed: ${seed:-unknown}"
+    echo "#    Replicate Number: ${run_num:-unknown}"
+    echo "#    Working Directory: ${workdir:-unknown}"
+    echo "#    Exit status: $rc"
+}
+trap write_footer EXIT
+
 # Number of forward/reverse switching pairs run per replicate.
 neq_reps=NEQ_REPS
 
@@ -48,12 +69,17 @@ seed=${seeds[$run_num-1]}
 ## Load modules
 MODULES
 
-## define the MPI equilibration engine (qdynp) and serial switching engine (qdyn_neq)
-QDYN
+## define the parallel equilibration engine (qdynp) and serial switching engine (qdyn)
+QDYNP
 QDYN_NEQ
 
-starttime=$(date +%s)
-starttime_readable=$(date)
+# Fail on the first error, so a failed mpirun can't exit 0 and be recorded COMPLETED. Set after
+# `module`, which returns non-zero on some stacks just for unload warnings.
+set -eo pipefail
+
+# Pin one rank per core: above two ranks mpirun maps by socket, which misplaces ranks when SLURM
+# splits the allocated cores unevenly across sockets.
+mpi_map="--map-by core --bind-to core"
 
 rundir=$workdir/$temperature/$run_num
 mkdir -p $rundir
@@ -76,7 +102,7 @@ sed -i "s/FEP_VAR/$fepfile/" *.inp
 
 # 1) Equilibration eq1 -> eq5 with the MPI engine across all cores (fixed lambda)
 for i in 1 2 3 4 5; do
-    time mpirun -np $ncores --bind-to core $qdyn eq$i.inp > eq$i.log
+    time mpirun -np $ncores $mpi_map $qdynp eq$i.inp > eq$i.log
 done
 
 # 2) Endpoint equilibration chain (MPI): one continuous trajectory per endpoint,
@@ -91,16 +117,16 @@ for rep in $(seq 0 $((neq_reps - 1))); do
         if [ "$rep" -eq 0 ]; then eqsrc=relax_${s}.inp; else eqsrc=eq6_${s}.inp; fi
         sed "s|RESTARTFILE|eq6_${s}_prev.re|; s|FINALFILE|eq6_${s}_${rep}.re|" \
             "$eqsrc" > eq6_${s}_run${rep}.inp
-        time mpirun -np $ncores --bind-to core $qdyn eq6_${s}_run${rep}.inp > eq6_${s}_${rep}.log
+        time mpirun -np $ncores $mpi_map $qdynp eq6_${s}_run${rep}.inp > eq6_${s}_${rep}.log
         cp eq6_${s}_${rep}.re eq6_${s}_prev.re
     done
 done
 
-# 3) Lambda switches with the serial engine, one per core. Each switch only needs
-# its own eq6 checkpoint, so they are independent and run concurrently: mpirun
-# launches one rank per switch, --bind-to core pins each to its own core, and the
-# launcher routes each rank to its (input, log) pair. State 1 (1->0) is the forward
-# work logged as neq_1; state 0 (0->1) is the reverse work logged as neq_0.
+# 3) Lambda switches with the serial engine, one per core. Each switch only needs its own eq6
+# checkpoint, so they are independent and run concurrently: mpirun launches one rank per switch,
+# pins each to its own core, and the launcher routes each rank to its (input, log) pair.
+# State 1 (1->0) is the forward work logged as neq_1; state 0 (0->1) is the reverse work
+# logged as neq_0.
 : > switch_list.txt
 for rep in $(seq 0 $((neq_reps - 1))); do
     for s in 0 1; do
@@ -117,31 +143,47 @@ idx=\$(( OMPI_COMM_WORLD_RANK + 1 + \${BATCH_OFFSET:-0} ))
 line=\$(sed -n "\${idx}p" switch_list.txt)
 [ -z "\$line" ] && exit 0
 set -- \$line
-$qdyn_neq "\$1" > "\$2"
+$qdyn "\$1" > "\$2"
 EOF
 chmod +x neq_launch.sh
 
 nsw=$(wc -l < switch_list.txt)
+# Setup sizes ncores to hold every switch in one wave (ncores >= nsw). If the node could not grant
+# that many and --ntasks-per-node was lowered, the switches run in successive waves; a count that
+# does not divide nsw leaves a partial last wave with billed cores idle. Prefer a divisor of nsw.
+if [ "$ncores" -gt 0 ] && [ "$ncores" -lt "$nsw" ] && [ $(( nsw % ncores )) -ne 0 ]; then
+    echo "WARNING: $nsw switches on $ncores cores run in uneven waves (partial last wave, cores idle)."
+    echo "         For a single wave request $nsw cores; otherwise set --ntasks-per-node to a divisor of $nsw."
+fi
 for (( off=0; off<nsw; off+=ncores )); do
     remaining=$(( nsw - off ))
     np=$(( remaining < ncores ? remaining : ncores ))
     export BATCH_OFFSET=$off
-    time mpirun -x BATCH_OFFSET --bind-to core -np $np ./neq_launch.sh
+    # A batch that loses ranks must not abort the others; the count below decides what to do.
+    time mpirun -x BATCH_OFFSET $mpi_map -np $np ./neq_launch.sh || true
 done
 
-#CLEANUP
+# qdyn's serial abort still exits 0, so a completed run is confirmed by "terminated normally" in
+# its log, not by the exit code.
 
-endtime=$(date +%s)
-endtime_readable=$(date)
-runtime=$((endtime - starttime))
-hours=$(($runtime / 3600))
-minutes=$(($runtime % 3600 / 60))
-seconds=$(($runtime % 60))
-
-echo "#    EXPRESS LOG for jobid: $SLURM_JOB_ID"
-echo "#    Starttime: $starttime_readable"
-echo "#    Endtime: $endtime_readable"
-echo "#    Runtime: ${hours}h:${minutes}m:${seconds}s"
-echo "#    Random seed: $seed"
-echo "#    Replicate Number: $run_num"
-echo "#    Working Directory: $workdir"
+# Equilibration must finish in full -- with no eq6 checkpoint there is nothing to switch from.
+eq_bad=""
+for log in eq*.log; do
+    grep -q "terminated normally" "$log" 2>/dev/null || eq_bad="$eq_bad $log"
+done
+if [ -n "$eq_bad" ]; then
+    echo "ERROR: replicate $run_num (T=$temperature, seed=$seed): equilibration incomplete:$eq_bad"
+    exit 1
+fi
+
+# A lost switch is tolerated -- analyze_neq discards it (analyze_neq._scan_leg). A partial loss is
+# just sampling noise; only a total loss (e.g. mpirun never launched) is fatal.
+nsw_ok=$(grep -l "terminated normally" neq_*.log 2>/dev/null | wc -l || true)
+if [ "$nsw_ok" -eq 0 ]; then
+    echo "ERROR: replicate $run_num (T=$temperature, seed=$seed): no switch produced work."
+    exit 1
+elif [ "$nsw_ok" -lt "$nsw" ]; then
+    echo "WARNING: replicate $run_num: $((nsw - nsw_ok))/$nsw switches lost (tolerated)."
+fi
+
+#CLEANUP

src/QligFEP/IO.py

@@ -4,6 +4,7 @@
 import stat
 import subprocess
 from pathlib import Path
+from typing import NamedTuple, Optional
 
 import numpy as np
 import pandas as pd
@@ -14,6 +15,65 @@
 
 qfep_error_regex = re.compile(r"ERROR:")
 
+# Failure markers -> the status label reported for a run, shared by both FEP analyzers
+# (analyze_FEP and analyze_neq) so equilibrium and NEQ runs are classified identically.
+#
+# Dict order decides the label when a log carries several markers.
+RUN_FAILURE_KEYWORDS = {
+    "A request was made to bind to": "MPI_LAUNCH_FAILED",
+    "There are not enough slots available": "MPI_LAUNCH_FAILED",
+    "mpirun was unable to": "MPI_LAUNCH_FAILED",
+    "DUE TO TIME LIMIT": "TIMEOUT",
+    "CANCELLED": "CANCELLED",
+    "Out Of Memory": "OOM",
+    "abnormally": "CRASHED",
+}
+
+
+class SlurmRunInfo(NamedTuple):
+    """Per-replicate run metadata parsed from one ``slurm*.out`` footer."""
+
+    runtime: str
+    seed: Optional[str]
+    replicate: Optional[str]
+    status: str
+
+
+def read_slurm_diagnostics(slurm_out) -> SlurmRunInfo:
+    """Parse one run.sh / run_neq.sh ``slurm*.out`` footer into a ``SlurmRunInfo``.
+
+    Runtime, seed and replicate come from the standardized ``#    Runtime:`` /
+    ``#    Random seed:`` / ``#    Replicate Number:`` footer the run scripts write. Reading the
+    footer -- rather than the ``Parameters`` line or the slurm filename -- means the true
+    ``run_num`` is used, so multi-temperature array jobs (where the filename's ``%a`` differs
+    from the replicate) are labeled correctly. ``status`` is ``"SUCCESS"`` unless a known failure
+    marker (``RUN_FAILURE_KEYWORDS``) appears anywhere in the log -- the first marker in dict order
+    wins -- or, failing that, the footer reports a non-zero ``#    Exit status:``. Logs written
+    before that footer line existed have none, and fall back to keyword matching alone.
+    """
+    with open(slurm_out, errors="ignore") as handle:
+        text = handle.read()
+    runtime, seed, replicate, status = "", None, None, "SUCCESS"
+    exit_status = None
+    for line in text.splitlines():
+        if line.startswith("#    Runtime:"):
+            runtime = line.split()[-1].strip()
+        elif line.startswith("#    Random seed:"):
+            seed = line.split()[-1].strip()
+        elif line.startswith("#    Replicate Number:"):
+            replicate = line.split()[-1].strip()
+        elif line.startswith("#    Exit status:"):
+            exit_status = line.split()[-1].strip()
+    for keyword, label in RUN_FAILURE_KEYWORDS.items():
+        if keyword in text:
+            status = label
+            break
+    # A keyword names a specific cause, so it is preferred; the reported exit code is the catch-all
+    # for failures we have no marker for.
+    if status == "SUCCESS" and exit_status not in (None, "0"):
+        status = "FAILED"
+    return SlurmRunInfo(runtime, seed, replicate, status)
+
 
 class QprepError(Exception):
     """Raised when qprep encounters a general error."""