A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 18:15 UTC
[CI][torch nightlies] Use main Dockerfile with flags for nightly torch tests (#30443)
Signed-off-by: Orion Reblitz-Richardson <orionr@meta.com> Signed-off-by: Orion Reblitz-Richardson <orionr@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Orion Reblitz-Richardson committed
Jan 23, 2026 at 08:22 UTC
68b0a6c1baf32221fa8ed98a941f611cbff2cdb1
4 files changed
+203
-32
docker/Dockerfile
+146
-19
@@ -148,12 +148,36 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
148
149
WORKDIR /workspace
150
151
-# install build and runtime dependencies
151
+# We can specify the standard or nightly build of PyTorch
152
+ARG PYTORCH_NIGHTLY
153
+
154
+# Install build and runtime dependencies, including PyTorch
155
+# Check whether to install torch nightly instead of release for this build
156
COPY requirements/common.txt requirements/common.txt
157
COPY requirements/cuda.txt requirements/cuda.txt
158
+COPY use_existing_torch.py use_existing_torch.py
159
+COPY pyproject.toml pyproject.toml
160
+RUN --mount=type=cache,target=/root/.cache/uv \
161
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
162
+ echo "Installing torch nightly..." \
163
+ && uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
164
+ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
165
+ && echo "Installing other requirements..." \
166
+ && /opt/venv/bin/python3 use_existing_torch.py --prefix \
167
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
168
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
169
+ else \
170
+ uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
171
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
172
+ fi
173
+
174
+# Track PyTorch lib versions used during build and match in downstream instances.
175
+# We do this for both nightly and release so we can strip dependencies/*.txt as needed.
176
+# Otherwise library dependencies can upgrade/downgrade torch incorrectly.
177
RUN --mount=type=cache,target=/root/.cache/uv \
155
- uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
156
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
178
+ uv pip freeze | grep -i "^torch=\|^torchvision=\|^torchaudio=" > torch_lib_versions.txt \
179
+ && TORCH_LIB_VERSIONS=$(cat torch_lib_versions.txt | xargs) \
180
+ && echo "Installed torch libs: ${TORCH_LIB_VERSIONS}"
181
182
# CUDA arch list used by torch
183
# Explicitly set the list to avoid issues with torch 2.2
@@ -171,8 +195,13 @@ ARG PIP_INDEX_URL UV_INDEX_URL
195
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
196
ARG PYTORCH_CUDA_INDEX_BASE_URL
197
174
-# install build dependencies
198
+# We can specify the standard or nightly build of PyTorch
199
+ARG PYTORCH_NIGHTLY
200
+
201
+# Install build dependencies
202
COPY requirements/build.txt requirements/build.txt
203
+COPY use_existing_torch.py use_existing_torch.py
204
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
205
206
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
207
# Reference: https://github.com/astral-sh/uv/pull/1694
@@ -182,8 +211,18 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
211
ENV UV_LINK_MODE=copy
212
213
RUN --mount=type=cache,target=/root/.cache/uv \
185
- uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
186
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
214
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
215
+ echo "Installing build requirements without torch..." \
216
+ && python3 use_existing_torch.py --prefix \
217
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
218
+ && echo "Installing torch nightly..." \
219
+ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \
220
+ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
221
+ else \
222
+ echo "Installing build requirements..." \
223
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
224
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
225
+ fi
226
227
WORKDIR /workspace
228
@@ -215,6 +254,13 @@ ARG VLLM_MAIN_CUDA_VERSION=""
254
# Use dummy version for csrc-build wheel (only .so files are extracted, version doesn't matter)
255
ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+csrc.build"
256
257
+# Use existing torch for nightly builds
258
+RUN --mount=type=cache,target=/root/.cache/uv \
259
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
260
+ python3 use_existing_torch.py --prefix; \
261
+ fi
262
+
263
+# Build the vLLM wheel
264
# if USE_SCCACHE is set, use sccache to speed up compilation
265
RUN --mount=type=cache,target=/root/.cache/uv \
266
if [ "$USE_SCCACHE" = "1" ]; then \
@@ -258,6 +304,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
304
export VLLM_DOCKER_BUILD_CONTEXT=1 && \
305
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
306
fi
307
+
308
#################### CSRC BUILD IMAGE ####################
309
310
#################### EXTENSIONS BUILD IMAGE ####################
@@ -314,8 +361,13 @@ ARG PIP_INDEX_URL UV_INDEX_URL
361
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
362
ARG PYTORCH_CUDA_INDEX_BASE_URL
363
317
-# install build dependencies
364
+# We can specify the standard or nightly build of PyTorch
365
+ARG PYTORCH_NIGHTLY
366
+
367
+# Install build dependencies
368
COPY requirements/build.txt requirements/build.txt
369
+COPY use_existing_torch.py use_existing_torch.py
370
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
371
372
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
373
# Reference: https://github.com/astral-sh/uv/pull/1694
@@ -325,14 +377,23 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
377
ENV UV_LINK_MODE=copy
378
379
RUN --mount=type=cache,target=/root/.cache/uv \
328
- uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
329
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
380
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
381
+ echo "Installing build requirements without torch..." \
382
+ && python3 use_existing_torch.py --prefix \
383
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
384
+ && echo "Installing torch nightly..." \
385
+ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \
386
+ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
387
+ else \
388
+ echo "Installing build requirements..." \
389
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \
390
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
391
+ fi
392
393
WORKDIR /workspace
394
395
# Copy pre-built csrc wheel directly
396
COPY --from=csrc-build /workspace/dist /precompiled-wheels
335
-
397
COPY . .
398
399
ARG GIT_REPO_CHECK=0
@@ -345,6 +406,13 @@ ENV VLLM_TARGET_DEVICE=${vllm_target_device}
406
# Skip adding +precompiled suffix to version (preserves git-derived version)
407
ENV VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX=1
408
409
+# Use existing torch for nightly builds
410
+RUN --mount=type=cache,target=/root/.cache/uv \
411
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
412
+ python3 use_existing_torch.py --prefix; \
413
+ fi
414
+
415
+# Build the vLLM wheel
416
RUN --mount=type=cache,target=/root/.cache/uv \
417
--mount=type=bind,source=.git,target=.git \
418
if [ "${vllm_target_device}" = "cuda" ]; then \
@@ -367,7 +435,8 @@ RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
435
else \
436
echo "Skipping wheel size check."; \
437
fi
370
-#################### EXTENSION Build IMAGE ####################
438
+
439
+#################### WHEEL BUILD IMAGE ####################
440
441
#################### DEV IMAGE ####################
442
FROM base AS dev
@@ -385,12 +454,34 @@ ENV UV_LINK_MODE=copy
454
455
# Install libnuma-dev, required by fastsafetensors (fixes #20384)
456
RUN apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*
457
+
458
+
459
+# We can specify the standard or nightly build of PyTorch
460
+ARG PYTORCH_NIGHTLY
461
+
462
+# Install development dependencies
463
COPY requirements/lint.txt requirements/lint.txt
464
+COPY requirements/test.in requirements/test.in
465
COPY requirements/test.txt requirements/test.txt
466
COPY requirements/dev.txt requirements/dev.txt
467
+COPY use_existing_torch.py use_existing_torch.py
468
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
469
RUN --mount=type=cache,target=/root/.cache/uv \
392
- uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \
393
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
470
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
471
+ echo "Installing dev requirements plus torch nightly..." \
472
+ && python3 use_existing_torch.py --prefix \
473
+ && cat torch_lib_versions.txt >> requirements/test.in \
474
+ && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \
475
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
476
+ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | xargs) --pre \
477
+ -r requirements/dev.txt \
478
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
479
+ else \
480
+ echo "Installing dev requirements..." \
481
+ && uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \
482
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
483
+ fi
484
+
485
#################### DEV IMAGE ####################
486
#################### vLLM installation IMAGE ####################
487
# image with vLLM installed
@@ -548,11 +639,26 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
639
ARG PYTORCH_CUDA_INDEX_BASE_URL
640
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
641
551
-# Install vllm wheel first, so that torch etc will be installed.
642
+# We can specify the standard or nightly build of PyTorch
643
+ARG PYTORCH_NIGHTLY
644
+
645
+# Install vLLM wheel first, so that torch etc will be installed.
646
+# Check whether to install torch nightly instead of release for this build.
647
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
648
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
649
--mount=type=cache,target=/root/.cache/uv \
554
- uv pip install --system dist/*.whl --verbose \
555
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
650
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
651
+ echo "Installing torch nightly..." \
652
+ && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
653
+ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
654
+ && echo "Installing vLLM..." \
655
+ && uv pip install --system dist/*.whl --verbose \
656
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
657
+ else \
658
+ echo "Installing vLLM..." \
659
+ && uv pip install --system dist/*.whl --verbose \
660
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
661
+ fi
662
663
RUN --mount=type=cache,target=/root/.cache/uv \
664
. /etc/environment && \
@@ -612,12 +718,33 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
718
&& apt-get update -y \
719
&& apt-get install -y git
720
615
-# install development dependencies (for testing)
721
+# We can specify the standard or nightly build of PyTorch
722
+ARG PYTORCH_NIGHTLY
723
+
724
+# Install development dependencies (for testing)
725
+COPY requirements/lint.txt requirements/lint.txt
726
+COPY requirements/test.in requirements/test.in
727
+COPY requirements/test.txt requirements/test.txt
728
+COPY requirements/dev.txt requirements/dev.txt
729
+COPY use_existing_torch.py use_existing_torch.py
730
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
731
RUN --mount=type=cache,target=/root/.cache/uv \
732
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
733
if [ "$CUDA_MAJOR" -ge 12 ]; then \
619
- uv pip install --system -r requirements/dev.txt \
620
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
734
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
735
+ echo "Installing dev requirements plus torch nightly..." \
736
+ && python3 use_existing_torch.py --prefix \
737
+ && cat torch_lib_versions.txt >> requirements/test.in \
738
+ && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \
739
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
740
+ && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
741
+ -r requirements/dev.txt \
742
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
743
+ else \
744
+ echo "Installing dev requirements..." \
745
+ && uv pip install --system -r requirements/dev.txt \
746
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
747
+ fi \
748
fi
749
750
# install development dependencies (for testing)
docker/Dockerfile.nightly_torch
+8
@@ -1,3 +1,11 @@
1
+#######
2
+#
3
+# THIS FILE IS DEPRECATED AND WILL BE REMOVED SHORTLY
4
+#
5
+# Please use the standard Dockerfile with PYTORCH_NIGHTLY=1 instead
6
+#
7
+#######
8
+
9
# The vLLM Dockerfile is used to construct vLLM image against torch nightly that can be directly used for testing
10
11
# for torch nightly, cuda >=12.6 is required,
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
use_existing_torch.py
+49
-13
@@ -1,18 +1,54 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import argparse
5
import glob
6
+import sys
7
6
-for file in (*glob.glob("requirements/*.txt"), "pyproject.toml"):
7
- print(f">>> cleaning {file}")
8
- with open(file) as f:
9
- lines = f.readlines()
10
- if "torch" in "".join(lines).lower():
11
- print("removed:")
12
- with open(file, "w") as f:
13
- for line in lines:
14
- if "torch" not in line.lower():
15
- f.write(line)
16
- else:
17
- print(line.strip())
18
- print(f"<<< done cleaning {file}\n")
8
+# Only strip targeted libraries when checking prefix
9
+TORCH_LIB_PREFIXES = (
10
+ # requirements/*.txt/in
11
+ "torch=",
12
+ "torchvision=",
13
+ "torchaudio=",
14
+ # pyproject.toml
15
+ '"torch =',
16
+ '"torchvision =',
17
+ '"torchaudio =',
18
+)
19
+
20
+
21
+def main(argv):
22
+ parser = argparse.ArgumentParser(
23
+ description="Strip torch lib requirements to use installed version."
24
+ )
25
+ parser.add_argument(
26
+ "--prefix",
27
+ action="store_true",
28
+ help="Strip prefix matches only (default: False)",
29
+ )
30
+ args = parser.parse_args(argv)
31
+
32
+ for file in (
33
+ *glob.glob("requirements/*.txt"),
34
+ *glob.glob("requirements/*.in"),
35
+ "pyproject.toml",
36
+ ):
37
+ with open(file) as f:
38
+ lines = f.readlines()
39
+ if "torch" in "".join(lines).lower():
40
+ with open(file, "w") as f:
41
+ for line in lines:
42
+ if (
43
+ args.prefix
44
+ and not line.lower().strip().startswith(TORCH_LIB_PREFIXES)
45
+ or not args.prefix
46
+ and "torch" not in line.lower()
47
+ ):
48
+ f.write(line)
49
+ else:
50
+ print(f">>> removed from {file}:", line.strip())
51
+
52
+
53
+if __name__ == "__main__":
54
+ main(sys.argv[1:])