Extract the untracked part of RunSavedModelWarmup to a separate function

PiperOrigin-RevId: 791455959

Author: AspirinSJL

tensorflow_serving/servables/tensorflow/BUILD

@@ -933,10 +933,13 @@ cc_library(
     deps = [
         ":session_bundle_config_cc_proto",
         "//tensorflow_serving/apis:prediction_log_cc_proto",
+        "//tensorflow_serving/util:executor",
         "//tensorflow_serving/util:threadpool_executor",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
         "@com_google_protobuf//:cc_wkt_protos",
         "@org_tensorflow//tensorflow/cc/saved_model:constants",
-        "@org_tensorflow//tensorflow/cc/saved_model:loader",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/kernels/batching_util:warmup",
@@ -979,14 +982,22 @@ cc_test(
         "//tensorflow_serving/apis:predict_cc_proto",
         "//tensorflow_serving/apis:prediction_log_cc_proto",
         "//tensorflow_serving/apis:regression_cc_proto",
-        "//tensorflow_serving/core/test_util:test_main",
+        "//tensorflow_serving/core/test_util:test_main",  # buildcleaner: keep
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_protobuf//:cc_wkt_protos",
+        "@local_xla//xla/tsl/lib/core:status_test_util",
         "@org_tensorflow//tensorflow/cc/saved_model:constants",
-        "@org_tensorflow//tensorflow/cc/saved_model:signature_constants",
+        "@org_tensorflow//tensorflow/cc/saved_model:loader",
+        "@org_tensorflow//tensorflow/core:framework_lite",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core:test",
+        "@org_tensorflow//tensorflow/core:tflite_portable_logging",
+        "@org_tensorflow//tensorflow/core/example:example_protos_cc",
         "@org_tensorflow//tensorflow/core/kernels/batching_util:warmup",
+        "@org_tensorflow//tensorflow/core/platform:errors",
+        "@org_tensorflow//tensorflow/core/platform:path",
     ],
 )
 

tensorflow_serving/servables/tensorflow/saved_model_warmup_util.cc

@@ -16,20 +16,30 @@ limitations under the License.
 #include "tensorflow_serving/servables/tensorflow/saved_model_warmup_util.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <utility>
 
 #include "google/protobuf/wrappers.pb.h"
+#include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "tensorflow/cc/saved_model/constants.h"
 #include "xla/tsl/platform/errors.h"
 #include "tensorflow/core/kernels/batching_util/warmup.h"
 #include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/io/record_reader.h"
 #include "tensorflow/core/lib/monitoring/sampler.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/env_time.h"
+#include "tensorflow/core/platform/file_system.h"
 #include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow_serving/util/executor.h"
 #include "tensorflow_serving/util/threadpool_executor.h"
 
 namespace tensorflow {
@@ -58,22 +68,9 @@ uint64_t GetLatencyMicroseconds(const uint64_t start_microseconds) {
 constexpr char WarmupConsts::kRequestsFileName[];
 constexpr int WarmupConsts::kMaxNumRecords;
 
-absl::Status RunSavedModelWarmup(
+absl::Status RunSavedModelWarmupUntracked(
     const ModelWarmupOptions& model_warmup_options, const string export_dir,
     std::function<absl::Status(PredictionLog)> warmup_request_executor) {
-  WarmupStateRegistry::Handle warmup_handle;
-  auto per_model_data = std::make_unique<WarmupStateRegistry::PerModelData>();
-  per_model_data->warmup_all_batch_sizes =
-      model_warmup_options.enable_all_batch_sizes_warmup();
-  if (!model_warmup_options.model_name().empty()) {
-    auto h = GetGlobalWarmupStateRegistry().Register(
-        {model_warmup_options.model_name(),
-         model_warmup_options.model_version()},
-        std::move(per_model_data));
-    TF_RETURN_IF_ERROR(h.status());
-    warmup_handle = std::move(h.value());
-  }
-
   const uint64_t start_microseconds = EnvTime::NowMicros();
   const string warmup_path =
       io::JoinPath(export_dir, kSavedModelAssetsExtraDirectory,
@@ -237,6 +234,26 @@ absl::Status RunSavedModelWarmup(
   return absl::OkStatus();
 }
 
+absl::Status RunSavedModelWarmup(
+    const ModelWarmupOptions& model_warmup_options, const string export_dir,
+    std::function<absl::Status(PredictionLog)> warmup_request_executor) {
+  WarmupStateRegistry::Handle warmup_handle;
+  auto per_model_data = std::make_unique<WarmupStateRegistry::PerModelData>();
+  per_model_data->warmup_all_batch_sizes =
+      model_warmup_options.enable_all_batch_sizes_warmup();
+  if (!model_warmup_options.model_name().empty()) {
+    auto h = GetGlobalWarmupStateRegistry().Register(
+        {model_warmup_options.model_name(),
+         model_warmup_options.model_version()},
+        std::move(per_model_data));
+    TF_RETURN_IF_ERROR(h.status());
+    warmup_handle = std::move(h.value());
+  }
+
+  return RunSavedModelWarmupUntracked(model_warmup_options, export_dir,
+                                      warmup_request_executor);
+}
+
 }  // namespace internal
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/servables/tensorflow/saved_model_warmup_util.h

@@ -16,7 +16,10 @@ limitations under the License.
 #ifndef THIRD_PARTY_TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SAVED_MODEL_WARMUP_UTIL_H_
 #define THIRD_PARTY_TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SAVED_MODEL_WARMUP_UTIL_H_
 
-#include "tensorflow/cc/saved_model/loader.h"
+#include <functional>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/saved_model.pb.h"
 #include "tensorflow_serving/apis/prediction_log.pb.h"
 #include "tensorflow_serving/servables/tensorflow/session_bundle_config.pb.h"
@@ -35,9 +38,18 @@ struct WarmupConsts {
 // to trigger lazy initializations (such as TF optimizations, XLA compilations)
 // at load time, and consequently improve first request latency.
 // Warmup is skipped if no warmup file present.
-Status RunSavedModelWarmup(
+absl::Status RunSavedModelWarmup(
+    const ModelWarmupOptions& model_warmup_options, const string export_dir,
+    std::function<absl::Status(PredictionLog)> warmup_request_executor);
+
+// Similar to `RunSavedModelWarmup()`, but does not track the warmup state.
+//
+// WARNING: Inside the function, multiple warmup threads might be dispatched to
+// run `warmup_request_executor`. Use with caution, especially when batching is
+// involved.
+absl::Status RunSavedModelWarmupUntracked(
     const ModelWarmupOptions& model_warmup_options, const string export_dir,
-    std::function<Status(PredictionLog)> warmup_request_executor);
+    std::function<absl::Status(PredictionLog)> warmup_request_executor);
 
 }  // namespace internal
 }  // namespace serving