[DeepRec] Fix warmup failed bug when use session_group. (#11)

DeepRec-AI · Nov 2, 2022 · c1bfe02 · c1bfe02
1 parent 98e8759
commit c1bfe02
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/tensorflow_serving/servables/tensorflow/saved_model_warmup.cc b/tensorflow_serving/servables/tensorflow/saved_model_warmup.cc
@@ -185,10 +185,20 @@ Status RunSavedModelWarmup(const ModelWarmupOptions& model_warmup_options,
 Status RunSavedModelWarmup(const ModelWarmupOptions& model_warmup_options,
                            const RunOptions& run_options,
                            const string& export_dir, SavedModelBundleV2* bundle) {
-  return InternalRunSavedModelWarmup(model_warmup_options,
-                                     run_options, export_dir,
-                                     bundle->meta_graph_def,
-                                     bundle->session_group->GetLeaderSession());
+  int N = bundle->session_group->GetSessionNum();
+  for (int i = 0; i < N; ++i) {
+    Status s = InternalRunSavedModelWarmup(
+        model_warmup_options,
+        run_options, export_dir,
+        bundle->meta_graph_def,
+        bundle->session_group->GetSession(i));
+    if (!s.ok()) {
+      LOG(ERROR) << "Warmup session group failed. " << s.error_message();
+      return s;
+    }
+  }
+
+  return Status::OK();
 }
 
 }  // namespace serving