splitting Dockerfile.rocm and updating defaults to upstream

docker/Dockerfile.rocm

@@ -28,14 +28,14 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0
 
 FROM ${BASE_IMAGE} AS base
 
-ARG ARG_PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250
+ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 
 # Install build dependencies and utilities
 RUN apt-get update -q -y && apt-get install -q -y \
     sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
     apt-transport-https ca-certificates wget curl \
-    build-essential libnuma-dev ccache mold
+    libnuma-dev ccache mold
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install --upgrade pip
 # Note: mold is installed but not set as the system default linker because
@@ -61,10 +61,6 @@ ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy
-# python binary fall back for non venv builds
-ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3
-# Expose paths from wheel installation
-ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH}
 # ccache directory - persisted across layer rebuilds via cache mounts.
 ENV CCACHE_DIR=/root/.cache/ccache
 ENV CCACHE_COMPILERCHECK=content
@@ -114,8 +110,8 @@ WORKDIR ${COMMON_WORKDIR}
 FROM base AS fetch_vllm_0
 ONBUILD COPY ./ vllm/
 FROM base AS fetch_vllm_1
-ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
-ARG VLLM_BRANCH="455_wip"
+ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
+ARG VLLM_BRANCH="main"
 ENV VLLM_REPO=${VLLM_REPO}
 ENV VLLM_BRANCH=${VLLM_BRANCH}
 ONBUILD RUN git clone ${VLLM_REPO} \
@@ -133,6 +129,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 # don't need the rust toolchain or protoc.
 FROM fetch_vllm AS rust-build
 ARG COMMON_WORKDIR
+ARG USE_SCCACHE
 
 # protoc is used by tonic-build/prost-build.
 RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
@@ -148,6 +145,10 @@ ENV CARGO_BUILD_JOBS=4
 ENV CARGO_NET_RETRY=10
 ENV RUSTUP_MAX_RETRIES=10
 
+# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on
+# a separate local sccache daemon while sharing the same remote cache backend.
+ENV SCCACHE_SERVER_PORT=4227
+
 RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
     cd ${COMMON_WORKDIR}/vllm \
     && uv pip install --system -r requirements/build/rust.txt
@@ -159,8 +160,15 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
 RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
     --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
     cd ${COMMON_WORKDIR}/vllm \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+        export RUSTC_WRAPPER=sccache \
+        && sccache --show-stats; \
+    fi \
     && bash build_rust.sh \
-    && test -x vllm/vllm-rs
+    && test -x vllm/vllm-rs \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+        sccache --show-stats; \
+    fi
 
 # -----------------------
 # vLLM native build stages
@@ -235,9 +243,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # RIXL/UCX build stages
@@ -246,7 +262,7 @@ ARG RIXL_BRANCH="39be1de8"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
 ARG UCX_BRANCH="bfb51733"
 ARG UCX_REPO="https://github.com/openucx/ucx.git"
-# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image
+ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
 ENV RIXL_HOME=/usr/local/rixl
 ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
@@ -302,7 +318,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
     git checkout ${RIXL_BRANCH} && \
     CC="ccache gcc" CXX="ccache g++" \
     meson setup build --prefix=${RIXL_HOME} \
-                     --force-fallback-for=abseil-cpp \
                      -Ducx_path=${UCX_HOME} \
                      -Drocm_path=${ROCM_PATH} && \
     cd build && \
@@ -315,10 +330,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 RUN cd /opt/rixl && \
     sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
         contrib/build-wheel.sh && \
-    # The wheel build re-runs meson via meson-python; force the bundled abseil
-    sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \
-        pyproject.toml && \
-    grep -q 'force-fallback-for' pyproject.toml && \
     mkdir -p /app/install && \
     _ucx_install_dir=${UCX_HOME} \
     ./contrib/build-wheel.sh \
@@ -335,7 +346,7 @@ ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
 # DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
 # it can be linked against DeepEP without arch mismatches.
 ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
-# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image
+ENV ROCM_PATH=/opt/rocm
 ENV ROCSHMEM_DIR=/opt/rocshmem
 
 RUN --mount=type=cache,target=/root/.cache/ccache \
@@ -525,9 +536,17 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmar
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
 
 # -----------------------
@@ -561,6 +580,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
     libibverbs1 \
     ibverbs-providers \
     ibverbs-utils \
+    unzip \
     pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
     libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
     && rm -rf /var/lib/apt/lists/*
@@ -584,6 +604,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV HF_XET_HIGH_PERFORMANCE=1
 ENV HF_HUB_DOWNLOAD_TIMEOUT=60
 
+# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
+ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
+
 # Pre-install vLLM test dependencies.
 COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -704,21 +727,16 @@ ENV SAFETENSORS_FAST_GPU=1
 # Performance environment variable.
 ENV HIP_FORCE_DEV_KERNARG=1
 
+# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py).
+ENV PYTORCH_NVML_BASED_CUDA_CHECK=1
+
 # Workaround for ROCm profiler limits
 RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
 ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
 RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
     && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
     && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
 
-
-### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better
-RUN pip3 uninstall -y triton && \
-    git clone https://github.com/triton-lang/triton.git && \
-    cd triton && \
-    git checkout c517f38c && \
-    TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install .
-
 CMD ["/bin/bash"]
 
 #Set entrypoint for vllm-openai official images