Move over recent changes to tensorflow_serving/session_bundle to tensorflow/contrib/session_bundle, including:

- Type and error checking when unpacking Any messages
- Increased session_bundle test coverage
Change: 127152552

4 files changed, 266 insertions, 16 deletions

diff --git a/tensorflow/contrib/session_bundle/session_bundle.cc b/tensorflow/contrib/session_bundle/session_bundle.cc
index 5029dcd7fe..6895c094be 100644
--- a/tensorflow/contrib/session_bundle/session_bundle.cc
+++ b/tensorflow/contrib/session_bundle/session_bundle.cc
@@ -91,7 +91,7 @@ string GetVariablesFilename(const StringPiece export_dir) {
   const char kVariablesFilename[] = "export";
   const char kVariablesFilenamePattern[] = "export-\?\?\?\?\?-of-\?\?\?\?\?";
   if (Env::Default()->FileExists(
-      tensorflow::io::JoinPath(export_dir, kVariablesFilename))) {
+          tensorflow::io::JoinPath(export_dir, kVariablesFilename))) {
     return tensorflow::io::JoinPath(export_dir, kVariablesFilename);
   } else {
     return tensorflow::io::JoinPath(export_dir, kVariablesFilenamePattern);
@@ -104,8 +104,8 @@ Status RunRestoreOp(const StringPiece export_dir,
                     const StringPiece variables_filename_const_op_name,
                     tensorflow::Session* session) {
   LOG(INFO) << "Running restore op for SessionBundle";
-  Tensor variables_tensor = CreateStringTensor(
-      GetVariablesFilename(export_dir));
+  Tensor variables_tensor =
+      CreateStringTensor(GetVariablesFilename(export_dir));
   std::vector<std::pair<string, Tensor>> inputs = {
       {variables_filename_const_op_name.ToString(), variables_tensor}};
   AddAssetsTensorsToInputs(export_dir, asset_files, &inputs);
@@ -137,11 +137,21 @@ tensorflow::Status LoadSessionBundleFromPath(
     // Use serving graph_def in MetaGraphDef collection_def.
     if (graph_collection_def.any_list().value_size() != 1) {
       return errors::FailedPrecondition(
-          strings::StrCat("Expected exactly one serving GraphDef in : ",
-                          bundle->meta_graph_def.DebugString()));
+          "Expected exactly one serving GraphDef in : ",
+          bundle->meta_graph_def.DebugString());
+    }
+    const auto& any = graph_collection_def.any_list().value(0);
+    if (!any.Is<GraphDef>()) {
+      return errors::FailedPrecondition(
+          "Expected Any type_url for: ",
+          tensorflow::GraphDef::default_instance().descriptor()->full_name(),
+          ". Got: ", string(any.type_url().data(), any.type_url().size()), ".");
     }
     tensorflow::GraphDef graph_def;
-    graph_collection_def.any_list().value(0).UnpackTo(&graph_def);
+    if (!any.UnpackTo(&graph_def)) {
+      return errors::FailedPrecondition("Failed to unpack: ",
+                                        any.DebugString());
+    }
     TF_RETURN_IF_ERROR(
         CreateSessionFromGraphDef(options, graph_def, &bundle->session));
   } else {
@@ -157,7 +167,17 @@ tensorflow::Status LoadSessionBundleFromPath(
     const auto& any_assets = assets_it->second.any_list().value();
     for (const auto& any_asset : any_assets) {
       AssetFile asset_file;
-      any_asset.UnpackTo(&asset_file);
+      if (!any_asset.Is<AssetFile>()) {
+        return errors::FailedPrecondition(
+            "Expected asset Any type_url for: ",
+            asset_file.descriptor()->full_name(), ". Got: ",
+            string(any_asset.type_url().data(), any_asset.type_url().size()),
+            ".");
+      }
+      if (!any_asset.UnpackTo(&asset_file)) {
+        return errors::FailedPrecondition("Failed to unpack: ",
+                                          any_asset.DebugString());
+      }
       asset_files.push_back(asset_file);
     }
   }
diff --git a/tensorflow/contrib/session_bundle/session_bundle_test.cc b/tensorflow/contrib/session_bundle/session_bundle_test.cc
index 5aeb344524..5ecbb32f88 100644
--- a/tensorflow/contrib/session_bundle/session_bundle_test.cc
+++ b/tensorflow/contrib/session_bundle/session_bundle_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "google/protobuf/any.pb.h"
+#include "tensorflow/contrib/session_bundle/signature.h"
 #include "tensorflow/contrib/session_bundle/test_util.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -36,9 +37,48 @@ namespace tensorflow {
 namespace serving {
 namespace {
 
-TEST(LoadSessionBundleFromPath, Basic) {
-  const string export_path = test_util::TestSrcDirPath(
+// Constants for the export file names.
+const char kVariablesFilename[] = "export-00000-of-00001";
+const char kMetaGraphDefFilename[] = "export.meta";
+
+// Function used to rewrite a MetaGraphDef.
+using MetaGraphDefTwiddler = std::function<void(MetaGraphDef*)>;
+
+// Copy the base half_plus_two to `export_path`.
+// Outputs the files using the passed names (typically the constants above).
+// The Twiddler can be used to update the MetaGraphDef before output.
+Status CopyExport(const string& export_path, const string& variables_filename,
+                  const string& meta_graph_def_filename,
+                  const MetaGraphDefTwiddler& twiddler) {
+  TF_RETURN_IF_ERROR(Env::Default()->CreateDir(export_path));
+  const string orig_path = test_util::TestSrcDirPath(
       "session_bundle/example/half_plus_two/00000123");
+  {
+    const string source =
+        tensorflow::io::JoinPath(orig_path, kVariablesFilename);
+    const string sink =
+        tensorflow::io::JoinPath(export_path, variables_filename);
+
+    string data;
+    TF_RETURN_IF_ERROR(ReadFileToString(Env::Default(), source, &data));
+    TF_RETURN_IF_ERROR(WriteStringToFile(Env::Default(), sink, data));
+  }
+  {
+    const string source =
+        tensorflow::io::JoinPath(orig_path, kMetaGraphDefFilename);
+    const string sink =
+        tensorflow::io::JoinPath(export_path, meta_graph_def_filename);
+
+    tensorflow::MetaGraphDef graph_def;
+    TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), source, &graph_def));
+    twiddler(&graph_def);
+    TF_RETURN_IF_ERROR(
+        WriteStringToFile(Env::Default(), sink, graph_def.SerializeAsString()));
+  }
+  return Status::OK();
+}
+
+void BasicTest(const string& export_path) {
   tensorflow::SessionOptions options;
   SessionBundle bundle;
   TF_ASSERT_OK(LoadSessionBundleFromPath(options, export_path, &bundle));
@@ -67,10 +107,8 @@ TEST(LoadSessionBundleFromPath, Basic) {
   Tensor input = test::AsTensor<float>({0, 1, 2, 3}, TensorShape({4, 1}));
 
   // Recover the Tensor names of our inputs and outputs.
-  auto collection_def = bundle.meta_graph_def.collection_def();
   Signatures signatures;
-  ASSERT_EQ(1, collection_def[kSignaturesKey].any_list().value_size());
-  collection_def[kSignaturesKey].any_list().value(0).UnpackTo(&signatures);
+  TF_ASSERT_OK(GetSignatures(bundle.meta_graph_def, &signatures));
   ASSERT_TRUE(signatures.default_signature().has_regression_signature());
   const tensorflow::serving::RegressionSignature regression_signature =
       signatures.default_signature().regression_signature();
@@ -86,6 +124,12 @@ TEST(LoadSessionBundleFromPath, Basic) {
       outputs[0], test::AsTensor<float>({2, 2.5, 3, 3.5}, TensorShape({4, 1})));
 }
 
+TEST(LoadSessionBundleFromPath, BasicTensorflowContrib) {
+  const string export_path = test_util::TestSrcDirPath(
+      "session_bundle/example/half_plus_two/00000123");
+  BasicTest(export_path);
+}
+
 TEST(LoadSessionBundleFromPath, BadExportPath) {
   const string export_path = test_util::TestSrcDirPath("/tmp/bigfoot");
   tensorflow::SessionOptions options;
@@ -97,6 +141,142 @@ TEST(LoadSessionBundleFromPath, BadExportPath) {
   EXPECT_TRUE(msg.find("Not found") != std::string::npos) << msg;
 }
 
+class SessionBundleTest : public ::testing::Test {
+ protected:
+  // Copy the half_plus_two graph and apply the twiddler to rewrite the
+  // MetaGraphDef.
+  // Returns the path of the export.
+  // ** Should only be called once per test **
+  string SetupExport(MetaGraphDefTwiddler twiddler) {
+    return SetupExport(twiddler, kVariablesFilename, kMetaGraphDefFilename);
+  }
+  // SetupExport that allows for the variables and meta_graph_def filenames
+  // to be overridden.
+  string SetupExport(MetaGraphDefTwiddler twiddler,
+                     const string& variables_filename,
+                     const string& meta_graph_def_filename) {
+    // Construct a unique path name based on the test name.
+    const ::testing::TestInfo* const test_info =
+        ::testing::UnitTest::GetInstance()->current_test_info();
+    const string export_path = tensorflow::io::JoinPath(
+        testing::TmpDir(),
+        strings::StrCat(test_info->test_case_name(), test_info->name()));
+    TF_CHECK_OK(CopyExport(export_path, variables_filename,
+                           meta_graph_def_filename, twiddler));
+    return export_path;
+  }
+
+  tensorflow::SessionOptions options_;
+  SessionBundle bundle_;
+  Status status_;
+};
+
+TEST_F(SessionBundleTest, Basic) {
+  const string export_path = SetupExport([](MetaGraphDef*) {});
+  BasicTest(export_path);
+}
+
+TEST_F(SessionBundleTest, UnshardedVariableFile) {
+  // Test that we can properly read the variables when exported
+  // without sharding.
+  const string export_path =
+      SetupExport([](MetaGraphDef*) {}, "export", kMetaGraphDefFilename);
+  BasicTest(export_path);
+}
+
+TEST_F(SessionBundleTest, ServingGraph_Empty) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    (*def->mutable_collection_def())[kGraphKey].clear_any_list();
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(StringPiece(status_.error_message())
+                  .contains("Expected exactly one serving GraphDef"))
+      << status_.error_message();
+}
+
+TEST_F(SessionBundleTest, ServingGraphAny_IncorrectType) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    // Pack an unexpected type in the GraphDef Any.
+    (*def->mutable_collection_def())[kGraphKey].clear_any_list();
+    auto* any = (*def->mutable_collection_def())[kGraphKey]
+                    .mutable_any_list()
+                    ->add_value();
+    any->PackFrom(AssetFile());
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(StringPiece(status_.error_message())
+                  .contains("Expected Any type_url for: tensorflow.GraphDef"))
+      << status_.error_message();
+}
+
+TEST_F(SessionBundleTest, ServingGraphAnyValue_Corrupted) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    // Pack an unexpected type in the GraphDef Any.
+    (*def->mutable_collection_def())[kGraphKey].clear_any_list();
+    auto* any = (*def->mutable_collection_def())[kGraphKey]
+                    .mutable_any_list()
+                    ->add_value();
+    any->PackFrom(GraphDef());
+    any->set_value("junk junk");
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(StringPiece(status_.error_message()).contains("Failed to unpack"))
+      << status_.error_message();
+}
+
+TEST_F(SessionBundleTest, AssetFileAny_IncorrectType) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    // Pack an unexpected type in the AssetFile Any.
+    (*def->mutable_collection_def())[kAssetsKey].clear_any_list();
+    auto* any = (*def->mutable_collection_def())[kAssetsKey]
+                    .mutable_any_list()
+                    ->add_value();
+    any->PackFrom(GraphDef());
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(
+      StringPiece(status_.error_message())
+          .contains(
+              "Expected asset Any type_url for: tensorflow.serving.AssetFile"))
+      << status_.error_message();
+}
+
+TEST_F(SessionBundleTest, AssetFileAny_ValueCorrupted) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    // Pack an unexpected type in the AssetFile Any.
+    (*def->mutable_collection_def())[kAssetsKey].clear_any_list();
+    auto* any = (*def->mutable_collection_def())[kAssetsKey]
+                    .mutable_any_list()
+                    ->add_value();
+    any->PackFrom(AssetFile());
+    any->set_value("junk junk");
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(StringPiece(status_.error_message()).contains("Failed to unpack"))
+      << status_.error_message();
+}
+
+TEST_F(SessionBundleTest, InitOp_TooManyValues) {
+  const string path = SetupExport([](MetaGraphDef* def) {
+    // Pack multiple init ops in to the collection.
+    (*def->mutable_collection_def())[kInitOpKey].clear_node_list();
+    auto* node_list =
+        (*def->mutable_collection_def())[kInitOpKey].mutable_node_list();
+    node_list->add_value("foo");
+    node_list->add_value("bar");
+  });
+  status_ = LoadSessionBundleFromPath(options_, path, &bundle_);
+  EXPECT_FALSE(status_.ok());
+  EXPECT_TRUE(StringPiece(status_.error_message())
+                  .contains("Expected exactly one serving init op"))
+      << status_.error_message();
+}
+
 }  // namespace
 }  // namespace serving
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/session_bundle/signature.cc b/tensorflow/contrib/session_bundle/signature.cc
index 50fea99ef2..04fb479523 100644
--- a/tensorflow/contrib/session_bundle/signature.cc
+++ b/tensorflow/contrib/session_bundle/signature.cc
@@ -48,14 +48,23 @@ Status BatchSizesMatch(const Tensor& input, const Tensor& output) {
 
 Status GetSignatures(const tensorflow::MetaGraphDef& meta_graph_def,
                      Signatures* signatures) {
-  auto collection_def = meta_graph_def.collection_def();
-  auto any_list = collection_def[kSignaturesKey].any_list();
-  if (any_list.value_size() != 1) {
+  const auto& collection_def = meta_graph_def.collection_def();
+  const auto it = collection_def.find(kSignaturesKey);
+  if (it == collection_def.end() || it->second.any_list().value_size() != 1) {
     return errors::FailedPrecondition(
         strings::StrCat("Expected exactly one signatures proto in : ",
                         meta_graph_def.DebugString()));
   }
-  any_list.value(0).UnpackTo(signatures);
+  const auto& any = it->second.any_list().value(0);
+  if (!any.Is<Signatures>()) {
+    return errors::FailedPrecondition(
+        "Expected signature Any type_url for: ",
+        signatures->descriptor()->full_name(), ". Got: ",
+        string(any.type_url().data(), any.type_url().size()), ".");
+  }
+  if (!any.UnpackTo(signatures)) {
+    return errors::FailedPrecondition("Failed to unpack: ", any.DebugString());
+  }
   return Status::OK();
 }
 
diff --git a/tensorflow/contrib/session_bundle/signature_test.cc b/tensorflow/contrib/session_bundle/signature_test.cc
index c890551217..10908cc649 100644
--- a/tensorflow/contrib/session_bundle/signature_test.cc
+++ b/tensorflow/contrib/session_bundle/signature_test.cc
@@ -478,6 +478,47 @@ TEST(SetAndGetSignatures, RoundTrip) {
                         .tensor_name());
 }
 
+TEST(GetSignatures, MissingSignature) {
+  tensorflow::MetaGraphDef meta_graph_def;
+  Signatures read_signatures;
+  const auto status = GetSignatures(meta_graph_def, &read_signatures);
+  EXPECT_FALSE(status.ok());
+  EXPECT_TRUE(
+      StringPiece(status.error_message()).contains("Expected exactly one"))
+      << status.error_message();
+}
+
+TEST(GetSignatures, WrongProtoInAny) {
+  tensorflow::MetaGraphDef meta_graph_def;
+  auto& collection_def = *(meta_graph_def.mutable_collection_def());
+  auto* any =
+      collection_def[kSignaturesKey].mutable_any_list()->mutable_value()->Add();
+  // Put an unexpected type into the Signatures Any.
+  any->PackFrom(TensorBinding());
+  Signatures read_signatures;
+  const auto status = GetSignatures(meta_graph_def, &read_signatures);
+  EXPECT_FALSE(status.ok());
+  EXPECT_TRUE(StringPiece(status.error_message())
+                  .contains("Expected signature Any type_url for: "
+                            "tensorflow.serving.Signatures"))
+      << status.error_message();
+}
+
+TEST(GetSignatures, JunkInAny) {
+  tensorflow::MetaGraphDef meta_graph_def;
+  auto& collection_def = *(meta_graph_def.mutable_collection_def());
+  auto* any =
+      collection_def[kSignaturesKey].mutable_any_list()->mutable_value()->Add();
+  // Create a valid Any then corrupt it.
+  any->PackFrom(Signatures());
+  any->set_value("junk junk");
+  Signatures read_signatures;
+  const auto status = GetSignatures(meta_graph_def, &read_signatures);
+  EXPECT_FALSE(status.ok());
+  EXPECT_TRUE(StringPiece(status.error_message()).contains("Failed to unpack"))
+      << status.error_message();
+}
+
 // GenericSignature test fixture that contains a signature initialized with two
 // bound Tensors.
 class GenericSignatureTest : public ::testing::Test {