Ensure that all RunManyGraphs calls complete before returning from cancellation.

This fixes a potential race condition and segfault, where the response
buffer could be deallocated with an RPC in flight. This would happen
if a step was cancelled before its master blocked on the RunGraph RPCs.
Change: 124255996

1 files changed, 6 insertions, 2 deletions

diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
index 9231fae362..870970b7ca 100644
--- a/tensorflow/core/distributed_runtime/master_session.cc
+++ b/tensorflow/core/distributed_runtime/master_session.cc
@@ -569,11 +569,15 @@ Status MasterSession::ReffedClientGraph::RunPartitions(
   bool success =
       cm->RegisterCallback(token, [&calls]() { calls.StartCancel(); });
   if (!success) {
-    return errors::Cancelled("Step was cancelled");
+    calls.StartCancel();
   }
   calls.Wait();
-  cm->DeregisterCallback(token);
   call_opts->ClearCancelCallback();
+  if (success) {
+    cm->DeregisterCallback(token);
+  } else {
+    return errors::Cancelled("Step was cancelled");
+  }
 
   // Collects fetches.
   Status status = calls.status();