File tree Expand file tree Collapse file tree 2 files changed +18
-20
lines changed Expand file tree Collapse file tree 2 files changed +18
-20
lines changed Original file line number Diff line number Diff line change @@ -78,26 +78,6 @@ RUN if [[ "${HUGECTR_DEV_MODE}" == "false" ]]; then \
7878 ; fi
7979
8080ENV PYTHONPATH=${PYTHONPATH}:${HUGECTR_HOME}/lib
81- ARG _HUGECTR_BACKEND_REPO="github.com/triton-inference-server/hugectr_backend.git"
82- ARG TRITON_VERSION
83- # Install Triton inference backend.
84- RUN if [ "${HUGECTR_DEV_MODE}" == "false" ]; then \
85- git clone --branch ${HUGECTR_BACKEND_VER} --depth 1 https://${_CI_JOB_TOKEN}${_HUGECTR_BACKEND_REPO} /repos/hugectr_triton_backend && \
86- mkdir /repos/hugectr_triton_backend/build && \
87- cd /repos/hugectr_triton_backend/build && \
88- cmake \
89- -DCMAKE_INSTALL_PREFIX:PATH=${HUGECTR_HOME} \
90- -DTRITON_COMMON_REPO_TAG="r${TRITON_VERSION}" \
91- -DTRITON_CORE_REPO_TAG="r${TRITON_VERSION}" \
92- -DTRITON_BACKEND_REPO_TAG="r${TRITON_VERSION}" .. && \
93- make -j$(nproc) && \
94- make install && \
95- cd ../.. && \
96- rm -rf hugectr_triton_backend && \
97- chmod +x ${HUGECTR_HOME}/lib/*.so ${HUGECTR_HOME}/backends/hugectr/*.so && \
98- rm -rf /repos \
99- ; fi
100- RUN ln -s ${HUGECTR_HOME}/backends/hugectr /opt/tritonserver/backends/hugectr
10181
10282# Clean up
10383RUN rm -rf /usr/local/share/jupyter/lab/staging/node_modules/marked
Original file line number Diff line number Diff line change @@ -38,6 +38,24 @@ COPY --chown=1000:1000 --from=dlfw /usr/local/lib/python${PYTHON_VERSION}/dist-p
3838COPY --chown=1000:1000 --from=dlfw /usr/local/lib/python${PYTHON_VERSION}/dist-packages/numpy-*.dist-info /usr/local/lib/python${PYTHON_VERSION}/dist-packages/numpy.dist-info/
3939COPY --chown=1000:1000 --from=dlfw /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch-*.egg-info /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch.egg-info/
4040
41+ # Argumeints "_XXXX" are only valid when $HUGECTR_DEV_MODE==false
42+ # Install hps_torch in merlin-pytorch
43+ ARG HUGECTR_DEV_MODE=false
44+ ARG _HUGECTR_REPO="github.com/NVIDIA-Merlin/HugeCTR.git"
45+ ARG _CI_JOB_TOKEN=""
46+ ARG HUGECTR_VER=main
47+
48+ RUN if [ "$HUGECTR_DEV_MODE" == "false" ]; then \
49+ export HUGECTR_HOME=/usr/local/hugectr && \
50+ rm -rf ${HUGECTR_HOME}/lib/libgmock* ${HUGECTR_HOME}/lib/pkgconfig/gmock* ${HUGECTR_HOME}/include/gmock && \
51+ rm -rf ${HUGECTR_HOME}/lib/libgtest* ${HUGECTR_HOME}/lib/pkgconfig/gtest* ${HUGECTR_HOME}/include/gtest && \
52+ git clone --branch ${HUGECTR_VER} --depth 1 --recurse-submodules --shallow-submodules https://${_CI_JOB_TOKEN}${_HUGECTR_REPO} /hugectr && \
53+ pushd /hugectr/hps_torch/ && \
54+ pip --no-cache-dir install ninja && \
55+ TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 9.0" python setup.py install && \
56+ popd && \
57+ rm -rf /hugectr \
58+ ; fi
4159
4260# Add all torch libraries to /usr/local
4361RUN ln -s /opt/tritonserver/backends/pytorch/* /usr/local/lib/
You can’t perform that action at this time.
0 commit comments