Handle partitioning of incomplete graphs gracefully.

The added test would lead to a SIGSEGV without the accompanying
changes in graph_partition.cc.

In practice, this code path can be triggered by invoking:
  Session::Extend(graph_def);
followed by:
  Session::Run()
when graph_def incompletely specifies a graph
(such as the one in the test, where not all inputs of the
target node have been specified).
Change: 138573807

2 files changed, 35 insertions, 0 deletions

diff --git a/tensorflow/core/graph/graph_partition.cc b/tensorflow/core/graph/graph_partition.cc
index 488fe47f38..8e9eceb699 100644
--- a/tensorflow/core/graph/graph_partition.cc
+++ b/tensorflow/core/graph/graph_partition.cc
@@ -857,6 +857,7 @@ Status Partition(const PartitionOptions& opts, Graph* g,
     ref_control_inputs.clear();
     const Edge* control_flow_edge = nullptr;
     int32 num_control_flow_edges = 0;
+    int32 num_input_edges = 0;
     for (const Edge* edge : dst->in_edges()) {
       if (edge->IsControlEdge()) {
         if (IsMerge(edge->src()) && IsControlLoop(edge->src())) {
@@ -871,9 +872,16 @@ Status Partition(const PartitionOptions& opts, Graph* g,
       } else {
         DCHECK(inputs[edge->dst_input()] == nullptr);
         inputs[edge->dst_input()] = edge;
+        ++num_input_edges;
       }
     }
 
+    if (num_input_edges != dst->num_inputs()) {
+      return errors::InvalidArgument("Incomplete graph, missing ",
+                                     (dst->num_inputs() - num_input_edges),
+                                     " inputs for ", dst->name());
+    }
+
     // Process in order so that all data edges are added as inputs to
     // dst in Edge::dst_input() order.
     for (const Edge* edge : inputs) {
diff --git a/tensorflow/core/graph/graph_partition_test.cc b/tensorflow/core/graph/graph_partition_test.cc
index 34342ff4f3..24bc8dd575 100644
--- a/tensorflow/core/graph/graph_partition_test.cc
+++ b/tensorflow/core/graph/graph_partition_test.cc
@@ -372,5 +372,32 @@ TEST_F(GraphPartitionTest, CrossDeviceLoop1) {
   }
 }
 
+TEST_F(GraphPartitionTest, PartitionIncompleteGraph) {
+  NodeDef ndef;
+  Graph g(OpRegistry::Global());
+  // Invalid graph since the Combine node requires an input.
+  bool parsed = protobuf::TextFormat::ParseFromString(
+      R"EOF(
+      name: "N"
+      op: "Combine"
+      )EOF",
+      &ndef);
+  ASSERT_TRUE(parsed);
+  Status status;
+  g.AddNode(ndef, &status);
+  TF_ASSERT_OK(status);
+
+  PartitionOptions popts;
+  popts.node_to_loc = SplitByDevice;
+  popts.new_name = [&g](const string& prefix) { return g.NewName(prefix); };
+  popts.get_incarnation = [](const string&) { return 1; };
+
+  std::unordered_map<string, GraphDef> partitions;
+  status = Partition(popts, &g, &partitions);
+  // Partitioning should fail, but not crash like it did before the
+  // changes that accompanied the addition of this test.
+  EXPECT_EQ(error::INVALID_ARGUMENT, status.code()) << status;
+}
+
 }  // namespace
 }  // namespace tensorflow