@@ -44,48 +44,41 @@ COPY docker/common/install.sh \
4444 docker/common/install_ucx.sh \
4545 docker/common/install_nixl.sh \
4646 docker/common/install_etcd.sh \
47- docker/common/install_mooncake.sh \
4847 ./
4948
50- RUN GITHUB_MIRROR=${GITHUB_MIRROR} \
51- PYTHON_VERSION=${PYTHON_VERSION} \
52- bash ./install.sh --base && rm install_base.sh
53-
54- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install.sh --cmake && rm install_cmake.sh
55-
56- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install.sh --ccache && rm install_ccache.sh
57-
58- RUN bash ./install.sh --cuda_toolkit && rm install_cuda_toolkit.sh
59-
6049ARG TRT_VER
6150ARG CUDA_VER
6251ARG CUDNN_VER
6352ARG NCCL_VER
6453ARG CUBLAS_VER
65- RUN TRT_VER=${TRT_VER} \
54+ ARG TORCH_INSTALL_TYPE="skip"
55+ RUN GITHUB_MIRROR=${GITHUB_MIRROR} \
56+ PYTHON_VERSION=${PYTHON_VERSION} \
57+ TRT_VER=${TRT_VER} \
6658 CUDA_VER=${CUDA_VER} \
6759 CUDNN_VER=${CUDNN_VER} \
6860 NCCL_VER=${NCCL_VER} \
6961 CUBLAS_VER=${CUBLAS_VER} \
70- bash ./install.sh --tensorrt && rm install_tensorrt.sh
71-
72- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install.sh --polygraphy && rm install_polygraphy.sh
73-
74- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install.sh --mpi4py && rm install_mpi4py.sh
75-
76- ARG TORCH_INSTALL_TYPE="skip"
77- RUN TORCH_INSTALL_TYPE=${TORCH_INSTALL_TYPE} bash ./install.sh --pytorch && rm install_pytorch.sh
78-
79- RUN bash ./install.sh --opencv && rm install.sh
80-
81- # Install UCX first
82- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_ucx.sh && rm install_ucx.sh
83-
84- # Install NIXL
85- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_nixl.sh && rm install_nixl.sh
86-
87- # Install etcd
88- RUN bash ./install_etcd.sh && rm install_etcd.sh
62+ TORCH_INSTALL_TYPE=${TORCH_INSTALL_TYPE} \
63+ bash ./install.sh --base --cmake --ccache --cuda_toolkit --tensorrt --polygraphy --mpi4py --pytorch --opencv && \
64+ rm install_base.sh && \
65+ rm install_cmake.sh && \
66+ rm install_ccache.sh && \
67+ rm install_cuda_toolkit.sh && \
68+ rm install_tensorrt.sh && \
69+ rm install_polygraphy.sh && \
70+ rm install_mpi4py.sh && \
71+ rm install_pytorch.sh && \
72+ rm install.sh
73+
74+ # Install UCX, NIXL, etcd
75+ # TODO: Combine these into the main install.sh script
76+ RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_ucx.sh && \
77+ GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_nixl.sh && \
78+ bash ./install_etcd.sh && \
79+ rm install_ucx.sh && \
80+ rm install_nixl.sh && \
81+ rm install_etcd.sh
8982
9083FROM ${TRITON_IMAGE}:${TRITON_BASE_TAG} AS triton
9184
@@ -100,16 +93,18 @@ COPY --from=triton /opt/tritonserver/caches /opt/tritonserver/caches
10093
10194# Copy all installation scripts at once to reduce layers
10295COPY docker/common/install_triton.sh \
96+ docker/common/install_mooncake.sh \
10397 ./
10498
105- RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_triton.sh && rm install_triton.sh
106-
10799# Install Mooncake, after triton handles boost requirement
108- RUN if [ -f /etc/redhat-release ]; then \
100+ RUN GITHUB_MIRROR=${GITHUB_MIRROR} bash ./install_triton.sh && \
101+ if [ -f /etc/redhat-release ]; then \
109102 echo "Rocky8 detected, skipping mooncake installation"; \
110103 else \
111104 bash ./install_mooncake.sh; \
112- fi && rm install_mooncake.sh
105+ fi && \
106+ rm install_triton.sh && \
107+ rm install_mooncake.sh
113108
114109FROM ${DEVEL_IMAGE} AS wheel
115110WORKDIR /src/tensorrt_llm
0 commit comments