Enable custom triton-rapids repo and tag (#409)

* enable changing of rapids triton branch and rapids backend triton remote

* add rapids triton stuff back

* address comments

* align naming in dockerfile

CMakeLists.txt

@@ -20,11 +20,12 @@ 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_REPO_ORGANIZATION "https://github.com/triton-inference-server" CACHE STRING "Git repository to pull triton libraries 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")
+set(RAPIDS_TRITON_REPO_PATH "https://github.com/rapidsai/rapids-triton.git" CACHE STRING "Git repository to pull rapids_triton from")
 
 if(TRITON_FIL_DOCKER_BUILD)
   project(RAPIDS_TRITON_BACKEND VERSION 22.10.00)
@@ -45,6 +46,7 @@ if(TRITON_FIL_DOCKER_BUILD)
       --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}
+      --build-arg RAPIDS_TRITON_REPO_PATH=${RAPIDS_TRITON_REPO_PATH}
       -f ${CMAKE_CURRENT_LIST_DIR}/ops/Dockerfile
       ${CMAKE_CURRENT_LIST_DIR}
     COMMAND docker rm triton_fil_builder || echo 'error ignored..' || true

build.sh

@@ -16,7 +16,6 @@
 set -e
 
 REPODIR=$(cd $(dirname $0); pwd)
-
 NUMARGS=$#
 ARGS=$*
 VALIDTARGETS="server tests"
@@ -49,10 +48,6 @@ 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
@@ -203,7 +198,6 @@ 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='https://github.com/triton-inference-server'
   [ ! -z $TRITON_REF ] || TRITON_REF="r${TRITON_VERSION}"
   [ ! -z $COMMON_REF ] || COMMON_REF="r${TRITON_VERSION}"
   [ ! -z $CORE_REF ] || CORE_REF="r${TRITON_VERSION}"
@@ -213,7 +207,6 @@ 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'
@@ -281,7 +274,7 @@ buildpy () {
   then
     rm -rf "$server_repo"
   fi
-  git clone $REPO_ORGANIZATION/server.git \
+  git clone https://github.com/triton-inference-server/server.git \
     "${server_repo}" \
     -b $TRITON_REF \
     --depth 1
@@ -294,7 +287,6 @@ 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 \
@@ -318,7 +310,6 @@ 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" \

ci/gitlab/build.sh

@@ -36,7 +36,6 @@ 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"

cmake/thirdparty/get_rapids-triton.cmake

@@ -39,14 +39,13 @@ endfunction()
 # CPM_raft_SOURCE=/path/to/local/raft
 set (RAPIDS_FORK https://github.com/rapidsai/rapids-triton.git)
 set (REPO_TAG branch-${RAPIDS_DEPENDENCIES_VERSION})
-message(STATUS "Setting repo tag to ${REPO_TAG}")
-# 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")
+message(STATUS "Setting repo tag to ${REPO_TAG} for rapids fork ${RAPIDS_FORK}")
+# if Triton tag and organization is non-default, change the fork and repo tag used
+# for rapids
+if (NOT RAPIDS_TRITON_REPO_PATH STREQUAL RAPIDS_FORK)
+  set (RAPIDS_FORK ${RAPIDS_TRITON_REPO_PATH})
   set (REPO_TAG ${RAPIDS_TRITON_REPO_TAG})
+  message(STATUS "Re-setting repo tag to ${REPO_TAG} for rapids fork ${RAPIDS_FORK}")
 endif()
 
 find_and_configure_rapids_triton(VERSION    ${RAPIDS_DEPENDENCIES_VERSION}

ops/Dockerfile

@@ -132,7 +132,9 @@ 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=main
-ENV RAPIDS_TRITON_REPO_TAG=${RAPIDS_TRITON_REPO_TAG}
+ENV RAPIDS_TRITON_REPO_TAG=$RAPIDS_TRITON_REPO_TAG
+ARG RAPIDS_TRITON_REPO_PATH=https://github.com/rapidsai/rapids-triton.git
+ENV RAPIDS_TRITON_REPO_PATH=$RAPIDS_TRITON_REPO_PATH
 
 ARG TRITON_ENABLE_GPU=ON
 ENV TRITON_ENABLE_GPU=$TRITON_ENABLE_GPU
@@ -162,6 +164,7 @@ RUN source /conda/dev/bin/activate \
       -DTRITON_COMMON_REPO_TAG="${TRITON_COMMON_REPO_TAG}" \
       -DTRITON_BACKEND_REPO_TAG="${TRITON_BACKEND_REPO_TAG}" \
       -DRAPIDS_TRITON_REPO_TAG="${RAPIDS_TRITON_REPO_TAG}" \
+      -DRAPIDS_TRITON_REPO_PATH="${RAPIDS_TRITON_REPO_PATH}" \
       -DTRITON_ENABLE_GPU="${TRITON_ENABLE_GPU}" \
       -DTRITON_ENABLE_STATS="${TRITON_ENABLE_STATS}" \
       -DRAPIDS_DEPENDENCIES_VERSION="${RAPIDS_DEPENDENCIES_VERSION}" \
@@ -171,7 +174,7 @@ RUN source /conda/dev/bin/activate \
 
 ENV CCACHE_DIR=/ccache
 
-RUN --mount=type=cache,target=/ccache/ source /conda/dev/bin/activate && ninja -j16 install
+RUN --mount=type=cache,target=/ccache/ source /conda/dev/bin/activate && ninja install
 
 
 # Stage for generating testing image