Moves model server TFRT integration code oss

PiperOrigin-RevId: 562885221
tensorflow · Sep 21, 2023 · 50ebab4 · 50ebab4
1 parent af1e870
commit 50ebab4
Show file tree

Hide file tree

Showing 58 changed files with 8,462 additions and 16 deletions.
diff --git a/tensorflow_serving/batching/BUILD b/tensorflow_serving/batching/BUILD
@@ -204,3 +204,39 @@ cc_test(
         "@org_tensorflow//tensorflow/core:test",
     ],
 )
+
+cc_library(
+    name = "tfrt_saved_model_with_batching",
+    srcs = ["tfrt_saved_model_with_batching.cc"],
+    hdrs = ["tfrt_saved_model_with_batching.h"],
+    deps = [
+        ":batching_options",
+        ":batching_session",
+        ":batching_util",
+        ":incremental_barrier",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@org_tensorflow//tensorflow/core:framework",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core/kernels/batching_util:batch_scheduler",
+        "@org_tensorflow//tensorflow/core/tfrt/saved_model",
+    ],
+)
+
+cc_test(
+    name = "tfrt_saved_model_with_batching_test",
+    srcs = ["tfrt_saved_model_with_batching_test.cc"],
+    deps = [
+        ":batching_util",
+        ":tfrt_saved_model_with_batching",
+        "//tensorflow_serving/core/test_util:test_main",
+        "//tensorflow_serving/servables/tensorflow/test_util:mock_tfrt_saved_model",
+        "@com_google_absl//absl/functional:bind_front",
+        "@com_google_googletest//:gtest",
+        "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:test",
+        "@org_tensorflow//tensorflow/core:testlib",
+        "@org_tensorflow//tensorflow/core/kernels/batching_util:basic_batch_scheduler",
+        "@org_tensorflow//tensorflow/core/tfrt/utils:tensor_util",
+    ],
+)