Dockerfile 22.6 KB
Newer Older
1

Simon Mo's avatar
Simon Mo committed
2
3
4
# The vLLM Dockerfile is used to construct vLLM image that can be directly used
# to run the OpenAI compatible server.

5
# Please update any changes made here to
6
7
# docs/contributing/dockerfile/dockerfile.md and
# docs/assets/contributing/dockerfile-stages-dependency.png
8

Huy Do's avatar
Huy Do committed
9
ARG CUDA_VERSION=12.8.1
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
ARG PYTHON_VERSION=3.12

# By parameterizing the base images, we allow third-party to use their own
# base images. One use case is hermetic builds with base images stored in
# private registries that use a different repository naming conventions.
#
# Example:
# docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04

# By parameterizing the Deadsnakes repository URL, we allow third-party to use
# their own mirror. When doing so, we don't benefit from the transparent
# installation of the GPG key of the PPA, as done by add-apt-repository, so we
# also need a URL for the GPG key.
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL

# The PyPA get-pip.py script is a self contained script+zip file, that provides
# both the installer script and the pip base85-encoded zip archive. This allows
# bootstrapping pip in environment where a dsitribution package does not exist.
#
# By parameterizing the URL for get-pip.py installation script, we allow
# third-party to use their own copy of the script stored in a private mirror.
# We set the default value to the PyPA owned get-pip.py script.
#
# Reference: https://pip.pypa.io/en/stable/installation/#get-pip-py
ARG GET_PIP_URL="https://bootstrap.pypa.io/get-pip.py"

# PIP supports fetching the packages from custom indexes, allowing third-party
# to host the packages in private mirrors. The PIP_INDEX_URL and
# PIP_EXTRA_INDEX_URL are standard PIP environment variables to override the
# default indexes. By letting them empty by default, PIP will use its default
# indexes if the build process doesn't override the indexes.
#
# Uv uses different variables. We set them by default to the same values as
# PIP, but they can be overridden.
ARG PIP_INDEX_URL
ARG PIP_EXTRA_INDEX_URL
ARG UV_INDEX_URL=${PIP_INDEX_URL}
ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}

# PyTorch provides its own indexes for standard and nightly builds
ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl
ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly

# PIP supports multiple authentication schemes, including keyring
# By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to
# disabled by default, we allow third-party to use keyring authentication for
# their private Python indexes, while not changing the default behavior which
# is no authentication.
#
# Reference: https://pip.pypa.io/en/stable/topics/authentication/#keyring-support
ARG PIP_KEYRING_PROVIDER=disabled
ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER}

66
67
68
# Flag enables build-in KV-connector dependency libs into docker images
ARG INSTALL_KV_CONNECTORS=false

Simon Mo's avatar
Simon Mo committed
69
#################### BASE BUILD IMAGE ####################
70
# prepare basic build environment
71
72
73
FROM ${BUILD_BASE_IMAGE} AS base
ARG CUDA_VERSION
ARG PYTHON_VERSION
74
ARG TARGETPLATFORM
75
ARG INSTALL_KV_CONNECTORS=false
76
77
ENV DEBIAN_FRONTEND=noninteractive

78
79
80
81
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL

82
83
84
85
86
# Install Python and other dependencies
RUN 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 \
87
88
89
90
91
92
93
94
95
96
97
98
99
    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
            mkdir -p -m 0755 /etc/apt/keyrings ; \
            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
        fi ; \
    else \
        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 ; \
    fi \
100
101
102
103
104
    && 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 \
105
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
106
    && python3 --version && python3 -m pip --version
107
108
109
110
111
112
113

ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL
ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

114
115
116
# Install uv for faster pip installs
RUN --mount=type=cache,target=/root/.cache/uv \
    python3 -m pip install uv
Stephen Krider's avatar
Stephen Krider committed
117

118
119
120
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
Huy Do's avatar
Huy Do committed
121
ENV UV_INDEX_STRATEGY="unsafe-best-match"
122

123
124
125
126
127
128
129
130
# 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
RUN apt-get install -y gcc-10 g++-10
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10
RUN <<EOF
gcc --version
EOF

131
132
133
134
# 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.
135
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
136

Stephen Krider's avatar
Stephen Krider committed
137
138
WORKDIR /workspace

139
140
# install build and runtime dependencies

141
142
143
144
# arm64 (GH200) build follows the practice of "use existing pytorch" build,
# we need to install torch and torchvision from the nightly builds first,
# pytorch will not appear as a vLLM dependency in all of the following steps
# after this step
145
RUN --mount=type=cache,target=/root/.cache/uv \
146
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
147
148
149
150
151
152
        uv pip install --system \
            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
            "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319";  \
        uv pip install --system \
            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
            --pre pytorch_triton==3.3.0+gitab727c40; \
