ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
FROM ${BASE_IMAGE} AS base

# Declare a variable to know if we want to use the nightly build or the stable build.
# NOTE: REMINDER to vLLM-Omni rebase maintainer
# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
# 1. If vLLM-omni maintainer is forced to use custom commits
#    during rebasing, they can change the two variables
# 2. Whenever vLLM upstream has released a stable version,
#    we should swap over to use stable release ASAP.
#    We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
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

# Step 3: Copy vllm-omni code and install without uv
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
COPY . ${COMMON_WORKDIR}/vllm-omni

# 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

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 ln -sf /usr/bin/python3 /usr/bin/python

FROM base AS test

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

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