[Docker] Cache test dependencies before vLLM install (#51184)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Michael Goin committed
Aug 10, 2026 at 13:34 UTC
21c667aa64d4023320fb608b81a343e7f9a47250
4 files changed
+61
-6
docker/Dockerfile
+55
-6
@@ -568,6 +568,7 @@ RUN dnf install -y numactl-devel \
568
ARG PYTORCH_NIGHTLY
569
570
# Install development dependencies
571
+COPY requirements/common.txt requirements/common.txt
572
COPY requirements/lint.txt requirements/lint.txt
573
COPY requirements/test/cuda.in requirements/test/cuda.in
574
COPY requirements/test/cuda.txt requirements/test/cuda.txt
@@ -593,7 +594,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
594
#################### DEV IMAGE ####################
595
#################### vLLM installation IMAGE ####################
596
# image with vLLM installed
596
-FROM ${FINAL_BASE_IMAGE} AS vllm-base
597
+FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
598
599
ARG CUDA_VERSION
600
ARG PYTHON_VERSION
@@ -790,6 +791,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
791
# VLLM INSTALLATION (depends on build stage)
792
# ============================================================
793
794
+FROM vllm-runtime-base AS vllm-base
795
+
796
ARG PIP_INDEX_URL UV_INDEX_URL
797
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
798
ARG PYTORCH_CUDA_INDEX_BASE_URL
@@ -849,11 +852,8 @@ COPY benchmarks benchmarks
852
COPY ./vllm/collect_env.py .
853
#################### vLLM installation IMAGE ####################
854
#################### TEST IMAGE ####################
852
-# image to run unit testing suite
853
-# note that this uses vllm installed by `pip`
854
-FROM vllm-base AS test
855
-
856
-ADD . /vllm-workspace/
855
+# Stable dependencies for the unit testing image
856
+FROM vllm-runtime-base AS test-deps
857
858
ARG PYTHON_VERSION
859
ARG TARGETPLATFORM
@@ -861,6 +861,7 @@ ARG TARGETPLATFORM
861
ARG PIP_INDEX_URL UV_INDEX_URL
862
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
863
ARG PYTORCH_CUDA_INDEX_BASE_URL
864
+ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
865
866
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
867
# Reference: https://github.com/astral-sh/uv/pull/1694
@@ -876,6 +877,7 @@ RUN apt-get update -y \
877
ARG PYTORCH_NIGHTLY
878
879
# Install development dependencies (for testing)
880
+COPY requirements/common.txt requirements/common.txt
881
COPY requirements/lint.txt requirements/lint.txt
882
COPY requirements/test/cuda.in requirements/test/cuda.in
883
COPY requirements/test/cuda.txt requirements/test/cuda.txt
@@ -906,6 +908,53 @@ RUN --mount=type=cache,target=/opt/uv/cache \
908
fi \
909
fi
910
911
+# image to run unit testing suite
912
+# note that this uses vllm installed by `pip`
913
+FROM test-deps AS test
914
+
915
+ARG PYTHON_VERSION
916
+
917
+ARG PIP_INDEX_URL UV_INDEX_URL
918
+ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
919
+ARG PYTORCH_CUDA_INDEX_BASE_URL
920
+ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
921
+
922
+# We can specify the standard or nightly build of PyTorch
923
+ARG PYTORCH_NIGHTLY
924
+
925
+# Install vLLM and EP kernel wheels after the stable test dependencies so
926
+# source-only changes do not invalidate the dependency layers above.
927
+COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
928
+COPY --from=build /workspace/dist/wheel.sha256 /tmp/vllm-wheel.sha256
929
+RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
930
+ --mount=type=cache,target=/opt/uv/cache \
931
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
932
+ echo "Installing torch nightly..." \
933
+ && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \
934
+ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
935
+ && echo "Installing vLLM..." \
936
+ && uv pip install --system dist/*.whl --verbose \
937
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
938
+ else \
939
+ echo "Installing vLLM..." \
940
+ && uv pip install --system dist/*.whl --verbose \
941
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
942
+ fi
943
+
944
+# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
945
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
946
+
947
+COPY --from=build /tmp/ep_kernels_workspace/dist/wheels.sha256 /tmp/ep-kernels-wheels.sha256
948
+RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
949
+ --mount=type=cache,target=/opt/uv/cache \
950
+ uv pip install --system ep_kernels/dist/*.whl --verbose \
951
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
952
+
953
+# See the matching comment in the vllm-base stage.
954
+ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH}
955
+
956
+ADD . /vllm-workspace/
957
+
958
# install development dependencies (for testing)
959
RUN --mount=type=cache,target=/opt/uv/cache \
960
uv pip install --system -e tests/vllm_test_utils
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
docs/contributing/dockerfile/dockerfile.md
+5
@@ -17,6 +17,9 @@ The edges of the build graph represent:
17
18
- `RUN --mount=(.\*)from=...` dependencies (with a dotted line and an empty diamond arrow head)
19
20
+The `test-deps` stage branches from `vllm-runtime-base` so Git and the test
21
+requirements remain cached independently of per-commit vLLM wheels and source.
22
+
23
> <figure markdown="span">
24
> { align="center" alt="query" width="100%" }
25
> </figure>
@@ -28,6 +31,7 @@ The edges of the build graph represent:
31
> ```bash
32
> dockerfilegraph \
33
> -o png \
34
+ > --concentrate \
35
> --legend \
36
> --dpi 200 \
37
> --max-label-length 50 \
@@ -47,6 +51,7 @@ The edges of the build graph represent:
51
> --dpi 200 \
52
> --max-label-length 50 \
53
> --filename docker/Dockerfile \
54
+ > --concentrate \
55
> --legend
56
> ```
57
>
tools/pre_commit/update-dockerfile-graph.sh
+1
@@ -47,6 +47,7 @@ if printf '%s\n' "${FILES[@]}" | grep -q "^docker/Dockerfile$"; then
47
--dpi 200 \
48
--max-label-length 50 \
49
--filename docker/Dockerfile \
50
+ --concentrate \
51
--legend
52
53
echo "Finding generated PNG file..."