3c545c0c3b · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -6,30 +6,106 @@

6 # docs/assets/contributing/dockerfile-stages-dependency.png

7

8 ARG CUDA_VERSION=12.8.1

9 +ARG PYTHON_VERSION=3.12

10 +

11 +# By parameterizing the base images, we allow third-party to use their own

12 +# base images. One use case is hermetic builds with base images stored in

13 +# private registries that use a different repository naming conventions.

14 +#

15 +# Example:

16 +# docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04

17 +ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04

18 +ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04

19 +

20 +# By parameterizing the Deadsnakes repository URL, we allow third-party to use

21 +# their own mirror. When doing so, we don't benefit from the transparent

22 +# installation of the GPG key of the PPA, as done by add-apt-repository, so we

23 +# also need a URL for the GPG key.

24 +ARG DEADSNAKES_MIRROR_URL

25 +ARG DEADSNAKES_GPGKEY_URL

26 +

27 +# The PyPA get-pip.py script is a self contained script+zip file, that provides

28 +# both the installer script and the pip base85-encoded zip archive. This allows

29 +# bootstrapping pip in environment where a dsitribution package does not exist.

30 +#

31 +# By parameterizing the URL for get-pip.py installation script, we allow

32 +# third-party to use their own copy of the script stored in a private mirror.

33 +# We set the default value to the PyPA owned get-pip.py script.

34 +#

35 +# Reference: https://pip.pypa.io/en/stable/installation/#get-pip-py

36 +ARG GET_PIP_URL="https://bootstrap.pypa.io/get-pip.py"

37 +

38 +# PIP supports fetching the packages from custom indexes, allowing third-party

39 +# to host the packages in private mirrors. The PIP_INDEX_URL and

40 +# PIP_EXTRA_INDEX_URL are standard PIP environment variables to override the

41 +# default indexes. By letting them empty by default, PIP will use its default

42 +# indexes if the build process doesn't override the indexes.

43 +#

44 +# Uv uses different variables. We set them by default to the same values as

45 +# PIP, but they can be overridden.

46 +ARG PIP_INDEX_URL

47 +ARG PIP_EXTRA_INDEX_URL

48 +ARG UV_INDEX_URL=${PIP_INDEX_URL}

49 +ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}

50 +

51 +# PyTorch provides its own indexes for standard and nightly builds

52 +ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl

53 +ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly

54 +

55 +# PIP supports multiple authentication schemes, including keyring

56 +# By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to

57 +# disabled by default, we allow third-party to use keyring authentication for

58 +# their private Python indexes, while not changing the default behavior which

59 +# is no authentication.

60 +#

61 +# Reference: https://pip.pypa.io/en/stable/topics/authentication/#keyring-support

62 +ARG PIP_KEYRING_PROVIDER=disabled

63 +ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER}

64 +

65 #################### BASE BUILD IMAGE ####################

66 # prepare basic build environment

11 -FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04 AS base

12 -ARG CUDA_VERSION=12.8.1

13 -ARG PYTHON_VERSION=3.12

67 +FROM ${BUILD_BASE_IMAGE} AS base

68 +ARG CUDA_VERSION

69 +ARG PYTHON_VERSION

70 ARG TARGETPLATFORM

71 ENV DEBIAN_FRONTEND=noninteractive

72

73 +ARG DEADSNAKES_MIRROR_URL

74 +ARG DEADSNAKES_GPGKEY_URL

75 +ARG GET_PIP_URL

76 +

77 # Install Python and other dependencies

78 RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \

79 && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \

80 && apt-get update -y \

81 && apt-get install -y ccache software-properties-common git curl sudo \

22 - && for i in 1 2 3; do \

23 - add-apt-repository -y ppa:deadsnakes/ppa && break || \

24 - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \

25 - done \

82 + && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \

83 + if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \

84 + mkdir -p -m 0755 /etc/apt/keyrings ; \

85 + curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \

86 + sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \

87 + echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \

88 + fi ; \

89 + else \

90 + for i in 1 2 3; do \

91 + add-apt-repository -y ppa:deadsnakes/ppa && break || \

92 + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \

93 + done ; \

94 + fi \

95 && apt-get update -y \

96 && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \

97 && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \

98 && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \

99 && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \

31 - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \

100 + && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \

101 && python3 --version && python3 -m pip --version

102 +

103 +ARG PIP_INDEX_URL UV_INDEX_URL

104 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

105 +ARG PYTORCH_CUDA_INDEX_BASE_URL

106 +ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL

107 +ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

108 +

109 # Install uv for faster pip installs

110 RUN --mount=type=cache,target=/root/.cache/uv \

