Server stop will wait until all background models are unloaded (#323)

* Unload all models wait until all models complete loading

* [NOT ROOT CAUSE] Revert "Unload all models wait until all models complete loading"

This reverts commit 62a4e99.

* Server shutdown must wait for background models

* Show bg model size when printing timeout info

Author: kthui

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,10 +341,11 @@ InferenceServer::Stop(const bool force)
       }
     } else {
       const auto& live_models = model_repository_manager_->LiveModelStates();
+      size_t bg_models_size = model_repository_manager_->BackgroundModelsSize();
+      size_t num_models = live_models.size() + bg_models_size;
 
-      LOG_INFO << "Timeout " << exit_timeout_iters << ": Found "
-               << live_models.size() << " live models and "
-               << inflight_request_counter_
+      LOG_INFO << "Timeout " << exit_timeout_iters << ": Found " << num_models
+               << " live models and " << inflight_request_counter_
                << " in-flight non-inference requests";
       if (LOG_VERBOSE_IS_ON(1)) {
         for (const auto& m : live_models) {
@@ -355,7 +356,7 @@ InferenceServer::Stop(const bool force)
         }
       }
 
-      if ((live_models.size() == 0) && (inflight_request_counter_ == 0)) {
+      if ((num_models == 0) && (inflight_request_counter_ == 0)) {
         return Status::Success;
       }
     }