Skip to content

Commit 21fdb59

Browse files
authored
[DeepRec] Fix warmup failed bug when use session_group. (#11)
1 parent 98e8759 commit 21fdb59

File tree

1 file changed

+14
-4
lines changed

1 file changed

+14
-4
lines changed

tensorflow_serving/servables/tensorflow/saved_model_warmup.cc

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -185,10 +185,20 @@ Status RunSavedModelWarmup(const ModelWarmupOptions& model_warmup_options,
185185
Status RunSavedModelWarmup(const ModelWarmupOptions& model_warmup_options,
186186
const RunOptions& run_options,
187187
const string& export_dir, SavedModelBundleV2* bundle) {
188-
return InternalRunSavedModelWarmup(model_warmup_options,
189-
run_options, export_dir,
190-
bundle->meta_graph_def,
191-
bundle->session_group->GetLeaderSession());
188+
int N = bundle->session_group->GetSessionNum();
189+
for (int i = 0; i < N; ++i) {
190+
Status s = InternalRunSavedModelWarmup(
191+
model_warmup_options,
192+
run_options, export_dir,
193+
bundle->meta_graph_def,
194+
bundle->session_group->GetSession(i));
195+
if (!s.ok()) {
196+
LOG(ERROR) << "Warmup session group failed. " << s.error_message();
197+
return s;
198+
}
199+
}
200+
201+
return Status::OK();
192202
}
193203

194204
} // namespace serving

0 commit comments

Comments
 (0)