fix: make cortex models stop works again (#1364)

Co-authored-by: vansangpfiev <[email protected]>
janhq · Sep 30, 2024 · 93bce74 · 93bce74
1 parent e3f2f46
commit 93bce74
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 36 deletions.
diff --git a/engine/commands/model_stop_cmd.cc b/engine/commands/model_stop_cmd.cc
@@ -1,31 +1,48 @@
 #include "model_stop_cmd.h"
+#include "config/yaml_config.h"
+#include "database/models.h"
 #include "httplib.h"
 #include "nlohmann/json.hpp"
+#include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
 
 namespace commands {
-ModelStopCmd::ModelStopCmd(std::string host, int port,
-                           const config::ModelConfig& mc)
-    : host_(std::move(host)), port_(port), mc_(mc) {}
 
-void ModelStopCmd::Exec() {
-  httplib::Client cli(host_ + ":" + std::to_string(port_));
-  nlohmann::json json_data;
-  json_data["model"] = mc_.name;
-  json_data["engine"] = mc_.engine;
+void ModelStopCmd::Exec(const std::string& host, int port,
+                        const std::string& model_handle) {
+  cortex::db::Models modellist_handler;
+  config::YamlHandler yaml_handler;
+  try {
+    auto model_entry = modellist_handler.GetModelInfo(model_handle);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
+    auto mc = yaml_handler.GetModelConfig();
+    httplib::Client cli(host + ":" + std::to_string(port));
+    nlohmann::json json_data;
+    json_data["model"] = mc.name;
+    json_data["engine"] = mc.engine;
 
-  auto data_str = json_data.dump();
+    auto data_str = json_data.dump();
 
-  auto res = cli.Post("/inferences/server/unloadmodel", httplib::Headers(),
-                      data_str.data(), data_str.size(), "application/json");
-  if (res) {
-    if (res->status == httplib::StatusCode::OK_200) {
-      // LOG_INFO << res->body;
-      CLI_LOG("Model unloaded!");
+    auto res = cli.Post("/inferences/server/unloadmodel", httplib::Headers(),
+                        data_str.data(), data_str.size(), "application/json");
+    if (res) {
+      if (res->status == httplib::StatusCode::OK_200) {
+        // LOG_INFO << res->body;
+        CLI_LOG("Model unloaded!");
+      } else {
+        CLI_LOG("Error: could not unload model - " << res->status);
+      }
+    } else {
+      auto err = res.error();
+      CTL_ERR("HTTP error: " << httplib::to_string(err));
     }
-  } else {
-    auto err = res.error();
-    CTL_ERR("HTTP error: " << httplib::to_string(err));
+  } catch (const std::exception& e) {
+    CLI_LOG("Fail to stop model information with ID '" + model_handle +
+            "': " + e.what());
   }
 }
 

diff --git a/engine/commands/model_stop_cmd.h b/engine/commands/model_stop_cmd.h
@@ -7,12 +7,6 @@ namespace commands {
 
 class ModelStopCmd {
  public:
-  ModelStopCmd(std::string host, int port, const config::ModelConfig& mc);
-  void Exec();
-
- private:
-  std::string host_;
-  int port_;
-  const config::ModelConfig& mc_;
+  void Exec(const std::string& host, int port, const std::string& model_handle);
 };
 }  // namespace commands
diff --git a/engine/controllers/command_line_parser.cc b/engine/controllers/command_line_parser.cc
@@ -220,17 +220,9 @@ void CommandLineParser::SetupModelCommands() {
       CLI_LOG(stop_model_cmd->help());
       return;
     };
-    commands::CmdInfo ci(cml_data_.model_id);
-    std::string model_file =
-        ci.branch == "main" ? ci.model_name : ci.model_name + "-" + ci.branch;
-    config::YamlHandler yaml_handler;
-    yaml_handler.ModelConfigFromFile(
-        file_manager_utils::GetModelsContainerPath().string() + "/" +
-        model_file + ".yaml");
-    commands::ModelStopCmd smc(cml_data_.config.apiServerHost,
-                               std::stoi(cml_data_.config.apiServerPort),
-                               yaml_handler.GetModelConfig());
-    smc.Exec();
+    commands::ModelStopCmd().Exec(cml_data_.config.apiServerHost,
+                                  std::stoi(cml_data_.config.apiServerPort),
+                                  cml_data_.model_id);
   });
 
   auto list_models_cmd =