monitor_batch2_speed_12400093.sh

#!/usr/bin/env bash
set -euo pipefail

JOBID=${JOBID:-12400105}
OUT_ROOT=${OUT_ROOT:-/home/jinliangl/home2/songby/20260514/ab-stacking/outputs_train_soc_abstacking_243_Bi_s3p2d2f1_gpu_lmax6_batch2}
SLURM_LOG_DIR=${SLURM_LOG_DIR:-/home/jinliangl/home2/songby/20260514/ab-stacking/train_slurm_logs}
MONITOR_LOG=${MONITOR_LOG:-${SLURM_LOG_DIR}/${JOBID}.batch2_speed_monitor.log}
POLL_SECONDS=${POLL_SECONDS:-300}
KEEP_THRESHOLD=${KEEP_THRESHOLD:-45.0}
CANCEL_THRESHOLD=${CANCEL_THRESHOLD:-50.0}

log() {
    printf '[%s] %s\n' "$(date '+%F %T')" "$*" | tee -a "${MONITOR_LOG}"
}

find_deepx_log() {
    find "${OUT_ROOT}" -maxdepth 3 -type f -name deepx.log -printf '%T@ %p\n' 2>/dev/null \
        | sort -n | tail -1 | awk '{print $2}'
}

check_epoch_speed() {
    local deepx_log=$1
    python - "$deepx_log" "$KEEP_THRESHOLD" "$CANCEL_THRESHOLD" <<'PY'
import re
import sys
from pathlib import Path

path = Path(sys.argv[1])
keep = float(sys.argv[2])
cancel = float(sys.argv[3])
times = []
epochs = []
for line in path.read_text(errors="replace").splitlines():
    m = re.search(r"Epoch\s+(\d+) \| Time ([0-9.]+) s", line)
    if not m:
        continue
    epoch = int(m.group(1))
    t = float(m.group(2))
    epochs.append(epoch)
    if epoch >= 2:
        times.append(t)

if not times:
    print("WAIT no_epoch")
    raise SystemExit

last = times[-5:]
avg = sum(last) / len(last)
print(f"INFO count={len(times)} last_epoch={max(epochs)} last_avg={avg:.2f} last={','.join(f'{x:.2f}' for x in last)}")

if len(times) >= 5 and avg <= keep:
    print("KEEP fast_enough")
elif len(times) >= 5 and avg >= cancel:
    print("CANCEL too_slow")
elif len(times) >= 8 and avg > keep:
    print("CANCEL not_in_30_40s_range")
else:
    print("WAIT need_more_epochs")
PY
}

log "Monitoring batch=2 DeepH job ${JOBID}; keep <= ${KEEP_THRESHOLD}s, cancel >= ${CANCEL_THRESHOLD}s."

while true; do
    sq=$(squeue -h -j "${JOBID}" -o '%T %M %R' 2>/dev/null || true)
    if [[ -z "${sq}" ]]; then
        log "Job ${JOBID} is no longer in queue."
        exit 0
    fi

    log "squeue: ${sq}"
    state=${sq%% *}

    deepx_log=$(find_deepx_log || true)
    if [[ -n "${deepx_log}" && -f "${deepx_log}" ]]; then
        result=$(check_epoch_speed "${deepx_log}" || true)
        while IFS= read -r line; do
            [[ -n "${line}" ]] && log "${line}"
        done <<<"${result}"
        if grep -q '^KEEP ' <<<"${result}"; then
            log "Batch=2 speed is good enough; leaving job ${JOBID} running."
            exit 0
        fi
        if grep -q '^CANCEL ' <<<"${result}"; then
            log "Canceling slow batch=2 test job ${JOBID}."
            scancel "${JOBID}" || true
            exit 0
        fi
    elif [[ "${state}" == "RUNNING" ]]; then
        log "Job is running but deepx.log has not appeared yet."
        if grep -RqiE 'traceback|error|oom|out of memory' "${SLURM_LOG_DIR}/${JOBID}."* 2>/dev/null; then
            log "Detected error text in Slurm logs; canceling ${JOBID}."
            scancel "${JOBID}" || true
            exit 1
        fi
    fi

    sleep "${POLL_SECONDS}"
done