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

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

Huy Do's avatar
Huy Do committed
8
ARG CUDA_VERSION=12.8.1
9
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
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}

Simon Mo's avatar
Simon Mo committed
65
#################### BASE BUILD IMAGE ####################
66
# prepare basic build environment
67
68
69
FROM ${BUILD_BASE_IMAGE} AS base
ARG CUDA_VERSION
ARG PYTHON_VERSION
70
ARG TARGETPLATFORM
71
72
ENV DEBIAN_FRONTEND=noninteractive

73
74
75
76
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL

77
78
79
80
81
# 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 \
82
83
84
85
86
87
88
89
90
91
92
93
94
    && 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 \
95
96
97
98
99
    && 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 \
100
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
101
    && python3 --version && python3 -m pip --version
102
103
104
105
106
107
108

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

109
110
111
# 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
112

113
114
115
# 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
116
ENV UV_INDEX_STRATEGY="unsafe-best-match"
117

118
119
120
121
122
123
124
125
# 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

126
127
128
129
# 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.
130
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
131

Stephen Krider's avatar
Stephen Krider committed
132
133
WORKDIR /workspace

134
135
# install build and runtime dependencies

136
137
138
139
# 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
140
RUN --mount=type=cache,target=/root/.cache/uv \
141
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
142
143
144
145
146
147
        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; \
148
    fi
Mor Zusman's avatar
Mor Zusman committed
149

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

156
157
158
159
# 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
160
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
161
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
162
163
164
# 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
165
166
#################### BASE BUILD IMAGE ####################

167
#################### WHEEL BUILD IMAGE ####################
168
FROM base AS build
169
ARG TARGETPLATFORM
170

171
172
173
174
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL

175
# install build dependencies
176
COPY requirements/build.txt requirements/build.txt
177

178
179
180
# 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
181
ENV UV_INDEX_STRATEGY="unsafe-best-match"
182

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

187
COPY . .
188
189
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
190
    if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
Stephen Krider's avatar
Stephen Krider committed
191
192

# max jobs used by Ninja to build extensions
193
194
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}
195
196
197
# number of threads used by nvcc
ARG nvcc_threads=8
ENV NVCC_THREADS=$nvcc_threads
198

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

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

236
# Check the size of the wheel if RUN_WHEEL_CHECK is true
237
COPY .buildkite/check-wheel-size.py check-wheel-size.py
238
# sync the default value with .buildkite/check-wheel-size.py
239
ARG VLLM_MAX_SIZE_MB=400
240
241
242
243
244
245
246
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
247
#################### EXTENSION Build IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
248

249
250
251
#################### DEV IMAGE ####################
FROM base as dev

252
253
254
255
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL

256
257
258
# 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
259
260
261
262
263
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"
264

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

273
274
#################### vLLM installation IMAGE ####################
# image with vLLM installed
275
# TODO: Restore to base image after FlashInfer AOT wheel fixed
276
277
278
FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
Simon Mo's avatar
Simon Mo committed
279
WORKDIR /vllm-workspace
280
ENV DEBIAN_FRONTEND=noninteractive
281
282
ARG TARGETPLATFORM

Kebe's avatar
Kebe committed
283
284
SHELL ["/bin/bash", "-c"]

285
286
287
288
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL

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

292
293
294
295
296
297
# 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 \
298
299
300
301
302
303
304
305
306
307
308
309
310
    && 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 \
311
312
313
314
315
    && 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 \
316
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
317
    && python3 --version && python3 -m pip --version
318
319
320
321
322
323
324

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

325
326
327
# Install uv for faster pip installs
RUN --mount=type=cache,target=/root/.cache/uv \
    python3 -m pip install uv
328

329
330
331
# 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
332
ENV UV_INDEX_STRATEGY="unsafe-best-match"
333

334
335
336
337
# 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.
338
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
339

340
341
342
343
# 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
344
RUN --mount=type=cache,target=/root/.cache/uv \
345
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
346
347
348
349
350
351
        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 ; \
352
353
    fi

354
# Install vllm wheel first, so that torch etc will be installed.
355
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
356
    --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
