refactor(docker): streamline Dockerfile configurations and enhance environment variable management

- Consolidated environment variable definitions in Dockerfile.ci for improved readability.
- Updated Dockerfile.fw_base to set default VLLM_VERSION and reorganized the build process for clarity.
- Refactored Dockerfile.fw_final to normalize COPY commands and streamline repository cloning and commit updates.
- Enhanced overall structure and maintainability of Dockerfiles across the project.

Signed-off-by: Ajay Balasa <abalasa@nvidia.com>

Author: balasaajay

docker/Dockerfile.ci

@@ -15,34 +15,26 @@
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
 FROM ${BASE_IMAGE} AS megatron_bridge
 WORKDIR /opt/Megatron-Bridge
-ENV PATH="/root/.local/bin:$PATH"
-ENV UV_PROJECT_ENVIRONMENT=/opt/venv
-ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
-ENV VIRTUAL_ENV=/opt/venv
-ENV PATH="$VIRTUAL_ENV/bin:$PATH"
-ENV UV_LINK_MODE=copy
-ENV UV_VERSION="0.7.2"
-ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3
-RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
-    uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages && \
-    # Address CVE-2025-68973
-    apt-get update && apt install -y --only-upgrade gnupg && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
-##############################################################################
-##
-## Install DeepEP and nvshmem
-##
-##############################################################################
 
 ARG INSTALL_DEEPEP=True
 ARG DEEPEP_COMMIT=34152ae28f80bcc3ee38d7a12cb2ad87cfd4ea72
 
-ENV HYBRID_EP_MULTINODE=1
-ENV RDMA_CORE_HOME=/opt/rdma-core/build
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64/:$LD_LIBRARY_PATH
+ENV PATH="/opt/venv/bin:/opt/venv/bin:/root/.local/bin:$PATH" \
+    UV_PROJECT_ENVIRONMENT=/opt/venv \
+    VIRTUAL_ENV=/opt/venv \
+    UV_LINK_MODE=copy \
+    UV_VERSION="0.7.2" \
+    NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \
+    HYBRID_EP_MULTINODE=1 \
+    RDMA_CORE_HOME=/opt/rdma-core/build \
+    LD_LIBRARY_PATH="/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
 RUN --mount=type=bind,source=docker/patches/deepep.patch,target=/opt/deepep.patch \
+    curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
+    uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages && \
+    # Address CVE-2025-68973
+    apt-get update && apt install -y --only-upgrade gnupg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* && \
     if [ "$INSTALL_DEEPEP" = "True" ]; then \
     # Upgrade system rdma-core to v60; libibverbs-dev supplies the unversioned
     # libibverbs.so symlink required at link time.
@@ -78,12 +70,7 @@ RUN --mount=type=bind,source=docker/patches/deepep.patch,target=/opt/deepep.patc
     popd; \
     fi
 
-COPY pyproject.toml uv.lock /opt/Megatron-Bridge/
-COPY src/megatron/bridge/__init__.py src/megatron/bridge/package_info.py /opt/Megatron-Bridge/src/megatron/bridge/
-COPY 3rdparty/Megatron-LM/pyproject.toml 3rdparty/Megatron-LM/setup.py /opt/Megatron-Bridge/3rdparty/Megatron-LM/
-COPY 3rdparty/Megatron-LM/megatron/training/__init__.py /opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/training/
-COPY 3rdparty/Megatron-LM/megatron/core/__init__.py 3rdparty/Megatron-LM/megatron/core/package_info.py /opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/
-COPY 3rdparty/Megatron-LM/megatron/core/datasets/Makefile 3rdparty/Megatron-LM/megatron/core/datasets/helpers.cpp /opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/datasets/
+COPY --chmod=644 . /opt/Megatron-Bridge
 
 # Build arg to skip --locked when testing with different MCore versions
 ARG MCORE_TRIGGERED_TESTING=false
@@ -109,5 +96,3 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     # Otherwise, the stub will be called in some cases and fail
     ln -sf "$(ldconfig -p | awk '/libcudart\.so\.[0-9]+ /{print $NF; exit}')" /opt/venv/lib/python3.12/site-packages/tilelang/lib/libcudart_stub.so && \
     uv cache prune ${UV_CACHE_PRUNE_ARGS}
