diff --git a/src/model_repository_manager/model_repository_manager.cc b/src/model_repository_manager/model_repository_manager.cc index ae982a1ad..8b06ad914 100644 --- a/src/model_repository_manager/model_repository_manager.cc +++ b/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> + models; + { + std::lock_guard 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)}; } diff --git a/src/model_repository_manager/model_repository_manager.h b/src/model_repository_manager/model_repository_manager.h index 60793500f..4d4e9531d 100644 --- a/src/model_repository_manager/model_repository_manager.h +++ b/src/model_repository_manager/model_repository_manager.h @@ -421,8 +421,11 @@ class ModelRepositoryManager { std::string, std::vector>& 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(); diff --git a/src/server.cc b/src/server.cc index e158df0c6..606ade8a3 100644 --- a/src/server.cc +++ b/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; }