# Multi-stage Dockerfile
ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
ARG BASE_TAG=23.08-py3

FROM ${BASE_IMAGE}:${BASE_TAG} as base

# `BASH_ENV` is only used for non-interactive bash.
ENV BASH_ENV="/tmp/bash_env"
ENV BUILD_TORCH_FROM_SOURCE=${BUILD_TORCH_FROM_SOURCE}
SHELL ["/bin/bash", "-c"]

FROM base as devel
ARG BUILD_TORCH_FROM_SOURCE=0

COPY docker/common/install_base.sh install_base.sh
RUN bash ./install_base.sh && rm install_base.sh

COPY docker/common/install_cmake.sh install_cmake.sh
RUN bash ./install_cmake.sh && rm install_cmake.sh

# Download & install internal TRT release
ARG RELEASE_URL_TRT
ARG TARGETARCH
ENV RELEASE_URL_TRT=$RELEASE_URL_TRT
ENV TRT_TARGETARCH=$TARGETARCH
COPY docker/common/install_tensorrt.sh install_tensorrt.sh
RUN bash ./install_tensorrt.sh && rm install_tensorrt.sh
ENV LD_LIBRARY_PATH=/usr/local/tensorrt/lib:${LD_LIBRARY_PATH}

# Install latest Polygraphy
COPY docker/common/install_polygraphy.sh install_polygraphy.sh
RUN bash ./install_polygraphy.sh && rm install_polygraphy.sh

# Install PyTorch from source (optional)
COPY docker/common/install_pytorch.sh install_pytorch.sh
RUN bash ./install_pytorch.sh $BUILD_TORCH_FROM_SOURCE && rm install_pytorch.sh

FROM devel as wheel
WORKDIR /src/tensorrt_llm
COPY benchmarks benchmarks
COPY cpp cpp
COPY benchmarks benchmarks
COPY scripts scripts
COPY tensorrt_llm tensorrt_llm
COPY 3rdparty 3rdparty
COPY setup.py requirements.txt ./

ARG BUILD_WHEEL_ARGS="--clean --trt_root /usr/local/tensorrt"
RUN python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS}

FROM devel as release

WORKDIR /app/tensorrt_llm
COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl .
COPY --from=wheel /src/tensorrt_llm/cpp/include/ include/
RUN pip install tensorrt_llm*.whl && \
    rm tensorrt_llm*.whl
COPY README.md ./
COPY examples examples