#!/bin/bash
# ==============================================================================
# submit.sh — Unified sbatch wrapper
#
# Injects --account and --partition from cluster_env.sh so that individual
# training scripts don't need to hard-code them in #SBATCH directives.
# (sbatch CLI args override #SBATCH directives in the script.)
#
# Usage:
#   bash scripts/submit.sh  <script.sh>  [extra sbatch args...]
#
# Examples:
#   # Use defaults from ~/.lmflow_cluster.conf
#   bash scripts/submit.sh chengyue_script/dvlm/training/multinode/mdm/j63/bd_size_4.sh
#
#   # Override account for this job only
#   LMFLOW_SLURM_ACCOUNT=nvr_lpr_llm bash scripts/submit.sh script.sh
#
#   # Pass extra sbatch args (e.g. change node count)
#   bash scripts/submit.sh script.sh -N 4
# ==============================================================================

set -eo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/cluster_env.sh"

if [ $# -lt 1 ]; then
    echo "Usage: bash scripts/submit.sh <script.sh> [extra sbatch args...]"
    echo ""
    echo "Current cluster config:"
    echo "  LMFLOW_SLURM_ACCOUNT   = ${LMFLOW_SLURM_ACCOUNT}"
    echo "  LMFLOW_SLURM_PARTITION = ${LMFLOW_SLURM_PARTITION}"
    echo "  LMFLOW_CONDA_ENV       = ${LMFLOW_CONDA_ENV}"
    echo "  VILA_DATASETS          = ${VILA_DATASETS}"
    echo ""
    echo "Override via ~/.lmflow_cluster.conf or env vars."
    exit 1
fi

TARGET_SCRIPT="$1"
shift

if [ ! -f "$TARGET_SCRIPT" ]; then
    echo "Error: script not found: ${TARGET_SCRIPT}" >&2
    exit 1
fi

echo "[submit] account=${LMFLOW_SLURM_ACCOUNT}  partition=${LMFLOW_SLURM_PARTITION}"
echo "[submit] sbatch ${TARGET_SCRIPT} $*"

sbatch \
    --account="${LMFLOW_SLURM_ACCOUNT}" \
    --partition="${LMFLOW_SLURM_PARTITION}" \
    "$@" \
    "${TARGET_SCRIPT}"
