slurm_train_deeph_soc_243_gpu.sh

#!/bin/bash
#SBATCH --account=coreai_devtech_all
#SBATCH --job-name=deeph_soc_243_gpu
#SBATCH --partition=interactive
#SBATCH --time=04:00:00
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=96
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH --signal=B:USR1@600
#SBATCH --output=/home/jinliangl/home2/songby/20260514/ab-stacking/train_slurm_logs/%j.gpu.out
#SBATCH --error=/home/jinliangl/home2/songby/20260514/ab-stacking/train_slurm_logs/%j.gpu.err

set -euo pipefail

VENV=/home/jinliangl/home2/songby/deepx-gpu/deeph
DATASET_DIR=/home/jinliangl/home2/songby/20260514/ab-stacking
BASE_CONFIG=${1:-${DATASET_DIR}/train_soc_243_gpu.toml}
CONFIG=${DATASET_DIR}/train_soc_243_gpu.autoresume.toml
SELF_SCRIPT=$(readlink -f "$0")
AUTO_RESUBMIT=${AUTO_RESUBMIT:-1}
RESUBMITTED_JOB_ID=""
RESUBMIT_SENT=0

THREADS=${SLURM_CPUS_PER_TASK:-96}

export OMP_NUM_THREADS="${THREADS}"
export MKL_NUM_THREADS="${THREADS}"
export OPENBLAS_NUM_THREADS="${THREADS}"
export NUMEXPR_NUM_THREADS="${THREADS}"
export TF_NUM_INTRAOP_THREADS="${THREADS}"
export TF_NUM_INTEROP_THREADS=1
export XLA_PYTHON_CLIENT_PREALLOCATE=false
export JAX_PLATFORMS=cuda,cpu

mkdir -p "${DATASET_DIR}/train_slurm_logs"

resubmit_next() {
    if [[ "${AUTO_RESUBMIT}" != "1" ]]; then
        echo "AUTO_RESUBMIT=${AUTO_RESUBMIT}; not submitting a follow-up job."
        return 0
    fi
    if [[ "${RESUBMIT_SENT}" == "1" ]]; then
        return 0
    fi
    RESUBMIT_SENT=1

    local dependency=()
    if [[ -n "${SLURM_JOB_ID:-}" ]]; then
        dependency=(--dependency="afterany:${SLURM_JOB_ID}")
    fi

    echo "Submitting follow-up training job at $(date)"
    local submit_output
    submit_output=$(sbatch "${dependency[@]}" "${SELF_SCRIPT}" "${BASE_CONFIG}")
    echo "${submit_output}"
    RESUBMITTED_JOB_ID=$(awk '{print $4}' <<<"${submit_output}")
}

handle_time_signal() {
    echo "Received pre-timeout signal at $(date)."
    resubmit_next
}

trap handle_time_signal USR1
trap handle_time_signal TERM

if [[ ! -f "${VENV}/bin/activate" ]]; then
    echo "ERROR: missing GPU venv activate script: ${VENV}/bin/activate" >&2
    exit 1
fi
if [[ ! -x "${VENV}/bin/deeph-train" ]]; then
    echo "ERROR: missing deeph-train: ${VENV}/bin/deeph-train" >&2
    exit 1
fi
if [[ ! -f "${BASE_CONFIG}" ]]; then
    echo "ERROR: missing base training config: ${BASE_CONFIG}" >&2
    exit 1
fi

source "${VENV}/bin/activate"
cd "${DATASET_DIR}"

prepare_output=$("${VENV}/bin/python" "${DATASET_DIR}/prepare_deeph_resume_config.py" \
    --base-config "${BASE_CONFIG}" \
    --generated-config "${CONFIG}")
prepare_status=${prepare_output%% *}
prepare_path=${prepare_output#* }

if [[ "${prepare_status}" == "COMPLETE" ]]; then
    echo "Latest DeepH output already completed: ${prepare_path}"
    echo "No follow-up job submitted."
    exit 0
fi
if [[ "${prepare_status}" != "CONFIG" ]]; then
    echo "ERROR: unexpected resume helper output: ${prepare_output}" >&2
    exit 1
fi
CONFIG="${prepare_path}"

echo "Job started: $(date)"
echo "Job id: ${SLURM_JOB_ID:-manual}"
echo "Node list: ${SLURM_JOB_NODELIST:-unknown}"
echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-unset}"
echo "Dataset dir: ${DATASET_DIR}"
echo "Base config: ${BASE_CONFIG}"
echo "Config: ${CONFIG}"
echo "GPU venv: ${VENV}"
echo "CPU threads: ${THREADS}"
echo "deeph-train: $(command -v deeph-train)"

nvidia-smi || true

python - <<'PY'
import jax

devices = jax.devices()
print("jax backend:", jax.default_backend())
print("jax devices:", devices)
gpu_devices = [d for d in devices if d.platform == "gpu"]
if not gpu_devices:
    raise SystemExit("ERROR: JAX did not detect any GPU devices")
print("jax gpu count:", len(gpu_devices))
PY

set +e
srun --ntasks=1 --cpus-per-task="${THREADS}" --cpu-bind=cores deeph-train "${CONFIG}" &
TRAIN_PID=$!
while true; do
    wait "${TRAIN_PID}"
    TRAIN_STATUS=$?
    if kill -0 "${TRAIN_PID}" 2>/dev/null; then
        continue
    fi
    break
done
set -e

if [[ "${TRAIN_STATUS}" -eq 0 && -n "${RESUBMITTED_JOB_ID}" ]]; then
    echo "Training finished normally after a follow-up job was queued; canceling ${RESUBMITTED_JOB_ID}."
    scancel "${RESUBMITTED_JOB_ID}" || true
fi

if [[ "${TRAIN_STATUS}" -ne 0 ]]; then
    echo "deeph-train exited with status ${TRAIN_STATUS}."
    exit "${TRAIN_STATUS}"
fi

echo "Job finished: $(date)"