forked from triton-inference-server/tensorrtllm_backend
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathDockerfile.trt_llm_backend
89 lines (71 loc) · 5.18 KB
/
Dockerfile.trt_llm_backend
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
ARG BASE_IMAGE=nvcr.io/nvidia/tritonserver
ARG BASE_TAG=25.01-py3
ARG PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:25.01-py3
FROM ${PYTORCH_IMAGE} as pytorch_image
FROM ${BASE_IMAGE}:${BASE_TAG} as base
# Copy PyTorch package from PyTorch image
COPY --from=pytorch_image /usr/local/lib/lib* /usr/local/lib/
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torch /usr/local/lib/python3.12/dist-packages/torch
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torch-2.6.0a0+ecf3bae40a.nv25.1.dist-info /usr/local/lib/python3.12/dist-packages/torch-2.6.0a0+ecf3bae40a.nv25.1.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torchgen /usr/local/lib/python3.12/dist-packages/torchgen
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torchvision /usr/local/lib/python3.12/dist-packages/torchvision
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torchvision-0.20.0a0.dist-info /usr/local/lib/python3.12/dist-packages/torchvision-0.20.0a0.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/torchvision.libs /usr/local/lib/python3.12/dist-packages/torchvision.libs
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/setuptools /usr/local/lib/python3.12/dist-packages/setuptools
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/setuptools-70.3.0.dist-info /usr/local/lib/python3.12/dist-packages/setuptools-70.3.0.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/functorch /usr/local/lib/python3.12/dist-packages/functorch
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/pytorch_triton-3.1.0+cf34004b8.internal.dist-info /usr/local/lib/python3.12/dist-packages/pytorch_triton-3.1.0+cf34004b8.internal.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/triton /usr/local/lib/python3.12/dist-packages/triton
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/jinja2 /usr/local/lib/python3.12/dist-packages/jinja2
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/jinja2-3.1.4.dist-info /usr/local/lib/python3.12/dist-packages/jinja2-3.1.4.dist-info
COPY --from=pytorch_image /usr/local/bin/ninja /usr/local/bin/
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/ninja /usr/local/lib/python3.12/dist-packages/ninja
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/ninja-1.11.1.3.dist-info /usr/local/lib/python3.12/dist-packages/ninja-1.11.1.3.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/networkx /usr/local/lib/python3.12/dist-packages/networkx
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/networkx-3.4.2.dist-info /usr/local/lib/python3.12/dist-packages/networkx-3.4.2.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/sympy /usr/local/lib/python3.12/dist-packages/sympy
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/sympy-1.13.1.dist-info /usr/local/lib/python3.12/dist-packages/sympy-1.13.1.dist-info
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/packaging /usr/local/lib/python3.12/dist-packages/packaging
COPY --from=pytorch_image /usr/local/lib/python3.12/dist-packages/packaging-23.2.dist-info /usr/local/lib/python3.12/dist-packages/packaging-23.2.dist-info
RUN apt-get update && apt-get install -y --no-install-recommends rapidjson-dev python-is-python3 ccache git-lfs
COPY requirements.txt /tmp/
RUN pip3 install -r /tmp/requirements.txt
# Remove previous TRT installation
RUN apt-get remove --purge -y tensorrt* libnvinfer*
RUN pip uninstall -y tensorrt
FROM base as dev
# Download & install internal TRT release
COPY tensorrt_llm/docker/common/install_tensorrt.sh /tmp/
RUN bash /tmp/install_tensorrt.sh && rm /tmp/install_tensorrt.sh
ENV LD_LIBRARY_PATH=/usr/local/tensorrt/lib:${LD_LIBRARY_PATH}
ENV TRT_ROOT=/usr/local/tensorrt
# Install latest Polygraphy
COPY tensorrt_llm/docker/common/install_polygraphy.sh /tmp/
RUN bash /tmp/install_polygraphy.sh && rm /tmp/install_polygraphy.sh
# CMake
COPY tensorrt_llm/docker/common/install_cmake.sh /tmp/
RUN bash /tmp/install_cmake.sh && rm /tmp/install_cmake.sh
ENV PATH="/usr/local/cmake/bin:${PATH}"
# Install mpi4py
COPY tensorrt_llm/docker/common/install_mpi4py.sh /tmp/
RUN bash /tmp/install_mpi4py.sh && rm /tmp/install_mpi4py.sh
# Use "skip" (default) for x86_64 arch and aarch64 arch
ARG TORCH_INSTALL_TYPE="skip"
COPY tensorrt_llm/docker/common/install_pytorch.sh install_pytorch.sh
RUN bash ./install_pytorch.sh $TORCH_INSTALL_TYPE && rm install_pytorch.sh
FROM dev as trt_llm_builder
WORKDIR /app
COPY scripts scripts
COPY tensorrt_llm tensorrt_llm
RUN cd tensorrt_llm && python3 scripts/build_wheel.py --trt_root="${TRT_ROOT}" -i -c && cd ..
FROM trt_llm_builder as trt_llm_backend_builder
WORKDIR /app/
COPY inflight_batcher_llm inflight_batcher_llm
RUN cd inflight_batcher_llm && bash scripts/build.sh && cd ..
FROM trt_llm_backend_builder as final
# Install TensorRT-LLM
WORKDIR /app/
COPY --from=trt_llm_builder /app/tensorrt_llm/build /app/tensorrt_llm/build
RUN cd /app/tensorrt_llm/build && pip3 install *.whl
# Install TensorRT-LLM backend
ENV LD_LIBRARY_PATH=/opt/tritonserver/backends/tensorrtllm:${LD_LIBRARY_PATH}