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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing/bin/python}"
CONFIG_PATH="${CONFIG_PATH:?CONFIG_PATH is required}"
OUTPUT_DIR="${OUTPUT_DIR:?OUTPUT_DIR is required}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
MASTER_PORT="${MASTER_PORT:-29741}"
USE_SWANLAB="${USE_SWANLAB:-true}"
MAX_STUDENT_STEPS="${MAX_STUDENT_STEPS:-}"
FORCED_EXIT_STEP="${FORCED_EXIT_STEP:-}"

cd "$ROOT_DIR"
IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
NPROC_PER_NODE="${#GPU_ARRAY[@]}"
if (( NPROC_PER_NODE != 8 )); then
  echo "错误：正式 Predictor-v4 DMD 训练需要 8 张 GPU，当前为 $NPROC_PER_NODE。" >&2
  exit 1
fi
for required in "$PYTHON_BIN" "$CONFIG_PATH"; do
  if [[ ! -e "$required" ]]; then
    echo "错误：找不到 $required" >&2
    exit 1
  fi
done
mkdir -p "$OUTPUT_DIR"
TRAIN_TMPDIR="${TRAIN_TMPDIR:-/mnt/local_nvme/zoubin/cz/tmp/predictor_v4_dmd}"
mkdir -p "$TRAIN_TMPDIR"

overrides=(
  --config_override "predictor_v4_dmd.output_dir=${OUTPUT_DIR}"
  --config_override "predictor_v4_dmd.use_swanlab=${USE_SWANLAB}"
)
if [[ -n "$MAX_STUDENT_STEPS" ]]; then
  overrides+=(
    --config_override "predictor_v4_dmd.max_student_steps=${MAX_STUDENT_STEPS}"
  )
fi
if [[ -n "$FORCED_EXIT_STEP" ]]; then
  overrides+=(
    --config_override "predictor_v4_dmd.forced_exit_step=${FORCED_EXIT_STEP}"
  )
fi

echo "[$(date -u +%FT%TZ)] 启动 Predictor-v4 Wan14B DMD：$CONFIG_PATH"
echo "[$(date -u +%FT%TZ)] 输出目录：$OUTPUT_DIR"
TMPDIR="$TRAIN_TMPDIR" \
CUDA_VISIBLE_DEVICES="$GPU_IDS" \
"$PYTHON_BIN" -m torch.distributed.run \
  --standalone \
  --nproc_per_node="$NPROC_PER_NODE" \
  --master_port="$MASTER_PORT" \
  train.py \
  --config_path "$CONFIG_PATH" \
  --no_visualize \
  --disable-wandb \
  "${overrides[@]}" \
  2>&1 | tee -a "$OUTPUT_DIR/launcher.log"