Let RunWithSortedInputsOutputs take graph_name.

SiqiaoWu1993 · tensorflower-gardener · commit 3bc0da7933cd · 2025-10-06T15:16:19.000-07:00
PiperOrigin-RevId: 815900754
diff --git a/tensorflow/core/tfrt/graph_executor/graph_executor.cc b/tensorflow/core/tfrt/graph_executor/graph_executor.cc
@@ -560,7 +560,7 @@ void CreateSortedNamesAndOriginalIndices(absl::Span<const std::string> names,
 }
 
 absl::Status GraphExecutor::RunWithSortedInputsOutputs(
-    const RunOptions& run_options,
+    const RunOptions& run_options, absl::string_view graph_name,
     absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs,
     absl::Span<const std::string> sorted_input_names,
     absl::Span<const tensorflow::DataType> sorted_input_dtypes,
@@ -570,12 +570,11 @@ absl::Status GraphExecutor::RunWithSortedInputsOutputs(
     absl::Span<const int> output_original_indices,
     std::vector<tensorflow::Tensor>* outputs) {
   // Load the client graph.
-  TF_ASSIGN_OR_RETURN(
-      LoadedClientGraph & loaded_client_graph,
-      GetOrCreateLoadedClientGraph(
-          run_options, sorted_input_names, sorted_input_dtypes,
-          sorted_output_names, sorted_target_node_names, run_options.work_queue,
-          /*graph_name=*/{}, inputs));
+  TF_ASSIGN_OR_RETURN(LoadedClientGraph & loaded_client_graph,
+                      GetOrCreateLoadedClientGraph(
+                          run_options, sorted_input_names, sorted_input_dtypes,
+                          sorted_output_names, sorted_target_node_names,
+                          run_options.work_queue, graph_name, inputs));
 
   // Get a shared_ptr of the executable so that during the current request the
   // executable to use is guaranteed to be alive.
@@ -683,9 +682,9 @@ absl::Status GraphExecutor::Run(
   std::sort(sorted_target_node_names.begin(), sorted_target_node_names.end());
 
   return RunWithSortedInputsOutputs(
-      run_options, inputs, sorted_input_names, sorted_input_dtypes,
-      sorted_output_names, sorted_target_node_names, input_original_indices,
-      output_original_indices, outputs);
+      run_options, /*graph_name=*/"", inputs, sorted_input_names,
+      sorted_input_dtypes, sorted_output_names, sorted_target_node_names,
+      input_original_indices, output_original_indices, outputs);
 }
 
 absl::Status GraphExecutor::Extend(const GraphDef& graph) {
diff --git a/tensorflow/core/tfrt/graph_executor/graph_executor.h b/tensorflow/core/tfrt/graph_executor/graph_executor.h
@@ -274,14 +274,17 @@ class GraphExecutor {
       std::vector<tensorflow::Tensor>* outputs);
 
   // Similar as `Run`, but it requires additional input parameters to specify
-  // the sorted input/output names and the original indices of the
-  // inputs/outputs. The caller must guarantee that inputs are in the same order
-  // as of `sorted_input_names`. The sorted input/output names are needed to
-  // consistently build the key for looking up the `LoadedClientGraph` in the
+  // the `graph_name`, the sorted input/output names and the original indices of
+  // the inputs/outputs. The caller must guarantee that inputs are in the same
+  // order as of `sorted_input_names`. The sorted input/output names are needed
+  // to consistently build the key for looking up the `LoadedClientGraph` in the
   // cache. The original indices are needed to map the results to the original
-  // inputs/outputs.
+  // inputs/outputs. The `graph_name` will be used to lookup the compiled graph
+  // in the cache. It is usually the signature name of the graph. If it is
+  // empty, a joined name will be constructed from the sorted input/output names
+  // to lookup the `LoadedClientGraph` in the cache.
   absl::Status RunWithSortedInputsOutputs(
-      const RunOptions& run_options,
+      const RunOptions& run_options, absl::string_view graph_name,
       absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs,
       absl::Span<const std::string> sorted_input_names,
       absl::Span<const tensorflow::DataType> sorted_input_dtypes,
