#!/bin/bash

# for rerun the task
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python

set -x

# ENV VARIABLES
export EVAL_LLM_URL=YOUR_EVAL_LLM_URL
export EVAL_LLM_KEY=YOUR_LLM_KEY
export EVAL_LLM_NAME=YOUR_LLM_NAME
export SANDBOX_FUSION_ENDPOINT=YOUR_SANDBOX_FUSION_ENDPOINT
export IDEALAB_SEARCH_KEY=YOUR_SEARCH_KEY
export JINA_API_KEY=YOUR_JINA_API
export SCRAPER_KEY=YOUR_SCRAPER_KEY


timestamp=$(date +%Y%m%d_%H%M%S)

NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"


source "./models/qwen3-30B-A3B.sh"

reward_shaping_gamma=0.995
SAVE_PATH=YOUR_MODEL_SAVE_PATH

CKPT_ARGS=(
   --hf-checkpoint YOUR_HF_MODEL_NAME
   --ref-load YOUR_HF_MODEL_NAME
   --load ${SAVE_PATH}
   --save ${SAVE_PATH}
   --save-interval 10
)

ROLLOUT_ARGS=(
   --prompt-data ./data/rl/rl_train.jsonl
   --input-key question
   --label-key label
   --rollout-shuffle

   --num-rollout 3000
   --rollout-batch-size 32
   --n-samples-per-prompt 16
   --rollout-max-response-len 40960
   --rollout-temperature 1.0
   --rollout-top-p 0.95

   --global-batch-size 128

   --rollout-function-path slime.rollout.context_rollout.context_rollout.context_generate_rollout
   --custom-rm-path slime.rollout.context_rollout.reward_utils.llm_utils.get_llm_reward
   --custom-reward-post-process-path slime.rollout.context_rollout.adv_utils.custom_reward_post_process_func

   --distributed-timeout-minutes 60

   --use-tis
   --iterResearch

   --step-reward-shaping
   --step-reward-shaping-gamma ${reward_shaping_gamma}
   --step-auxiliary-reward 0.0
   --step-reward-shaping-clip

)

EVAL_ARGS=(
   # --eval-interval 10
   --eval-prompt-data iterResearch EVAL_DATA_PATH
   --n-samples-per-eval-prompt 1
   --eval-max-response-len 40960
   --eval-temperature 0.6
   --eval-top-p 0.95
)

PERF_ARGS=(
   --tensor-model-parallel-size 4
   --sequence-parallel
   --pipeline-model-parallel-size 1
   --context-parallel-size 1
   --expert-model-parallel-size 8
   --expert-tensor-parallel-size 1

   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   # --micro-batch-size 1
   --use-dynamic-batch-size
   --max-tokens-per-gpu 20480
)

GRPO_ARGS=(
   --advantage-estimator gspo
   --use-kl-loss
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --kl-coef 0.00
   --entropy-coef 0.00
   --eps-clip 4e-4
   --eps-clip-high 3e-4
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-6
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98

   # --optimizer-cpu-offload
   # --overlap-cpu-optimizer-d2h-h2d
   # --use-precision-aware-optimizer
)

WANDB_ARGS=(
   --use-wandb
   --wandb-project YOUR_PROJECT_NAME
   --wandb-group YOUR_EXPERIMENT_NAME
   --wandb-key ${WANDB_API_KEY}
)

SGLANG_ARGS=(
   --sglang-server-concurrency 6
   --sglang-context-length 40960
   --rollout-num-gpus-per-engine 1
   --sglang-mem-fraction-static 0.8
   --sglang-max-running-requests 32
   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
)

MISC_ARGS=(
   # default dropout in megatron is 0.1
   --attention-dropout 0.0
   --hidden-dropout 0.0
   # should be good for model performance
   --accumulate-allreduce-grads-in-fp32
   --attention-softmax-in-fp32
   # need to comment this when using model with MLA
   --attention-backend flash
)

if [ $RANK -eq 0 ]; then
   mkdir -p ${SAVE_PATH}
   # launch the master node of ray in container
    export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
    export no_proxy="127.0.0.1,${MASTER_ADDR}"
    ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NPROC_PER_NODE} --port=6379 --disable-usage-stats

    # Build the runtime environment JSON with proper variable substitution
    RUNTIME_ENV_JSON="{
    \"env_vars\": {
        \"PYTHONPATH\": \"/root/Megatron-LM/\",
        \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
        \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
         \"no_proxy\": \"${no_proxy}\",
         \"MASTER_ADDR\": \"${MASTER_ADDR}\"
    }
    }"

    ray job submit --address="http://127.0.0.1:8265" \
        --runtime-env-json="${RUNTIME_ENV_JSON}" \
        -- python3 ./slime/train.py \
        --actor-num-nodes ${WORLD_SIZE} \
        --actor-num-gpus-per-node ${NPROC_PER_NODE} \
        --colocate \
        ${MODEL_ARGS[@]} \
        ${CKPT_ARGS[@]} \
        ${ROLLOUT_ARGS[@]} \
        ${OPTIMIZER_ARGS[@]} \
        ${GRPO_ARGS[@]} \
        ${DISTRIBUTED_ARGS[@]} \
        ${WANDB_ARGS[@]} \
        ${PERF_ARGS[@]} \
        ${EVAL_ARGS[@]} \
        ${SGLANG_ARGS[@]} \
        ${MISC_ARGS[@]} 2>&1 | tee ${SAVE_PATH}/log_${timestamp}.log

else
   sleep 60

   pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --block --address=${MASTER_ADDR}:6379 --num-gpus ${NPROC_PER_NODE} --disable-usage-stats
fi
