#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib.sh
source "${SCRIPT_DIR}/lib.sh"

SUBMISSION_ID="${SUBMISSION_ID:-mvp_ppo_$(timestamp)_$RANDOM}"
# 为了让“输出目录 = submission id”，默认把 JOB_TAG 也设成 SUBMISSION_ID（可手动覆盖）。
JOB_TAG="${JOB_TAG:-${SUBMISSION_ID}}"
JOB_DIR="/mnt/shared/jobs/${SUBMISSION_ID}"

MODEL_ID="Qwen/Qwen2.5-0.5B-Instruct"
TRAIN_FILE="/mnt/shared/datasets/gsm8k/train.parquet"
VAL_FILE="/mnt/shared/datasets/gsm8k/test.parquet"

echo "[head] create job dir: ${JOB_DIR}"
dexec "${HEAD_CONTAINER}" bash -lc "mkdir -p ${JOB_DIR}/logs ${JOB_DIR}/checkpoints ${JOB_DIR}/config"

SUBMIT_CMD="python3 -m verl.trainer.main_ppo \
data.train_files=${TRAIN_FILE} \
data.val_files=${VAL_FILE} \
data.train_batch_size=256 \
data.max_prompt_length=512 \
data.max_response_length=512 \
actor_rollout_ref.model.path=${MODEL_ID} \
actor_rollout_ref.actor.optim.lr=1e-6 \
actor_rollout_ref.actor.ppo_mini_batch_size=64 \
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
actor_rollout_ref.rollout.name=sglang \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
critic.optim.lr=1e-5 \
critic.model.path=${MODEL_ID} \
critic.ppo_micro_batch_size_per_gpu=4 \
algorithm.kl_ctrl.kl_coef=0.001 \
trainer.logger=console \
trainer.val_before_train=False \
trainer.n_gpus_per_node=4 \
trainer.nnodes=2 \
trainer.save_freq=10 \
trainer.test_freq=29 \
trainer.total_epochs=1 \
trainer.total_training_steps=29 \
trainer.resume_mode=disable \
trainer.default_local_dir=${JOB_DIR}/checkpoints \
+ray_kwargs.ray_init.address=auto \
hydra.run.dir=${JOB_DIR}/logs/hydra"

printf "%s\n" "${SUBMIT_CMD}" | dexec "${HEAD_CONTAINER}" bash -lc "cat > ${JOB_DIR}/config/submit_cmd.txt"

echo "[head] submit PPO via ray job submit (force driver on worker via entrypoint resources)"

SUBMIT_OUT="$(dexec "${HEAD_CONTAINER}" bash -lc "export HF_HOME=/mnt/shared/hf HUGGINGFACE_HUB_CACHE=/mnt/shared/hf/hub TRANSFORMERS_CACHE=/mnt/shared/hf/transformers HF_ENDPOINT=https://hf-mirror.com PYTHONUNBUFFERED=1; ray job submit --address=http://127.0.0.1:8265 --submission-id='${SUBMISSION_ID}' --entrypoint-num-cpus=1 --entrypoint-resources='{\"worker_node\": 1}' --runtime-env-json='{\"env_vars\":{\"HF_HOME\":\"/mnt/shared/hf\",\"HUGGINGFACE_HUB_CACHE\":\"/mnt/shared/hf/hub\",\"TRANSFORMERS_CACHE\":\"/mnt/shared/hf/transformers\",\"HF_ENDPOINT\":\"https://hf-mirror.com\",\"PYTHONUNBUFFERED\":\"1\"}}' --no-wait -- ${SUBMIT_CMD}")"

printf "%s\n" "${SUBMIT_OUT}"
printf "%s\n" "${SUBMIT_OUT}" | dexec "${HEAD_CONTAINER}" bash -lc "cat > ${JOB_DIR}/logs/ray_job_submit.out"

PARSED_SUBMISSION_ID="$(printf "%s\n" "${SUBMIT_OUT}" | sed -r 's/\x1b\\[[0-9;]*m//g' | grep -Eo "raysubmit_[A-Za-z0-9_-]+" | head -n 1 || true)"
if [[ -n "${PARSED_SUBMISSION_ID}" && "${PARSED_SUBMISSION_ID}" != "${SUBMISSION_ID}" ]]; then
  echo "WARN: submission id mismatch: expected=${SUBMISSION_ID} parsed=${PARSED_SUBMISSION_ID}" >&2
fi

echo "${SUBMISSION_ID}" | dexec "${HEAD_CONTAINER}" bash -lc "cat > ${JOB_DIR}/config/ray_submission_id.txt"
echo "ray submission id: ${SUBMISSION_ID}"

echo "submitted. track via Ray dashboard: http://<host>:8265 (driver should be scheduled on a worker due to entrypoint resources)"
echo "job dir: ${JOB_DIR}"
