Make graph construction work while graph is being concurrently run.

The overall approach is to use Graph._lock to synchronize Session.run calls and construction methods that rely on graph mutation. We don't want to synchronize the actual running of the graph, only the Extend call, so this change exposes an ExtendSession method to the Python API and disables extending automatically in TF_SessionRun.

PiperOrigin-RevId: 188106818

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h
index 91667056e0..027e2d2b15 100644
--- a/tensorflow/c/c_api_internal.h
+++ b/tensorflow/c/c_api_internal.h
@@ -133,6 +133,12 @@ struct TF_Session {
   // buffers of a TF_Tensor pinned in device memory.
   const tensorflow::DeviceMgr* device_mgr;   // Owned by session.
   std::vector<tensorflow::Device*> devices;  // Owned by device_mgr.
+
+  // If true, TF_SessionRun and similar methods will call
+  // ExtendSessionGraphHelper before running the graph (this is the default
+  // public behavior). Can be set to false if the caller needs to call
+  // ExtendSessionGraphHelper manually.
+  bool extend_before_run GUARDED_BY(mu);
 };
 
 struct TF_ImportGraphDefOptions {
@@ -212,6 +218,8 @@ void TF_GraphSetOutputHandleShapesAndTypes(TF_Graph* graph, TF_Output output,
 void RecordMutation(TF_Graph* graph, const TF_Operation& op,
                     const char* mutation_type);
 
+bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status);
+
 }  // end namespace tensorflow
 
 #endif  // TENSORFLOW_C_C_API_INTERNAL_H_