ci(docker): make the fw_base vLLM wheel build optional (#5725)

Signed-off-by: Dong Hyuk Chang <9426164+thomasdhc@users.noreply.github.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Author: thomasdhc

docker/Dockerfile.fw_base

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 # Build layer variables
-# Supported mode: FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder
+# Supported mode: FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder,
+# VLLM_WHEEL_SRC=vllm_wheel_build|vllm_wheel_none
 # TRT-LLM stages are not present in this Dockerfile.
 ARG FW_DEP_BUILDER
 ARG FW_BASE_FINAL
+ARG VLLM_WHEEL_SRC=vllm_wheel_build
 
 ARG NEMO_FW_BASE_IMAGE
 FROM ${NEMO_FW_BASE_IMAGE} AS base
@@ -31,7 +33,7 @@ ENV NVIDIA_PRODUCT_NAME="NeMo Framework" \
 ##
 ##############################################################################
 
-FROM base AS vllm_wheel
+FROM base AS vllm_wheel_build
 
 ARG VLLM_VERSION=v0.14.1
 ARG MAX_JOBS=4
@@ -47,6 +49,13 @@ RUN mkdir -p /src/vllm && \
     pip install -r requirements/build/cuda.txt && \
     pip wheel --no-deps --no-build-isolation -v .
 
+# /src/vllm must exist for the install step's bind mount.
+FROM base AS vllm_wheel_none
+
+RUN mkdir -p /src/vllm
+
+FROM ${VLLM_WHEEL_SRC} AS vllm_wheel
+
 ##############################################################################
 ##
 ## Top layer of FW_BASE container
@@ -127,10 +136,14 @@ RUN --mount=type=bind,source=docker/common/install_nsys.sh,target=/opt/install_n
 # Copy Wheels
 RUN --mount=type=bind,from=vllm_wheel,source=/src/vllm/,target=/tmp/vllm/ \
     --mount=type=bind,source=docker/patches/vllm.patch,target=/opt/vllm.patch \
+    if ! compgen -G "/tmp/vllm/vllm*.whl" >/dev/null; then \
+    echo "No vLLM wheel produced; skipping vLLM install."; \
+    else \
     pip install /tmp/vllm/vllm*.whl && \
     pushd /usr/local/lib/python3.12/dist-packages/vllm && \
     patch -p1 < /opt/vllm.patch && \
-    popd
+    popd; \
+    fi
 
 ##############################################################################
 ##

docker/README.md

@@ -132,6 +132,7 @@ docker build \
 | `FW_BASE_FINAL` | Output stage. `trtllm_install` (with TRT-LLM) or `fw_toolkit_builder` (without) |
 | `UV_VERSION` | uv version to install |
 | `VLLM_VERSION` | vLLM git tag to build |
+| `VLLM_WHEEL_SRC` | Stage supplying the vLLM wheel. `vllm_wheel_build` (default) builds it from source; `vllm_wheel_none` skips both the build and the install |
 | `TRT_LLM_COMMIT` | TensorRT-LLM git commit or tag |
 | `TRT_LLM_VERSION` | TensorRT-LLM version string embedded as an image environment variable |
 | `TRT_VER` | TensorRT version for the TRT-LLM install scripts |