ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT}"

export TMPDIR=/pfss/mlde/workspaces/mlde_wsp_MGPATH/phuc/tmp
export PIP_CACHE_DIR=/pfss/mlde/workspaces/mlde_wsp_MGPATH/phuc/.cache/pip
mkdir -p "$TMPDIR" "$PIP_CACHE_DIR"

CACHE_ROOT="${ROOT}/.cache"
export HF_HOME="${CACHE_ROOT}/huggingface"
export HUGGINGFACE_HUB_CACHE="${HF_HOME}/hub"
export HF_DATASETS_CACHE="${HF_HOME}/datasets"
export TRANSFORMERS_CACHE="${HF_HOME}/hub"
export VLLM_CACHE_ROOT="${CACHE_ROOT}/vllm"
export FLASHINFER_CACHE_DIR="${CACHE_ROOT}/flashinfer"
export FLASHINFER_JIT_DIR="${CACHE_ROOT}/flashinfer/jit"
export TRITON_CACHE_DIR="${CACHE_ROOT}/triton"
export TORCHINDUCTOR_CACHE_DIR="${CACHE_ROOT}/torchinductor"
export XDG_CACHE_HOME="${CACHE_ROOT}"
mkdir -p "${HUGGINGFACE_HUB_CACHE}" "${HF_DATASETS_CACHE}" "${VLLM_CACHE_ROOT}" \
  "${FLASHINFER_JIT_DIR}" "${TRITON_CACHE_DIR}" "${TORCHINDUCTOR_CACHE_DIR}"

# DeepSpeed imports torch.utils.cpp_extension and requires CUDA_HOME/nvcc.
# Same pattern as rubric_med: conda cuda-nvcc 12.8 + CUDA_HOME=$CONDA_PREFIX.
if [ -z "${CUDA_HOME:-}" ] && [ -n "${CONDA_PREFIX:-}" ] && [ -x "${CONDA_PREFIX}/bin/nvcc" ]; then
  export CUDA_HOME="${CONDA_PREFIX}"
  export PATH="${CUDA_HOME}/bin:${PATH}"
fi

# source ~/.bashrc
# conda activate pre_rlvr
# # cd verl/
# # pip install -e .
# pip install lmdeploy==0.15.0
# pip install deepspeed==0.16.4
# # Official Dao-AILab linux_x86_64 wheels need GLIBC 2.32 (Ubuntu 22.04).
# # This host is Ubuntu 20.04 / GLIBC 2.31, so use Astral's manylinux_2_24 build.
# pip install --no-deps \
#   "https://wheels.astral.sh/artifacts/d4ffd81f93ca34e2e3f8b93499cd62e26faed24250657e7efcc3afd5a9b479df/flash_attn-2.8.3+cu.12.8.torch.2.8-cp311-cp311-manylinux_2_24_x86_64.whl"


deepspeed --num_gpus=4 src/run_dlms.py \
    --model_path models/dLMs-0.999zoo_op2-20+0.001teacher_op2 \
    --dataset context \
    --learning_rate 1e-6 \
    --model_name SDAR-100M-0.999zoo_op2-20+0.001teacher_op2-process \
    --num_iterations 200 \
    --mini_batch_size 1536\
    --clip_ratio_high 0.28 \
    --per_device_batch_size 64 \
    --use_step_process_reward \
    --num_samples 1024 \
    --num_generations_per_sample 6 \

deepspeed --num_gpus=4 src/run_ar.py \
    --model_path models/AR-0.999zoo_op2-20+0.001teacher_op2 \
    --dataset context\
    --learning_rate 1e-6 \
    --model_name Qwen3-100M-AR-0.999zoo_op2-20+0.001teacher_op2-process\
    --num_iterations 200 \
    --mini_batch_size 1536\
    --clip_ratio_high 0.28 \
    --per_device_batch_size 64 \
    --use_step_process_reward \
    --num_samples 1024 \
    --num_generations_per_sample 6 \