# Keep the ROCm runtime aligned with the vLLM release targeted by
# docker/Dockerfile.ci. A stale base image can build successfully while failing
# at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0

# Keep the dependency-heavy portion independent from the source tree. The
# registry cache is keyed on these inputs, while the final source copy can be
# linked onto the cached image without materializing its large ROCm layers.
FROM ${BASE_IMAGE} AS deps

# NOTE: REMINDER to vLLM-Omni rebase maintainer
# BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
# VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when
# USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's
# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
# intentional custom source pin must update that regression in the same change,
# and should return to the stable release as soon as one is available.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}

ARG COMMON_WORKDIR=/app
WORKDIR ${COMMON_WORKDIR}

# Step 1: Setup - Install system dependencies.
# Need to include ffmpeg because vllm rocm upstream docker image
# does not include it.
RUN apt-get update && \
    apt-get install -y espeak-ng ffmpeg git jq && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Step 2: Conditionally reinstall vllm from source for nightly builds
RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
        python3 -m pip uninstall -y vllm && rm -rf vllm && \
        git clone https://github.com/vllm-project/vllm.git && \
        cd vllm && \
        git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
        python3 -m pip install -r requirements/rocm.txt && \
        python3 setup.py clean --all && \
        PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
        cd ../ && \
        rm -rf vllm/.git; \
    fi

# Fail during image construction instead of letting every AMD test fail at
# runtime when the base image and vLLM-Omni target incompatible vLLM APIs.
RUN python3 -c "import vllm; from vllm.v1.kv_cache_interface import compute_layout_strides; from vllm.v1.kv_cache_layout import KVCacheLayout; print(f'Validated vLLM {vllm.__version__}')"

# 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: Install project dependencies against a placeholder package tree.
# The editable install keeps resolving imports from this directory after the
# final COPY --link replaces the placeholder with the real source tree.
WORKDIR ${COMMON_WORKDIR}/vllm-omni
RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
COPY pyproject.toml setup.py ./
COPY requirements/ requirements/
ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0

RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
        --no-cache-dir -e ".[dev]" --no-build-isolation && \
    python3 -m pip uninstall -y onnxruntime && \
    python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
        "onnxruntime-rocm>=1.22.2" && \
    python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
    python3 -c "from torchcodec.decoders import VideoDecoder"

RUN ln -sf /usr/bin/python3 /usr/bin/python

# This is a workaround to ensure pytest exits with the correct status code in CI tests.
RUN printf '%s\n' \
    'import os' \
    '' \
    '_exit_code = 1' \
    '' \
    'def pytest_sessionfinish(session, exitstatus):' \
    '    global _exit_code' \
    '    _exit_code = int(exitstatus)' \
    '' \
    'def pytest_unconfigure(config):' \
    '    import sys' \
    '    sys.stdout.flush()' \
    '    sys.stderr.flush()' \
    '    os._exit(_exit_code)' \
    > ${COMMON_WORKDIR}/vllm-omni/conftest.py

# Keep source changes in one lightweight linked layer. Do not add RUN, COPY, or
# ADD instructions that depend on the resulting filesystem after this point;
# doing so would force BuildKit to download and extract the cached ROCm layers.
FROM deps AS base
ARG COMMON_WORKDIR=/app
WORKDIR ${COMMON_WORKDIR}/vllm-omni
COPY --link . .

FROM base AS test

CMD ["/bin/bash"]
ENTRYPOINT []

# Set entrypoint for vllm-openai official images
FROM base AS vllm-openai
ENTRYPOINT ["vllm", "serve", "--omni"]