Build uv venv for slime

gistfile1.txt

#!/bin/bash
set -ex

###############################################################################
# Config
###############################################################################
export BASE_DIR=$(pwd)

export SGLANG_COMMIT=${SGLANG_COMMIT:-"24c91001cf99ba642be791e099d358f4dfe955f5"}
export MEGATRON_COMMIT=${MEGATRON_COMMIT:-"3714d81d418c9f1bca4594fc35f9e8289f652862"}
export CUDA_VERSION=${CUDA_VERSION:-"129"}

# Toggle apex build (needs nvcc)
export WITH_APEX=${WITH_APEX:-1}

# flash-attn prebuilt wheel (torch2.9 + py312 + cu12)
FLASH_ATTN_WHEEL_URL=${FLASH_ATTN_WHEEL_URL:-"https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"}

export CUDA_HOME=${CUDA_HOME:-"/usr/local/cuda"}
export PATH="$CUDA_HOME/bin:$PATH"
export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"

cd "$BASE_DIR"

###############################################################################
# System deps (Ubuntu/Debian)
###############################################################################
apt-get update
apt-get install -y --no-install-recommends \
  git curl ca-certificates build-essential \
  cmake ninja-build

###############################################################################
# Install uv + create venv
###############################################################################
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"

uv venv --python python3.12 .venv
source .venv/bin/activate

uv pip install -U pip setuptools wheel

###############################################################################
# Torch (cu129)
###############################################################################
uv pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 \
  --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}

###############################################################################
# Core + extra deps (你新增的都在这里)
###############################################################################
uv pip install \
  accelerate \
  blobfile \
  datasets \
  "httpx[http2]" \
  "mcp[cli]" \
  memray \
  numba \
  omegaconf \
  pillow \
  pylatexenc \
  pyyaml \
  qwen_vl_utils \
  "ray[default]" \
  ring_flash_attn \
  "sglang-router>=0.2.3" \
  tensorboard \
  transformers \
  wandb\
  poetry\
  nvidia-resiliency-ext\
  pybind11\
  nvidia-ml-py

# 保守策略（numba / torch 生态更稳）
uv pip install "numpy<2"

###############################################################################
# flash-attn (prebuilt wheel)
###############################################################################
uv pip uninstall flash-attn flash_attn
uv pip install --no-deps "${FLASH_ATTN_WHEEL_URL}"

###############################################################################
# Transformer Engine (wheel-first)
###############################################################################
uv pip install --extra-index-url https://pypi.nvidia.com \
  transformer-engine-cu12==2.10.0 transformer-engine-torch==2.10.0

###############################################################################
# Other git deps
###############################################################################
uv pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
uv pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b \
  --no-cache-dir --force-reinstall
uv pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
uv pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation

###############################################################################
# Apex (optional)
###############################################################################
NVCC_APPEND_FLAGS="--threads 4" \
  ".venv/bin/python" -m pip install --disable-pip-version-check --no-cache-dir \
  --no-build-isolation \
  --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4


###############################################################################
# sglang
###############################################################################
rm -rf sglang
git clone https://github.com/sgl-project/sglang.git 
cd sglang
git checkout "${SGLANG_COMMIT}"
uv pip install -e "python[all]"
cd "$BASE_DIR"

###############################################################################
# Megatron-LM
###############################################################################
rm -rf Megatron-LM
git clone https://github.com/NVIDIA/Megatron-LM.git --recursive
cd Megatron-LM
git checkout "${MEGATRON_COMMIT}"
uv pip install -e .
cd "$BASE_DIR"

###############################################################################
# slime
###############################################################################
cd "$BASE_DIR"
uv pip install -e .

###############################################################################
# cuDNN runtime pin
###############################################################################
uv pip install nvidia-cudnn-cu12==9.16.0.29

###############################################################################
# Apply patches
###############################################################################
cd "$BASE_DIR/sglang"
git apply "$BASE_DIR/docker/patch/v0.5.7/sglang.patch"
cd "$BASE_DIR/Megatron-LM"
git apply "$BASE_DIR/docker/patch/v0.5.7/megatron.patch"

echo "ALL DONE"