Dockerfile 22.4 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
382
# 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"
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
415
416
417
418
419
420
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
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 ####################