111 python3 -m pip install uv

@@ -63,15 +139,19 @@ WORKDIR /workspace

139 # after this step

140 RUN --mount=type=cache,target=/root/.cache/uv \

141 if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \

66 - uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319"; \

67 - uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \

142 + uv pip install --system \

143 + --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \

144 + "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319"; \

145 + uv pip install --system \

146 + --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \

147 + --pre pytorch_triton==3.3.0+gitab727c40; \

148 fi

149

150 COPY requirements/common.txt requirements/common.txt

151 COPY requirements/cuda.txt requirements/cuda.txt

152 RUN --mount=type=cache,target=/root/.cache/uv \

153 uv pip install --system -r requirements/cuda.txt \

74 - --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

154 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

155

156 # cuda arch list used by torch

157 # can be useful for both `dev` and `test`

@@ -88,6 +168,10 @@ ENV VLLM_FA_CMAKE_GPU_ARCHES=${vllm_fa_cmake_gpu_arches}

168 FROM base AS build

169 ARG TARGETPLATFORM

170

171 +ARG PIP_INDEX_URL UV_INDEX_URL

172 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

173 +ARG PYTORCH_CUDA_INDEX_BASE_URL

174 +

175 # install build dependencies

176 COPY requirements/build.txt requirements/build.txt

177

@@ -98,7 +182,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"

182

183 RUN --mount=type=cache,target=/root/.cache/uv \

184 uv pip install --system -r requirements/build.txt \

101 - --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

185 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

186

187 COPY . .

188 ARG GIT_REPO_CHECK=0

@@ -113,6 +197,8 @@ ARG nvcc_threads=8

197 ENV NVCC_THREADS=$nvcc_threads

198

199 ARG USE_SCCACHE

200 +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

201 +ARG SCCACHE_ENDPOINT

202 ARG SCCACHE_BUCKET_NAME=vllm-build-sccache

203 ARG SCCACHE_REGION_NAME=us-west-2

204 ARG SCCACHE_S3_NO_CREDENTIALS=0

@@ -121,10 +207,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \

207 --mount=type=bind,source=.git,target=.git \

208 if [ "$USE_SCCACHE" = "1" ]; then \

209 echo "Installing sccache..." \

124 - && curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \

210 + && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \

211 && tar -xzf sccache.tar.gz \

212 && sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \

213 && 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 && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \

216 && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \

217 && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \

@@ -162,6 +249,10 @@ RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \

249 #################### DEV IMAGE ####################

250 FROM base as dev

251

252 +ARG PIP_INDEX_URL UV_INDEX_URL

253 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

254 +ARG PYTORCH_CUDA_INDEX_BASE_URL

255 +

256 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out

257 # Reference: https://github.com/astral-sh/uv/pull/1694

258 ENV UV_HTTP_TIMEOUT=500

@@ -176,21 +267,25 @@ COPY requirements/test.txt requirements/test.txt

267 COPY requirements/dev.txt requirements/dev.txt

268 RUN --mount=type=cache,target=/root/.cache/uv \

269 uv pip install --system -r requirements/dev.txt \

179 - --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

270 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

271 #################### DEV IMAGE ####################

272

273 #################### vLLM installation IMAGE ####################

274 # image with vLLM installed

275 # TODO: Restore to base image after FlashInfer AOT wheel fixed

185 -FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS vllm-base

186 -ARG CUDA_VERSION=12.8.1

187 -ARG PYTHON_VERSION=3.12

276 +FROM ${FINAL_BASE_IMAGE} AS vllm-base

277 +ARG CUDA_VERSION

278 +ARG PYTHON_VERSION

279 WORKDIR /vllm-workspace

280 ENV DEBIAN_FRONTEND=noninteractive

281 ARG TARGETPLATFORM

282

283 SHELL ["/bin/bash", "-c"]

284

285 +ARG DEADSNAKES_MIRROR_URL

286 +ARG DEADSNAKES_GPGKEY_URL

287 +ARG GET_PIP_URL

288 +

289 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \

290 echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment

291

@@ -200,17 +295,33 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \

295 && apt-get update -y \

296 && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \

297 && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \

203 - && for i in 1 2 3; do \

204 - add-apt-repository -y ppa:deadsnakes/ppa && break || \

205 - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \

206 - done \

298 + && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \

299 + if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \

300 + mkdir -p -m 0755 /etc/apt/keyrings ; \

301 + curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \

302 + sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \

303 + echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \

304 + fi ; \

305 + else \

306 + for i in 1 2 3; do \

