#!/usr/bin/env bash
# Cluster launcher for one V4.4 ablation arm: builds (or verifies) the arm's frozen
# mixture, then trains, auto-resuming from last.pt when resubmitted.
# Usage: cluster_v44.sh <arm> [gpu]; without [gpu] the scheduler's CUDA_VISIBLE_DEVICES is kept.
set -euo pipefail
cd "$(dirname "$0")/../.."
arm="${1:?usage: cluster_v44.sh <s1|s2|s3|chat0|chat10|crop60> [gpu]}"
config="configs/train-v44-${arm}.toml"
[[ -f "$config" ]] || { echo "No config for arm $arm" >&2; exit 2; }
if [[ -n "${2:-}" ]]; then export CUDA_VISIBLE_DEVICES="$2"; fi
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 PYTHONUNBUFFERED=1
# The image exports PYTHONPATH to its own Python 3.10 tree, which shadows
# the training venv's triton (cp312) and breaks torch.compile.
unset PYTHONPATH
# Per-arm compile caches so six concurrent compilations never share a directory.
export TORCHINDUCTOR_CACHE_DIR="/tmp/vimeml-v44-${arm}/torchinductor-cache"
export TRITON_CACHE_DIR="/tmp/vimeml-v44-${arm}/triton-cache"
mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR"
output_dir=$(grep '^output_dir' "$config" | head -1 | sed 's/.*= *"\(.*\)"/\1/')
# W&B never uses the machine's login: drop every inherited WANDB_* variable (key,
# entity, service socket), point config/cache/data/netrc into this arm's own
# directory, then read only the private credential file.
while read -r name; do unset "$name"; done < <(compgen -e | grep '^WANDB_' || true)
tracking="$PWD/artifacts/tracking/$(basename "$output_dir")"
mkdir -p "$tracking"/{config,cache,data}
export WANDB_DIR="$tracking" WANDB_CONFIG_DIR="$tracking/config" WANDB_CACHE_DIR="$tracking/cache"
export WANDB_DATA_DIR="$tracking/data" NETRC="$tracking/netrc"
credentials="${VIMEML_WANDB_ENV:-$HOME/.vimeml-wandb.env}"
if [[ -f "$credentials" ]]; then
  if [[ -n "$(find "$credentials" -perm /077)" ]]; then
    echo "Refusing $credentials: readable by others; chmod 600 it" >&2; exit 2
  fi
  set -a; source "$credentials"; set +a
else
  echo "No W&B credential file at $credentials; this arm logs offline" >&2
fi
python_bin="${VIMEML_PYTHON:-python3}"
"$python_bin" -X utf8 -u scripts/corpus/prepare_v44.py --arm "$arm"
resume=()
if [[ -f "$output_dir/last.pt" ]]; then
  resume=(--resume)
fi
exec "$python_bin" -X utf8 -u scripts/training/train_v44.py --config "$config" "${resume[@]}"