Dockerfile 23.6 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
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
18
# TODO: Restore to base image after FlashInfer AOT wheel fixed
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 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
# Flag enables built-in KV-connector dependency libs into docker images
67
68
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
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
124

125
126
127
128
129
130
131
132
# 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

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

Stephen Krider's avatar
Stephen Krider committed
139
140
WORKDIR /workspace

141
# install build and runtime dependencies
142
143
COPY requirements/common.txt requirements/common.txt
COPY requirements/cuda.txt requirements/cuda.txt
144
RUN --mount=type=cache,target=/root/.cache/uv \
Huy Do's avatar
Huy Do committed
145
    uv pip install --system -r requirements/cuda.txt \
146
    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
147

148
149
150
151
# 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
152
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
153
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
Simon Mo's avatar
Simon Mo committed
154
155
#################### BASE BUILD IMAGE ####################

156
#################### WHEEL BUILD IMAGE ####################
157
FROM base AS build
158
ARG TARGETPLATFORM
159

160
161
162
163
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ARG PYTORCH_CUDA_INDEX_BASE_URL

164
# install build dependencies
165
COPY requirements/build.txt requirements/build.txt
166

167
168
169
# 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
170
ENV UV_INDEX_STRATEGY="unsafe-best-match"
171
172
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
173

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

178
COPY . .
179
180
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
181
    if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
Stephen Krider's avatar
Stephen Krider committed
182
183

# max jobs used by Ninja to build extensions
184
185
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}
186
187
188
# number of threads used by nvcc
ARG nvcc_threads=8
ENV NVCC_THREADS=$nvcc_threads
189

190
ARG USE_SCCACHE
191
192
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
193
194
ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
ARG SCCACHE_REGION_NAME=us-west-2
195
ARG SCCACHE_S3_NO_CREDENTIALS=0
196
197

# Flag to control whether to use pre-built vLLM wheels
198
ARG VLLM_USE_PRECOMPILED=""
199

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

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

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

250
#################### DEV IMAGE ####################
251
FROM base AS dev
252

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

257
258
259
# 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
260
ENV UV_INDEX_STRATEGY="unsafe-best-match"
261
262
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
Huy Do's avatar
Huy Do committed
263

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

272
273
#################### vLLM installation IMAGE ####################
# image with vLLM installed
274
275
276
FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
277
ARG INSTALL_KV_CONNECTORS=false
Simon Mo's avatar
Simon Mo committed
278
WORKDIR /vllm-workspace
279
ENV DEBIAN_FRONTEND=noninteractive
280
281
ARG TARGETPLATFORM

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

284
285
286
287
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL

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

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

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

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

328
329
330
# 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
331
ENV UV_INDEX_STRATEGY="unsafe-best-match"
332
333
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
334

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

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

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

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

373
# Install FlashInfer from source
374
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
375
376
377
378
379
380
# Keep this in sync with "flashinfer" extra in setup.py
ARG FLASHINFER_GIT_REF="v0.2.12"
# Flag to control whether to compile FlashInfer AOT kernels
# Set to "true" to enable AOT compilation:
# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
ARG FLASHINFER_AOT_COMPILE=false
381
382
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
  . /etc/environment
