#!/bin/bash
# The v2 society campaign (prereg tasks/prereg-llm-society-v2.md §9): one (seed, arm) per array
# element on one L40S each, 16 elements = 4 seeds x 4 arms. Each element is self-contained: founders
# are trained inline and cached per seed (the four arm-elements of a seed share them via the
# filesystem; the first to arrive trains, the others wait on the adapter_config.json check), the loop
# checkpoints every generation and resumes, so a killed element is re-queued with the same index and
# picks up where it stopped. ~6 h per element at k_inherit=300 (prereg §9); 8 h walltime.
#   submit:  qsub hpc/llm_society_v2.pbs      status: qstat -u $USER -t
#   index -> seed = 1 + (i-1) // 4, arm = (full no_grounding no_sex no_diversity)[(i-1) % 4]
#PBS -l select=1:ncpus=8:mem=64gb:ngpus=1:gpu_type=L40S
#PBS -l walltime=08:00:00
#PBS -N lam_society_v2
#PBS -J 1-16

cd "$PBS_O_WORKDIR"
export HF_HOME="$EPHEMERAL/hf_cache"
export TOKENIZERS_PARALLELISM=false
export UV_CACHE_DIR="$EPHEMERAL/uvcache"

ARMS=(full no_grounding no_sex no_diversity)
I=$((PBS_ARRAY_INDEX - 1))
SEED=$((1 + I / 4))
ARM=${ARMS[$((I % 4))]}

source .venv/bin/activate
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
echo "seed=$SEED arm=$ARM start=$(date)"

# One config per (seed, arm): generated from the seed template so the resolved config is exact.
CFG="configs/llm/_gen/society_v2_s${SEED}_${ARM}.yaml"
mkdir -p configs/llm/_gen
python - "$SEED" "$ARM" "$CFG" <<'EOF'
import sys, yaml
seed, arm, out = int(sys.argv[1]), sys.argv[2], sys.argv[3]
cfg = yaml.safe_load(open(f"configs/llm/society_v2_s{seed}.yaml"))
cfg["arms"] = [arm]
cfg["output"] = {"dir": f"results/llm_society_v2/s{seed}_{arm}"}
yaml.safe_dump(cfg, open(out, "w"), sort_keys=False)
EOF

python -m llm.experiment "$CFG"
echo "done: $(date)"