Compare commits

...

40 Commits

Author SHA1 Message Date
1184e6e0a6 fix linter 2025-08-06 10:46:09 -07:00
584fec0443 fix linter 2025-08-06 10:42:47 -07:00
8b16c9c19e fix linter 2025-08-06 10:31:45 -07:00
0223eafe54 fix linter 2025-08-05 23:39:12 -07:00
39b5bfdfb0 fix linter 2025-08-05 23:31:41 -07:00
7b4010f804 fix linter 2025-08-05 19:29:43 -07:00
1bebb2eec8 Merge branch 'main' into vllmbuildci 2025-08-05 19:22:57 -07:00
5109207d7e fix linter 2025-08-05 19:21:29 -07:00
b30a421aeb fix linter 2025-08-05 19:19:56 -07:00
af066d8818 setup 2025-08-05 19:10:41 -07:00
75a3bb0ea1 setup 2025-08-05 19:10:12 -07:00
ab198d1f12 setup 2025-08-05 17:27:28 -07:00
8cd9fe649f setup 2025-08-05 17:19:50 -07:00
615939e060 setup 2025-08-05 17:19:13 -07:00
5cba6d3a69 setup 2025-08-05 15:57:14 -07:00
ce9d1261f5 setup 2025-08-05 15:50:41 -07:00
c655d9c118 setup 2025-08-05 15:48:48 -07:00
2ef6e44e68 setup 2025-08-05 15:25:43 -07:00
4e900806cd setup 2025-08-05 14:09:33 -07:00
79413329c2 setup 2025-08-05 14:07:41 -07:00
b418ad7984 setup 2025-08-05 14:05:58 -07:00
d8cf04aa4e setup 2025-08-05 14:04:48 -07:00
c24b7b6104 setup 2025-08-05 13:57:47 -07:00
c6672d6508 setup 2025-08-05 13:57:29 -07:00
77aee73526 setup 2025-08-05 13:57:15 -07:00
6ff6516ddd setup 2025-08-05 13:49:21 -07:00
b952bca99a setup 2025-08-05 12:16:10 -07:00
00271474df setup 2025-08-05 12:12:47 -07:00
2757b05889 setup 2025-08-05 12:08:33 -07:00
790ba79fca setup 2025-08-05 11:28:23 -07:00
9a327e4ba5 setup 2025-08-05 11:25:51 -07:00
11b7f5277a setup 2025-08-04 17:31:12 -07:00
0ba75476df setup 2025-08-04 17:28:30 -07:00
e6abe22ce0 setup 2025-08-04 16:57:03 -07:00
c6ef94ec4c setup 2025-08-04 16:55:38 -07:00
218a90bd3b setup 2025-08-04 16:55:12 -07:00
c0470be279 setup 2025-08-04 16:33:27 -07:00
f4230ace9a setup 2025-08-04 16:31:53 -07:00
a8428ca203 setup 2025-08-04 16:28:52 -07:00
f9dbbdecb2 setup 2025-08-04 16:27:42 -07:00
13 changed files with 1305 additions and 0 deletions

22
.github/ci_configs/vllm_ci_config.yaml vendored Normal file
View File

@ -0,0 +1,22 @@
# Todo: potentially add build config
external_build:
build_target: vllm
artifact_dir: shared
torch_whl_dir: dist
# if set, overrides the base image used in vllm build
base_image: ${DOCKER_IMAGE}
# if set, replaces the vllm dockerfile.torch_nightly with the local dockerfile.
dockerfile_path: "./.github/docker/external/vllm/Dockerfile.base"
test:
- name: Basic Correctness Test # name of the test
id: vllm_basic_correctness_test # the unique id to identify the test config, it must be unique in the ci config yml file
env_vars: # global env vars
- VLLM_WORKER_MULTIPROC_METHOD=spawn
preset: # this can be bash, python etc anything you want to run to set the proper test env
run:
- test: pytest -v -s basic_correctness/test_cumem.py
- test: pytest -v -s basic_correctness/test_basic_correctness.py
- test: pytest -v -s basic_correctness/test_cpu_offload.py
- test: pytest -v -s basic_correctness/test_preemption.py
env_vars:
- VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT=1 # test-level only env var

View File

