#!/usr/bin/env bash
# 在联网环境下：用与离线机一致的 PyTorch 官方镜像拉齐 vLLM 依赖 wheel，并做离线 dry-run 自检。
# 离线目标：pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel 内执行 pip install -e .
#
# 资源分两类（缺一不可）：
#   1) Python wheel/sdist：pip download 到 OUT（默认 …/vllm_wheelhouse_pytorch210_cu128）
#   2) CMake FetchContent 用的 Git 源码：在宿主机 git clone 到 CMAKE_SRC_DIR（默认 …/vllm_offline_cmake_src）
#      离线编译时若不设置 VLLM_CUTLASS_SRC_DIR / VLLM_FLASH_ATTN_SRC_DIR / QUTLASS_SRC_DIR 等，CMake 仍会访问 GitHub 并失败。
#
# 版本须与当前 vLLM 树中 cmake 配置一致；升级 vLLM 后请重新运行本脚本并核对下面 *_TAG / *_REV。
#
# 仅做离线 dry-run（跳过 git 克隆、docker pull、pip download）：OFFLINE_DRY_RUN_ONLY=1 bash vllm_offline_fetch.sh
# 要求 OUT 指向的 wheelhouse 已存在且非空。

set -euo pipefail

IMAGE="${IMAGE:-pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel}"
# 默认与本仓库同级目录 vllm/（含 pyproject.toml）；也可 export VLLM_SRC=… 覆盖
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VLLM_SRC="${VLLM_SRC:-$SCRIPT_DIR/vllm}"
OUT="${OUT:-$SCRIPT_DIR/vllm_wheelhouse_pytorch210_cu128}"
CMAKE_SRC_DIR="${CMAKE_SRC_DIR:-$SCRIPT_DIR/vllm_offline_cmake_src}"
# pip 解析本地 vllm 时会走 PEP517 隔离环境，镜像内通常无 git，setuptools-scm 会报错；用伪装版本即可拉依赖
SCM_PRETEND="${SETUPTOOLS_SCM_PRETEND_VERSION:-0.0.0+wheelfetch}"
# 源码挂载不可为只读：setup 会向 vllm/_version.py 写版本；跑完可 git checkout -- vllm/_version.py
FETCH_CMAKE_SOURCES="${FETCH_CMAKE_SOURCES:-1}"
# 顶层 CMakeLists.txt 在 CUDA 下无条件 include triton_kernels.cmake / flashmla.cmake，
# 配置阶段就会 FetchContent；离线无网时必须提供本地目录，否则 cmake 仍会访问 GitHub。
FETCH_TRITON_KERNELS_SRC="${FETCH_TRITON_KERNELS_SRC:-1}"
FETCH_FLASH_MLA_SRC="${FETCH_FLASH_MLA_SRC:-1}"
OFFLINE_DRY_RUN_ONLY="${OFFLINE_DRY_RUN_ONLY:-0}"

# --- 与 vllm/CMakeLists.txt、cmake/external_projects/*.cmake 保持同步 ---
CUTLASS_TAG="${CUTLASS_TAG:-v4.2.1}"
FLASH_ATTN_REV="${FLASH_ATTN_REV:-140c00c0241bb60cc6e44e7c1be9998d4b20d8d2}"
QUTLASS_REV="${QUTLASS_REV:-830d2c4537c7396e14a02a46fbddd18b5d107c65}"
TRITON_KERNELS_TAG="${TRITON_KERNELS_TAG:-v3.6.0}"
FLASH_MLA_REV="${FLASH_MLA_REV:-692917b1cda61b93ac9ee2d846ec54e75afe87b1}"

if [[ ! -f "$VLLM_SRC/pyproject.toml" ]]; then
  echo "错误: 找不到 vLLM 根目录（需要 pyproject.toml）: $VLLM_SRC" >&2
  echo "请设置环境变量 VLLM_SRC 指向你的 vllm 源码根目录。" >&2
  exit 1
fi

if [[ "$OFFLINE_DRY_RUN_ONLY" == "1" ]]; then
  if [[ ! -d "$OUT" ]] || [[ -z "$(ls -A "$OUT" 2>/dev/null || true)" ]]; then
    echo "错误: OFFLINE_DRY_RUN_ONLY=1 需要已有非空的 wheel 目录 OUT=$OUT" >&2
    exit 1
  fi
  echo "[vllm-offline] 仅运行离线 dry-run（跳过 CMake 克隆 / docker pull / pip download）"
else
  mkdir -p "$OUT"
fi

# rev：commit SHA、分支名或 tag（与 git fetch 的 refspec 兼容）
# 使用 --progress：避免 -q 时长时间无输出，看起来像卡死
clone_shallow_at() {
  local url="$1" rev="$2" dest="$3"
  rm -rf "$dest"
  mkdir -p "$dest"
  git -C "$dest" init -q
  git -C "$dest" remote add origin "$url"
  git -C "$dest" fetch --progress --depth 1 origin "$rev"
  git -C "$dest" checkout -q FETCH_HEAD
}

