diff --git a/.github/scripts/verify_cuda_wheel_artifacts.py b/.github/scripts/verify_cuda_wheel_artifacts.py new file mode 100755 index 0000000000..0c7c8ad046 --- /dev/null +++ b/.github/scripts/verify_cuda_wheel_artifacts.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Verify native artifacts required by a CUDA wheel.""" + +from __future__ import annotations + +import argparse +import stat +import sys +import zipfile +from pathlib import Path + +import tomllib + + +def supported_python_minors(pyproject: Path) -> list[int]: + with pyproject.open("rb") as file: + spec = tomllib.load(file)["project"]["requires-python"] + bounds = spec.removeprefix(">=3.").split(",<3.") + if len(bounds) != 2 or not all(bound.isdigit() for bound in bounds): + raise ValueError(f"Cannot parse requires-python from {pyproject}") + lower, upper = map(int, bounds) + return list(range(lower, upper)) + + +def verify(wheel: Path, pyproject: Path) -> None: + with zipfile.ZipFile(wheel) as archive: + entries = {entry.filename: entry for entry in archive.infolist()} + + missing: list[str] = [] + rust_binary = "aphrodite/aphrodite-rs" + if rust_binary not in entries: + missing.append(rust_binary) + else: + mode = entries[rust_binary].external_attr >> 16 + if not mode & stat.S_IXUSR: + missing.append(f"{rust_binary} (not executable)") + + if not any( + name.startswith("aphrodite/_rust_tool_parser") + and name.endswith(".so") + and "/" not in name.removeprefix("aphrodite/") + for name in entries + ): + missing.append("aphrodite/_rust_tool_parser*.so") + + for minor in supported_python_minors(pyproject): + prefix = f"aphrodite/third_party/deep_gemm/_C.cpython-3{minor}-" + if not any(name.startswith(prefix) and name.endswith(".so") for name in entries): + missing.append(f"{prefix}*.so") + + if missing: + formatted = "\n".join(f" - {item}" for item in missing) + raise RuntimeError(f"{wheel} is missing required native artifacts:\n{formatted}") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("wheel", type=Path) + parser.add_argument("--pyproject", type=Path, default=Path("pyproject.toml")) + args = parser.parse_args() + verify(args.wheel, args.pyproject) + print(f"Verified CUDA wheel native artifacts: {args.wheel}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/Dockerfile b/docker/Dockerfile index 5f60f9c3cb..15a271c60d 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -227,6 +227,72 @@ ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} #################### BUILD BASE IMAGE #################### +#################### RUST BUILD IMAGE #################### +# Build the Rust frontend separately so the CUDA wheel build does not need a +# Rust toolchain and can reuse Cargo artifacts across wheel commits. +FROM scratch AS rust-input +COPY requirements/build/rust.txt /rust-input/aphrodite/requirements/build/rust.txt +COPY rust/Cargo.lock rust/Cargo.toml /rust-input/aphrodite/rust/ +COPY rust/proto /rust-input/aphrodite/rust/proto +COPY rust/src /rust-input/aphrodite/rust/src +COPY tools/build_rust.py /rust-input/aphrodite/tools/build_rust.py + +FROM base AS rust-build +ARG USE_SCCACHE +ARG SCCACHE_ENDPOINT + +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 +ENV CARGO_HOME=/root/.cargo +ENV RUSTUP_HOME=/root/.rustup +ENV PATH=${CARGO_HOME}/bin:${PATH} +ENV SCCACHE_SERVER_PORT=4227 + +WORKDIR /workspace + +COPY --from=rust-input \ + /rust-input/aphrodite/requirements/build/rust.txt \ + requirements/build/rust.txt + +RUN apt-get update -y \ + && apt-get install -y --no-install-recommends \ + make \ + unzip \ + protobuf-compiler \ + libprotobuf-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal --default-toolchain stable \ + && rustc --version \ + && cargo --version + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install \ + --python /opt/venv/bin/python3 \ + -r requirements/build/rust.txt + +COPY --from=rust-input /rust-input/aphrodite/rust rust +COPY --from=rust-input \ + /rust-input/aphrodite/tools/build_rust.py \ + tools/build_rust.py + +RUN --mount=type=cache,id=aphrodite-cuda-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=aphrodite-cuda-cargo-git,target=/root/.cargo/git,sharing=locked \ + --mount=type=cache,id=aphrodite-cuda-cargo-target,target=/workspace/rust/target,sharing=locked \ + --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ + if [ "$USE_SCCACHE" = "1" ]; then \ + if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \ + export RUSTC_WRAPPER=sccache; \ + sccache --show-stats; \ + fi \ + && python3 tools/build_rust.py --release \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + sccache --show-stats; \ + fi +#################### RUST BUILD IMAGE #################### + #################### EXTENSIONS BUILD IMAGE #################### # Build DeepEP wheels FROM base AS extensions-build @@ -301,8 +367,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \ WORKDIR /workspace +COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/ +ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv +RUN --mount=type=cache,target=/root/.cache/uv \ + tools/setup_deepgemm_pythons.sh > /tmp/deepgemm-pythons.txt + COPY . . +# setup.py packages these separately built Rust artifacts as-is. +RUN rm -f aphrodite/aphrodite-rs aphrodite/_rust_*.so +COPY --from=rust-build /workspace/aphrodite/aphrodite-rs aphrodite/aphrodite-rs +COPY --from=rust-build /workspace/aphrodite/_rust_*.so aphrodite/ + ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi @@ -318,6 +394,7 @@ ENV APHRODITE_VERSION_OVERRIDE=${APHRODITE_VERSION_OVERRIDE} ENV MAX_JOBS=${max_jobs} ENV NVCC_THREADS=${nvcc_threads} ENV CMAKE_BUILD_TYPE=Release +ENV APHRODITE_USE_PRECOMPILED_RUST=1 # Use existing torch for nightly builds RUN --mount=type=cache,target=/root/.cache/uv \ @@ -336,11 +413,15 @@ RUN --mount=type=cache,id=aphrodite-wheel-uv,target=/root/.cache/uv \ export RUSTC_WRAPPER=sccache; \ sccache --show-stats; \ fi \ + && export DEEPGEMM_PYTHON_INTERPRETERS="$(cat /tmp/deepgemm-pythons.txt)" \ && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \ && if [ "$USE_SCCACHE" = "1" ]; then \ sccache --show-stats; \ fi +RUN python3 .github/scripts/verify_cuda_wheel_artifacts.py \ + "$(find dist -maxdepth 1 -type f -name '*.whl' -print -quit)" + # Check the size of the CUDA 13 wheel uploaded to PyPI COPY .buildkite/check-wheel-size.py check-wheel-size.py # sync the default value with .buildkite/check-wheel-size.py diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 370c77a828..55135d12cf 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -193,7 +193,7 @@ RUN --mount=type=cache,id=aphrodite-rocm-cargo-registry,target=/root/.cargo/regi export RUSTC_WRAPPER=sccache; \ sccache --show-stats; \ fi \ - && python3 tools/build_rust.py \ + && python3 tools/build_rust.py --release \ && if [ "$USE_SCCACHE" = "1" ]; then \ sccache --show-stats; \ fi diff --git a/tools/setup_deepgemm_pythons.sh b/tools/setup_deepgemm_pythons.sh new file mode 100755 index 0000000000..6ff91a6a3c --- /dev/null +++ b/tools/setup_deepgemm_pythons.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Provision one bare Python environment per supported CPython version and +# print their interpreter paths as a colon-separated list for CMake. +# +# Usage: +# export DEEPGEMM_PYTHON_INTERPRETERS="$( +# tools/setup_deepgemm_pythons.sh +# )" +# +# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv). +set -euo pipefail + +if [ "$#" -eq 0 ]; then + pyproject="$(dirname "$0")/../pyproject.toml" + spec="$( + grep -E '^requires-python' "$pyproject" | + grep -oE '>=3\.[0-9]+,<3\.[0-9]+' + )" + lo="${spec#>=3.}" + lo="${lo%%,*}" + hi="${spec##*<3.}" + set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./') +fi + +prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}" +mkdir -p "$prefix" + +paths="" +for version in "$@"; do + venv="${prefix}/${version}" + if [ ! -x "${venv}/bin/python" ]; then + uv venv \ + --python "$version" \ + "$venv" \ + --python-preference only-managed \ + --seed \ + >/dev/null + fi + paths="${paths}:${venv}/bin/python" +done + +echo "${paths#:}"