@ -0,0 +1,422 @@
# The vLLM Dockerfile is used to construct vLLM image against torch nightly and torch main that can be directly used for testing
ARG CUDA_VERSION=12.8.1
ARG PYTHON_VERSION=3.12
# Build_BASE_IMAGE: used to build xformers, and vllm wheels, it can be replaced with a different base image from local machine,
# by default, it uses the torch-nightly-base stage from this docker image
ARG BUILD_BASE_IMAGE=torch-nightly-base
# FINAL_BASE_IMAGE: used to set up vllm-instaled environment and build flashinfer,
# by default, it uses devel-ubuntu22.04 official image.
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
#################### TORCH NIGHTLY BASE IMAGE ####################
# A base image for building vLLM with devel ubuntu 22.04, this is mainly used to build vllm in vllm builtkite ci
From nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 as torch-nightly-base
ARG CUDA_VERSION=12.8.1
ARG PYTHON_VERSION=3.12
ARG TARGETPLATFORM
ENV DEBIAN_FRONTEND=noninteractive
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
# Install Python and other dependencies if it does not existed
RUN if ! command -v python3 >/dev/null || ! python3 --version | grep -q "${PYTHON_VERSION}"; then \
echo "Installing Python ${PYTHON_VERSION}..." && \
echo 'tzdata tzdata/Areas select America' | debconf-set-selections && \
echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections && \
apt-get update -y && \
apt-get install -y ccache software-properties-common git curl sudo && \
for i in 1 2 3; do \
add-apt-repository -y ppa:deadsnakes/ppa && break || \
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
done && \
apt-get update -y && \
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && \
ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config && \
curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION}; \
else \
echo "Python ${PYTHON_VERSION} already present, skipping setup."; \
fi \
&& python3 --version && python3 -m pip --version
# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
# as it was causing spam when compiling the CUTLASS kernels
# Ensure gcc >= 10 to avoid CUTLASS issues (bug 92519)
RUN current_gcc_version=$(gcc -dumpversion | cut -f1 -d.) && \
if [ "$current_gcc_version" -lt 10 ]; then \
echo "GCC version is $current_gcc_version, installing gcc-10..."; \
apt-get update && \
apt-get install -y gcc-10 g++-10 && \
update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 100 && \
update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-10 100; \
else \
echo "GCC version is $current_gcc_version, no need to install gcc-10."; \
fi && \
gcc --version && g++ --version
# install uv for faster pip installs
RUN --mount=type=cache,target=/root/.cache/uv \
python3 -m pip install uv==0.8.4
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
#################### TORCH NIGHTLY BASE IMAGE ####################
#################### BASE BUILD IMAGE ####################
# A base image for building vLLM with torch nightly or torch wheels
# prepare basic build environment
FROM ${BUILD_BASE_IMAGE} AS base
USER root
# Workaround for https://github.com/openai/triton/issues/2507 and
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
# this won't be needed for future versions of this docker image
# or future versions of triton.
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
# Install uv for faster pip installs if not existed
RUN --mount=type=cache,target=/root/.cache/uv \
if ! python3 -m uv --version >/dev/null 2>&1; then \
python3 -m pip install uv==0.8.4; \
fi
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
WORKDIR /workspace
# install build and runtime dependencies
COPY requirements/common.txt requirements/common.txt
COPY use_existing_torch.py use_existing_torch.py
COPY pyproject.toml pyproject.toml
# install build and runtime dependencies without stable torch version
RUN python3 use_existing_torch.py
# default mount file as placeholder, this just avoid the mount error
# change to a different vllm folder if this does not exist anymore
ARG TORCH_WHEELS_PATH="./requirements"
ARG PINNED_TORCH_VERSION
# Install torch, torchaudio and torchvision based on the input
# if TORCH_WHEELS_PATH is default "./requirements", it will pull thethe nightly versions using pip
# otherwise, it will use the whls from TORCH_WHEELS_PATH from the host machine
RUN --mount=type=bind,source=${TORCH_WHEELS_PATH},target=/dist \
--mount=type=cache,target=/root/.cache/uv \
if [ -n "$TORCH_WHEELS_PATH" ] && [ "$TORCH_WHEELS_PATH" != "./requirements" ] && [ -d "/dist" ] && ls /dist/torch*.whl >/dev/null 2>&1; then \
torch_whl=$(find /dist -maxdepth 1 -name 'torch-*.whl' -print -quit); \
vision_whl=$(find /dist/vision -name 'torchvision*.whl' | head -n1 | xargs); \
audio_whl=$(find /dist/audio -name 'torchaudio*.whl' | head -n1 | xargs); \
uv pip install --system "${torch_whl}[opt-einsum]"; \
uv pip install --system "${vision_whl}"; \
uv pip install --system "${audio_whl}"; \
elif [ -n "$PINNED_TORCH_VERSION" ]; then \
echo "[INFO] Installing pinned torch nightly version: $PINNED_TORCH_VERSION"; \
uv pip install --system "$PINNED_TORCH_VERSION" --index-url https://download.pytorch.org/whl/nightly/cu128; \
else \
echo "[INFO] Installing torch nightly with latest one"; \
uv pip install --system torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128; \
fi
# Install numba 0.61.2 for cuda environment
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system numba==0.61.2
# Install common dependencies from vllm common.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/common.txt
# Must put before installing xformers, so it can install the correct version of xfomrers.
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
ARG max_jobs=16
ENV MAX_JOBS=${max_jobs}
# Build xformers with cuda and torch nightly/wheel
# following official xformers guidance: https://github.com/facebookresearch/xformers#build
ARG XFORMERS_COMMIT=f2de641ef670510cadab099ce6954031f52f191c
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/uv \
echo 'git clone xformers...' \
&& git clone https://github.com/facebookresearch/xformers.git --recursive \
&& cd xformers \
&& git checkout ${XFORMERS_COMMIT} \
&& git submodule update --init --recursive \
&& echo 'finish git clone xformers...' \
&& rm -rf build \
&& python3 setup.py bdist_wheel --dist-dir=../xformers-dist --verbose \
&& cd .. \
&& rm -rf xformers
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system xformers-dist/*.whl --verbose
# Build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
# track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
RUN cat torch_build_versions.txt
RUN pip freeze | grep -E 'torch|xformers|torchvision|torchaudio'
# Cuda arch list used by torch
# can be useful for `test`
# explicitly set the list to avoid issues with torch 2.2
# see https://github.com/pytorch/pytorch/pull/123243
# Override the arch list for flash-attn to reduce the binary size
ARG vllm_fa_cmake_gpu_arches='80-real;90-real'
ENV VLLM_FA_CMAKE_GPU_ARCHES=${vllm_fa_cmake_gpu_arches}
#################### BASE BUILD IMAGE ####################
#################### WHEEL BUILD IMAGE ####################
# Image used to build vllm wheel
FROM base AS build
ARG TARGETPLATFORM
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
COPY . .
RUN python3 use_existing_torch.py
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/build.txt
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
# Max jobs used by Ninja to build extensions
ARG max_jobs=16
ENV MAX_JOBS=${max_jobs}
ARG nvcc_threads=2
ENV NVCC_THREADS=$nvcc_threads
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
ARG USE_SCCACHE
ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
ARG SCCACHE_REGION_NAME=us-west-2
ARG SCCACHE_S3_NO_CREDENTIALS=0
# if USE_SCCACHE is set, use sccache to speed up compilation
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=.git,target=.git \
if [ "$USE_SCCACHE" = "1" ]; then \
echo "Installing sccache..." \
&& curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
&& tar -xzf sccache.tar.gz \
&& sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \
&& rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \
&& export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
&& export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
&& export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
&& export SCCACHE_IDLE_TIMEOUT=0 \
&& export CMAKE_BUILD_TYPE=Release \
&& sccache --show-stats \
&& python3 setup.py bdist_wheel --dist-dir=vllm-dist --py-limited-api=cp38 \
&& sccache --show-stats; \
fi
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=.git,target=.git \
if [ "$USE_SCCACHE" != "1" ]; then \
# Clean any existing CMake artifacts
rm -rf .deps && \
mkdir -p .deps && \
python3 setup.py bdist_wheel --dist-dir=vllm-dist --py-limited-api=cp38; \
fi
RUN echo "[DEBUG] Listing current directory:" && \
ls -al && \
echo "[DEBUG] Showing torch_build_versions.txt content:" && \
cat torch_build_versions.txt
#################### WHEEL BUILD IMAGE ####################
################### VLLM INSTALLED IMAGE ####################
# Setup clean environment for vLLM for test and api server using ubuntu22.04 with AOT flashinfer
FROM ${FINAL_BASE_IMAGE} AS vllm-base
USER root
# prepare for environment starts
WORKDIR /workspace
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
# Install Python and other dependencies if it does not existed
RUN if ! command -v python3 >/dev/null || ! python3 --version | grep -q "${PYTHON_VERSION}"; then \
echo "Installing Python ${PYTHON_VERSION}..." && \
echo 'tzdata tzdata/Areas select America' | debconf-set-selections && \
echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections && \
apt-get update -y && \
apt-get install -y ccache software-properties-common git curl sudo && \
for i in 1 2 3; do \
add-apt-repository -y ppa:deadsnakes/ppa && break || \
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
done && \
apt-get update -y && \
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && \
ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config && \
curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION}; \
else \
echo "Python ${PYTHON_VERSION} already present, skipping setup."; \
fi \
&& python3 --version && python3 -m pip --version
# Get the torch versions, and whls used in previous stagtes for consistency
COPY --from=base /workspace/torch_build_versions.txt ./torch_build_versions.txt
COPY --from=base /workspace/xformers-dist /wheels/xformers
COPY --from=build /workspace/vllm-dist /wheels/vllm
RUN echo "[DEBUG] Listing current directory before torch install step:" && \
ls -al && \
echo "[DEBUG] Showing torch_build_versions.txt content:" && \
cat torch_build_versions.txt
# Workaround for https://github.com/openai/triton/issues/2507 and
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
# this won't be needed for future versions of this docker image
# or future versions of triton.
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
# Install uv for faster pip installs if not existed
RUN --mount=type=cache,target=/root/.cache/uv \
if ! python3 -m uv --version > /dev/null 2>&1; then \
python3 -m pip install uv==0.8.4; \
fi
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
# Default mount file as placeholder, this just avoid the mount error
ARG TORCH_WHEELS_PATH="./requirements"
# Install torch, torchaudio and torchvision
# if TORCH_WHEELS_PATH is default "./requirements", it will pull the nightly versions using pip using torch_build_versions.txt
# otherwise, it will use the whls from TORCH_WHEELS_PATH from the host machine
RUN --mount=type=bind,source=${TORCH_WHEELS_PATH},target=/dist \
--mount=type=cache,target=/root/.cache/uv \
if [ -n "$TORCH_WHEELS_PATH" ] && [ "$TORCH_WHEELS_PATH" != "./requirements" ] && [ -d "/dist" ] && ls /dist/torch*.whl >/dev/null 2>&1; then \
torch_whl=$(find /dist -maxdepth 1 -name 'torch-*.whl' -print -quit); \
vision_whl=$(find /dist/vision -name 'torchvision*.whl' | head -n1 | xargs); \
audio_whl=$(find /dist/audio -name 'torchaudio*.whl' | head -n1 | xargs); \
echo "Found: '${torch_whl}' '${audio_whl}' '${vision_whl}'"; \
uv pip install --system "${torch_whl}[opt-einsum]"; \
uv pip install --system "${vision_whl}"; \
uv pip install --system "${audio_whl}"; \
else \
echo "[INFO] Installing torch versions from torch_build_versions.txt"; \
uv pip install --system $(cat torch_build_versions.txt | xargs) --index-url https://download.pytorch.org/whl/nightly/cu128; \
fi
# Install the vllm wheel from previous stage
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system /wheels/vllm/*.whl --verbose
# Install xformers wheel from previous stage
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system /wheels/xformers/*.whl --verbose
# Build flashinfer from source.
ARG torch_cuda_arch_list='8.0;8.9;9.0a'
# install package for build flashinfer
# see issue: https://github.com/flashinfer-ai/flashinfer/issues/738
RUN pip install build==1.3.0
RUN pip freeze | grep -E 'setuptools|packaging|build'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
# Build flashinfer for torch nightly from source around 10 mins
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
# Keep this in sync with https://github.com/vllm-project/vllm/blob/main/requirements/cuda.txt
ARG FLASHINFER_GIT_REF="v0.2.9rc2"
RUN --mount=type=cache,target=/root/.cache/uv \
git clone --depth 1 --recursive --shallow-submodules \
--branch ${FLASHINFER_GIT_REF} \
${FLASHINFER_GIT_REPO} flashinfer \
&& echo "Building FlashInfer with AOT for arches: ${torch_cuda_arch_list}" \
&& cd flashinfer \
&& python3 -m flashinfer.aot \
&& python3 -m build --no-isolation --wheel --outdir ../wheels/flashinfer \
&& cd .. \
&& rm -rf flashinfer
# install flashinfer python
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system wheels/flashinfer/*.whl --verbose
# Logging to confirm the torch versions
RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
################### VLLM INSTALLED IMAGE ####################
#################### UNITTEST IMAGE #############################
FROM vllm-base as test
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
COPY tests/ tests/
COPY examples examples
COPY benchmarks benchmarks
COPY ./vllm/collect_env.py .
COPY requirements/common.txt requirements/common.txt
COPY use_existing_torch.py use_existing_torch.py
COPY pyproject.toml pyproject.toml
# Install build and runtime dependencies without stable torch version
COPY requirements/nightly_torch_test.txt requirements/nightly_torch_test.txt
RUN python3 use_existing_torch.py
# install packages
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/common.txt
# enable fast downloads from hf (for testing)
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system hf_transfer
ENV HF_HUB_ENABLE_HF_TRANSFER 1
# install development dependencies (for testing)
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -e tests/vllm_test_utils
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/nightly_torch_test.txt
# Workaround for #17068
# pinned commit for v2.2.4
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@95d8aba8a8c75aedcaa6143713b11e745e7cd0d9#egg=mamba-ssm"
# Logging to confirm the torch versions
RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'
# Logging to confirm all the packages are installed
RUN pip freeze
#################### UNITTEST IMAGE #############################
#################### EXPORT STAGE ####################
FROM scratch as export-wheels
# Just copy the wheels we prepared in previous stages
COPY --from=base /workspace/xformers-dist /wheels/xformers
COPY --from=build /workspace/vllm-dist /wheels/vllm
COPY --from=vllm-base /workspace/wheels/flashinfer /wheels/flashinfer-python

View File

@ -5,6 +5,7 @@ ciflow_push_tags:
- ciflow/binaries_libtorch
- ciflow/binaries_wheel
- ciflow/triton_binaries
- ciflow/vllm
- ciflow/inductor
- ciflow/inductor-periodic
- ciflow/inductor-rocm

View File

@ -0,0 +1,323 @@
name: linux-external-build
on:
workflow_call:
inputs:
build-environment:
required: true
type: string
description: Top-level label for what's being built/tested.
build-target:
required: true
type: string
description: target library to build
ci-config:
required: false
type: string
description: CI config to use for the build and test.
use-gha:
required: false
type: string
default: ""
description: If set to any value, upload to GHA. Otherwise upload to S3.
build-generates-artifacts:
required: false
type: boolean
default: true
description: If set, upload generated build artifacts.
artifacts-folder-name:
required: false
type: string
description: must be different from build-environment
default: ""
docker-image:
required: true
type: string
description: Docker image to run in or replace the external base image.
cuda-arch-list:
required: false
type: string
default: "8.0"
description: |
List of CUDA architectures CI build should target.
max_jobs:
required: false
type: number
description:
Maximum number of jobs to run the external build
default: 16
runner_prefix:
required: false
default: ""
type: string
description: Prefix for runner label
runner:
required: false
type: string
default: "linux.2xlarge"
description: |
Label of the runner this job should run on.
s3-bucket:
description: S3 bucket to download artifact
required: false
type: string
default: "gha-artifacts"
aws-role-to-assume:
description: Role to assume for downloading artifacts
required: false
type: string
default: ""
disable-monitor:
description: |
Disable utilization monitoring for build job
required: false
type: boolean
default: false
monitor-log-interval:
description: |
Set the interval for the monitor script to log utilization.
required: false
type: number
default: 5
monitor-data-collect-interval:
description: |
Set the interval for the monitor script to collect data.
required: false
type: number
default: 1
build-additional-packages:
description: |
If set, the build job will also builds these packages and saves their
wheels as artifacts
required: false
type: string
default: ""
secrets:
HUGGING_FACE_HUB_TOKEN:
required: false
description: |
HF Auth token to avoid rate limits when downloading models or datasets from hub
SCRIBE_GRAPHQL_ACCESS_TOKEN:
required: false
description: |
FB app token to write to scribe endpoint
jobs:
build-external-lib:
environment: ${{ github.ref == 'refs/heads/main' && 'scribe-protected' || startsWith(github.ref, 'refs/heads/release/') && 'scribe-protected' || contains(github.event.pull_request.labels.*.name, 'ci-scribe') && 'scribe-pr' || '' }}
# Don't run on forked repos
if: github.repository_owner == 'pytorch'
runs-on: ${{ inputs.runner_prefix}}${{ inputs.runner }}
timeout-minutes: 240
steps:
- name: Setup SSH (Click me for login details)
uses: pytorch/test-infra/.github/actions/setup-ssh@main
with:
github-secret: ${{ secrets.GITHUB_TOKEN }}
instructions: |
Build is done inside the container, to start an interactive session run:
docker exec -it $(docker container ps --format '{{.ID}}') bash
# [pytorch repo ref]
# Use a pytorch/pytorch reference instead of a reference to the local
# checkout because when we run this action we don't *have* a local
# checkout. In other cases you should prefer a local checkout.
- name: Checkout PyTorch
uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
with:
no-sudo: true
- name: Get workflow job id
id: get-job-id
uses: ./.github/actions/get-workflow-job-id
if: always()
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
- name: configure aws credentials
uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0
if: ${{ inputs.aws-role-to-assume != ''}}
with:
role-to-assume: ${{ inputs.aws-role-to-assume }}
role-session-name: gha-linux-build
aws-region: us-east-1
- name: Setup Linux
uses: ./.github/actions/setup-linux
if: inputs.build-environment != 'linux-s390x-binary-manywheel'
- name: Login to Amazon ECR
if: ${{ inputs.aws-role-to-assume != ''}}
id: login-ecr
continue-on-error: true
uses: aws-actions/amazon-ecr-login@062b18b96a7aff071d4dc91bc00c4c1a7945b076 # v2.0.1
- name: Parse ref
id: parse-ref
run: .github/scripts/parse_ref.py
- name: Start monitoring script
id: monitor-script
if: ${{ !inputs.disable-monitor }}
shell: bash
continue-on-error: true
env:
JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
WORKFLOW_NAME: ${{ github.workflow }}
WORKFLOW_RUN_ID: ${{github.run_id}}
MONITOR_LOG_INTERVAL: ${{ inputs.monitor-log-interval }}
MONITOR_DATA_COLLECT_INTERVAL: ${{ inputs.monitor-data-collect-interval }}
run: |
mkdir -p ../../usage_logs
python3 -m pip install psutil==5.9.8 dataclasses_json==0.6.7
python3 -m tools.stats.monitor \
--log-interval "$MONITOR_LOG_INTERVAL" \
--data-collect-interval "$MONITOR_DATA_COLLECT_INTERVAL" \
> "../../usage_logs/usage_log_build_${JOB_ID}.txt" 2>&1 &
echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}"
- name: Calculate docker image
id: calculate-docker-image
uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
if: inputs.build-environment != 'linux-s390x-binary-manywheel'
with:
docker-image-name: ${{ inputs.docker-image }}
- name: Use following to pull public copy of the image
id: print-ghcr-mirror
if: inputs.build-environment != 'linux-s390x-binary-manywheel'
env:
ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
shell: bash
run: |
tag=${ECR_DOCKER_IMAGE##*:}
echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
- name: Pull docker image
uses: pytorch/test-infra/.github/actions/pull-docker-image@main
with:
docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
- name: Download pytorch build artifacts
uses: ./.github/actions/download-build-artifacts
with:
name: ${{ inputs.build-environment }}
s3-bucket: ${{ inputs.s3-bucket }}
use-gha: ${{ inputs.use-gha }}
- name: Download TD artifacts
continue-on-error: true
uses: ./.github/actions/download-td-artifacts
- name: Calculate Max Jobs
id: set_max_jobs
run: |
if [[ -n "${{ inputs.max_jobs }}" ]]; then
echo "Using input max_jobs: ${{ inputs.max_jobs }}"
echo "MAX_JOBS=${{ inputs.max_jobs }}" >> "$GITHUB_OUTPUT"
else
DEFAULT_JOBS=$(nproc --ignore=6)
echo "Fallback to nproc: $DEFAULT_JOBS"
echo "MAX_JOBS=$DEFAULT_JOBS" >> "$GITHUB_OUTPUT"
fi
- name: Build external project
id: build
env:
BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
BRANCH: ${{ steps.parse-ref.outputs.branch }}
PR_NUMBER: ${{ github.event.pull_request.number }}
SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
# Do not set SCCACHE_S3_KEY_PREFIX to share the cache between all build jobs
SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
SCCACHE_REGION: us-east-1
PR_LABELS: ${{ toJson(github.event.pull_request.labels.*.name) }}
TORCH_CUDA_ARCH_LIST: ${{ inputs.cuda-arch-list }}
OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
MAX_JOBS: {{ steps.set_max_jobs.outputs.MAX_JOBS }}
DOCKER_IMAGE: ${{ inputs.docker-image }}
BUILD_TARGET: ${{ inputs.build-target }}
CI_CONFIG: ${{ inputs.ci-config }}
run: |
set -euo pipefail
python3 --version
docker images
START_TIME=$(date +%s)
(
cd scripts/torch_cli
python3 -m pip install -e .
)
python3 -m cli.run --config "$CI_CONFIG" build external "$BUILD_TARGET"
END_TIME=$(date +%s)
echo "build_time=$((END_TIME - START_TIME))" >> "$GITHUB_OUTPUT"
- name: Archive artifacts into zip
if: ${{ inputs.build-generates-artifacts && steps.build.outcome && steps.build.outcome != 'skipped'}}
run: |
zip -1 -r artifacts.zip shared/
# By default it will upload the artifacts to <github_org>/<github_repo>/<workflow_id>/<name>-<target>-additional-build/
# to avoid override the pytorch build artifacts
- name: Store External Build Artifacts on S3
if: ${{ inputs.build-generates-artifacts }}
uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0
with:
name: ${{ inputs.artifacts-folder-name || format('{0}-{1}-additional-build', inputs.build-environment, inputs.build-target) }}
retention-days: 14
if-no-files-found: warn
path: artifacts.zip
s3-bucket: ${{ inputs.s3-bucket }}
- name: Stop monitoring script
if: ${{ always() && steps.monitor-script.outputs.monitor-script-pid }}
shell: bash
continue-on-error: true
env:
MONITOR_SCRIPT_PID: ${{ steps.monitor-script.outputs.monitor-script-pid }}
run: |
kill "$MONITOR_SCRIPT_PID"
- name: Copy logs
shell: bash
if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}}
continue-on-error: true
run: |
rm -f ./usage_logs
mkdir -p ./usage_logs
cp ../../usage_logs/usage_log_build_*.txt ./usage_logs/
- name: Upload raw usage log to s3
if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}}
uses: seemethere/upload-artifact-s3@v5
with:
s3-prefix: |
${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
retention-days: 14
if-no-files-found: warn
path: usage_logs/usage_log_build_*.txt
- name: Upload utilization stats
if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel' }}
continue-on-error: true
uses: ./.github/actions/upload-utilization-stats
with:
job_id: ${{ steps.get-job-id.outputs.job-id }}
job_name: ${{ steps.get-job-id.outputs.job-name }}
workflow_name: ${{ github.workflow }}
workflow_run_id: ${{github.run_id}}
workflow_attempt: ${{github.run_attempt}}
artifact_prefix: usage_log_build_${{ steps.get-job-id.outputs.job-id }}
- name: Teardown Linux
uses: pytorch/test-infra/.github/actions/teardown-linux@main
if: always() && inputs.build-environment != 'linux-s390x-binary-manywheel'
- name: Cleanup docker
if: always() && inputs.build-environment == 'linux-s390x-binary-manywheel'
shell: bash
run: |
# on s390x stop the container for clean worker stop
docker stop -a || true
docker kill -a || true

94
.github/workflows/vllm.yml vendored Normal file
View File

@ -0,0 +1,94 @@
name: vllm-test
on:
push:
tags:
- ciflow/vllm/*
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
cancel-in-progress: true
permissions:
id-token: write
contents: read
jobs:
get-label-type:
name: get-label-type
uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }}
with:
triggering_actor: ${{ github.triggering_actor }}
issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
curr_branch: ${{ github.head_ref || github.ref_name }}
curr_ref_type: ${{ github.ref_type }}
opt_out_experiments: lf
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-80:
name: cuda12.8-py3.12-gcc11-sm80-vllm-test
uses: ./.github/workflows/_linux-build.yml
needs: get-label-type
with:
build-additional-packages: "vision audio torchao xformers"
build-environment: linux-jammy-cuda12.8-py3.12-gcc11-sm80
docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.12-gcc11-vllm
cuda-arch-list: '8.0'
test-matrix: |
{ include: [
{ config: "vllm_basic_correctness_test", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
{ config: "vllm_regression_test", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
{ config: "vllm_entrypoints_test_llm", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
{ config: "vllm_basic_models_test", shard: 1, num_shards: 1, runner: "linux.aws.a100" },
]}
secrets: inherit
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-external-build-sm80:
name: cuda12.8-py3.12-gcc11-sm80-vllm-test
uses: ./.github/workflows/_linux-external-build-main.yml
needs: [
get-label-type,
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-80
]
with:
build-additional-packages: "vision audio"
build-environment: linux-jammy-cuda12.8-py3.12-gcc11-sm80
build-target: vllm
ci-config: ".github/ci_configs/vllm_ci_config.yaml"
runner_prefix: ${{ needs.get-label-type.outputs.label-type }}
docker-image: ${{ needs.linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-80.outputs.docker-image }}
runner: linux.24xlarge.memory
secrets: inherit
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-90:
name: cuda12.8-py3.12-gcc11-sm90-vllm-test
uses: ./.github/workflows/_linux-build.yml
needs: get-label-type
with:
build-additional-packages: "vision audio torchao xformers"
build-environment: linux-jammy-cuda12.8-py3.12-gcc11-sm90
docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.12-gcc11-vllm
cuda-arch-list: '9.0'
test-matrix: |
{ include: [
{ config: "vllm_basic_correctness_test", shard: 1, num_shards: 1, runner: "linux.aws.h100" },
]}
secrets: inherit
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-external-build-sm90:
name: cuda12.8-py3.12-gcc11-sm90-vllm-test
uses: ./.github/workflows/_linux-external-build-main.yml
needs: [
get-label-type,
linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-90
]
with:
build-additional-packages: "vision audio"
build-environment: linux-jammy-cuda12.8-py3.12-gcc11-sm90
build-target: vllm
ci-config: ".github/ci_configs/vllm_ci_config.yaml"
runner_prefix: ${{ needs.get-label-type.outputs.label-type }}
docker-image: ${{ needs.linux-jammy-cuda12_8-cudnn9-py3_12-gcc11-vllm-build-90.outputs.docker-image }}
runner: linux.24xlarge.memory
secrets: inherit

View File

@ -0,0 +1,5 @@
from cli.build_cli.external.controller import ExternalBuildController
from cli.build_cli.main import BuildController
BUILD_CONTROLLERS = [BuildController, ExternalBuildController]

View File

@ -0,0 +1,23 @@
from cement import Controller, ex
from cli.lib.vllm_utils import build_vllm
from cli.lib.utils import generate_dataclass_help, read_yaml_file
class ExternalBuildController(Controller):
class Meta:
label = "build_external"
aliases = ["external"]
stacked_on = "build"
stacked_type = "nested"
help = "Build external libraries with pytorch ci"
@ex(
help="Build vllm",
)
def vllm(self):
config = self.app.pargs.config
config_map = {}
if config:
print("please input config file, otherwise use default config")
config_map = read_yaml_file(config)
build_vllm(config=config_map)

View File

@ -0,0 +1,9 @@
from cement import Controller
class BuildController(Controller):
class Meta:
label = "build"
stacked_on = "base"
stacked_type = "nested"
description = "Build CLI group"

View File

View File

@ -0,0 +1,185 @@
import os
import shlex
import shutil
import subprocess
import sys
from dataclasses import fields
from pathlib import Path
from textwrap import indent
from typing import Optional
import yaml
def run_shell(
cmd: str,
logging: bool = True,
cwd: Optional[str] = None,
env: Optional[dict] = None,
):
if logging:
print(f"[shell] {cmd}", flush=True)
try:
subprocess.run(
cmd,
shell=True,
executable="/bin/bash",
stdout=sys.stdout,
stderr=sys.stderr,
check=True,
env=env,
cwd=cwd,
)
except subprocess.CalledProcessError as e:
print("docker build failed]")
print("Command:", " ".join(cmd))
print("Exit code:", e.returncode)
print("STDOUT:\n", e.stdout)
print("STDERR:\n", e.stderr)
def run_cmd(
cmd: str,
logging: bool = True,
cwd: Optional[str] = None,
env: Optional[dict] = None,
):
args = shlex.split(cmd)
if logging:
print(f"[cmd] {' '.join(args)}", flush=True)
try:
subprocess.run(
args,
shell=False,
stdout=sys.stdout,
stderr=sys.stderr,
check=True,
env=env,
cwd=cwd,
)
except subprocess.CalledProcessError as e:
print("[❌ docker build failed]")
print("Command:", " ".join(cmd))
print("Exit code:", e.returncode)
print("STDOUT:\n", e.stdout)
print("STDERR:\n", e.stderr)
# eliainwy
def get_post_build_pinned_commit(name: str, prefix=".github/ci_commit_pins") -> str:
path = Path(prefix) / f"{name}.txt"
if not path.exists():
raise FileNotFoundError(f"Pin file not found: {path}")
return path.read_text(encoding="utf-8").strip()
def get_env(name: str, default: str = "") -> str:
return os.environ.get(name, default)
def force_create_dir(path: str):
"""
Forcefully create a fresh directory.
If the directory exists, it will be removed first.
"""
remove_dir(path)
ensure_dir_exists(path)
def ensure_dir_exists(path: str):
"""
Ensure the directory exists. Create it if necessary.
"""
if not os.path.exists(path):
print(f"[INFO] Creating directory: {path}", flush=True)
os.makedirs(path, exist_ok=True)
else:
print(f"[INFO] Directory already exists: {path}", flush=True)
def remove_dir(path: str):
"""
Remove a directory if it exists.
"""
if os.path.exists(path):
print(f"[INFO] Removing directory: {path}", flush=True)
shutil.rmtree(path)
else:
print(f"[INFO] Directory not found (skipped): {path}", flush=True)
def get_abs_path(path: str):
return os.path.abspath(path)
def generate_dataclass_help(cls) -> str:
"""Auto-generate help text for dataclass default values."""
lines = []
for field in fields(cls):
default = field.default
if default is not None and default != "":
lines.append(f"{field.name:<22} = {repr(default)}")
else:
lines.append(f"{field.name:<22} = ''")
return indent("\n".join(lines), " ")
def get_existing_abs_path(path: str) -> str:
path = os.path.abspath(path)
if not os.path.exists(path):
raise FileNotFoundError(f"Path does not exist: {path}")
return path
def clone_vllm(commit: str):
"""
cloning vllm and checkout pinned commit
"""
print(f"clonening vllm....", flush=True)
cwd = "vllm"
# delete the directory if it exists
remove_dir(cwd)
# Clone the repo & checkout commit
run_cmd("git clone https://github.com/vllm-project/vllm.git")
run_cmd(f"git checkout {commit}", cwd=cwd)
run_cmd("git submodule update --init --recursive", cwd=cwd)
def pip_install(package: str):
cmd = f"python3 -m pip install {package}"
subprocess.run(shlex.split(cmd), check=True)
def uv_pip_install(package: str):
cmd = f"python3 -m uv pip install --system {package}"
subprocess.run(shlex.split(cmd), check=True)
def read_yaml_file(file_path: str) -> dict:
p = get_abs_path(file_path)
if not os.path.exists(p):
raise FileNotFoundError(f"YAML file not found: {file_path}")
try:
with open(p, "r", encoding="utf-8") as f:
raw_content = f.read()
# Replace environment variables with env var such as ${DOCKER_IMAGE}
expanded_content = os.path.expandvars(raw_content)
data = yaml.safe_load(expanded_content)
if data is None:
return {}
if not isinstance(data, dict):
raise ValueError(
f"YAML content must be a dictionary, got {type(data).__name__}"
)
return data
except yaml.YAMLError as e:
raise ValueError(f"Failed to parse YAML file '{file_path}': {e}") from e
except Exception as e:
raise RuntimeError(
f"Unexpected error while reading YAML file '{file_path}': {e}"
) from e

View File

@ -0,0 +1,166 @@
import os
import subprocess
import textwrap
from dataclasses import dataclass
from typing import Any, Dict
from cli.lib.utils import (
clone_vllm,
ensure_dir_exists,
force_create_dir,
get_abs_path,
get_env,
get_existing_abs_path,
get_post_build_pinned_commit,
run_cmd,
)
@dataclass
class VllmBuildConfig:
tag_name: str = get_env("TAG", "vllm-wheels")
cuda: str = get_env("CUDA_VERSION", "12.8.1")
py: str = get_env("PYTHON_VERSION", "3.12")
max_jobs: str = get_env("MAX_JOBS", "64")
target: str = get_env("TARGET", "export-wheels")
sccache_bucket: str = get_env("SCCACHE_BUCKET", "")
sccache_region: str = get_env("SCCACHE_REGION", "")
torch_cuda_arch_list: str = get_env("TORCH_CUDA_ARCH_LIST", "8.0")
vllm_fa_cmake_gpu_arches = get_env("VLLM_FA_CMAKE_GPU_ARCHES", "80-real")
dev = get_env("DEV")
artifact_dir: str = ""
torch_whl_dir: str = ""
base_image: str = ""
dockerfile_path: str = ""
_DEFAULT_RESULT_PATH = "./results"
_VLLM_TEMP_FOLDER = "tmp"
def local_image_exists(image: str):
try:
subprocess.check_output(
["docker", "image", "inspect", image], stderr=subprocess.DEVNULL
)
return True
except subprocess.CalledProcessError:
return False
def _prepare_artifact_dir(path: str):
if not path:
path = _DEFAULT_RESULT_PATH
abs_path = get_abs_path(path)
ensure_dir_exists(abs_path)
return abs_path
def getVllmBuildConfig(config: Dict[str, Any]):
print("config", config)
external_build_config = config.get("external_build", {})
return VllmBuildConfig(
artifact_dir=external_build_config.get("artifact_dir", ""),
torch_whl_dir=external_build_config.get("torch_whl_dir", ""),
base_image=external_build_config.get("base_image", ""),
dockerfile_path=external_build_config.get("dockerfile_path", ""),
)
def build_vllm(config: Dict[str, Any]):
cfg = getVllmBuildConfig(config)
print(f"Target artifact dir path is {cfg.artifact_dir}", flush=True)
print("config peek", cfg)
if cfg.dev:
vllm_commit = get_post_build_pinned_commit("vllm", ".")
else:
vllm_commit = get_post_build_pinned_commit("vllm")
clone_vllm(vllm_commit)
# replace dockerfile
if cfg.dockerfile_path:
abs_file_path = get_existing_abs_path(cfg.dockerfile_path)
print(
f"use user provided dockerfile {cfg.dockerfile_path} with path {abs_file_path}"
)
run_cmd(
f"cp {abs_file_path} ./vllm/docker/Dockerfile.nightly_torch",
)
else:
print("using vllm default dockerfile.torch_nightly for build")
reault_path = _prepare_artifact_dir(cfg.artifact_dir)
torch_arg, _ = _prepare_torch_wheels(cfg.torch_whl_dir)
base_arg, final_base_img, pull_flag = _get_base_image_args(cfg.base_image)
cmd = _generate_docker_build_cmd(
cfg, reault_path, torch_arg, base_arg, final_base_img, pull_flag
)
print("Running docker build", flush=True)
print(cmd, flush=True)
run_cmd(cmd, cwd="vllm", env=os.environ.copy())
def _prepare_torch_wheels(torch_whl_dir: str) -> tuple[str, str]:
if not torch_whl_dir:
return "", ""
abs_whl_dir = get_existing_abs_path(torch_whl_dir)
tmp_dir = f"./vllm/{_VLLM_TEMP_FOLDER}"
force_create_dir(tmp_dir)
run_cmd(f"cp -a {abs_whl_dir}/. {tmp_dir}", logging=True)
return f"--build-arg TORCH_WHEELS_PATH={_VLLM_TEMP_FOLDER}", tmp_dir
def _get_base_image_args(base_image: str) -> tuple[str, str, str]:
"""
Returns:
- base_image_arg: docker buildx arg string for base image
- pull_flag: --pull=true or --pull=false depending on whether the image exists locally
"""
pull_flag = ""
if not base_image:
return "", "", ""
base_image_arg = f"--build-arg BUILD_BASE_IMAGE={base_image}"
final_base_image_arg = f"--build-arg FINAL_BASE_IMAGE={base_image}"
if local_image_exists(base_image):
print(f"[INFO] Found local image: {base_image}", flush=True)
pull_flag = "--pull=false"
return base_image_arg, final_base_image_arg, pull_flag
print(
f"[INFO] Local image not found: {base_image}, will try to pull from remote",
flush=True,
)
return base_image_arg, final_base_image_arg, ""
def _generate_docker_build_cmd(
cfg: VllmBuildConfig,
result_path: str,
torch_arg: str,
base_image_arg: str,
final_base_image_arg: str,
pull_flag: str,
) -> str:
return textwrap.dedent(
f"""
docker buildx build \
--output type=local,dest={result_path} \
-f docker/Dockerfile.nightly_torch \
{pull_flag} \
{torch_arg} \
{base_image_arg} \
{final_base_image_arg} \
--build-arg max_jobs={cfg.max_jobs} \
--build-arg CUDA_VERSION={cfg.cuda} \
--build-arg PYTHON_VERSION={cfg.py} \
--build-arg USE_SCCACHE={int(bool(cfg.sccache_bucket and cfg.sccache_region))} \
--build-arg SCCACHE_BUCKET_NAME={cfg.sccache_bucket} \
--build-arg SCCACHE_REGION_NAME={cfg.sccache_region} \
--build-arg torch_cuda_arch_list={cfg.torch_cuda_arch_list} \
--build-arg vllm_fa_cmake_gpu_arches={cfg.vllm_fa_cmake_gpu_arches}\
--target {cfg.target} \
-t {cfg.tag_name} \
--progress=plain .
"""
).strip()

View File

@ -0,0 +1,35 @@
from cement import App, Controller, ex
from cli.build_cli import BUILD_CONTROLLERS
class MainController(Controller):
class Meta:
label = "base"
description = "Main entry point"
arguments = [
(
["--config"],
{
"help": "Path to config file for ci build and test",
"dest": "config",
"default": "",
"required": False,
},
),
]
class TorchCli(App):
class Meta:
label = "cli"
base_controller = "base"
handlers = [MainController] + BUILD_CONTROLLERS
def main():
with TorchCli() as app:
app.run()
if __name__ == "__main__":
main()

View File

@ -0,0 +1,20 @@
[project]
name = "torch-ci"
version = "0.1.0"
dependencies = [
"cement>=3.0.0",
"pyyaml>=6.0"
]
[tool.setuptools]
packages = ["cli", "cli.build_cli", "cli.lib"]
[tool.setuptools.package-dir]
cli = "cli"
[tool.ruff.lint]
# Enable preview mode for linting
preview = true
# Now you can select your preview rules, like RUF048
extend-select = ["RUF048"]