From d6240742f365cd813991b540b8ea797347aac72e Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Fri, 24 Jul 2026 23:48:32 +0530 Subject: [PATCH] [BugFix][CI] Upgrade ROCm and CUDA docker base images to v0.25.0 (#5115) Signed-off-by: akshatvishu --- docker/Dockerfile.ci | 5 +- docker/Dockerfile.rocm | 27 ++++++- tools/install_torchcodec_rocm.sh | 125 +++++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+), 5 deletions(-) create mode 100755 tools/install_torchcodec_rocm.sh diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 84ce5aa41..caae57abf 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,6 +1,5 @@ -# Target vllm version: v0.24.0 (releases/v0.24.0 branch). -# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must -# be updated on every rebase to match the new upstream vllm commit/tag. +# Target vllm version: v0.25.0 (releases/v0.25.0 branch). +# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. ARG VLLM_BASE_IMAGE=vllm/vllm-openai ARG VLLM_BASE_TAG=v0.25.0 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 49c6f02aa..9be3e5fb5 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -17,7 +17,7 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} ARG COMMON_WORKDIR=/app WORKDIR ${COMMON_WORKDIR} -# Step 1: Setup - Install system dependencies +# Step 1: Setup - Install system dependencies. # Need to include ffmpeg because vllm rocm upstream docker image # does not include it. RUN apt-get update && \ @@ -38,6 +38,26 @@ RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \ rm -rf vllm/.git; \ fi +# Build TorchCodec after any optional nightly vLLM reinstall so it targets the +# final torch/ROCm ABI. Keep this layer before COPY . so source changes do not +# invalidate the expensive native build. FFmpeg development headers are +# removed after the wheel is built. +COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec_rocm.sh +RUN --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + apt-get update && \ + apt-get install -y --no-install-recommends \ + pkg-config \ + libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev && \ + bash /tmp/install_torchcodec_rocm.sh && \ + apt-get purge -y \ + libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev && \ + rm -f /tmp/install_torchcodec_rocm.sh && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + # Step 3: Copy vllm-omni code and install without uv RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni COPY . ${COMMON_WORKDIR}/vllm-omni @@ -59,7 +79,10 @@ RUN printf '%s\n' \ ' os._exit(_exit_code)' \ > ${COMMON_WORKDIR}/vllm-omni/conftest.py -RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation +RUN cd ${COMMON_WORKDIR}/vllm-omni && \ + uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \ + --no-cache-dir ".[dev]" --no-build-isolation && \ + python3 -c "from torchcodec.decoders import VideoDecoder" RUN ln -sf /usr/bin/python3 /usr/bin/python diff --git a/tools/install_torchcodec_rocm.sh b/tools/install_torchcodec_rocm.sh new file mode 100755 index 000000000..787079603 --- /dev/null +++ b/tools/install_torchcodec_rocm.sh @@ -0,0 +1,125 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Build TorchCodec against the ROCm PyTorch installation. The PyPI wheel is +# built against upstream PyTorch and is not ABI-compatible with ROCm builds. + +set -euo pipefail + +TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}" +# Keep this pinned for ROCm/PyTorch ABI compatibility. When rebasing the ROCm +# base image or vLLM, revalidate this version against the active torch build. +TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}" +TORCHCODEC_WHEEL_CACHE="${TORCHCODEC_WHEEL_CACHE:-/root/.cache/torchcodec-wheels}" + +echo "=== TorchCodec ROCm installation ===" + +verify_torchcodec() { + python3 -c "from torchcodec.decoders import VideoDecoder" +} + +if verify_torchcodec >/dev/null 2>&1; then + echo "TorchCodec is already installed and working." + exit 0 +fi + +install_system_deps() { + if ! command -v apt-get >/dev/null 2>&1; then + echo "TorchCodec build dependencies are missing and apt-get is unavailable." + return 1 + fi + apt-get update + apt-get install -y --no-install-recommends \ + pkg-config \ + ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev +} + +if ! command -v pkg-config >/dev/null 2>&1 || \ + ! pkg-config --exists libavcodec libavformat libavutil libswscale \ + libavdevice libavfilter libswresample 2>/dev/null; then + echo "Installing missing FFmpeg development dependencies..." + install_system_deps +fi + +if ! python3 -c "import packaging, pybind11" >/dev/null 2>&1; then + python3 -m pip install packaging pybind11 setuptools wheel +fi + +PYBIND11_DIR=$(python3 -c "import pybind11; print(pybind11.get_cmake_dir())") +export pybind11_DIR="$PYBIND11_DIR" +export CMAKE_PREFIX_PATH="${PYBIND11_DIR}:${CMAKE_PREFIX_PATH:-}" + +safe_component() { + printf '%s' "$1" | sed 's/[^A-Za-z0-9_.-]/_/g' +} + +PYTHON_SOABI=$(python3 -c \ + "import sysconfig; print(sysconfig.get_config_var('SOABI') or 'unknown')") +TORCH_VERSION=$(python3 -c "import torch; print(torch.__version__)") +ROCM_VERSION=$(python3 -c "import torch; print(torch.version.hip or 'none')") +FFMPEG_VERSION=$(pkg-config --modversion libavcodec 2>/dev/null || echo unknown) +ARCH_TAG="${PYTORCH_ROCM_ARCH:-all}" + +CACHE_DIR="${TORCHCODEC_WHEEL_CACHE}/$(safe_component "$TORCHCODEC_BRANCH")/$(safe_component "$PYTHON_SOABI")/$(safe_component "$TORCH_VERSION")/$(safe_component "$ROCM_VERSION")/$(safe_component "$FFMPEG_VERSION")/$(safe_component "$ARCH_TAG")" +CACHED_WHEEL="" +if [ -d "$CACHE_DIR" ]; then + CACHED_WHEEL=$(find "$CACHE_DIR" -maxdepth 1 -type f -name 'torchcodec-*.whl' -print -quit) +fi + +validate_wheel_filename() { + WHEEL_PATH="$1" python3 - <<'PY' +from os import environ +from pathlib import Path + +from packaging.utils import parse_wheel_filename + +wheel = Path(environ["WHEEL_PATH"]) +parse_wheel_filename(wheel.name) +PY +} + +if [ -n "$CACHED_WHEEL" ] && validate_wheel_filename "$CACHED_WHEEL"; then + echo "Trying cached wheel: $CACHED_WHEEL" + if python3 -m pip install --no-deps --force-reinstall "$CACHED_WHEEL" && \ + verify_torchcodec; then + echo "Installed TorchCodec from cache." + exit 0 + fi + echo "Cached wheel is incompatible; rebuilding TorchCodec." + rm -f "$CACHED_WHEEL" +fi + +BUILD_DIR=$(mktemp -d -t torchcodec-XXXXXX) +trap 'rm -rf "$BUILD_DIR"' EXIT + +cd "$BUILD_DIR" +echo "Cloning TorchCodec $TORCHCODEC_BRANCH..." +git clone --depth 1 --branch "$TORCHCODEC_BRANCH" "$TORCHCODEC_REPO" torchcodec +cd torchcodec + +export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build" +export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1 +export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1 +export CMAKE_GENERATOR=Ninja +export MAX_JOBS="${MAX_JOBS:-$(nproc)}" + +if command -v ccache >/dev/null 2>&1; then + export CMAKE_C_COMPILER_LAUNCHER=ccache + export CMAKE_CXX_COMPILER_LAUNCHER=ccache +fi + +echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..." +python3 -m pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist" +BUILT_WHEEL=$(find "$BUILD_DIR/dist" -maxdepth 1 -type f -name 'torchcodec-*.whl' -print -quit) +test -n "$BUILT_WHEEL" +validate_wheel_filename "$BUILT_WHEEL" + +python3 -m pip install --no-deps --force-reinstall "$BUILT_WHEEL" +verify_torchcodec + +mkdir -p "$CACHE_DIR" +cp "$BUILT_WHEEL" "$CACHE_DIR/$(basename "$BUILT_WHEEL")" +echo "Cached TorchCodec wheel in $CACHE_DIR" +echo "=== TorchCodec ROCm installation complete ==="