#!/bin/bash
#SBATCH --account=reformo
#SBATCH --partition=booster
#SBATCH --nodes=1
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=16
#SBATCH --time=01:00:00
#SBATCH --job-name=m2-cont-large-smoke
#SBATCH --output=/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/logs/continuous_large_%j.out
#SBATCH --error=/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/logs/continuous_large_%j.err
set -euo pipefail
PLAN=$(realpath "$1")
SC=/e/scratch/reformo/schuhmann1_moss
CODE=/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/continuous_large_talker_v1
source "$SC/code/fastgen.sh"
export HF_HOME="$SC/hfcache"
export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 TOKENIZERS_PARALLELISM=false
export PYTHONPATH="$CODE:/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/cascade_v3_timed:$SC/code/small_tts:$SC/code/vc_env_pkgs:$PYTHONPATH"
export OMP_NUM_THREADS=4 PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=$((21000 + SLURM_JOB_ID % 19000))
export M2_INVOCATION_SUFFIX="${SLURM_RESTART_COUNT:-}"
export M2_CONTINUOUS_PLAN="$PLAN"
WORLD=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["world_size"])' "$PLAN")
RANKS_PER_NODE=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["ranks_per_node"])' "$PLAN")

srun --ntasks="$WORLD" --ntasks-per-node="$RANKS_PER_NODE" --gpus-per-task=1 \
  --kill-on-bad-exit=1 --cpu-bind=none bash -c '
  cd /e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/continuous_large_talker_v1
  python3 -u continuous_train.py "$M2_CONTINUOUS_PLAN"
  RC=$?
  echo "M2_CONTINUOUS_TASK RANK=$SLURM_PROCID EXIT=$RC"
  exit $RC
'