Internal change.

PiperOrigin-RevId: 686677820

Author: tensorflower-gardener

tensorflow_serving/servables/tensorflow/BUILD

@@ -159,15 +159,17 @@ cc_library(
     ],
     deps = [
         ":bundle_factory_util",
+        ":saved_model_config_cc_proto",
+        ":saved_model_config_util",
         ":session_bundle_config_cc_proto",
         ":tflite_session_lib",
         "//tensorflow_serving/batching:batching_session",
         "//tensorflow_serving/core:loader",
         "//tensorflow_serving/resources:resources_cc_proto",
-        "//tensorflow_serving/session_bundle:session_bundle_util",
+        "//tensorflow_serving/session_bundle:session_bundle_util",  # buildcleaner: keep
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
-        "@com_google_protobuf//:cc_wkt_protos",
         "@org_tensorflow//tensorflow/cc/saved_model:loader",
         "@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
         "@org_tensorflow//tensorflow/core:core_cpu",
@@ -190,14 +192,20 @@ cc_test(
         ":bundle_factory_test",
         ":bundle_factory_test_util",
         ":saved_model_bundle_factory",
+        ":saved_model_config_cc_proto",
         ":session_bundle_config_cc_proto",
         "//tensorflow_serving/core/test_util:session_test_util",
         "//tensorflow_serving/core/test_util:test_main",
+        "@com_google_absl//absl/status",
         "@com_google_protobuf//:cc_wkt_protos",
+        "@org_tensorflow//tensorflow/cc/saved_model:constants",
         "@org_tensorflow//tensorflow/cc/saved_model:loader",
         "@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
         "@org_tensorflow//tensorflow/core:core_cpu",
+        "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:lib",
+        "@org_tensorflow//tensorflow/core:portable_tensorflow_test_lib",
+        "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core:test",
     ],
 )

tensorflow_serving/servables/tensorflow/saved_model_bundle_factory.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/log/log.h"
 #include "absl/strings/string_view.h"
 #include "tensorflow/cc/saved_model/tag_constants.h"
 #include "tensorflow/core/framework/tensor.pb.h"
@@ -31,6 +32,8 @@ limitations under the License.
 #include "tensorflow/core/protobuf/rewriter_config.pb.h"
 #include "tensorflow/core/public/session_options.h"
 #include "tensorflow_serving/servables/tensorflow/bundle_factory_util.h"
+#include "tensorflow_serving/servables/tensorflow/saved_model_config.pb.h"
+#include "tensorflow_serving/servables/tensorflow/saved_model_config_util.h"
 #include "tensorflow_serving/servables/tensorflow/tflite_session.h"
 #include "tensorflow_serving/session_bundle/session_bundle_util.h"
 
