mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
59 lines
1.7 KiB
Docker
59 lines
1.7 KiB
Docker
# 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"
|
|
SHELL ["/bin/bash", "-c"]
|
|
|
|
FROM base as devel
|
|
|
|
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
|
|
ARG TORCH_INSTALL_TYPE="skip"
|
|
COPY docker/common/install_pytorch.sh install_pytorch.sh
|
|
RUN bash ./install_pytorch.sh $TORCH_INSTALL_TYPE && 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
|