Internal change only

PiperOrigin-RevId: 792359286

Author: SiqiaoWu1993

tensorflow_serving/servables/tensorflow/BUILD

@@ -1254,6 +1254,8 @@ cc_library(
         "//tensorflow_serving/resources:resource_values",
         "//tensorflow_serving/resources:resources_cc_proto",
         "//tensorflow_serving/session_bundle:graph_rewriter",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -1266,7 +1268,9 @@ cc_library(
         "@org_tensorflow//tensorflow/core:core_cpu",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
+        "@org_tensorflow//tensorflow/core/kernels/batching_util:batch_scheduler_hdrs",
         "@org_tensorflow//tensorflow/core/kernels/batching_util:shared_batch_scheduler_hdrs",
+        "@org_tensorflow//tensorflow/core/tfrt/graph_executor:graph_execution_options",
         "@org_tensorflow//tensorflow/core/tfrt/runtime",
         "@org_tensorflow//tensorflow/core/tfrt/saved_model:saved_model_cpu",
     ],
@@ -1320,18 +1324,22 @@ cc_library(
     ],
     deps = [
         ":bundle_factory_util",
+        ":file_acl",
         ":machine_learning_metadata",
         ":servable",
         ":tfrt_saved_model_factory",
         ":tfrt_saved_model_source_adapter_cc_proto",
         ":tfrt_servable",
+        "//tensorflow_serving/core:loader",
         "//tensorflow_serving/core:simple_loader",
         "//tensorflow_serving/core:source_adapter",
         "//tensorflow_serving/core:storage_path",
         "//tensorflow_serving/resources:resource_util",
         "//tensorflow_serving/resources:resource_values",
         "//tensorflow_serving/resources:resources_cc_proto",
-        "//tensorflow_serving/servables/tensorflow:file_acl",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
         "@org_tensorflow//tensorflow/cc/saved_model:loader",
         "@org_tensorflow//tensorflow/core:lib",
     ],

tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc

@@ -17,13 +17,15 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 #include <unordered_set>
 #include <utility>
 #include <vector>
 
 #include "google/protobuf/wrappers.pb.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -33,8 +35,11 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
 #include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
 #include "tensorflow/core/public/session_options.h"
@@ -401,5 +406,14 @@ CreateThreadPoolFactoryFromConfig(const TfrtSavedModelConfig& config) {
   return thread_pool_factory;
 }
 
+// copybara:strip_begin (Do not leak in tesorflow serving OSS.)
+absl::Status TfrtSavedModelFactory::CreateOrbaxServable(
+    const Loader::Metadata& metadata, const string& path,
+    std::unique_ptr<Servable>* servable) {
+  return absl::UnimplementedError(
+      "CreateOrbaxServable is not implemented yet.");
+}
+// copybara:strip_end
+
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.h

@@ -16,21 +16,30 @@ limitations under the License.
 #ifndef TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_TFRT_SAVED_MODEL_FACTORY_H_
 #define TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_TFRT_SAVED_MODEL_FACTORY_H_
 
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <string>
+#include <unordered_set>
+#include <utility>
 
+#include "absl/base/attributes.h"
+#include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
-#include "absl/types/optional.h"
+#include "xla/tsl/platform/macros.h"
 #include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h"
-#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
 #include "tensorflow/core/tfrt/runtime/runtime.h"
 #include "tensorflow/core/tfrt/saved_model/saved_model.h"
 #include "tensorflow_serving/batching/tfrt_saved_model_with_batching.h"
 #include "tensorflow_serving/core/loader.h"
 #include "tensorflow_serving/resources/resources.pb.h"
+#include "tensorflow_serving/servables/tensorflow/servable.h"
 #include "tensorflow_serving/servables/tensorflow/session_bundle_config.pb.h"
 #include "tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.pb.h"
 #include "tensorflow_serving/servables/tensorflow/tfrt_servable.h"
@@ -106,6 +115,12 @@ class TfrtSavedModelFactory {
   absl::Status EstimateResourceRequirement(const string& path,
                                            ResourceAllocation* estimate) const;
 
+  // copybara:strip_begin (Do not leak in tesorflow serving OSS.)
+  virtual absl::Status CreateOrbaxServable(const Loader::Metadata& metadata,
+                                           const string& path,
+                                           std::unique_ptr<Servable>* servable);
+  // copybara:strip_end
+
   const TfrtSavedModelConfig& config() const { return config_; }
   TfrtSavedModelConfig& mutable_config() { return config_; }
   absl::string_view GetServingResourceType() const;

tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.cc

@@ -16,11 +16,21 @@ limitations under the License.
 #include "tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.h"
 
 #include <memory>
-
+#include <string>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/saved_model/loader.h"
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
-#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow_serving/core/loader.h"
 #include "tensorflow_serving/core/simple_loader.h"
+#include "tensorflow_serving/core/source_adapter.h"
+#include "tensorflow_serving/core/storage_path.h"
 #include "tensorflow_serving/resources/resource_util.h"
 #include "tensorflow_serving/resources/resource_values.h"
 #include "tensorflow_serving/resources/resources.pb.h"
@@ -29,11 +39,31 @@ limitations under the License.
 #include "tensorflow_serving/servables/tensorflow/machine_learning_metadata.h"
 #include "tensorflow_serving/servables/tensorflow/servable.h"
 #include "tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.h"
-#include "tensorflow_serving/servables/tensorflow/tfrt_servable.h"
 
 namespace tensorflow {
 namespace serving {
 
+// copybara:strip_begin (Do not leak in tensorflow serving OSS.)
+namespace {
+// Orbax manifest file name.
+inline constexpr char kOrbaxModelManifestPb[] = "manifest.pb";
+// Orbax manifest version file name.
+inline constexpr char kOrbaxModelManifestVersionTxt[] = "manifest_version.txt";
+
+absl::Status IsOrbaxModelDirectory(absl::string_view path) {
+  const std::string orbax_model_manifest_pb_path =
+      tensorflow::io::JoinPath(path, kOrbaxModelManifestPb);
+  const std::string orbax_model_manifest_version_path =
+      tensorflow::io::JoinPath(path, kOrbaxModelManifestVersionTxt);
+  tsl::Env* env = tsl::Env::Default();
+  TF_RETURN_IF_ERROR(env->FileExists(orbax_model_manifest_pb_path));
+  TF_RETURN_IF_ERROR(env->FileExists(orbax_model_manifest_version_path));
+  return absl::OkStatus();
+}
+
+}  // namespace
+// copybara:strip_end
+
 absl::Status TfrtSavedModelSourceAdapter::Create(
     const TfrtSavedModelSourceAdapterConfig& config,
     std::unique_ptr<TfrtSavedModelSourceAdapter>* adapter) {
@@ -57,9 +87,20 @@ TfrtSavedModelSourceAdapter::GetServableCreator(
   return [factory, path](const Loader::Metadata& metadata,
                          std::unique_ptr<Servable>* servable) {
     TF_RETURN_IF_ERROR(RegisterModelRoot(metadata.servable_id, path));
-    TF_RETURN_IF_ERROR(
-        factory->CreateTfrtSavedModelWithMetadata(metadata, path, servable));
-    return absl::OkStatus();
+    if (MaybeSavedModelDirectory(path)) {
+      return factory->CreateTfrtSavedModelWithMetadata(metadata, path,
+                                                       servable);
+    }
+
+    // copybara:strip_begin (Do not leak in tesorflow serving OSS.)
+    if (IsOrbaxModelDirectory(path).ok()) {
+      return factory->CreateOrbaxServable(metadata, path, servable);
+    }
+    // copybara:strip_end
+
+    return absl::InvalidArgumentError(
+        absl::StrCat("Unsupported model directory: ", path,
+                     ". Only SavedModel and Orbax Model are supported."));
   };
 }
 

tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto

@@ -238,6 +238,12 @@ message TfrtSavedModelConfig {
   // If non-zero, all models on this server are switched to use a prioritized
   // batching function using this number of global threads.
   int64 tfrt_batch_queue_global_prioritization_num_threads = 2029;
+
+  // copybara:strip_begin (Do not leak in tesorflow serving OSS.)
+  // If true, allow Orbax as an additional format for loading models.
+  bool allow_orbax = 2030;
+  // copybara:strip_end
+  bool allow_xla_cpu = 2031;
 }
 
 // Config proto for TfrtSavedModelSourceAdapter.