153
    fi
Mor Zusman's avatar
Mor Zusman committed
154

155
156
COPY requirements/common.txt requirements/common.txt
COPY requirements/cuda.txt requirements/cuda.txt
157
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
158
    uv pip install --system -r requirements/cuda.txt \
159
    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
160

161
162
163
164
# cuda arch list used by torch
# can be useful for both `dev` and `test`
# explicitly set the list to avoid issues with torch 2.2
# see https://github.com/pytorch/pytorch/pull/123243
165
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
166
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
167
168
169
# 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}
Simon Mo's avatar
Simon Mo committed
170
171
#################### BASE BUILD IMAGE ####################

172
#################### WHEEL BUILD IMAGE ####################
173
FROM base AS build
174
ARG TARGETPLATFORM
175

176
177
178
179
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL

180
# install build dependencies
181
COPY requirements/build.txt requirements/build.txt
182

183
184
185
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
Huy Do's avatar
Huy Do committed
186
ENV UV_INDEX_STRATEGY="unsafe-best-match"
187

188
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
189
    uv pip install --system -r requirements/build.txt \
190
    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
191

192
COPY . .
193
194
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
195
    if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
Stephen Krider's avatar
Stephen Krider committed
196
197

# max jobs used by Ninja to build extensions
198
199
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}
200
201
202
# number of threads used by nvcc
ARG nvcc_threads=8
ENV NVCC_THREADS=$nvcc_threads
203

204
ARG USE_SCCACHE
205
206
ARG SCCACHE_DOWNLOAD_URL=https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz
ARG SCCACHE_ENDPOINT
207
208
ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
ARG SCCACHE_REGION_NAME=us-west-2
209
ARG SCCACHE_S3_NO_CREDENTIALS=0
210
# if USE_SCCACHE is set, use sccache to speed up compilation
211
RUN --mount=type=cache,target=/root/.cache/uv \
212
    --mount=type=bind,source=.git,target=.git \
213
214
    if [ "$USE_SCCACHE" = "1" ]; then \
        echo "Installing sccache..." \
215
        && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \
216
217
218
        && 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 \
219
        && if [ ! -z ${SCCACHE_ENDPOINT} ] ; then export SCCACHE_ENDPOINT=${SCCACHE_ENDPOINT} ; fi \
220
221
        && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
        && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
222
        && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
223
        && export SCCACHE_IDLE_TIMEOUT=0 \
224
        && export CMAKE_BUILD_TYPE=Release \
225
        && sccache --show-stats \
226
        && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
227
228
229
        && sccache --show-stats; \
    fi

230
231
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
232
    --mount=type=cache,target=/root/.cache/uv \
233
    --mount=type=bind,source=.git,target=.git  \
234
    if [ "$USE_SCCACHE" != "1" ]; then \
235
236
237
        # Clean any existing CMake artifacts
        rm -rf .deps && \
        mkdir -p .deps && \
238
        python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
239
    fi
240

241
# Check the size of the wheel if RUN_WHEEL_CHECK is true
242
COPY .buildkite/check-wheel-size.py check-wheel-size.py
243
# sync the default value with .buildkite/check-wheel-size.py
244
ARG VLLM_MAX_SIZE_MB=400
245
246
247
248
249
250
251
ENV VLLM_MAX_SIZE_MB=$VLLM_MAX_SIZE_MB
ARG RUN_WHEEL_CHECK=true
RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
        python3 check-wheel-size.py dist; \
    else \
        echo "Skipping wheel size check."; \
    fi
Simon Mo's avatar
Simon Mo committed
252
#################### EXTENSION Build IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
253

254
255
256
#################### DEV IMAGE ####################
FROM base as dev

257
258
259
260
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL

261
262
263
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
Huy Do's avatar
Huy Do committed
264
265
266
267
268
ENV UV_INDEX_STRATEGY="unsafe-best-match"

# Workaround for #17068
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
269

270
271
272
COPY requirements/lint.txt requirements/lint.txt
COPY requirements/test.txt requirements/test.txt
COPY requirements/dev.txt requirements/dev.txt
273
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
274
    uv pip install --system -r requirements/dev.txt \
275
    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
276
#################### DEV IMAGE ####################
277