383
384
385
386
    git clone --depth 1 --recursive --shallow-submodules \
        --branch ${FLASHINFER_GIT_REF} \
        ${FLASHINFER_GIT_REPO} flashinfer
    pushd flashinfer
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
        if [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
            # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
            # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
            if [[ "${CUDA_VERSION}" == 11.* ]]; then
                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"
            fi
            echo "🏗️  Installing FlashInfer with AOT compilation for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
            # Build AOT kernels
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
                python3 -m flashinfer.aot
            # Install with no-build-isolation since we already built AOT kernels
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
                uv pip install --system --no-build-isolation . \
                --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
            # Download pre-compiled cubins
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
                python3 -m flashinfer --download-cubin || echo "WARNING: Failed to download flashinfer cubins."
        else
            echo "🏗️  Installing FlashInfer without AOT compilation in JIT mode"
            uv pip install --system . \
                --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
        fi
414
415
    popd
    rm -rf flashinfer
416
BASH
417
COPY examples examples
418
COPY benchmarks benchmarks
419
COPY ./vllm/collect_env.py .
420

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

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

434
435
# Install DeepGEMM from source
ARG DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
436
ARG DEEPGEMM_GIT_REF="7b6b5563b9d4c1ae07ffbce7f78ad3ac9204827c"
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
  . /etc/environment
    CUDA_MAJOR="${CUDA_VERSION%%.*}"
    CUDA_MINOR="${CUDA_VERSION#${CUDA_MAJOR}.}"
    CUDA_MINOR="${CUDA_MINOR%%.*}"
    if [ "$CUDA_MAJOR" -ge 12 ] && [ "$CUDA_MINOR" -ge 8 ]; then
        git clone --recursive --shallow-submodules \
            ${DEEPGEMM_GIT_REPO} deepgemm
        echo "🏗️  Building DeepGEMM"
        pushd deepgemm
            git checkout ${DEEPGEMM_GIT_REF}
            # Build DeepGEMM
            # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh)
            rm -rf build dist
            rm -rf *.egg-info
            python3 setup.py bdist_wheel
            uv pip install --system dist/*.whl
        popd
        rm -rf deepgemm
    else
        echo "Skipping DeepGEMM installation (requires CUDA 12.8+ but got ${CUDA_VERSION})"
    fi
BASH

461
462
463
464
465
466
467
468
# Install EP kernels(pplx-kernels and DeepEP), NixL
COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
COPY tools/install_nixl.sh install_nixl.sh
ENV CUDA_HOME=/usr/local/cuda
RUN export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0a+PTX}" \
    && bash install_python_libraries.sh \
    && bash install_nixl.sh --force

469
#################### vLLM installation IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
470

471
472
473
474
#################### 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
475

476
ADD . /vllm-workspace/
Stephen Krider's avatar
Stephen Krider committed
477

478
479
480
481
482
ARG PYTHON_VERSION

ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

483
484
485
# 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
486
ENV UV_INDEX_STRATEGY="unsafe-best-match"
487
488
# Use copy mode to avoid hardlink failures with Docker cache mounts
ENV UV_LINK_MODE=copy
489

Huy Do's avatar
Huy Do committed
490
# install development dependencies (for testing)
491
RUN --mount=type=cache,target=/root/.cache/uv \
492
493
494
495
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
    if [ "$CUDA_MAJOR" -ge 12 ]; then \
        uv pip install --system -r requirements/dev.txt; \
    fi
496

youkaichao's avatar
youkaichao committed
497
# install development dependencies (for testing)
498
RUN --mount=type=cache,target=/root/.cache/uv \
499
    uv pip install --system -e tests/vllm_test_utils
youkaichao's avatar
youkaichao committed
500

501
# enable fast downloads from hf (for testing)
502
RUN --mount=type=cache,target=/root/.cache/uv \
503
    uv pip install --system hf_transfer
504
505
ENV HF_HUB_ENABLE_HF_TRANSFER 1

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

509
510
# Source code is used in the `python_only_compile.sh` test
# We hide it inside `src/` so that this source code
511
# will not be imported by other tests
512
513
RUN mkdir src
RUN mv vllm src/vllm
514
#################### TEST IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
515

Simon Mo's avatar
Simon Mo committed
516
#################### OPENAI API SERVER ####################
517
518
# base openai image with additional requirements, for any subsequent openai-style images
FROM vllm-base AS vllm-openai-base
519
ARG TARGETPLATFORM
520
ARG INSTALL_KV_CONNECTORS=false
521

522
523
524
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

525
526
527
528
# 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

529
530
COPY requirements/kv_connectors.txt requirements/kv_connectors.txt

531
# install additional dependencies for openai api server
532
RUN --mount=type=cache,target=/root/.cache/uv \
533
534
535
    if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
        uv pip install --system -r requirements/kv_connectors.txt; \
    fi; \
536
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
537
        BITSANDBYTES_VERSION="0.42.0"; \
538
    else \
539
540
        BITSANDBYTES_VERSION="0.46.1"; \
    fi; \
541
    uv pip install --system accelerate hf_transfer modelscope "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm==0.9.10' boto3 runai-model-streamer runai-model-streamer[s3]
542

yhu422's avatar
yhu422 committed
543
544
ENV VLLM_USAGE_SOURCE production-docker-image

545
546
547
# define sagemaker first, so it is not default from `docker build`
FROM vllm-openai-base AS vllm-sagemaker

548
COPY examples/online_serving/sagemaker-entrypoint.sh .
549
550
551
552
553
RUN chmod +x sagemaker-entrypoint.sh
ENTRYPOINT ["./sagemaker-entrypoint.sh"]

FROM vllm-openai-base AS vllm-openai

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