Add more tracer

PiperOrigin-RevId: 716347614
tensorflow · Jan 16, 2025 · 51f1cf6 · 51f1cf6
1 parent ebb71bf
commit 51f1cf6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/tensorflow_serving/servables/tensorflow/BUILD b/tensorflow_serving/servables/tensorflow/BUILD
@@ -1181,6 +1181,7 @@ cc_library(
         "@org_tensorflow//tensorflow/cc/saved_model:util",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
+        "@org_tensorflow//tensorflow/core/platform:tracing",
         "@org_tensorflow//tensorflow/core/tfrt/runtime:tf_threadpool_concurrent_work_queue",
         "@org_tensorflow//tensorflow/core/tfrt/saved_model:saved_model_cpu",
     ],

diff --git a/tensorflow_serving/servables/tensorflow/tfrt_predict_util.cc b/tensorflow_serving/servables/tensorflow/tfrt_predict_util.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/tracing.h"  // NOLINT
 #include "tensorflow/core/protobuf/error_codes.pb.h"
 #include "tensorflow/core/protobuf/named_tensor.pb.h"
 #include "tensorflow/core/tfrt/runtime/tf_threadpool_concurrent_work_queue.h"
@@ -190,12 +191,14 @@ absl::Status RunPredict(
   }
 
   if (IsOutputFilterEmptyOrFullSet(request, function_metadata.value())) {
+    TRACELITERAL("Pre process prediction without output filter");
     // Pre-processing.
     std::vector<Tensor> input_tensors;
     TF_RETURN_IF_ERROR(PreProcessPredictionWithoutOutputFilter(
         function_metadata.value(), request, &input_tensors));
 
     // Executes requests.
+    TRACELITERAL("Execute prediction without output filter");
     std::vector<Tensor> outputs;
     const uint64_t start_microseconds = EnvTime::NowMicros();
     if (const auto status =
@@ -213,6 +216,7 @@ absl::Status RunPredict(
                          end_microseconds - start_microseconds);
 
     // Post-processing.
+    TRACELITERAL("Post process prediction without output filter");
     return PostProcessPredictionResultWithoutOutputFilter(
         function_metadata->GetOutputNames(), outputs, option, request,
         response);
@@ -229,13 +233,15 @@ absl::Status RunPredict(
     }
     const SignatureDef& signature = iter->second;
 
+    TRACELITERAL("Pre process prediction with output filter");
     std::vector<std::pair<string, Tensor>> input_tensors;
     std::vector<string> output_tensor_names;
     std::vector<string> output_tensor_aliases;
     TF_RETURN_IF_ERROR(PreProcessPrediction(signature, request, &input_tensors,
                                             &output_tensor_names,
                                             &output_tensor_aliases));
 
+    TRACELITERAL("Execute prediction with output filter");
     const uint64_t start_microseconds = EnvTime::NowMicros();
     std::vector<Tensor> outputs;
     if (const auto status = saved_model->RunByTensorNames(
@@ -253,6 +259,7 @@ absl::Status RunPredict(
                          /*runtime=*/"TFRT",
                          end_microseconds - start_microseconds);
 
+    TRACELITERAL("Post process prediction with output filter");
     return PostProcessPredictionResult(output_tensor_aliases, outputs, option,
                                        response);
   }