if [[ "$OFFLINE_DRY_RUN_ONLY" != "1" ]] && [[ "$FETCH_CMAKE_SOURCES" == "1" ]]; then
  if ! command -v git >/dev/null 2>&1; then
    echo "错误: 需要 git 以拉取 CUTLASS / flash-attention / QuTLASS 源码（或设置 FETCH_CMAKE_SOURCES=0 并自行准备目录）。" >&2
    exit 1
  fi
  echo "[vllm-offline] 拉取 CMake 用 Git 依赖到: $CMAKE_SRC_DIR"
  echo "[vllm-offline] （CUTLASS / flash-attention 等仓库较大，接下来数分钟仅有 Git 进度输出属正常，并非卡死）"
  mkdir -p "$CMAKE_SRC_DIR"
  # cutlass：官方 tag 浅克隆最稳
  echo "[vllm-offline] 浅克隆 cutlass ($CUTLASS_TAG) …"
  rm -rf "$CMAKE_SRC_DIR/cutlass"
  git clone --progress --depth 1 --branch "$CUTLASS_TAG" "https://github.com/nvidia/cutlass.git" "$CMAKE_SRC_DIR/cutlass"
  echo "[vllm-offline] 浅克隆 flash-attention ($FLASH_ATTN_REV) …"
  clone_shallow_at "https://github.com/vllm-project/flash-attention.git" "$FLASH_ATTN_REV" "$CMAKE_SRC_DIR/flash-attention"
  # flash-attn 的 #include <cute/tensor.hpp> 依赖子仓库 csrc/cutlass；若未 init，该目录为空导致 nvcc 报错。
  # 只拉 csrc/cutlass，跳过 composable_kernel（ROCm）。
  echo "[vllm-offline] 初始化 flash-attention 子模块 csrc/cutlass（CuTe / CUTLASS 头文件）…"
  git -C "$CMAKE_SRC_DIR/flash-attention" submodule update --init --progress --depth 1 csrc/cutlass
  echo "[vllm-offline] 浅克隆 qutlass ($QUTLASS_REV) …"
  clone_shallow_at "https://github.com/IST-DASLab/qutlass.git" "$QUTLASS_REV" "$CMAKE_SRC_DIR/qutlass"
  if [[ "$FETCH_TRITON_KERNELS_SRC" == "1" ]]; then
    echo "[vllm-offline] 浅克隆 triton（仓库大，CMake 配置阶段即 FetchContent；离线必带或设 TRITON_KERNELS_SRC_DIR）…"
    rm -rf "$CMAKE_SRC_DIR/triton"
    git clone --progress --depth 1 --branch "$TRITON_KERNELS_TAG" "https://github.com/triton-lang/triton.git" "$CMAKE_SRC_DIR/triton"
  fi
  if [[ "$FETCH_FLASH_MLA_SRC" == "1" ]]; then
    echo "[vllm-offline] 浅克隆 FlashMLA（CMake 无条件 include；离线必带或设 FLASH_MLA_SRC_DIR，与是否编出 .so 无关）…"
    clone_shallow_at "https://github.com/vllm-project/FlashMLA.git" "$FLASH_MLA_REV" "$CMAKE_SRC_DIR/FlashMLA"
    # flashmla.cmake 将 ${FLASH_MLA_SRC_DIR}/csrc/cutlass/include 加入 nvcc；子模块未 init 则目录为空。
    echo "[vllm-offline] 初始化 FlashMLA 子模块 csrc/cutlass…"
    git -C "$CMAKE_SRC_DIR/FlashMLA" submodule update --init --progress --depth 1 csrc/cutlass
  fi
  echo "[vllm-offline] CMake 源码目录就绪。"
fi

if [[ "$OFFLINE_DRY_RUN_ONLY" == "1" ]]; then
  :
else
echo "[vllm-offline] 拉取镜像: $IMAGE"
docker pull "$IMAGE"

echo "[vllm-offline] 下载依赖到: $OUT"
docker run --rm \
  -e PIP_DISABLE_PIP_VERSION_CHECK=1 \
  -e PIP_BREAK_SYSTEM_PACKAGES=1 \
  -e "SETUPTOOLS_SCM_PRETEND_VERSION=$SCM_PRETEND" \
  -v "$VLLM_SRC:/src" \
  -v "$OUT:/out" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    echo "[vllm-offline] 镜像内 Python: $(python -V 2>&1)"
    echo "[vllm-offline] 注意: wheel 的 cp3xx 标签必须与离线安装时一致（同一镜像即可）。"
    echo "[vllm-offline] SETUPTOOLS_SCM_PRETEND_VERSION=$SETUPTOOLS_SCM_PRETEND_VERSION（供 pip 解析 /src 元数据，避免隔离环境无 git）"
    python -m pip install -U pip setuptools wheel
    export PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu128"
    # 构建期 + 运行期（cuda.txt 已包含 -r common.txt）
    python -m pip download -d /out \
      --extra-index-url https://download.pytorch.org/whl/cu128 \
      -r /src/requirements/build.txt \
      -r /src/requirements/cuda.txt
    # 按本地源码元数据再拉一层（在源码根目录执行，利于解析动态依赖）
    cd /src
    python -m pip download -d /out \
      --extra-index-url https://download.pytorch.org/whl/cu128 \
      .
    echo "[vllm-offline] wheel 文件数量: $(ls -1 /out | wc -l)"
  '