278
279
#################### vLLM installation IMAGE ####################
# image with vLLM installed
280
# TODO: Restore to base image after FlashInfer AOT wheel fixed
281
282
283
FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
284
ARG INSTALL_KV_CONNECTORS=false
Simon Mo's avatar
Simon Mo committed
285
WORKDIR /vllm-workspace
286
ENV DEBIAN_FRONTEND=noninteractive
287
288
ARG TARGETPLATFORM

Kebe's avatar
Kebe committed
289
290
SHELL ["/bin/bash", "-c"]

291
292
293
294
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL

295
296
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
    echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
297

298
299
300
301
302
303
# Install Python and other dependencies
RUN 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 wget sudo vim python3-pip \
    && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
304
305
306
307
308
309
310
311
312
313
314
315
316
    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
            mkdir -p -m 0755 /etc/apt/keyrings ; \
            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
        fi ; \
    else \
        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 ; \
    fi \
317
318
319
320
321
    && apt-get update -y \
    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \
    && 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 \
322
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
323
    && python3 --version && python3 -m pip --version
324
325
326
327
328
329
330

ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL
ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

331
332
333
# Install uv for faster pip installs
RUN --mount=type=cache,target=/root/.cache/uv \
    python3 -m pip install uv
334

335
336
337
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
Huy Do's avatar
Huy Do committed
338
ENV UV_INDEX_STRATEGY="unsafe-best-match"
339

340
341
342
343
# 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.
344
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
345

346
347
348
349
# arm64 (GH200) build follows the practice of "use existing pytorch" build,
# we need to install torch and torchvision from the nightly builds first,
# pytorch will not appear as a vLLM dependency in all of the following steps
# after this step
350
RUN --mount=type=cache,target=/root/.cache/uv \
351
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
352
353
354
355
356
357
        uv pip install --system \
            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
            "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \
        uv pip install --system \
            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
            --pre pytorch_triton==3.3.0+gitab727c40 ; \
358
359
    fi

360
# Install vllm wheel first, so that torch etc will be installed.
361
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
362
    --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
