Use non default git repositories for triton dependencies (#404)

triton-inference-server · Sep 20, 2024 · f68a695 · f68a695
1 parent 88ac83e
commit f68a695
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 5 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -20,9 +20,11 @@ option(TRITON_ENABLE_STATS "Include statistics collections in backend" ON)
 option(TRITON_FIL_DOCKER_BUILD OFF)
 option(BUILD_SHARED_LIBS OFF)
 
+set(TRITON_REPO_ORGANIZATION "https://github.com/triton-inference-server" CACHE STRING "Git repository to pull from")
 set(TRITON_COMMON_REPO_TAG "main" CACHE STRING "Tag for triton-inference-server/common repo")
 set(TRITON_CORE_REPO_TAG "main" CACHE STRING "Tag for triton-inference-server/core repo")
 set(TRITON_BACKEND_REPO_TAG "main" CACHE STRING "Tag for triton-inference-server/backend repo")
+set(RAPIDS_TRITON_REPO_TAG "main" CACHE STRING "Tag for rapidsai/rapids-triton repo")
 
 if(TRITON_FIL_DOCKER_BUILD)
   project(RAPIDS_TRITON_BACKEND VERSION 22.10.00)
@@ -38,9 +40,11 @@ if(TRITON_FIL_DOCKER_BUILD)
       --build-arg BASE_IMAGE=${TRITON_BUILD_CONTAINER}
       --build-arg TRITON_ENABLE_GPU=${TRITON_ENABLE_GPU}
       --build-arg TRITON_ENABLE_STATS=${TRITON_ENABLE_GPU}
+      --build-arg TRITON_REPO_ORGANIZATION=${TRITON_REPO_ORGANIZATION}
       --build-arg TRITON_COMMON_REPO_TAG=${TRITON_COMMON_REPO_TAG}
       --build-arg TRITON_CORE_REPO_TAG=${TRITON_CORE_REPO_TAG}
       --build-arg TRITON_BACKEND_REPO_TAG=${TRITON_BACKEND_REPO_TAG}
+      --build-arg RAPIDS_TRITON_REPO_TAG=${RAPIDS_TRITON_REPO_TAG}
       -f ${CMAKE_CURRENT_LIST_DIR}/ops/Dockerfile
       ${CMAKE_CURRENT_LIST_DIR}
     COMMAND docker rm triton_fil_builder || echo 'error ignored..' || true

diff --git a/build.sh b/build.sh
@@ -49,6 +49,10 @@ HELP="$0 [<target> ...] [<flag> ...]
                       the build the test image
    PREBUILT_IMAGE   - A server image to be tested (used as base of test image)
    TRITON_REF       - Commit ref for Triton when using build.py
+   REPO_ORGANIZATION
+                    - Repository organization to pull triton repositories from
+                      Default is \"https://github.com/triton-inference-server\"
+
    COMMON_REF       - Commit ref for Triton common repo when using build.py
    CORE_REF         - Commit ref for Triton core repo when using build.py
    BACKEND_REF      - Commit ref for Triton backend repo when using build.py
@@ -199,6 +203,7 @@ then
   # If the user has specified a TRITON_VERSION (or if we are performing a host
   # build), set the upstream repo references to the corresponding branches
   # (unless otherwise specified by the user)
+  [ ! -z $REPO_ORGANIZATION ] || REPO_ORGANIZATION="${REPO_ORGANIZATION}"
   [ ! -z $TRITON_REF ] || TRITON_REF="r${TRITON_VERSION}"
   [ ! -z $COMMON_REF ] || COMMON_REF="r${TRITON_VERSION}"
   [ ! -z $CORE_REF ] || CORE_REF="r${TRITON_VERSION}"
@@ -208,6 +213,7 @@ else
   # If TRITON_VERSION has not been set, these values will only be used for a
   # full build.py build, so it is safe to default to main rather than a release
   # branch.
+  [ ! -z $REPO_ORGANIZATION ] || REPO_ORGANIZATION='https://github.com/triton-inference-server'
   [ ! -z $TRITON_REF ] || TRITON_REF='main'
   [ ! -z $COMMON_REF ] || COMMON_REF='main'
   [ ! -z $CORE_REF ] || CORE_REF='main'
@@ -275,7 +281,7 @@ buildpy () {
   then
     rm -rf "$server_repo"
   fi
-  git clone https://github.com/triton-inference-server/server.git \
+  git clone $REPO_ORGANIZATION/server.git \
     "${server_repo}" \
     -b $TRITON_REF \
     --depth 1
@@ -288,6 +294,7 @@ buildpy () {
     --enable-stats \
     --endpoint=http \
     --endpoint=grpc \
+    --github-organization=$REPO_ORGANIZATION \
     --repo-tag=common:$COMMON_REF \
     --repo-tag=core:$CORE_REF \
     --repo-tag=backend:$BACKEND_REF \
@@ -311,6 +318,7 @@ hostbuild () {
     -GNinja \
     -DCMAKE_BUILD_TYPE="$BUILD_TYPE" \
     -DBUILD_TESTS="OFF" \
+    -DTRITON_REPO_ORGANIZATION="$REPO_ORGANIZATION" \
     -DTRITON_CORE_REPO_TAG="$CORE_REF" \
     -DTRITON_COMMON_REPO_TAG="$COMMON_REF" \
     -DTRITON_BACKEND_REPO_TAG="$BACKEND_REF" \

diff --git a/ci/gitlab/build.sh b/ci/gitlab/build.sh
@@ -36,6 +36,7 @@ else
   export PREBUILT_IMAGE="$PREBUILT_SERVER_TAG"
   export SERVER_TAG="$PREBUILT_SERVER_TAG"
 fi
+[ -z $TRITON_REPO_ORGANIZATION ] || export REPO_ORGANIZATION="$TRITON_REPO_ORGANIZATION"
 [ -z $TRITON_SERVER_REPO_TAG ] || export TRITON_REF="$TRITON_SERVER_REPO_TAG"
 [ -z $TRITON_COMMON_REPO_TAG ] || export COMMON_REF="$TRITON_COMMON_REPO_TAG"
 [ -z $TRITON_CORE_REPO_TAG ] || export CORE_REF="$TRITON_CORE_REPO_TAG"

diff --git a/cmake/thirdparty/get_rapids-triton.cmake b/cmake/thirdparty/get_rapids-triton.cmake
@@ -25,7 +25,7 @@ function(find_and_configure_rapids_triton)
       BUILD_EXPORT_SET    rapids_triton-exports
       INSTALL_EXPORT_SET  rapids_triton-exports
         CPM_ARGS
-            GIT_REPOSITORY https://github.com/${PKG_FORK}/rapids-triton.git
+            GIT_REPOSITORY ${PKG_FORK}
             GIT_TAG        ${PKG_PINNED_TAG}
             SOURCE_SUBDIR  cpp
             OPTIONS
@@ -40,7 +40,19 @@ endfunction()
 # Change pinned tag here to test a commit in CI
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
+set (RAPIDS_FORK https://github.com/rapidsai/rapids-triton.git)
+set (REPO_TAG branch-${RAPIDS_DEPENDENCIES_VERSION})
+# if Triton tag and organization is specified, change the fork and the repo
+if (NOT TRITON_REPO_ORGANIZATION STREQUAL "https://github.com/triton-inference-server")
+  set (RAPIDS_FORK ${TRITON_REPO_ORGANIZATION}/rapids_triton.git)
+  message(STATUS "Setting repo fork to ${RAPIDS_FORK}")
+endif()
+if (NOT RAPIDS_TRITON_REPO_TAG STREQUAL "main")
+  set (REPO_TAG ${RAPIDS_TRITON_REPO_TAG})
+  message(STATUS "Setting repo tag to ${REPO_TAG}")
+endif()
+
 find_and_configure_rapids_triton(VERSION    ${RAPIDS_TRITON_MIN_VERSION_rapids_projects}
-                                 FORK       rapidsai
-                                 PINNED_TAG branch-${RAPIDS_DEPENDENCIES_VERSION}
+                                 FORK       ${RAPIDS_FORK}
+                                 PINNED_TAG ${REPO_TAG}
                                  )
diff --git a/ops/Dockerfile b/ops/Dockerfile
@@ -123,12 +123,16 @@ ENV BUILD_TESTS=$BUILD_TESTS
 ARG BUILD_EXAMPLE
 ENV BUILD_EXAMPLE=$BUILD_EXAMPLE
 
+ARG TRITON_REPO_ORGANIZATION=${TRITON_REPO_ORGANIZATION}
+ENV TRITON_REPO_ORGANIZATION=$TRITON_REPO_ORGANIZATION
 ARG TRITON_CORE_REPO_TAG=r${TRITON_VERSION}
 ENV TRITON_CORE_REPO_TAG=$TRITON_CORE_REPO_TAG
 ARG TRITON_COMMON_REPO_TAG=r${TRITON_VERSION}
 ENV TRITON_COMMON_REPO_TAG=$TRITON_COMMON_REPO_TAG
 ARG TRITON_BACKEND_REPO_TAG=r${TRITON_VERSION}
 ENV TRITON_BACKEND_REPO_TAG=$TRITON_BACKEND_REPO_TAG
+ARG RAPIDS_TRITON_REPO_TAG=${RAPIDS_TRITON_REPO_TAG}
+ENV RAPIDS_TRITON_REPO_TAG=${RAPIDS_TRITON_REPO_TAG}
 
 ARG TRITON_ENABLE_GPU=ON
 ENV TRITON_ENABLE_GPU=$TRITON_ENABLE_GPU
@@ -153,9 +157,11 @@ RUN source /conda/dev/bin/activate \
       -GNinja \
       -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \
       -DBUILD_TESTS="${BUILD_TESTS}" \
+      -DTRITON_REPO_ORGANIZATION="${TRITON_REPO_ORGANIZATION}" \
       -DTRITON_CORE_REPO_TAG="${TRITON_CORE_REPO_TAG}" \
       -DTRITON_COMMON_REPO_TAG="${TRITON_COMMON_REPO_TAG}" \
       -DTRITON_BACKEND_REPO_TAG="${TRITON_BACKEND_REPO_TAG}" \
+      -DRAPIDS_TRITON_REPO_TAG="${RAPIDS_TRITON_REPO_TAG}" \
       -DTRITON_ENABLE_GPU="${TRITON_ENABLE_GPU}" \
       -DTRITON_ENABLE_STATS="${TRITON_ENABLE_STATS}" \
       -DRAPIDS_DEPENDENCIES_VERSION="${RAPIDS_DEPENDENCIES_VERSION}" \
@@ -165,7 +171,7 @@ RUN source /conda/dev/bin/activate \
 
 ENV CCACHE_DIR=/ccache
 
-RUN --mount=type=cache,target=/ccache/ source /conda/dev/bin/activate && ninja install
+RUN --mount=type=cache,target=/ccache/ source /conda/dev/bin/activate && ninja -j16 install
 
 
 # Stage for generating testing image