#!/bin/bash
#SBATCH --job-name=mxf_rl_drgrpo
#SBATCH --partition=general
#SBATCH --qos=high
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=8
#SBATCH --mem=256G
#SBATCH --time=48:00:00
#SBATCH --no-requeue
#SBATCH --output=/workspace-vast/celeste/maxact-fast/logs/%x_%j.out

# Stage 6: Dr. GRPO RL (scripts/rl.py) — colocated single task: vLLM TP across all GPUs at
# gpu_memory_utilization=0.35, HF actor+LoRA on cuda:0. `mkdir -p logs` once before first submit.
#
#   sbatch scripts/sbatch_rl.sh                                       # RLConfig defaults
#   sbatch scripts/sbatch_rl.sh --groups-per-step 64 --sync-every 4   # extra args pass through
#
# Venv (dedicated — vllm pins are hard, see pyproject [vllm] extra). Build once with:
#   uv venv /workspace-vast/celeste/envs/mxf-vllm --python 3.12
#   uv pip install --python /workspace-vast/celeste/envs/mxf-vllm/bin/python \
#     "vllm==0.19.0" "vllm-lens==1.1.0" "transformers==4.57.1" peft wandb --torch-backend=cu128
# Pins are load-bearing: vllm-lens 1.1.0 is built against vLLM 0.19.0 — on vLLM>=0.22 the hook
# crashes then SILENTLY skips injection; cu128 matches the cluster's driver 570 (cu130 wheels
# fail at import). transformers must stay <5 (apply_chat_template API break).
# /workspace-vast/celeste/envs/vllm-lens (nla-experiments) is a known-good fallback with the same pins.

set -euo pipefail
ROOT=/workspace-vast/celeste/maxact-fast
VENV=${VENV:-/workspace-vast/celeste/envs/mxf-vllm}
source "$VENV/bin/activate"

export VLLM_ALLOW_INSECURE_SERIALIZATION=1   # pickle for apply_model(partial) weight sync
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export HF_HOME=/workspace-vast/pretrained_ckpts
export HF_TOKEN_PATH=/workspace-vast/celeste/.cache/huggingface/token
: "${WANDB_API_KEY:?set WANDB_API_KEY in your shell}"
export PYTHONUNBUFFERED=1
export PYTHONPATH=$ROOT/src:${PYTHONPATH:-}
# NEVER set CUDA_VISIBLE_DEVICES here — SLURM does it.

cd "$ROOT"
python scripts/rl.py --tp "${SLURM_GPUS_ON_NODE:-1}" "$@"