357
    uv pip install --system dist/*.whl --verbose \
358
        --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
359

360
# If we need to build FlashInfer wheel before its release:
361
# $ # Note we remove 7.0 from the arch list compared to the list below, since FlashInfer only supports sm75+
362
# $ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0'
363
364
# $ git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
# $ cd flashinfer
365
366
367
368
369
370
# $ 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
371

372
373
374
375
376
# 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"
ARG FLASHINFER_GIT_REF="v0.2.6.post1"
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
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
      if [[ "$CUDA_VERSION" == 12.8* ]]; then
          uv pip install --system ${FLASHINFER_CUDA128_INDEX_URL}/${FLASHINFER_CUDA128_WHEEL}
      else
          export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0'
          git clone ${FLASHINFER_GIT_REPO} --single-branch --branch ${FLASHINFER_GIT_REF} --recursive
          # Needed to build AOT kernels
          (cd flashinfer && \
              python3 -m flashinfer.aot && \
              uv pip install --system --no-build-isolation . \
          )
          rm -rf flashinfer

          # Default arches (skipping 10.0a and 12.0 since these need 12.8)
          # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
          TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
          if [[ "${CUDA_VERSION}" == 11.* ]]; then
              TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
          fi
          echo "🏗️  Building FlashInfer for arches: ${TORCH_CUDA_ARCH_LIST}"

          git clone --depth 1 --recursive --shallow-submodules \
            --branch v0.2.6.post1 \
            https://github.com/flashinfer-ai/flashinfer.git flashinfer

          pushd flashinfer
            python3 -m flashinfer.aot
            TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
              uv pip install --system --no-build-isolation .
          popd

          rm -rf flashinfer
      fi \
  fi
BASH
415
COPY examples examples
416
COPY benchmarks benchmarks
417
COPY ./vllm/collect_env.py .
418

Huy Do's avatar
Huy Do committed
419
420
421
422
RUN --mount=type=cache,target=/root/.cache/uv \
. /etc/environment && \
uv pip list

423
# Even when we build Flashinfer with AOT mode, there's still
424
425
426
# 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
427
COPY requirements/build.txt requirements/build.txt
428
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
429
    uv pip install --system -r requirements/build.txt \
430
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
431

432
#################### vLLM installation IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
433

434
435
436
437
#################### 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
438

439
ADD . /vllm-workspace/
Stephen Krider's avatar
Stephen Krider committed
440

441
442
443
444
445
ARG PYTHON_VERSION

ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

446
447
448
# 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
449
ENV UV_INDEX_STRATEGY="unsafe-best-match"
450

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

# install development dependencies (for testing)
456
RUN --mount=type=cache,target=/root/.cache/uv \
457
458
459
460
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
    if [ "$CUDA_MAJOR" -ge 12 ]; then \
        uv pip install --system -r requirements/dev.txt; \
    fi
461

youkaichao's avatar
youkaichao committed
462
# install development dependencies (for testing)
463
RUN --mount=type=cache,target=/root/.cache/uv \
464
    uv pip install --system -e tests/vllm_test_utils
youkaichao's avatar
youkaichao committed
465

466
# enable fast downloads from hf (for testing)
467
RUN --mount=type=cache,target=/root/.cache/uv \
468
    uv pip install --system hf_transfer
469
470
ENV HF_HUB_ENABLE_HF_TRANSFER 1

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

474
475
476
477
478
# 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/
479
RUN cp -r examples test_docs/
480
RUN mv vllm test_docs/
481
RUN mv mkdocs.yaml test_docs/
482
#################### TEST IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
483

Simon Mo's avatar
Simon Mo committed
484
#################### OPENAI API SERVER ####################
485
486
# base openai image with additional requirements, for any subsequent openai-style images
FROM vllm-base AS vllm-openai-base
487
ARG TARGETPLATFORM
488

489
490
491
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

492
493
494
495
# 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

496
# install additional dependencies for openai api server
497
RUN --mount=type=cache,target=/root/.cache/uv \
498
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
499
        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]; \
500
    else \
501
        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]; \
502
    fi
503

yhu422's avatar
yhu422 committed
504
505
ENV VLLM_USAGE_SOURCE production-docker-image

506
507
508
# define sagemaker first, so it is not default from `docker build`
FROM vllm-openai-base AS vllm-sagemaker

509
COPY examples/online_serving/sagemaker-entrypoint.sh .
510
511
512
513
514
RUN chmod +x sagemaker-entrypoint.sh
ENTRYPOINT ["./sagemaker-entrypoint.sh"]

FROM vllm-openai-base AS vllm-openai

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