363
    uv pip install --system dist/*.whl --verbose \
364
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
Mor Zusman's avatar
Mor Zusman committed
365

366
# If we need to build FlashInfer wheel before its release:
367
# $ # Note we remove 7.0 from the arch list compared to the list below, since FlashInfer only supports sm75+
368
# $ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0'
369
370
# $ git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
# $ cd flashinfer
371
372
373
374
375
376
# $ git checkout v0.2.6.post1
# $ python -m flashinfer.aot
# $ python -m build --no-isolation --wheel
# $ ls -la dist
# -rw-rw-r-- 1 mgoin mgoin 205M Jun  9 18:03 flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
# $ # upload the wheel to a public location, e.g. https://wheels.vllm.ai/flashinfer/v0.2.6.post1/flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
377

378
379
380
381
# Allow specifying a version, Git revision or local .whl file
ARG FLASHINFER_CUDA128_INDEX_URL="https://download.pytorch.org/whl/cu128/flashinfer"
ARG FLASHINFER_CUDA128_WHEEL="flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl"
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
382
383
384
385
ARG FLASHINFER_GIT_REF="v0.2.8rc1"
# Flag to control whether to use pre-built FlashInfer wheels (set to false to force build from source)
# TODO: Currently disabled because the pre-built wheels are not available for FLASHINFER_GIT_REF
ARG USE_FLASHINFER_PREBUILT_WHEEL=false
386
387
388
389
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
  . /etc/environment
  if [ "$TARGETPLATFORM" != "linux/arm64" ]; then
      # FlashInfer already has a wheel for PyTorch 2.7.0 and CUDA 12.8. This is enough for CI use
390
      if [[ "$CUDA_VERSION" == 12.8* ]] && [[ "$USE_FLASHINFER_PREBUILT_WHEEL" == "true" ]]; then
391
392
          uv pip install --system ${FLASHINFER_CUDA128_INDEX_URL}/${FLASHINFER_CUDA128_WHEEL}
      else
393
          # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
394
395
          # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
          if [[ "${CUDA_VERSION}" == 11.* ]]; then
396
397
398
399
400
401
              FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
          elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
              FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
          else
              # CUDA 12.8+ supports 10.0a and 12.0
              FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
402
          fi
403
          echo "🏗️  Building FlashInfer for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
404
405

          git clone --depth 1 --recursive --shallow-submodules \
406
407
            --branch ${FLASHINFER_GIT_REF} \
            ${FLASHINFER_GIT_REPO} flashinfer
408

409
          # Needed to build AOT kernels
410
          pushd flashinfer
411
412
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
              python3 -m flashinfer.aot
413
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
414
415
416
417
418
419
420
              uv pip install --system --no-build-isolation .
          popd

          rm -rf flashinfer
      fi \
  fi
BASH
421
COPY examples examples
422
COPY benchmarks benchmarks
423
COPY ./vllm/collect_env.py .
424

Huy Do's avatar
Huy Do committed
425
426
427
428
RUN --mount=type=cache,target=/root/.cache/uv \
. /etc/environment && \
uv pip list

429
# Even when we build Flashinfer with AOT mode, there's still
430
431
432
# some issues w.r.t. JIT compilation. Therefore we need to
# install build dependencies for JIT compilation.
# TODO: Remove this once FlashInfer AOT wheel is fixed
433
COPY requirements/build.txt requirements/build.txt
434
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
435
    uv pip install --system -r requirements/build.txt \
436
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
437

438
#################### vLLM installation IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
439

440
441
442
443
#################### TEST IMAGE ####################
# image to run unit testing suite
# note that this uses vllm installed by `pip`
FROM vllm-base AS test
Stephen Krider's avatar
Stephen Krider committed
444

445
ADD . /vllm-workspace/
Stephen Krider's avatar
Stephen Krider committed
446

447
448
449
450
451
ARG PYTHON_VERSION

ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

452
453
454
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500
Huy Do's avatar
Huy Do committed
455
ENV UV_INDEX_STRATEGY="unsafe-best-match"
456

457
458
# Workaround for #17068
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
459
460
461
    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"

# install development dependencies (for testing)
462
RUN --mount=type=cache,target=/root/.cache/uv \
463
464
465
466
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
    if [ "$CUDA_MAJOR" -ge 12 ]; then \
        uv pip install --system -r requirements/dev.txt; \
    fi
467

youkaichao's avatar
youkaichao committed
468
# install development dependencies (for testing)
469
RUN --mount=type=cache,target=/root/.cache/uv \
470
    uv pip install --system -e tests/vllm_test_utils
youkaichao's avatar
youkaichao committed
471

472
# enable fast downloads from hf (for testing)
473
RUN --mount=type=cache,target=/root/.cache/uv \
474
    uv pip install --system hf_transfer
475
476
ENV HF_HUB_ENABLE_HF_TRANSFER 1

Joe Runde's avatar
Joe Runde committed
477
# Copy in the v1 package for testing (it isn't distributed yet)
478
COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
Joe Runde's avatar
Joe Runde committed
479

480
481
482
483
484
# doc requires source code
# we hide them inside `test_docs/` , so that this source code
# will not be imported by other tests
RUN mkdir test_docs
RUN mv docs test_docs/
485
RUN cp -r examples test_docs/
486
RUN mv vllm test_docs/
487
RUN mv mkdocs.yaml test_docs/
488
#################### TEST IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
489

Simon Mo's avatar
Simon Mo committed
490
#################### OPENAI API SERVER ####################
491
492
# base openai image with additional requirements, for any subsequent openai-style images
FROM vllm-base AS vllm-openai-base
493
ARG TARGETPLATFORM
494
ARG INSTALL_KV_CONNECTORS=false
495

496
497
498
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

499
500
501
502
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
ENV UV_HTTP_TIMEOUT=500

503
504
COPY requirements/kv_connectors.txt requirements/kv_connectors.txt

505
# install additional dependencies for openai api server
506
RUN --mount=type=cache,target=/root/.cache/uv \
507
508
509
    if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
        uv pip install --system -r requirements/kv_connectors.txt; \
    fi; \
510
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
511
        uv pip install --system accelerate hf_transfer 'modelscope!=1.15.0' 'bitsandbytes>=0.42.0' 'timm==0.9.10' boto3 runai-model-streamer runai-model-streamer[s3]; \
512
    else \
513
        uv pip install --system accelerate hf_transfer 'modelscope!=1.15.0' 'bitsandbytes>=0.46.1' 'timm==0.9.10' boto3 runai-model-streamer runai-model-streamer[s3]; \
514
    fi
515

yhu422's avatar
yhu422 committed
516
517
ENV VLLM_USAGE_SOURCE production-docker-image

518
519
520
# define sagemaker first, so it is not default from `docker build`
FROM vllm-openai-base AS vllm-sagemaker

521
COPY examples/online_serving/sagemaker-entrypoint.sh .
522
523
524
525
526
RUN chmod +x sagemaker-entrypoint.sh
ENTRYPOINT ["./sagemaker-entrypoint.sh"]

FROM vllm-openai-base AS vllm-openai

Stephen Krider's avatar
Stephen Krider committed
527
ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
Simon Mo's avatar
Simon Mo committed
528
#################### OPENAI API SERVER ####################