-
-COPY --chmod=644 . /opt/Megatron-Bridge

docker/Dockerfile.fw_base

@@ -13,19 +13,17 @@
 # limitations under the License.
 
 # Build layer variables
-# To install trtllm set FW_DEP_BUILDER=trtllm_builder, FW_BASE_FINAL=trtllm_install
-# To exclude trtllm set FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder
+# Supported mode: FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder
+# TRT-LLM stages are not present in this Dockerfile.
 ARG FW_DEP_BUILDER
 ARG FW_BASE_FINAL
 
 ARG NEMO_FW_BASE_IMAGE
 FROM ${NEMO_FW_BASE_IMAGE} AS base
 
-ENV NVIDIA_PRODUCT_NAME="NeMo Framework"
-ENV PIP_NO_CACHE_DIR=1
-ENV DEBIAN_FRONTEND=noninteractive
-
-WORKDIR /opt
+ENV NVIDIA_PRODUCT_NAME="NeMo Framework" \
+    PIP_NO_CACHE_DIR=1 \
+    DEBIAN_FRONTEND=noninteractive
 
 ##############################################################################
 ##
@@ -35,14 +33,14 @@ WORKDIR /opt
 
 FROM base AS vllm_wheel
 
-ARG VLLM_VERSION
+ARG VLLM_VERSION=v0.14.1
 ARG MAX_JOBS=4
 ENV VLLM_VERSION=$VLLM_VERSION
 
-WORKDIR /src/vllm
-
 # Build vllm
-RUN git clone https://github.com/vllm-project/vllm.git . && \
+RUN mkdir -p /src/vllm && \
+    git clone https://github.com/vllm-project/vllm.git /src/vllm && \
+    cd /src/vllm && \
     echo "Building vLLM version: $VLLM_VERSION" && \
     git checkout $VLLM_VERSION && \
     python use_existing_torch.py && \
@@ -53,7 +51,6 @@ RUN git clone https://github.com/vllm-project/vllm.git . && \
 ##
 ## Top layer of FW_BASE container
 ##
-## When TRTLLM is re-installed, FW_DEP_BUILDER=trtllm_builder
 ## default, FW_DEP_BUILDER=base
 ##
 ##############################################################################
@@ -69,11 +66,11 @@ FROM ${FW_DEP_BUILDER} AS fw_dep_builder
 ##
 ##############################################################################
 
-ARG UV_VERSION
-ENV UV_PROJECT_ENVIRONMENT=/opt/venv
-ENV UV_CACHE_DIR=/opt/uv_cache
-ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
-ENV UV_LINK_MODE=copy
+ARG UV_VERSION=0.9.17
+ENV UV_PROJECT_ENVIRONMENT=/opt/venv \
+    UV_CACHE_DIR=/opt/uv_cache \
+    PATH="/opt/venv/bin:$PATH" \
+    UV_LINK_MODE=copy
 RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | XDG_BIN_HOME=/usr/local/bin sh && \
     uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages && \
     python -m ipykernel install --name=python3 --display-name="Python 3 (ipykernel)" && \
@@ -92,7 +89,6 @@ FROM fw_dep_builder AS fw_toolkit_builder
 ##############################################################################
 ##
 ## Reinstall CUDA toolkit packages
-##
 ##############################################################################
 
 ARG REINSTALL_NSYS
@@ -146,8 +142,6 @@ RUN --mount=type=bind,from=vllm_wheel,source=/src/vllm/,target=/tmp/vllm/ \
 
 FROM ${FW_BASE_FINAL} AS nemo_fw_base_final
 
-WORKDIR /opt
-
 ##############################################################################
 ##
 ## Remove nvidia-modelopt
@@ -157,5 +151,4 @@ WORKDIR /opt
 ##
 ##############################################################################
 
-RUN pip uninstall -y nvidia-modelopt || true
-
+RUN cd /opt && pip uninstall -y nvidia-modelopt || true