Server stop will wait until all loading models to complete before unload all models (#324)

* Unload all models must wait until all models complete loading or unloading

* Use LoadUnloadModels() in UnloadAllModels()

* Erase removed nodes only after affected nodes is completed

* Respect server shutdown timeout when waiting for load

Author: kthui

src/model_repository_manager/model_repository_manager.cc

@@ -957,16 +957,21 @@ ModelRepositoryManager::PollModels(
 Status
 ModelRepositoryManager::UnloadAllModels()
 {
-  Status status;
-  for (const auto& name_info : infos_) {
-    Status unload_status = model_life_cycle_->AsyncUnload(name_info.first);
-    if (!unload_status.IsOk()) {
-      status = Status(
-          unload_status.ErrorCode(),
-          "Failed to gracefully unload models: " + unload_status.Message());
+  // Find all the models to be unloaded.
+  std::unordered_map<std::string, std::vector<const InferenceParameter*>>
+      models;
+  {
+    std::lock_guard<std::mutex> lock(mu_);
+    for (const auto& pair : infos_) {
+      // Only the model name is needed.
+      models[pair.first.name_];
     }
   }
-  return Status::Success;
+  // Unload all models found.
+  bool polled;
+  return LoadUnloadModels(
+      models, ActionType::UNLOAD, true /* unload_dependents */, &polled,
+      nullptr /* no_parallel_conflict */);
 }
 
 Status
@@ -1821,12 +1826,14 @@ ModelRepositoryManager::DependencyGraph::RemoveNodes(
       }
 
       all_removed_nodes.emplace(model_id);
-      // Exclude removed node from affected nodes to skip some evaluations.
-      all_affected_nodes.erase(model_id);
     }
 
     curr_removal.swap(next_removal);
   }
+  // Exclude removed nodes from affected nodes.
+  for (const auto& removed_node : all_removed_nodes) {
+    all_affected_nodes.erase(removed_node);
+  }
   return {std::move(all_affected_nodes), std::move(all_removed_nodes)};
 }
 

src/model_repository_manager/model_repository_manager.h

@@ -421,8 +421,11 @@ class ModelRepositoryManager {
           std::string, std::vector<const InferenceParameter*>>& models,
       const ActionType type, const bool unload_dependents);
 
-  /// Unload all models. This function should be called before shutting down
-  /// the model repository manager.
+  /// Unload all models that are tracked by the model repository manager. If a
+  /// model is loading or unloading when this function is called, or a model
+  /// failed to unload, an error is returned. New models may be loaded while
+  /// this function is unloading models. This function should be called before
+  /// shutting down the model repository manager.
   /// \return error status.
   Status UnloadAllModels();
 

src/server.cc

@@ -330,11 +330,11 @@ InferenceServer::Stop(const bool force)
       }
 
       if (inflight_status.size() == 0) {
-        unloading_model = true;
         status = model_repository_manager_->UnloadAllModels();
         if (!status.IsOk()) {
-          LOG_ERROR << status.Message();
+          LOG_WARNING << status.Message();
         } else {
+          unloading_model = true;
           LOG_INFO << "All models are stopped, unloading models";
           continue;
         }