[Ready to be reviewed] Support Customized HCTR Repo in the unified containers

docker/dockerfile.ctr

@@ -69,7 +69,7 @@ RUN git clone --branch apache-arrow-4.0.1 --recurse-submodules https://github.co
               -DCMAKE_LIBRARY_PATH=${CUDA_CUDA_LIBRARY} \
               -DARROW_FLIGHT=ON \
               -DARROW_GANDIVA=OFF \
-              -DARROW_ORC=ON \
+              -DARROW_ORC=OFF \
               -DARROW_WITH_BZ2=ON \
               -DARROW_WITH_ZLIB=ON \
               -DARROW_WITH_ZSTD=ON \
@@ -89,7 +89,7 @@ RUN git clone --branch apache-arrow-4.0.1 --recurse-submodules https://github.co
       pushd python && \
         export PYARROW_WITH_PARQUET=ON && \
         export PYARROW_WITH_CUDA=ON && \
-        export PYARROW_WITH_ORC=ON && \
+        export PYARROW_WITH_ORC=OFF && \
         export PYARROW_WITH_DATASET=ON && \
         python setup.py build_ext --build-type=release bdist_wheel && \
         pip install dist/*.whl && \
@@ -215,6 +215,12 @@ FROM phase3 AS phase4
 ARG RELEASE=false
 ARG HUGECTR_VER=vnightly
 
+# Arguments "_XXXX" are only valid when $HUGECTR_DEV_MODE==false
+ARG HUGECTR_DEV_MODE=false
+ARG _HUGECTR_BRANCH=master
+ARG _HUGECTR_REPO="github.com/NVIDIA-Merlin/HugeCTR.git"
+ARG _CI_JOB_TOKEN=""
+
 RUN pip3 install --no-cache-dir mpi4py ortools sklearn onnx onnxruntime
 
 ENV OMPI_MCA_plm_rsh_agent=sh
@@ -232,20 +238,26 @@ RUN rm -rf /usr/lib/x86_64-linux-gnu/libibverbs.so && \
     ln -s /usr/lib/x86_64-linux-gnu/libibverbs.so.1.14.36.0 /usr/lib/x86_64-linux-gnu/libibverbs.so
 
 # Install hugectr
-RUN git clone https://github.com/NVIDIA-Merlin/HugeCTR.git /hugectr && \
-      cd /hugectr && if [ "$RELEASE" == "true" ] && [ ${HUGECTR_VER} != "vnightly" ]; then git fetch --all --tags && git checkout tags/${HUGECTR_VER}; else git checkout master; fi && \
-    git submodule update --init --recursive && \
-    mkdir build && cd build && \
-    LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs/:$LD_LIBRARY_PATH && \
-    export PATH=$PATH:/usr/local/cuda-${CUDA_SHORT_VERSION}/compat/ && \
-    cmake -DCMAKE_CXX_COMPILER=/usr/bin/g++ -DCMAKE_C_COMPILER=/usr/bin/gcc -DCMAKE_BUILD_TYPE=Release -DSM="60;61;70;75;80" \
-          -DENABLE_MULTINODES=ON .. && \
-    make -j$(nproc) && make install && \
-    chmod +x /usr/local/hugectr/bin/* && \
-    chmod +x /usr/local/hugectr/lib/* && \
-    cd /hugectr/onnx_converter && \
-    python3 setup.py install && \
-    rm -rf /hugectr/build
+RUN if [ "$HUGECTR_DEV_MODE" == "false" ]; then \
+        git clone https://${_CI_JOB_TOKEN}${_HUGECTR_REPO} /hugectr && cd /hugectr; \
+        if [ "$RELEASE" == "true" ] && [ "$HUGECTR_VER" != "vnightly" ]; then \
+            git fetch --all --tags && git checkout tags/${HUGECTR_VER}; \
+        else \
+            git checkout ${_HUGECTR_BRANCH}; \
+        fi; \
+        git submodule update --init --recursive && \
+        mkdir build && cd build && \
+        LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs/:$LD_LIBRARY_PATH && \
+        export PATH=$PATH:/usr/local/cuda-${CUDA_SHORT_VERSION}/compat/ && \
+        cmake -DCMAKE_CXX_COMPILER=/usr/bin/g++ -DCMAKE_C_COMPILER=/usr/bin/gcc -DCMAKE_BUILD_TYPE=Release -DSM="60;61;70;75;80" \
+            -DENABLE_MULTINODES=ON .. && \
+        make -j$(nproc) && make install && \
+        chmod +x /usr/local/hugectr/bin/* && \
+        chmod +x /usr/local/hugectr/lib/* && \
+        cd /hugectr/onnx_converter && \
+        python3 setup.py install && \
+        rm -rf /hugectr; \
+    fi
 
 ENV PATH=/usr/local/hugectr/bin:$PATH
 ENV LIBRARY_PATH=/usr/local/hugectr/lib:$LIBRARY_PATH

docker/dockerfile.tf

@@ -95,19 +95,29 @@ ARG HUGECTR_VER=vnightly
 ARG SM="60;61;70;75;80"
 ARG USE_NVTX=OFF
 
+# Arguments "_XXXX" are only valid when $HUGECTR_DEV_MODE==false
+ARG HUGECTR_DEV_MODE=false
+ARG _HUGECTR_BRANCH=master
+ARG _HUGECTR_REPO="github.com/NVIDIA-Merlin/HugeCTR.git"
+ARG _CI_JOB_TOKEN=""
+
 RUN mkdir -p /usr/local/nvidia/lib64 && \
     ln -s /usr/local/cuda/lib64/libcusolver.so /usr/local/nvidia/lib64/libcusolver.so.10
 
 RUN ln -s /usr/lib/x86_64-linux-gnu/libibverbs.so.1 /usr/lib/x86_64-linux-gnu/libibverbs.so
 
-RUN git clone https://github.com/NVIDIA-Merlin/HugeCTR.git build-env && \
-    pushd build-env && \
-      if [ "$RELEASE" == "true" ] && [ ${HUGECTR_VER} != "vnightly" ] ; then git fetch --all --tags && git checkout tags/${HUGECTR_VER}; else echo ${HUGECTR_VER} && git checkout master; fi && \
-      cd sparse_operation_kit && \
-      python setup.py install && \
-    popd && \
-    rm -rf build-env && \
-    rm -rf /var/tmp/HugeCTR
+RUN if [ "$HUGECTR_DEV_MODE" == "false" ]; then \
+        git clone https://${_CI_JOB_TOKEN}${_HUGECTR_REPO} build-env && pushd build-env && git fetch --all; \
+        if [ "$RELEASE" == "true" ] && [ ${HUGECTR_VER} != "vnightly" ]; then \
+            git fetch --all --tags && git checkout tags/${HUGECTR_VER}; \
+        else \
+            git checkout ${_HUGECTR_BRANCH}; \
+        fi; \
+        cd sparse_operation_kit && \
+        python setup.py install && \
+        popd && \
+        rm -rf build-env; \
+    fi
 
 RUN pip install pybind11
 RUN pip install numba numpy --upgrade

docker/dockerfile.tri

@@ -80,7 +80,7 @@ RUN git clone --branch apache-arrow-4.0.1 --recurse-submodules https://github.co
               -DCMAKE_LIBRARY_PATH=${CUDA_CUDA_LIBRARY} \
               -DARROW_FLIGHT=ON \
               -DARROW_GANDIVA=OFF \
-              -DARROW_ORC=ON \
+              -DARROW_ORC=OFF \
               -DARROW_WITH_BZ2=ON \
               -DARROW_WITH_ZLIB=ON \
               -DARROW_WITH_ZSTD=ON \
@@ -100,7 +100,7 @@ RUN git clone --branch apache-arrow-4.0.1 --recurse-submodules https://github.co
       pushd python && \
         export PYARROW_WITH_PARQUET=ON && \
         export PYARROW_WITH_CUDA=ON && \
-        export PYARROW_WITH_ORC=ON && \
+        export PYARROW_WITH_ORC=OFF && \
         export PYARROW_WITH_DATASET=ON && \
         python setup.py build_ext --build-type=release bdist_wheel && \
         pip install dist/*.whl && \
@@ -215,31 +215,51 @@ RUN apt-get update -y && \
 
 ENV CPATH=/usr/local/include:$CPATH
 
+# Arguments "_XXXX" are only valid when $HUGECTR_DEV_MODE==false
+ARG HUGECTR_DEV_MODE=false
+ARG _HUGECTR_BRANCH=master
+ARG _HUGECTR_REPO="github.com/NVIDIA-Merlin/HugeCTR.git"
+ARG _HUGECTR_BACKEND_BRANCH=main
+ARG _HUGECTR_BACKEND_REPO="github.com/triton-inference-server/hugectr_backend"
+
+ARG _CI_JOB_TOKEN=""
+
 # Install HugeCTR
 RUN apt update -y && apt install rapidjson-dev -y
-RUN git clone https://github.com/NVIDIA-Merlin/HugeCTR.git /hugectr && \
-      cd /hugectr && if [ "$RELEASE" == "true" ] && [ ${HUGECTR_VER} != "vnightly" ]; then git fetch --all --tags && git checkout tags/${HUGECTR_VER}; else git checkout master; fi && \
+RUN if [ "$HUGECTR_DEV_MODE" == "false" ]; then \
+      git clone https://${_CI_JOB_TOKEN}${_HUGECTR_REPO} /hugectr && cd /hugectr && git fetch --all; \
+      if [ "$RELEASE" == "true" ] && [ ${HUGECTR_VER} != "vnightly" ]; then \
+        git fetch --all --tags && git checkout tags/${HUGECTR_VER}; \
+      else \
+        git checkout ${_HUGECTR_BRANCH}; \
+      fi; \
       git submodule update --init --recursive && \
       mkdir -p build && cd build &&\
       cmake -DCMAKE_BUILD_TYPE=Release -DSM=$SM -DENABLE_INFERENCE=ON .. && \
       make -j$(nproc) && make install && \
-      chmod +x /usr/local/hugectr/bin/* &&\
-      export CPATH=/usr/local/hugectr/include:$CPATH && \
-      export LIBRARY_PATH=/usr/local/hugectr/lib:$LIBRARY_PATH && \
-      git clone https://github.com/triton-inference-server/hugectr_backend /repos/hugectr_inference_backend && \
-      cd /repos/hugectr_inference_backend && if [ "$RELEASE" == "true" ] && [ ${HUGECTR_BACKEND_VER} != "vnightly" ] ; then git fetch --all --tags && git checkout tags/${HUGECTR_BACKEND_VER}; else git checkout main; fi && \
+      chmod +x /usr/local/hugectr/bin/*; \
+    fi
+
+ENV CPATH=/usr/local/hugectr/include:$CPATH
+ENV LIBRARY_PATH=/usr/local/hugectr/lib:$LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/usr/local/hugectr/lib:$LD_LIBRARY_PATH
+ENV PATH=/usr/local/hugectr/bin:$PATH
+
+RUN if [ "$HUGECTR_DEV_MODE" == "false" ]; then \
+      git clone https://${_CI_JOB_TOKEN}${_HUGECTR_BACKEND_REPO} /repos/hugectr_inference_backend && cd /repos/hugectr_inference_backend &&  \
+      if [ "$RELEASE" == "true" ] && [ "$HUGECTR_BACKEND_VER" != "vnightly" ]; then \
+        git fetch --all --tags && git checkout tags/${HUGECTR_BACKEND_VER}; \
+      else \
+        git checkout ${_HUGECTR_BACKEND_BRANCH}; \
+      fi && \
       mkdir -p build && cd build && \
       cmake -DCMAKE_INSTALL_PREFIX:PATH=/usr/local/hugectr \
             -DTRITON_COMMON_REPO_TAG="r$TRITON_VERSION" \
             -DTRITON_CORE_REPO_TAG="r$TRITON_VERSION" \
             -DTRITON_BACKEND_REPO_TAG="r$TRITON_VERSION" .. && \
       make -j$(nproc) && make install && \
-      rm -rf /hugectr/build
-
-ENV CPATH=/usr/local/hugectr/include:$CPATH
-ENV LIBRARY_PATH=/usr/local/hugectr/lib:$LIBRARY_PATH
-ENV LD_LIBRARY_PATH=/usr/local/hugectr/lib:$LD_LIBRARY_PATH
-ENV PATH=/usr/local/hugectr/bin:$PATH
+      rm -rf /repos/hugectr_inference_backend; \
+    fi
 
 RUN ln -s /usr/local/hugectr/backends/hugectr /opt/tritonserver/backends/