Add some internal change.

PiperOrigin-RevId: 633410091
tensorflow · May 14, 2024 · 820f077 · 820f077
1 parent 3e4ed32
commit 820f077
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
@@ -490,6 +490,7 @@ tf_cc_test(
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

diff --git a/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc
@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
@@ -118,6 +119,44 @@ TEST(IfrtExecutableRegistry, Basic) {
   ASSERT_EQ(executable_ptr, raw_ptr);
 }
 
+TEST(IfrtExecutableRegistry, DuplicateRegistrationFails) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  EXPECT_THAT(
+      ServingExecutableRegistry::Register(program_id, std::move(executable)),
+      testing::StatusIs(absl::StatusCode::kAlreadyExists));
+}
+
+TEST(IfrtExecutableRegistry, ReleaseOk) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_EQ(ServingExecutableRegistry::Lookup(program_id), nullptr);
+}
+
 TEST(IfrtExecutableRegistry, FreezeOk) {
   mlir::DialectRegistry registry;
   mlir::registerAllDialects(registry);
@@ -142,6 +181,27 @@ TEST(IfrtExecutableRegistry, FreezeOk) {
   ASSERT_EQ(executable_ptr, raw_ptr);
 }
 
+TEST(IfrtExecutableRegistry, FreezeFailedProgramNotRegistered) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_THAT(handle.Freeze(),
+              testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+}
+
 TEST(IfrtExecutableRegistry, InvalidProgramIdShallReturnNull) {
   int64_t program_id = 1234;