|
| 1 | +# Multi-stage Dockerfile |
| 2 | +ARG BASE_IMAGE=nvcr.io/nvidia/pytorch |
| 3 | +ARG BASE_TAG=24.02-py3 |
| 4 | +ARG DEVEL_IMAGE=devel |
| 5 | + |
| 6 | +FROM ${BASE_IMAGE}:${BASE_TAG} as base |
| 7 | + |
| 8 | +# https://www.gnu.org/software/bash/manual/html_node/Bash-Startup-Files.html |
| 9 | +# The default values come from `nvcr.io/nvidia/pytorch` |
| 10 | +ENV BASH_ENV=${BASH_ENV:-/etc/bash.bashrc} |
| 11 | +ENV ENV=${ENV:-/etc/shinit_v2} |
| 12 | +SHELL ["/bin/bash", "-c"] |
| 13 | + |
| 14 | +FROM base as devel |
| 15 | + |
| 16 | +COPY docker/common/install_base.sh install_base.sh |
| 17 | +RUN bash ./install_base.sh && rm install_base.sh |
| 18 | + |
| 19 | +COPY docker/common/install_cmake.sh install_cmake.sh |
| 20 | +RUN bash ./install_cmake.sh && rm install_cmake.sh |
| 21 | + |
| 22 | +COPY docker/common/install_ccache.sh install_ccache.sh |
| 23 | +RUN bash ./install_ccache.sh && rm install_ccache.sh |
| 24 | + |
| 25 | +# Download & install internal TRT release |
| 26 | +ARG TRT_VER |
| 27 | +ARG CUDA_VER |
| 28 | +ARG CUDNN_VER |
| 29 | +ARG NCCL_VER |
| 30 | +ARG CUBLAS_VER |
| 31 | +COPY docker/common/install_tensorrt.sh install_tensorrt.sh |
| 32 | +RUN bash ./install_tensorrt.sh \ |
| 33 | + --TRT_VER=${TRT_VER} \ |
| 34 | + --CUDA_VER=${CUDA_VER} \ |
| 35 | + --CUDNN_VER=${CUDNN_VER} \ |
| 36 | + --NCCL_VER=${NCCL_VER} \ |
| 37 | + --CUBLAS_VER=${CUBLAS_VER} && \ |
| 38 | + rm install_tensorrt.sh |
| 39 | + |
| 40 | +# Install latest Polygraphy |
| 41 | +COPY docker/common/install_polygraphy.sh install_polygraphy.sh |
| 42 | +RUN bash ./install_polygraphy.sh && rm install_polygraphy.sh |
| 43 | + |
| 44 | +# Install mpi4py |
| 45 | +COPY docker/common/install_mpi4py.sh install_mpi4py.sh |
| 46 | +RUN bash ./install_mpi4py.sh && rm install_mpi4py.sh |
| 47 | + |
| 48 | +# Install PyTorch |
| 49 | +ARG TORCH_INSTALL_TYPE="skip" |
| 50 | +COPY docker/common/install_pytorch.sh install_pytorch.sh |
| 51 | +RUN bash ./install_pytorch.sh $TORCH_INSTALL_TYPE && rm install_pytorch.sh |
| 52 | +COPY setup.py requirements.txt requirements-dev.txt ./ |
| 53 | + |
| 54 | +RUN pip install --no-cache-dir -r requirements-dev.txt |
| 55 | + |
| 56 | +RUN apt update && echo Y | apt install sudo cargo -y |
| 57 | + |
| 58 | +RUN cargo install sccache --locked |
| 59 | + |
| 60 | +ARG RUNNER_VERSION=2.317.0 |
| 61 | + |
| 62 | +ARG USER_ID=1000 |
| 63 | +ARG USER_NAME=runner |
| 64 | +ARG GROUP_ID=1000 |
| 65 | +ARG GROUP_NAME=runner |
| 66 | + |
| 67 | +RUN (getent group ${GROUP_ID} || groupadd --gid ${GROUP_ID} ${GROUP_NAME}) && \ |
| 68 | + (getent passwd ${USER_ID} || useradd --gid ${GROUP_ID} --uid ${USER_ID} --create-home --no-log-init --shell /bin/bash ${USER_NAME}) |
| 69 | + |
| 70 | + |
| 71 | +RUN usermod -aG sudo ${USER_NAME} \ |
| 72 | + && echo "%sudo ALL=(ALL:ALL) NOPASSWD:ALL" > /etc/sudoers \ |
| 73 | + && echo "Defaults env_keep += \"DEBIAN_FRONTEND\"" >> /etc/sudoers |
| 74 | + |
| 75 | +ENV HOME=/home/runner |
| 76 | + |
| 77 | +# cd into the user directory, download and unzip the github actions runner |
| 78 | +RUN cd /home/runner && mkdir actions-runner && cd actions-runner \ |
| 79 | + && curl -O -L https://github.com/actions/runner/releases/download/v${RUNNER_VERSION}/actions-runner-linux-x64-${RUNNER_VERSION}.tar.gz \ |
| 80 | + && tar xzf ./actions-runner-linux-x64-${RUNNER_VERSION}.tar.gz |
| 81 | + |
| 82 | +RUN chown -R runner:runner /home/runner && /home/runner/actions-runner/bin/installdependencies.sh |
| 83 | + |
| 84 | +ADD docker/common/start.sh start.sh |
| 85 | + |
| 86 | +RUN chmod +x start.sh |
| 87 | + |
| 88 | +RUN sudo chmod +x /root/.cargo/bin/sccache |
| 89 | + |
| 90 | +USER runner |
| 91 | + |
| 92 | +ENV PATH=/root/.cargo/bin${PATH:+:${PATH}} |
| 93 | + |
| 94 | +ENTRYPOINT ["./start.sh"] |
| 95 | + |
| 96 | +FROM ${DEVEL_IMAGE} as wheel |
| 97 | +WORKDIR /src/tensorrt_llm |
| 98 | +COPY benchmarks benchmarks |
| 99 | +COPY cpp cpp |
| 100 | +COPY benchmarks benchmarks |
| 101 | +COPY scripts scripts |
| 102 | +COPY tensorrt_llm tensorrt_llm |
| 103 | +COPY 3rdparty 3rdparty |
| 104 | +COPY setup.py requirements.txt requirements-dev.txt ./ |
| 105 | + |
| 106 | +ARG BUILD_WHEEL_ARGS="--clean --trt_root /usr/local/tensorrt --python_bindings --benchmarks" |
| 107 | +RUN python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS} |
| 108 | + |
| 109 | +FROM ${DEVEL_IMAGE} as release |
| 110 | + |
| 111 | +WORKDIR /app/tensorrt_llm |
| 112 | +COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl . |
| 113 | +RUN pip install tensorrt_llm*.whl --extra-index-url https://pypi.nvidia.com && \ |
| 114 | + rm tensorrt_llm*.whl |
| 115 | +COPY README.md ./ |
| 116 | +COPY docs docs |
| 117 | +COPY cpp/include include |
| 118 | +RUN ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/libs")') lib && \ |
| 119 | + test -f lib/libnvinfer_plugin_tensorrt_llm.so && \ |
| 120 | + ln -sv lib/libnvinfer_plugin_tensorrt_llm.so lib/libnvinfer_plugin_tensorrt_llm.so.9 && \ |
| 121 | + echo "/app/tensorrt_llm/lib" > /etc/ld.so.conf.d/tensorrt_llm.conf && \ |
| 122 | + ldconfig |
| 123 | +ARG SRC_DIR=/src/tensorrt_llm |
| 124 | +COPY --from=wheel ${SRC_DIR}/benchmarks benchmarks |
| 125 | +ARG CPP_BUILD_DIR=${SRC_DIR}/cpp/build |
| 126 | +COPY --from=wheel \ |
| 127 | + ${CPP_BUILD_DIR}/benchmarks/bertBenchmark \ |
| 128 | + ${CPP_BUILD_DIR}/benchmarks/gptManagerBenchmark \ |
| 129 | + ${CPP_BUILD_DIR}/benchmarks/gptSessionBenchmark \ |
| 130 | + benchmarks/cpp/ |
| 131 | +COPY examples examples |
| 132 | +RUN chmod -R a+w examples && \ |
| 133 | + rm -v \ |
| 134 | + benchmarks/cpp/bertBenchmark.cpp \ |
| 135 | + benchmarks/cpp/gptManagerBenchmark.cpp \ |
| 136 | + benchmarks/cpp/gptSessionBenchmark.cpp \ |
| 137 | + benchmarks/cpp/CMakeLists.txt |
| 138 | +ARG GIT_COMMIT |
| 139 | +ARG TRT_LLM_VER |
| 140 | +ENV TRT_LLM_GIT_COMMIT=${GIT_COMMIT} \ |
| 141 | + TRT_LLM_VERSION=${TRT_LLM_VER} |
| 142 | + |
0 commit comments