Fixes a bug in the locking discipline in `DirectSession`.

The `execution_state_` member should be protected by
`graph_def_lock_`, because it can be replaced at the same times that
the `graph_def_` is mutated.
Change: 123701264

2 files changed, 13 insertions, 13 deletions

diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
index 2882158315..1f864b018b 100644
--- a/tensorflow/core/common_runtime/direct_session.cc
+++ b/tensorflow/core/common_runtime/direct_session.cc
@@ -208,7 +208,7 @@ DirectSession::~DirectSession() {
 }
 
 void DirectSession::MaybeInitializeExecutionState(const GraphDef& graph) {
-  // If already initialied, do nothing.
+  // If already initialized, do nothing.
   if (flib_def_ && execution_state_) {
     return;
   }
@@ -839,6 +839,7 @@ Status DirectSession::GetOrCreateExecutors(
 Status DirectSession::CreateGraphs(const BuildGraphOptions& options,
                                    std::unordered_map<string, Graph*>* outputs,
                                    RunStateArgs* run_state_args) {
+  mutex_lock l(graph_def_lock_);
   std::unique_ptr<SimpleClientGraph> client_graph;
   SimpleClientGraph* cgraph = nullptr;
 
@@ -947,16 +948,13 @@ Status DirectSession::CreateGraphs(const BuildGraphOptions& options,
     Device* d;
     s = device_mgr_->LookupDevice(partition_name, &d);
     if (!s.ok()) break;
-    {
-      mutex_lock l(graph_def_lock_);
-      // TODO(pbar) The library is currently shared and immutable. There
-      // may be possible use cases where a device may want to modify
-      // function definitions - in which case the library would need to be
-      // replicated per device.
-      s = d->MaybeRewriteGraph(flib_def_->ToProto(), graph_def);
-      if (!s.ok()) {
-        break;
-      }
+    // TODO(pbar) The library is currently shared and immutable. There
+    // may be possible use cases where a device may want to modify
+    // function definitions - in which case the library would need to be
+    // replicated per device.
+    s = d->MaybeRewriteGraph(flib_def_->ToProto(), graph_def);
+    if (!s.ok()) {
+      break;
     }
     Graph* device_graph = new Graph(flib_def_.get());
     GraphConstructorOptions device_opts;
diff --git a/tensorflow/core/common_runtime/direct_session.h b/tensorflow/core/common_runtime/direct_session.h
index 0c7f4103e4..13adce2ab9 100644
--- a/tensorflow/core/common_runtime/direct_session.h
+++ b/tensorflow/core/common_runtime/direct_session.h
@@ -162,7 +162,8 @@ class DirectSession : public Session {
 
   // Initializes the base execution state given the 'graph',
   // if not already initialized.
-  void MaybeInitializeExecutionState(const GraphDef& graph);
+  void MaybeInitializeExecutionState(const GraphDef& graph)
+      EXCLUSIVE_LOCKS_REQUIRED(graph_def_lock_);
 
   // Retrieves an already existing set of executors to run 'inputs' and
   // 'outputs', or creates and caches them for future use.
@@ -249,7 +250,8 @@ class DirectSession : public Session {
   std::unordered_map<string, string> stateful_placements_ GUARDED_BY(mu_);
 
   // Execution_state; used when placing the entire graph.
-  std::unique_ptr<SimpleGraphExecutionState> execution_state_;
+  std::unique_ptr<SimpleGraphExecutionState> execution_state_
+      GUARDED_BY(graph_def_lock_);
   std::unique_ptr<FunctionLibraryDefinition> flib_def_;
 
   // For generating unique names.