Make FunctionLibraryDefinition thread-safe.

The eager runtime mutates the FunctionLibraryRuntime's FunctionLibraryDefinition, which is shared across threads; at the same time, OpKernels might read the FunctionLibraryDefinition. This is not thread-safe unless FunctionLibraryDefinition is thread-safe.

This change makes FunctionLibraryDefinition, which is basically a map from function names to FunctionDefs, thread-safe. This is almost entirely accomplished by guarding the map with a mutex. There is however one complication: Find and RemoveFunction cannot be made thread-safe in a straightforward way (Find returns a raw pointer to a FunctionDef while Remove can delete the corresponding FunctionDef). In light of the fact that clients only ever call RemoveFunction when they in fact want to replace an existing function with a new one, we make the following modifications to FunctionLibraryDefinition's API:

1. A Contains method is added to check for the existence of a function.
2. A ReplaceFunction method is added.
3. RemoveFunction and RemoveGradient are made private.

We also update clients of the FunctionLibraryDefinition to use Contains & ReplaceFunction instead of Find and RemoveFunction.

PiperOrigin-RevId: 208271076

1 files changed, 2 insertions, 3 deletions

diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc
index 9c9eacb5b5..c23b7d3699 100644
--- a/tensorflow/core/common_runtime/graph_execution_state.cc
+++ b/tensorflow/core/common_runtime/graph_execution_state.cc
@@ -643,10 +643,9 @@ Status GraphExecutionState::OptimizeGraph(
     for (const FunctionDef& fdef : new_graph.library().function()) {
       const string& func_name = fdef.signature().name();
 
-      if ((*optimized_flib)->Find(func_name)) {
+      if ((*optimized_flib)->Contains(func_name)) {
         VLOG(3) << "Replace function: name=" << func_name;
-        TF_RETURN_IF_ERROR((*optimized_flib)->RemoveFunction(func_name));
-        TF_RETURN_IF_ERROR((*optimized_flib)->AddFunctionDef(fdef));
+        TF_RETURN_IF_ERROR((*optimized_flib)->ReplaceFunction(func_name, fdef));
       } else {
         VLOG(3) << "Add new function: name=" << func_name;
         TF_RETURN_IF_ERROR((*optimized_flib)->AddFunctionDef(fdef));