#!/usr/bin/env bash
# verl char_count recipe | Qwen3-1.7B | 直接 GRPO（无 SFT）| 单卡
# 基于 verl-recipe@6d8d5c1 的 char_count/train_grpo.sh，走本镜像已验证的 uv 链路
set -euo pipefail

ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
REPO=${VERL_DIR:-$ROOT/.cache_verl}
cd "$REPO"

# ---- MLflow ----
# Console only by default; optional MLflow via LOGGER and MLFLOW_TRACKING_URI.
LOGGER=${LOGGER:-'["console"]'}

# ---- 模型 / 数据 ----
export MODEL_PATH=${MODEL_PATH:-$ROOT/models/Qwen3-1.7B}
export TRAIN_FILE=${TRAIN_FILE:-$ROOT/data/rl/train.parquet}
export TEST_FILE=${TEST_FILE:-$ROOT/data/rl/test.parquet}

# ---- 单卡拓扑 ----
export NGPUS_PER_NODE=1
export NNODES=1
export ROLLOUT_TP=1

# ---- 批次（recipe：128 prompts × n=8 = 1024 seqs；mini 16，整除 1024）----
export TRAIN_BATCH_SIZE=128
export PPO_MINI_BATCH_SIZE=16
export PPO_MICRO_BATCH_SIZE_PER_GPU=8
export LOG_PROB_MICRO_BATCH_SIZE_PER_GPU=8
export MAX_PROMPT_LENGTH=128
export MAX_RESPONSE_LENGTH=128

# ---- 优化 ----
export ACTOR_LR=1e-6
export ENTROPY_COEFF=0
export ROLLOUT_N=8
export ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.7}

# ---- 实验 ----
export PROJECT_NAME=verl_char_count
export EXPERIMENT_NAME=qwen3-1.7b-char-count-grpo-direct-1gpu
export SAVE_FREQ=-1
export TEST_FREQ=5
export TOTAL_EPOCHS=2

# Setup is explicit; training never mutates the dependency environment.
if [ ! -f "$MODEL_PATH/config.json" ] || [ ! -f "$TRAIN_FILE" ] || [ ! -f "$TEST_FILE" ]; then
    echo "Missing model or dataset. Follow README.md setup first." >&2
    exit 1
fi
export UV_NO_SYNC=1
export PATH="$REPO/.venv/bin:$PATH"
if [[ "$LOGGER" == *mlflow* ]] && [[ -z "${MLFLOW_TRACKING_URI:-}" ]]; then
    echo "Set MLFLOW_TRACKING_URI when enabling MLflow." >&2
    exit 1
fi

exec bash "$ROOT/scripts/run_grpo_base.sh" \
    actor_rollout_ref.actor.strategy=fsdp2 \
    actor_rollout_ref.actor.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \
    '++data.apply_chat_template_kwargs={enable_thinking: false}' \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.kl_loss_coef=0.0 \
    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \
    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \
    actor_rollout_ref.rollout.enforce_eager=True \
    "reward.custom_reward_function.path=$ROOT/src/reward_tolerant.py" \
    reward.custom_reward_function.name=compute_score \
    ++trainer.val_before_train=True \
    "trainer.logger=$LOGGER" \
    "trainer.default_local_dir=$ROOT/checkpoints/char-count-qwen1.7b" \
    "$@"