Server shutdown must wait for background models

src/model_repository_manager/model_lifecycle.cc

@@ -259,6 +259,14 @@ ModelLifeCycle::InflightStatus()
   return inflight_status;
 }
 
+size_t
+ModelLifeCycle::BackgroundModelsSize()
+{
+  LOG_VERBOSE(2) << "BackgroundModelsSize()";
+  std::lock_guard<std::mutex> map_lock(map_mtx_);
+  return background_models_.size();
+}
+
 const ModelStateMap
 ModelLifeCycle::ModelStates()
 {

src/model_repository_manager/model_lifecycle.h

@@ -221,6 +221,9 @@ class ModelLifeCycle {
   // that don't have in-flight inferences will not be included.
   const std::set<std::tuple<ModelIdentifier, int64_t, size_t>> InflightStatus();
 
+  // Return the number of model(s) in the background.
+  size_t BackgroundModelsSize();
+
  private:
   struct ModelInfo {
     ModelInfo(

src/model_repository_manager/model_repository_manager.cc

@@ -981,6 +981,12 @@ ModelRepositoryManager::InflightStatus()
   return model_life_cycle_->InflightStatus();
 }
 
+size_t
+ModelRepositoryManager::BackgroundModelsSize()
+{
+  return model_life_cycle_->BackgroundModelsSize();
+}
+
 const ModelStateMap
 ModelRepositoryManager::LiveModelStates(bool strict_readiness)
 {

src/model_repository_manager/model_repository_manager.h

@@ -438,6 +438,9 @@ class ModelRepositoryManager {
   /// if it doesn't have in-flight inferences.
   const std::set<std::tuple<ModelIdentifier, int64_t, size_t>> InflightStatus();
 
+  /// \return the number of model(s) in the background.
+  size_t BackgroundModelsSize();
+
   /// \param strict_readiness If true, only models that have at least one
   /// ready version will be considered as live. Otherwise, the models that
   /// have loading / unloading versions will also be live.

src/server.cc

@@ -341,6 +341,7 @@ InferenceServer::Stop(const bool force)
       }
     } else {
       const auto& live_models = model_repository_manager_->LiveModelStates();
+      size_t bg_models_size = model_repository_manager_->BackgroundModelsSize();
 
       LOG_INFO << "Timeout " << exit_timeout_iters << ": Found "
                << live_models.size() << " live models and "
@@ -355,7 +356,8 @@ InferenceServer::Stop(const bool force)
         }
       }
 
-      if ((live_models.size() == 0) && (inflight_request_counter_ == 0)) {
+      if ((live_models.size() == 0) && (bg_models_size == 0) &&
+          (inflight_request_counter_ == 0)) {
         return Status::Success;
       }
     }