#!/bin/bash

USE_MEGATRON=${USE_MEGATRON:-1}
USE_SGLANG=${USE_SGLANG:-1}

export MAX_JOBS=32

echo "0. Install uv (The fast package installer)"
pip install uv

echo "1. install inference frameworks and pytorch they need"
if [ $USE_SGLANG -eq 1 ]; then
    # --system is needed if not running inside an active virtual environment
    uv pip install --system "sglang[all]==0.5.2" --no-cache-dir && uv pip install --system torch-memory-saver --no-cache-dir
fi
uv pip install --system --no-cache-dir "vllm==0.11.0"

echo "2. install basic packages"
uv pip install --system "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \
    "numpy<2.0.0" "pyarrow>=15.0.0" pandas "tensordict>=0.8.0,<=0.10.0,!=0.9.0" torchdata \
    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \
    pytest py-spy pre-commit ruff tensorboard 

echo "pyext is lack of maintainace and cannot work with python 3.12."
echo "if you need it for prime code rewarding, please install using patched fork:"
echo "uv pip install --system git+https://github.com/ShaohonChen/PyExt.git@py311support"

uv pip install --system "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1"


echo "3. install FlashAttention and FlashInfer"
# Install flash-attn-2.8.1 (cxx11abi=False)
# uv can install directly from the file after wget
wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl && \
    uv pip install --system --no-cache-dir flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl

uv pip install --system --no-cache-dir flashinfer-python==0.3.1


if [ $USE_MEGATRON -eq 1 ]; then
    echo "4. install TransformerEngine and Megatron"
    echo "Notice that TransformerEngine installation can take very long time, please be patient"
    uv pip install --system "onnxscript==0.3.1"
    
    # Keeping no-deps here as per original script logic
    NVTE_FRAMEWORK=pytorch uv pip install --system --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@v2.6
    uv pip install --system --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.1
fi


echo "5. May need to fix opencv"
uv pip install --system opencv-python
uv pip install --system opencv-fixer && \
    python -c "from opencv_fixer import AutoFix; AutoFix()"


if [ $USE_MEGATRON -eq 1 ]; then
    echo "6. Install cudnn python package (avoid being overridden)"
    uv pip install --system nvidia-cudnn-cu12==9.10.2.21
fi

echo "Successfully installed all packages"