diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index 7cca9624d..17577aff5 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -168,9 +168,32 @@ void Models::AbortPullModel( } } -void Models::ListModel( - const HttpRequestPtr& req, - std::function&& callback) const { +void Models::ListModel(const HttpRequestPtr& req, + std::function&& callback, + std::optional author) const { + if (author.has_value()) { + auto res = model_src_svc_->GetRepositoryList(author.value()); + if (res.has_error()) { + Json::Value ret; + ret["message"] = res.error(); + auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); + resp->setStatusCode(k400BadRequest); + callback(resp); + } else { + auto& info = res.value(); + Json::Value ret; + Json::Value arr(Json::arrayValue); + for (auto const& i : info) { + arr.append(i); + } + ret["data"] = arr; + auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); + resp->setStatusCode(k200OK); + callback(resp); + } + return; + } + namespace fs = std::filesystem; namespace fmu = file_manager_utils; Json::Value ret; diff --git a/engine/controllers/models.h b/engine/controllers/models.h index 8f6af4280..a905cc843 100644 --- a/engine/controllers/models.h +++ b/engine/controllers/models.h @@ -30,7 +30,7 @@ class Models : public drogon::HttpController { ADD_METHOD_TO(Models::PullModel, "/v1/models/pull", Options, Post); ADD_METHOD_TO(Models::AbortPullModel, "/v1/models/pull", Options, Delete); - ADD_METHOD_TO(Models::ListModel, "/v1/models", Get); + ADD_METHOD_TO(Models::ListModel, "/v1/models?author={author}", Get); ADD_METHOD_TO(Models::GetModel, "/v1/models/{1}", Get); ADD_METHOD_TO(Models::UpdateModel, "/v1/models/{1}", Options, Patch); ADD_METHOD_TO(Models::ImportModel, "/v1/models/import", Options, Post); @@ -63,7 +63,8 @@ class Models : public drogon::HttpController { void AbortPullModel(const HttpRequestPtr& req, std::function&& callback); void ListModel(const HttpRequestPtr& req, - std::function&& callback) const; + std::function&& callback, + std::optional author) const; void GetModel(const HttpRequestPtr& req, std::function&& callback, const std::string& model_id) const; diff --git a/engine/services/database_service.cc b/engine/services/database_service.cc index 695e36f72..6cf16ec69 100644 --- a/engine/services/database_service.cc +++ b/engine/services/database_service.cc @@ -6,25 +6,30 @@ std::optional DatabaseService::UpsertEngine( const std::string& api_key, const std::string& url, const std::string& version, const std::string& variant, const std::string& status, const std::string& metadata) { + std::lock_guard l(mtx_); return cortex::db::Engines().UpsertEngine(engine_name, type, api_key, url, version, variant, status, metadata); } std::optional> DatabaseService::GetEngines() const { + std::lock_guard l(mtx_); return cortex::db::Engines().GetEngines(); } std::optional DatabaseService::GetEngineById(int id) const { + std::lock_guard l(mtx_); return cortex::db::Engines().GetEngineById(id); } std::optional DatabaseService::GetEngineByNameAndVariant( const std::string& engine_name, const std::optional variant) const { + std::lock_guard l(mtx_); return cortex::db::Engines().GetEngineByNameAndVariant(engine_name, variant); } std::optional DatabaseService::DeleteEngineById(int id) { + std::lock_guard l(mtx_); return cortex::db::Engines().DeleteEngineById(id); } // end engines @@ -32,21 +37,25 @@ std::optional DatabaseService::DeleteEngineById(int id) { // begin file cpp::result, std::string> DatabaseService::GetFileList() const { + std::lock_guard l(mtx_); return cortex::db::File().GetFileList(); } cpp::result DatabaseService::GetFileById( const std::string& file_id) const { + std::lock_guard l(mtx_); return cortex::db::File().GetFileById(file_id); } cpp::result DatabaseService::AddFileEntry( OpenAi::File& file) { + std::lock_guard l(mtx_); return cortex::db::File().AddFileEntry(file); } cpp::result DatabaseService::DeleteFileEntry( const std::string& file_id) { + std::lock_guard l(mtx_); return cortex::db::File().DeleteFileEntry(file_id); } // end file @@ -54,77 +63,104 @@ cpp::result DatabaseService::DeleteFileEntry( // begin hardware cpp::result, std::string> DatabaseService::LoadHardwareList() const { + std::lock_guard l(mtx_); return cortex::db::Hardware().LoadHardwareList(); } cpp::result DatabaseService::AddHardwareEntry( const HardwareEntry& new_entry) { + std::lock_guard l(mtx_); return cortex::db::Hardware().AddHardwareEntry(new_entry); } +bool DatabaseService::HasHardwareEntry(const std::string& id) { + std::lock_guard l(mtx_); + return cortex::db::Hardware().HasHardwareEntry(id); +} + cpp::result DatabaseService::UpdateHardwareEntry( const std::string& id, const HardwareEntry& updated_entry) { + std::lock_guard l(mtx_); return cortex::db::Hardware().UpdateHardwareEntry(id, updated_entry); } cpp::result DatabaseService::DeleteHardwareEntry( const std::string& id) { + std::lock_guard l(mtx_); return cortex::db::Hardware().DeleteHardwareEntry(id); } + +cpp::result DatabaseService::UpdateHardwareEntry( + const std::string& id, int hw_id, int sw_id) const { + std::lock_guard l(mtx_); + return cortex::db::Hardware().UpdateHardwareEntry(id, hw_id, sw_id); +} // end hardware // begin models cpp::result, std::string> DatabaseService::LoadModelList() const { + std::lock_guard l(mtx_); return cortex::db::Models().LoadModelList(); } cpp::result DatabaseService::GetModelInfo( const std::string& identifier) const { + std::lock_guard l(mtx_); return cortex::db::Models().GetModelInfo(identifier); } cpp::result DatabaseService::AddModelEntry( ModelEntry new_entry) { + std::lock_guard l(mtx_); return cortex::db::Models().AddModelEntry(new_entry); } cpp::result DatabaseService::UpdateModelEntry( const std::string& identifier, const ModelEntry& updated_entry) { + std::lock_guard l(mtx_); return cortex::db::Models().UpdateModelEntry(identifier, updated_entry); } cpp::result DatabaseService::DeleteModelEntry( const std::string& identifier) { + std::lock_guard l(mtx_); return cortex::db::Models().DeleteModelEntry(identifier); } cpp::result DatabaseService::DeleteModelEntryWithOrg( const std::string& src) { + std::lock_guard l(mtx_); return cortex::db::Models().DeleteModelEntryWithOrg(src); } cpp::result DatabaseService::DeleteModelEntryWithRepo( const std::string& src) { + + std::lock_guard l(mtx_); return cortex::db::Models().DeleteModelEntryWithRepo(src); } cpp::result, std::string> DatabaseService::FindRelatedModel(const std::string& identifier) const { + std::lock_guard l(mtx_); return cortex::db::Models().FindRelatedModel(identifier); } bool DatabaseService::HasModel(const std::string& identifier) const { + std::lock_guard l(mtx_); return cortex::db::Models().HasModel(identifier); } cpp::result, std::string> DatabaseService::GetModels( const std::string& model_src) const { + std::lock_guard l(mtx_); return cortex::db::Models().GetModels(model_src); } cpp::result, std::string> DatabaseService::GetModelSources() const { + std::lock_guard l(mtx_); return cortex::db::Models().GetModelSources(); } // end models \ No newline at end of file diff --git a/engine/services/database_service.h b/engine/services/database_service.h index e45fc57b2..660c57cb7 100644 --- a/engine/services/database_service.h +++ b/engine/services/database_service.h @@ -1,4 +1,5 @@ #pragma once +#include #include "database/engines.h" #include "database/file.h" #include "database/hardware.h" @@ -39,9 +40,13 @@ class DatabaseService { cpp::result, std::string> LoadHardwareList() const; cpp::result AddHardwareEntry( const HardwareEntry& new_entry); + bool HasHardwareEntry(const std::string& id); cpp::result UpdateHardwareEntry( const std::string& id, const HardwareEntry& updated_entry); cpp::result DeleteHardwareEntry(const std::string& id); + cpp::result UpdateHardwareEntry(const std::string& id, + int hw_id, + int sw_id) const; // models cpp::result, std::string> LoadModelList() const; @@ -62,8 +67,8 @@ class DatabaseService { bool HasModel(const std::string& identifier) const; cpp::result, std::string> GetModels( const std::string& model_src) const; - cpp::result, std::string> GetModelSources() - const; + cpp::result, std::string> GetModelSources() const; private: + mutable std::mutex mtx_; }; \ No newline at end of file diff --git a/engine/services/hardware_service.cc b/engine/services/hardware_service.cc index 6b19a0119..56ecadd6d 100644 --- a/engine/services/hardware_service.cc +++ b/engine/services/hardware_service.cc @@ -296,8 +296,7 @@ void HardwareService::UpdateHardwareInfos() { using HwEntry = cortex::db::HardwareEntry; CheckDependencies(); auto gpus = cortex::hw::GetGPUInfo(); - cortex::db::Hardware hw_db; - auto b = hw_db.LoadHardwareList(); + auto b = db_service_->LoadHardwareList(); // delete if not exists auto exists = [&gpus](const std::string& uuid) { for (auto const& g : gpus) { @@ -308,12 +307,12 @@ void HardwareService::UpdateHardwareInfos() { }; for (auto const& he : b.value()) { if (!exists(he.uuid)) { - hw_db.DeleteHardwareEntry(he.uuid); + db_service_->DeleteHardwareEntry(he.uuid); } } // Get updated list - b = hw_db.LoadHardwareList(); + b = db_service_->LoadHardwareList(); std::vector> activated_gpu_bf; std::string debug_b; for (auto const& he : b.value()) { @@ -326,15 +325,15 @@ void HardwareService::UpdateHardwareInfos() { for (auto const& gpu : gpus) { // ignore error // Note: only support NVIDIA for now, so hardware_id = software_id - if (hw_db.HasHardwareEntry(gpu.uuid)) { - auto res = hw_db.UpdateHardwareEntry(gpu.uuid, std::stoi(gpu.id), + if (db_service_->HasHardwareEntry(gpu.uuid)) { + auto res = db_service_->UpdateHardwareEntry(gpu.uuid, std::stoi(gpu.id), std::stoi(gpu.id)); if (res.has_error()) { CTL_WRN(res.error()); } } else { auto res = - hw_db.AddHardwareEntry(HwEntry{.uuid = gpu.uuid, + db_service_->AddHardwareEntry(HwEntry{.uuid = gpu.uuid, .type = "gpu", .hardware_id = std::stoi(gpu.id), .software_id = std::stoi(gpu.id), diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index a9131c886..979cf9342 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -858,8 +858,7 @@ cpp::result ModelService::StartModel( } else { // only report to user the error for (auto& depend : depends) { - - StopModel(depend); + (void)StopModel(depend); } } CTL_ERR("Model failed to start with status code: " << status); @@ -869,7 +868,7 @@ cpp::result ModelService::StartModel( // Running remote model if (engine_svc_->IsRemoteEngine(mc.engine)) { - engine_svc_->LoadEngine(mc.engine); + (void)engine_svc_->LoadEngine(mc.engine); config::RemoteModelConfig remote_mc; remote_mc.LoadFromYamlFile( fmu::ToAbsoluteCortexDataPath( @@ -1035,7 +1034,7 @@ cpp::result ModelService::StopModel( // Stop all depends model auto depends = python_model_config.depends; for (auto& depend : depends) { - StopModel(depend); + (void)StopModel(depend); } } @@ -1233,7 +1232,7 @@ cpp::result, std::string> ModelService::MayFallbackToCpu(const std::string& model_path, int ngl, int ctx_len, int n_batch, int n_ubatch, const std::string& kv_cache_type) { - // TODO(sang) temporary disable this function + // TODO(sang) temporary disable this function return std::nullopt; assert(hw_service_); auto hw_info = hw_service_->GetHardwareInfo(); diff --git a/engine/services/model_source_service.cc b/engine/services/model_source_service.cc index f0e29b8bc..5298aa75c 100644 --- a/engine/services/model_source_service.cc +++ b/engine/services/model_source_service.cc @@ -1,5 +1,6 @@ #include "model_source_service.h" #include +#include #include #include "database/models.h" #include "json/json.h" @@ -199,6 +200,26 @@ cpp::result ModelSourceService::GetModelSource( return ms; } +cpp::result, std::string> +ModelSourceService::GetRepositoryList(std::string_view author) { + std::string as(author); + if (cortexso_repos_.find(as) != cortexso_repos_.end() && + !cortexso_repos_.at(as).empty()) + return cortexso_repos_.at(as); + const auto begin = std::chrono::high_resolution_clock::now(); + auto res = + curl_utils::SimpleGet("https://huggingface.co/api/models?author=" + as); + if (res.has_value()) { + auto repos = ParseJsonString(res.value()); + for (auto& r : repos) { + cortexso_repos_[as].push_back(r.id); + } + return cortexso_repos_.at(as); + } else { + return cpp::fail(res.error()); + } +} + cpp::result ModelSourceService::AddHfOrg( const std::string& model_source, const std::string& author) { auto res = curl_utils::SimpleGet("https://huggingface.co/api/models?author=" + @@ -361,6 +382,7 @@ cpp::result ModelSourceService::AddCortexsoOrg( cpp::result ModelSourceService::AddCortexsoRepo( const std::string& model_source, const std::string& author, const std::string& model_name) { + auto begin = std::chrono::system_clock::now(); auto branches = huggingface_utils::GetModelRepositoryBranches("cortexso", model_name); if (branches.has_error()) { @@ -377,19 +399,25 @@ cpp::result ModelSourceService::AddCortexsoRepo( if (!readme.has_error()) { desc = readme.value(); } - // Get models from db + // Get models from db auto model_list_before = db_service_->GetModels(model_source) .value_or(std::vector{}); std::unordered_set updated_model_list; - + std::vector> tasks; for (auto const& [branch, _] : branches.value()) { - CTL_INF(branch); - auto add_res = AddCortexsoRepoBranch(model_source, author, model_name, - branch, repo_info->metadata, desc) - .value_or(std::unordered_set{}); - for (auto const& a : add_res) { - updated_model_list.insert(a); + CTL_DBG(branch); + tasks.push_back(std::async(std::launch::async, [&, branch = branch] { + return AddCortexsoRepoBranch(model_source, author, model_name, branch, + repo_info->metadata, desc) + .value_or(std::string{}); + })); + } + + for (auto& task : tasks) { + auto add_res = task.get(); + if (!add_res.empty()) { + updated_model_list.insert(add_res); } } @@ -402,18 +430,19 @@ cpp::result ModelSourceService::AddCortexsoRepo( } } } + + auto end = std::chrono::system_clock::now(); + CTL_INF( + "Duration ms: " << std::chrono::duration_cast( + end - begin) + .count()); return true; } -cpp::result, std::string> -ModelSourceService::AddCortexsoRepoBranch(const std::string& model_source, - const std::string& author, - const std::string& model_name, - const std::string& branch, - const std::string& metadata, - const std::string& desc) { - std::unordered_set res; - +cpp::result ModelSourceService::AddCortexsoRepoBranch( + const std::string& model_source, const std::string& author, + const std::string& model_name, const std::string& branch, + const std::string& metadata, const std::string& desc) { url_parser::Url url = { .protocol = "https", .host = kHuggingFaceHost, @@ -468,9 +497,9 @@ ModelSourceService::AddCortexsoRepoBranch(const std::string& model_source, } } } - res.insert(model_id); + return model_id; } - return res; + return {}; } void ModelSourceService::SyncModelSource() { diff --git a/engine/services/model_source_service.h b/engine/services/model_source_service.h index 606d25d4f..40ca1f887 100644 --- a/engine/services/model_source_service.h +++ b/engine/services/model_source_service.h @@ -53,6 +53,9 @@ class ModelSourceService { cpp::result GetModelSource(const std::string& src); + cpp::result, std::string> GetRepositoryList( + std::string_view author); + private: cpp::result AddHfOrg(const std::string& model_source, const std::string& author); @@ -72,7 +75,7 @@ class ModelSourceService { const std::string& model_source, const std::string& author, const std::string& model_name); - cpp::result, std::string> + cpp::result AddCortexsoRepoBranch(const std::string& model_source, const std::string& author, const std::string& model_name, @@ -85,4 +88,6 @@ class ModelSourceService { std::shared_ptr db_service_ = nullptr; std::thread sync_db_thread_; std::atomic running_; + + std::unordered_map> cortexso_repos_; }; \ No newline at end of file