Merge commit for internal changes

WORKSPACE

@@ -23,6 +23,6 @@ load("//tensorflow_serving:workspace.bzl", "tf_serving_workspace")
 tf_serving_workspace()
 
 # Specify the minimum required bazel version.
-load("@org_tensorflow//tensorflow:workspace.bzl", "check_version")
+load("@org_tensorflow//tensorflow:workspace.bzl", "check_bazel_version_at_least")
 
-check_version("0.5.4")
+check_bazel_version_at_least("0.5.4")

tensorflow_serving/g3doc/serving_inception.md

@@ -188,8 +188,9 @@ $ gcloud auth login --project tensorflow-serving
 
 ### Create a container cluster
 
-First we create a [Google Container Engine](https://cloud.google.com/container-engine/)
-cluster for service deployment.
+First we create a
+[Google Kubernetes Engine](https://cloud.google.com/container-engine/) cluster
+for service deployment.
 
 ```shell
 $ gcloud container clusters create inception-serving-cluster --num-nodes 5

tensorflow_serving/model_servers/main.cc

@@ -57,7 +57,9 @@ limitations under the License.
 #include "grpc++/support/status.h"
 #include "grpc++/support/status_code_enum.h"
 #include "grpc/grpc.h"
+#include "tensorflow/cc/saved_model/tag_constants.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/protobuf.h"
@@ -311,6 +313,7 @@ int main(int argc, char** argv) {
   bool flush_filesystem_caches = true;
   tensorflow::string model_base_path;
   const bool use_saved_model = true;
+  tensorflow::string saved_model_tags = tensorflow::kSavedModelTagServe;
   // Tensorflow session parallelism of zero means that both inter and intra op
   // thread pools will be auto configured.
   tensorflow::int64 tensorflow_session_parallelism = 0;
@@ -364,7 +367,10 @@ int main(int argc, char** argv) {
           "Fraction that each process occupies of the GPU memory space "
           "the value is between 0.0 and 1.0 (with 0.0 as the default) "
           "If 1.0, the server will allocate all the memory when the server "
-          "starts, If 0.0, Tensorflow will automatically select a value.")};
+          "starts, If 0.0, Tensorflow will automatically select a value."),
+      tensorflow::Flag("saved_model_tags", &saved_model_tags,
+                       "Comma-separated set of tags corresponding to the meta "
+                       "graph def to load from SavedModel.")};
 
   string usage = tensorflow::Flags::Usage(argv[0], flag_list);
   const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
@@ -416,6 +422,11 @@ int main(int argc, char** argv) {
         ->set_intra_op_parallelism_threads(tensorflow_session_parallelism);
     session_bundle_config.mutable_session_config()
         ->set_inter_op_parallelism_threads(tensorflow_session_parallelism);
+    const std::vector<string> tags =
+        tensorflow::str_util::Split(saved_model_tags, ",");
+    for (const string& tag : tags) {
+      *session_bundle_config.add_saved_model_tags() = tag;
+    }
     options.platform_config_map = CreateTensorFlowPlatformConfigMap(
         session_bundle_config, use_saved_model);
   } else {

tensorflow_serving/model_servers/server_core.cc

@@ -388,20 +388,25 @@ Status ServerCore::AddModelsViaCustomModelConfig() {
       config_.custom_model_config(), servable_event_bus_.get(), &manager_);
 }
 
-Status ServerCore::MaybeUpdateServerRequestLogger() {
+Status ServerCore::MaybeUpdateServerRequestLogger(
+    const ModelServerConfig::ConfigCase config_case) {
   if (options_.server_request_logger_updater) {
     return options_.server_request_logger_updater(
         config_, options_.server_request_logger.get());
   }
 
-  std::map<string, LoggingConfig> logging_config_map;
-  for (const auto& model_config : config_.model_config_list().config()) {
-    if (model_config.has_logging_config()) {
-      logging_config_map.insert(
-          {model_config.name(), model_config.logging_config()});
+  if (config_case == ModelServerConfig::kModelConfigList) {
+    std::map<string, LoggingConfig> logging_config_map;
+    for (const auto& model_config : config_.model_config_list().config()) {
+      if (model_config.has_logging_config()) {
+        logging_config_map.insert(
+            {model_config.name(), model_config.logging_config()});
+      }
     }
+    return options_.server_request_logger->Update(logging_config_map);
   }
-  return options_.server_request_logger->Update(logging_config_map);
+
+  return Status::OK();
 }
 
 Status ServerCore::ReloadConfig(const ModelServerConfig& new_config) {
@@ -457,7 +462,7 @@ Status ServerCore::ReloadConfig(const ModelServerConfig& new_config) {
     default:
       return errors::InvalidArgument("Invalid ServerModelConfig");
   }
-  TF_RETURN_IF_ERROR(MaybeUpdateServerRequestLogger());
+  TF_RETURN_IF_ERROR(MaybeUpdateServerRequestLogger(config_.config_case()));
 
   if (options_.flush_filesystem_caches) {
     return Env::Default()->FlushFileSystemCaches();

tensorflow_serving/model_servers/server_core.h

@@ -312,7 +312,9 @@ class ServerCore : public Manager {
   Status AddModelsViaCustomModelConfig() EXCLUSIVE_LOCKS_REQUIRED(config_mu_);
 
   // Updates the ServerRequestLogger based on the ModelConfigList.
-  Status MaybeUpdateServerRequestLogger() EXCLUSIVE_LOCKS_REQUIRED(config_mu_);
+  Status MaybeUpdateServerRequestLogger(
+      ModelServerConfig::ConfigCase config_case)
+      EXCLUSIVE_LOCKS_REQUIRED(config_mu_);
 
   // ************************************************************************
   // Request Processing.

tensorflow_serving/servables/tensorflow/BUILD

@@ -110,6 +110,7 @@ cc_library(
         ":session_bundle_config_proto",
         "//tensorflow_serving/resources:resources_proto",
         "//tensorflow_serving/test_util",
+        "@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
         "@org_tensorflow//tensorflow/core:core_cpu",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:test",
@@ -197,6 +198,7 @@ cc_test(
         ":session_bundle_config_proto",
         "//tensorflow_serving/core/test_util:test_main",
         "@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:lib",
         "@org_tensorflow//tensorflow/core:test",

tensorflow_serving/servables/tensorflow/saved_model_bundle_factory.cc

@@ -78,9 +78,16 @@ Status SavedModelBundleFactory::EstimateResourceRequirement(
 Status SavedModelBundleFactory::CreateSavedModelBundle(
     const string& path, std::unique_ptr<SavedModelBundle>* bundle) {
   bundle->reset(new SavedModelBundle);
+  std::unordered_set<string> saved_model_tags(
+      config_.saved_model_tags().begin(), config_.saved_model_tags().end());
+  // Defaults to loading the meta graph def corresponding to the `serve` tag if
+  // no `saved_model_tags` are specified.
+  if (saved_model_tags.empty()) {
+    saved_model_tags.insert(kSavedModelTagServe);
+  }
   TF_RETURN_IF_ERROR(LoadSessionBundleOrSavedModelBundle(
       GetSessionOptions(config_), GetRunOptions(config_), path,
-      {kSavedModelTagServe}, bundle->get()));
+      saved_model_tags, bundle->get()));
   if (!config_.experimental_fixed_input_tensors().empty()) {
     LOG(INFO) << "Wrapping session to inject fixed input tensors";
     std::vector<std::pair<string, Tensor>> fixed_input_tensors;

tensorflow_serving/servables/tensorflow/saved_model_bundle_factory_test.cc

@@ -24,6 +24,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "tensorflow/cc/saved_model/loader.h"
+#include "tensorflow/cc/saved_model/tag_constants.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
@@ -74,6 +75,7 @@ TEST_F(SavedModelBundleFactoryTest, FixedInputTensors) {
   fixed_input.AsProtoField(fixed_input_proto.mutable_tensor());
 
   SessionBundleConfig config;
+  *config.add_saved_model_tags() = kSavedModelTagServe;
   *config.add_experimental_fixed_input_tensors() = fixed_input_proto;
   std::unique_ptr<Session> session;
   TF_ASSERT_OK(CreateSession(config, &session));

tensorflow_serving/servables/tensorflow/session_bundle_config.proto

@@ -51,6 +51,10 @@ message SessionBundleConfig {
   // Remove it once resource estimates are moved inside SavedModel.
   uint64 experimental_transient_ram_bytes_during_load = 5;
 
+  // Set of SavedModel tags identifying the specific meta graph def to be
+  // loaded.
+  repeated string saved_model_tags = 6;
+
   // EXPERIMENTAL. THIS FIELD MAY CHANGE OR GO AWAY. USE WITH CAUTION.
   //
   // Input tensors to append to every Session::Run() call.