@@ -125,15 +128,16 @@ Status SavedModelBundleFactory::InternalCreateSavedModelBundle(
   if (saved_model_tags.empty()) {
     saved_model_tags.insert(kSavedModelTagServe);
   }
+  bool is_tflite = config_.prefer_tflite_model() && TfLiteModelFound(path);
   const auto& session_options = [&]() {
     auto result = GetSessionOptions(config_);
     string mixed_precision_value = config_.mixed_precision();
+    tensorflow::ConfigProto& config = result.config;
+    GraphOptions* gopt = config.mutable_graph_options();
+    RewriterConfig* rwcfg = gopt->mutable_rewrite_options();
     if (!mixed_precision_value.empty()) {
       if (mixed_precision_value == "bfloat16") {
         LOG(INFO) << "Running inference with bfloat16 auto mixed precision";
-        tensorflow::ConfigProto& config = result.config;
-        GraphOptions* gopt = config.mutable_graph_options();
-        RewriterConfig* rwcfg = gopt->mutable_rewrite_options();
         rwcfg->set_auto_mixed_precision_onednn_bfloat16(RewriterConfig::ON);
       } else {
         LOG(WARNING)
@@ -147,10 +151,20 @@ Status SavedModelBundleFactory::InternalCreateSavedModelBundle(
       session_metadata->set_name(metadata->servable_id.name);
       session_metadata->set_version(metadata->servable_id.version);
     }
+    // Set other options from saved model config.
+    if (!is_tflite) {
+      absl::StatusOr<SavedModelConfig> model_config =
+          LoadSavedModelConfigOrDefault(path);
+      if (!model_config.ok()) {
+        LOG(WARNING) << "Failed to load saved model config: "
+                     << model_config.status();
+      } else if (model_config->has_session_overrides()) {
+        UpdateRewriterConfig(model_config->session_overrides(), rwcfg);
+      }
+    }
     return result;
   }();
 
-  bool is_tflite = config_.prefer_tflite_model() && TfLiteModelFound(path);
   if (is_tflite) {
     int num_tflite_pools = config_.num_tflite_pools();
     if (num_tflite_pools == 0 && config_.num_tflite_interpreters() > 0) {

tensorflow_serving/servables/tensorflow/saved_model_bundle_factory_test.cc

@@ -25,19 +25,28 @@ limitations under the License.
 #include "google/protobuf/wrappers.pb.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
 #include "tensorflow/cc/saved_model/constants.h"
 #include "tensorflow/cc/saved_model/loader.h"
 #include "tensorflow/cc/saved_model/tag_constants.h"
+#include "xla/tsl/lib/core/status_test_util.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/named_tensor.pb.h"
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"
 #include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/public/version.h"
+#include "tsl/platform/errors.h"
 #include "tensorflow_serving/core/test_util/session_test_util.h"
 #include "tensorflow_serving/servables/tensorflow/bundle_factory_test.h"
 #include "tensorflow_serving/servables/tensorflow/bundle_factory_test_util.h"
+#include "tensorflow_serving/servables/tensorflow/saved_model_config.pb.h"
 #include "tensorflow_serving/servables/tensorflow/session_bundle_config.pb.h"
 
 namespace tensorflow {
@@ -56,10 +65,11 @@ Status CreateBundleFromPath(const CreationType creation_type,
                             const string& path,
                             std::unique_ptr<SavedModelBundle>* bundle) {
   std::unique_ptr<SavedModelBundleFactory> factory;
-  TF_RETURN_IF_ERROR(SavedModelBundleFactory::Create(config, &factory));
   auto config_with_session_hook = config;
   config_with_session_hook.set_session_target(
       test_util::kNewSessionHookSessionTargetPrefix);
+  TF_RETURN_IF_ERROR(
+      SavedModelBundleFactory::Create(config_with_session_hook, &factory));
   test_util::SetNewSessionHook([&](const SessionOptions& session_options) {
     const bool enable_session_metadata =
         creation_type == CreationType::kWithMetadata;
@@ -288,6 +298,48 @@ TEST_P(SavedModelBundleFactoryTest, RunOptions) { TestRunOptions(); }
 
 TEST_P(SavedModelBundleFactoryTest, RunOptionsError) { TestRunOptionsError(); }
 
+TEST_P(SavedModelBundleFactoryTest, SetGraphOptionsFromSavedModelConfig) {
+  const std::string dst_dir = io::JoinPath(testing::TmpDir(), "model");
+  test_util::CopyDirOrDie(export_dir_, dst_dir);
+  tensorflow::Env* env = tensorflow::Env::Default();
+  TF_ASSERT_OK(env->CreateDir(io::JoinPath(dst_dir, "assets.extra")));
+  SavedModelConfig saved_model_config;
+  saved_model_config.mutable_session_overrides()->set_disable_meta_optimizer(
+      true);
+  TF_ASSERT_OK(tensorflow::WriteBinaryProto(
+      env, io::JoinPath(dst_dir, "assets.extra", "saved_model_config.pb"),
+      saved_model_config));
+  export_dir_ = dst_dir;
+
+  SessionBundleConfig config = GetSessionBundleConfig();
+  config.set_session_target(test_util::kNewSessionHookSessionTargetPrefix);
+  std::unique_ptr<SavedModelBundle> bundle;
+  if (ExpectCreateBundleFailure()) {
+    EXPECT_FALSE(CreateBundleFromPath(GetParam().creation_type, config,
+                                      export_dir_, &bundle)
+                     .ok());
+    return;
+  }
+  std::unique_ptr<SavedModelBundleFactory> factory;
+  TF_ASSERT_OK(SavedModelBundleFactory::Create(config, &factory));
+  test_util::SetNewSessionHook([&](const SessionOptions& session_options) {
+    EXPECT_TRUE(session_options.config.graph_options()
+                    .rewrite_options()
+                    .disable_meta_optimizer());
+    return absl::OkStatus();
+  });
+
+  switch (GetParam().creation_type) {
+    case CreationType::kWithoutMetadata:
+      TF_ASSERT_OK(factory->CreateSavedModelBundle(export_dir_, &bundle));
+      break;
+    case CreationType::kWithMetadata:
+      TF_ASSERT_OK(factory->CreateSavedModelBundleWithMetadata(
+          CreateMetadata(), export_dir_, &bundle));
+      break;
+  }
+}
+
 }  // namespace
 }  // namespace serving
 }  // namespace tensorflow