307 + add-apt-repository -y ppa:deadsnakes/ppa && break || \

308 + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \

309 + done ; \

310 + fi \

311 && apt-get update -y \

312 && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \

313 && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \

314 && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \

315 && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \

212 - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \

316 + && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \

317 && python3 --version && python3 -m pip --version

318 +

319 +ARG PIP_INDEX_URL UV_INDEX_URL

320 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

321 +ARG PYTORCH_CUDA_INDEX_BASE_URL

322 +ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL

323 +ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

324 +

325 # Install uv for faster pip installs

326 RUN --mount=type=cache,target=/root/.cache/uv \

327 python3 -m pip install uv

@@ -232,15 +343,19 @@ RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/

343 # after this step

344 RUN --mount=type=cache,target=/root/.cache/uv \

345 if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \

235 - uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319"; \

236 - uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \

346 + uv pip install --system \

347 + --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \

348 + "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \

349 + uv pip install --system \

350 + --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \

351 + --pre pytorch_triton==3.3.0+gitab727c40 ; \

352 fi

353

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 \

357 uv pip install --system dist/*.whl --verbose \

243 - --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

358 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

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+

@@ -254,15 +369,20 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist

369 # -rw-rw-r-- 1 mgoin mgoin 205M Jun 9 18:03 flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl

370 # $ # 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 +# Allow specifying a version, Git revision or local .whl file

373 +ARG FLASHINFER_CUDA128_INDEX_URL="https://download.pytorch.org/whl/cu128/flashinfer"

374 +ARG FLASHINFER_CUDA128_WHEEL="flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl"

375 +ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"

376 +ARG FLASHINFER_GIT_REF="v0.2.6.post1"

377 RUN --mount=type=cache,target=/root/.cache/uv \

378 . /etc/environment && \

379 if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \

380 # FlashInfer already has a wheel for PyTorch 2.7.0 and CUDA 12.8. This is enough for CI use

381 if [[ "$CUDA_VERSION" == 12.8* ]]; then \

262 - uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl; \

382 + uv pip install --system ${FLASHINFER_CUDA128_INDEX_URL}/${FLASHINFER_CUDA128_WHEEL} ; \

383 else \

384 export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0' && \

265 - git clone https://github.com/flashinfer-ai/flashinfer.git --single-branch --branch v0.2.6.post1 --recursive && \

385 + git clone ${FLASHINFER_GIT_REPO} --single-branch --branch ${FLASHINFER_GIT_REF} --recursive && \

386 # Needed to build AOT kernels

387 (cd flashinfer && \

388 python3 -m flashinfer.aot && \

@@ -286,7 +406,7 @@ uv pip list

406 COPY requirements/build.txt requirements/build.txt

407 RUN --mount=type=cache,target=/root/.cache/uv \

408 uv pip install --system -r requirements/build.txt \

289 - --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

409 + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

410

411 #################### vLLM installation IMAGE ####################

412

@@ -297,6 +417,11 @@ FROM vllm-base AS test

417

418 ADD . /vllm-workspace/

419

420 +ARG PYTHON_VERSION

421 +

422 +ARG PIP_INDEX_URL UV_INDEX_URL

423 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

424 +

425 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out

426 # Reference: https://github.com/astral-sh/uv/pull/1694

427 ENV UV_HTTP_TIMEOUT=500

@@ -307,7 +432,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \

432 uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"

433

434 # install development dependencies (for testing)

310 -RUN --mount=type=cache,target=/root/.cache/uv \

435 +RUN --mount=type=cache,target=/root/.cache/uv \

436 CUDA_MAJOR="${CUDA_VERSION%%.*}"; \

437 if [ "$CUDA_MAJOR" -ge 12 ]; then \

438 uv pip install --system -r requirements/dev.txt; \

@@ -323,7 +448,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \

448 ENV HF_HUB_ENABLE_HF_TRANSFER 1

449

450 # Copy in the v1 package for testing (it isn't distributed yet)

326 -COPY vllm/v1 /usr/local/lib/python3.12/dist-packages/vllm/v1

451 +COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1

452

453 # doc requires source code

454 # we hide them inside `test_docs/` , so that this source code

@@ -340,6 +465,9 @@ RUN mv mkdocs.yaml test_docs/

465 FROM vllm-base AS vllm-openai-base

466 ARG TARGETPLATFORM

467

468 +ARG PIP_INDEX_URL UV_INDEX_URL

469 +ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL

470 +

471 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out

472 # Reference: https://github.com/astral-sh/uv/pull/1694

473 ENV UV_HTTP_TIMEOUT=500

[CI/Build] Allow hermetic builds (#18064)