fi

echo "[vllm-offline] 离线 dry-run：requirements（不应再尝试访问 PyPI）"
docker run --rm \
  -e PIP_NO_INDEX=1 \
  -e PIP_BREAK_SYSTEM_PACKAGES=1 \
  -e "SETUPTOOLS_SCM_PRETEND_VERSION=$SCM_PRETEND" \
  -v "$VLLM_SRC:/src" \
  -v "$OUT:/out:ro" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    python -m pip install -U pip setuptools wheel
    python -m pip install --dry-run --no-index --find-links=/out \
      -r /src/requirements/build.txt -r /src/requirements/cuda.txt
  '

echo "[vllm-offline] 离线 dry-run：pip install -e .（不真正编译，仅解析/校验依赖闭包）"
docker run --rm \
  -e PIP_NO_INDEX=1 \
  -e PIP_BREAK_SYSTEM_PACKAGES=1 \
  -e "SETUPTOOLS_SCM_PRETEND_VERSION=$SCM_PRETEND" \
  -v "$VLLM_SRC:/src" \
  -v "$OUT:/out:ro" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    python -m pip install -U pip setuptools wheel
    python -m pip install --dry-run --no-index --find-links=/out -e /src
  '

echo "[vllm-offline] 离线 dry-run：先装 build.txt 再 -e .（--no-build-isolation，更贴近真实编译安装的前半段）"
docker run --rm \
  -e PIP_NO_INDEX=1 \
  -e PIP_BREAK_SYSTEM_PACKAGES=1 \
  -e "SETUPTOOLS_SCM_PRETEND_VERSION=$SCM_PRETEND" \
  -v "$VLLM_SRC:/src" \
  -v "$OUT:/out:ro" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    python -m pip install -U pip setuptools wheel
    python -m pip install --no-index --find-links=/out -r /src/requirements/build.txt
    python -m pip install --dry-run --no-build-isolation --no-index --find-links=/out -e /src
  '

echo ""
echo "完成。"
echo "  wheel 目录:        $OUT"
echo "  CMake Git 源码:    $CMAKE_SRC_DIR  (若 FETCH_CMAKE_SOURCES=1)"
echo ""
echo "拷贝到离线机时建议带上："
echo "  - 完整 vLLM 源码目录（保留 .git 以便 setuptools-scm 取版本；否则安装前 export SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0+offline）"
echo "  - 整个 $OUT"
echo "  - 整个 $CMAKE_SRC_DIR"
echo ""
echo "离线安装示例（在目标镜像内，先指向本地源码，避免 CMake 访问外网）："
echo "  export VLLM_CUTLASS_SRC_DIR=$CMAKE_SRC_DIR/cutlass"
echo "  export VLLM_FLASH_ATTN_SRC_DIR=$CMAKE_SRC_DIR/flash-attention"
echo "  export QUTLASS_SRC_DIR=$CMAKE_SRC_DIR/qutlass"
if [[ "$FETCH_TRITON_KERNELS_SRC" == "1" ]]; then
  echo "  export TRITON_KERNELS_SRC_DIR=$CMAKE_SRC_DIR/triton/python/triton_kernels/triton_kernels"
fi
if [[ "$FETCH_FLASH_MLA_SRC" == "1" ]]; then
  echo "  export FLASH_MLA_SRC_DIR=$CMAKE_SRC_DIR/FlashMLA"
fi
echo "  export PIP_NO_INDEX=1"
echo "  export PIP_FIND_LINKS=$OUT"
echo "  export PIP_BREAK_SYSTEM_PACKAGES=1   # 该 PyTorch 镜像是 PEP 668 受管环境，pip install 需要此变量或 --break-system-packages"
echo "  export TORCH_CUDA_ARCH_LIST=9.0   # 按你的 GPU 修改，例如 8.9 或 9.0+PTX"
echo "  # 目标镜像常无 git：apt-get install -y git   或   export SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0+offline"
echo "  python -m pip install -U pip setuptools wheel"
echo "  python -m pip install --no-index --find-links=$OUT -r $VLLM_SRC/requirements/build.txt"
echo "  python -m pip install --no-build-isolation --no-index --find-links=$OUT -e $VLLM_SRC"
echo ""
echo "说明：镜像里已含 torch 2.10 + CUDA 12.8 时，一般不必再从 wheelhouse 重装 torch；"
echo "若 pip 仍尝试联网，请确认 PIP_NO_INDEX=1 且 PIP_FIND_LINKS 指向离线包目录。"
