Add plumbing for git_repo/repo_tag for tensorflow

triton-inference-server · Dec 6, 2023 · fd4b24d · fd4b24d
1 parent acc4893
commit fd4b24d
Showing 1 changed file with 35 additions and 3 deletions.
diff --git a/build.py b/build.py
@@ -593,7 +593,9 @@ def backend_cmake_args(
     elif be == "openvino":
         args = openvino_cmake_args()
     elif be == "tensorflow":
-        args = tensorflow_cmake_args(images, library_paths)
+        args = tensorflow_cmake_args(
+            images, library_paths, components, components_git_orgs
+        )
     elif be == "python":
         args = []
     elif be == "dali":
@@ -848,11 +850,39 @@ def tensorrt_cmake_args():
     return cargs
 
 
-def tensorflow_cmake_args(images, library_paths):
+def tensorflow_cmake_args(images, library_paths, components, components_git_orgs):
     backend_name = "tensorflow"
 
     # If platform is jetpack do not use docker images
-    extra_args = []
+    extra_args = [
+        cmake_backend_arg(
+            "tensorflow", "TRITON_COMMON_REPO_TAG", "STRING", components["common"]
+        ),
+        cmake_backend_arg(
+            "tensorflow",
+            "TRITON_COMMON_GIT_REPO",
+            "STRING",
+            components_git_orgs["common"] + "common.git",
+        ),
+        cmake_backend_arg(
+            "tensorflow", "TRITON_CORE_REPO_TAG", "STRING", components["core"]
+        ),
+        cmake_backend_arg(
+            "tensorflow",
+            "TRITON_CORE_GIT_REPO",
+            "STRING",
+            components_git_orgs["core"] + "core.git",
+        ),
+        cmake_backend_arg(
+            "tensorflow", "TRITON_BACKEND_REPO_TAG", "STRING", components["backend"]
+        ),
+        cmake_backend_arg(
+            "tensorflow",
+            "TRITON_BACKEND_GIT_REPO",
+            "STRING",
+            components_git_orgs["backend"] + "backend.git",
+        ),
+    ]
     if target_platform() == "jetpack":
         if backend_name in library_paths:
             extra_args = [
@@ -1878,6 +1908,7 @@ def backend_build(
     install_dir,
     images,
     components,
+    components_git_orgs,
     github_org,
     library_paths,
 ):
@@ -2883,6 +2914,7 @@ def enable_all():
                     script_install_dir,
                     images,
                     components,
+                    components_git_orgs,
                     github_org,
                     library_paths,
                 )