feat: cortex pull and cortex engines install CLI uses API server

engine/cli/CMakeLists.txt

@@ -80,6 +80,8 @@ add_executable(${TARGET_NAME} main.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/../services/engine_service.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/../services/model_service.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/../services/inference_service.cc
+    ${CMAKE_CURRENT_SOURCE_DIR}/utils/easywsclient.cc
+    ${CMAKE_CURRENT_SOURCE_DIR}/utils/download_progress.cc
   )
 
 target_link_libraries(${TARGET_NAME} PRIVATE httplib::httplib)

engine/cli/command_line_parser.cc

@@ -130,7 +130,9 @@ void CommandLineParser::SetupCommonCommands() {
       return;
     }
     try {
-      commands::ModelPullCmd(download_service_).Exec(cml_data_.model_id);
+      commands::ModelPullCmd(download_service_)
+          .Exec(cml_data_.config.apiServerHost,
+                std::stoi(cml_data_.config.apiServerPort), cml_data_.model_id);
     } catch (const std::exception& e) {
       CLI_LOG(e.what());
     }
@@ -462,7 +464,9 @@ void CommandLineParser::EngineInstall(CLI::App* parent,
     if (std::exchange(executed_, true))
       return;
     try {
-      commands::EngineInstallCmd(download_service_)
+      commands::EngineInstallCmd(download_service_,
+                                 cml_data_.config.apiServerHost,
+                                 std::stoi(cml_data_.config.apiServerPort))
           .Exec(engine_name, version, src);
     } catch (const std::exception& e) {
       CTL_ERR(e.what());

engine/cli/commands/engine_install_cmd.cc

@@ -1,16 +1,68 @@
 #include "engine_install_cmd.h"
+#include "server_start_cmd.h"
+#include "utils/download_progress.h"
+#include "utils/engine_constants.h"
+#include "utils/json_helper.h"
 #include "utils/logging_utils.h"
 
 namespace commands {
-
-void EngineInstallCmd::Exec(const std::string& engine,
+bool EngineInstallCmd::Exec(const std::string& engine,
                             const std::string& version,
                             const std::string& src) {
-  auto result = engine_service_.InstallEngine(engine, version, src);
-  if (result.has_error()) {
-    CLI_LOG(result.error());
-  } else if(result && result.value()){
-    CLI_LOG("Engine " << engine << " installed successfully!");
+  // Handle local install, if fails, fallback to remote install
+  if (!src.empty()) {
+    auto res = engine_service_.UnzipEngine(engine, version, src);
+    if (res.has_error()) {
+      CLI_LOG(res.error());
+      return false;
+    }
+    if (res.value()) {
+      CLI_LOG("Engine " << engine << " installed successfully!");
+      return true;
+    }
+  }
+
+  // Start server if server is not started yet
+  if (!commands::IsServerAlive(host_, port_)) {
+    CLI_LOG("Starting server ...");
+    commands::ServerStartCmd ssc;
+    if (!ssc.Exec(host_, port_)) {
+      return false;
+    }
+  }
+
+  httplib::Client cli(host_ + ":" + std::to_string(port_));
+  Json::Value json_data;
+  auto data_str = json_data.toStyledString();
+  cli.set_read_timeout(std::chrono::seconds(60));
+  auto res = cli.Post("/v1/engines/install/" + engine, httplib::Headers(),
+                      data_str.data(), data_str.size(), "application/json");
+
+  if (res) {
+    if (res->status != httplib::StatusCode::OK_200) {
+      auto root = json_helper::ParseJsonString(res->body);
+      CLI_LOG(root["message"].asString());
+      return false;
+    }
+  } else {
+    auto err = res.error();
+    CTL_ERR("HTTP error: " << httplib::to_string(err));
+    return false;
   }
+
+  CLI_LOG("Start downloading ...")
+  DownloadProgress dp;
+  dp.Connect(host_, port_);
+  if (!dp.Handle(engine))
+    return false;
+
+  bool check_cuda_download = !system_info_utils::GetCudaVersion().empty();
+  if (check_cuda_download) {
+    if (!dp.Handle("cuda"))
+      return false;
+  }
+
+  CLI_LOG("Engine " << engine << " downloaded successfully!")
+  return true;
 }
 };  // namespace commands

engine/cli/commands/engine_install_cmd.h

@@ -7,13 +7,15 @@ namespace commands {
 
 class EngineInstallCmd {
  public:
-  explicit EngineInstallCmd(std::shared_ptr<DownloadService> download_service)
-      : engine_service_{EngineService(download_service)} {};
+  explicit EngineInstallCmd(std::shared_ptr<DownloadService> download_service, const std::string& host, int port)
+      : engine_service_{EngineService(download_service)}, host_(host), port_(port) {};
 
-  void Exec(const std::string& engine, const std::string& version = "latest",
+  bool Exec(const std::string& engine, const std::string& version = "latest",
             const std::string& src = "");
 
  private:
   EngineService engine_service_;
+  std::string host_;
+  int port_;
 };
 }  // namespace commands

engine/cli/commands/model_pull_cmd.cc

@@ -1,11 +1,183 @@
 #include "model_pull_cmd.h"
+#include <memory>
+#include "cli/utils/easywsclient.hpp"
+#include "cli/utils/indicators.hpp"
+#include "common/event.h"
+#include "database/models.h"
+#include "server_start_cmd.h"
+#include "utils/cli_selection_utils.h"
+#include "utils/download_progress.h"
+#include "utils/format_utils.h"
+#include "utils/huggingface_utils.h"
+#include "utils/json_helper.h"
 #include "utils/logging_utils.h"
+#include "utils/scope_exit.h"
+#include "utils/string_utils.h"
+#if defined(_WIN32)
+#include <signal.h>
+#endif
 
 namespace commands {
-void ModelPullCmd::Exec(const std::string& input) {
-  auto result = model_service_.DownloadModel(input);
-  if (result.has_error()) {
-    CLI_LOG(result.error());
+std::function<void(int)> shutdown_handler;
+inline void signal_handler(int signal) {
+  if (shutdown_handler) {
+    shutdown_handler(signal);
+  }
+}
+std::optional<std::string> ModelPullCmd::Exec(const std::string& host, int port,
+                                              const std::string& input) {
+
+  // model_id: use to check the download progress
+  // model: use as a parameter for pull API
+  std::string model_id = input;
+  std::string model = input;
+
+  // Start server if server is not started yet
+  if (!commands::IsServerAlive(host, port)) {
+    CLI_LOG("Starting server ...");
+    commands::ServerStartCmd ssc;
+    if (!ssc.Exec(host, port)) {
+      return std::nullopt;
+    }
+  }
+
+  // Get model info from Server
+  httplib::Client cli(host + ":" + std::to_string(port));
+  cli.set_read_timeout(std::chrono::seconds(60));
+  Json::Value j_data;
+  j_data["model"] = input;
+  auto d_str = j_data.toStyledString();
+  auto res = cli.Post("/models/pull/info", httplib::Headers(), d_str.data(),
+                      d_str.size(), "application/json");
+
+  if (res) {
+    if (res->status == httplib::StatusCode::OK_200) {
+      // CLI_LOG(res->body);
+      auto root = json_helper::ParseJsonString(res->body);
+      std::string id = root["id"].asString();
+      bool is_cortexso = root["modelSource"].asString() == "cortexso";
+      std::string default_branch = root["defaultBranch"].asString();
+      std::vector<std::string> downloaded;
+      for (auto const& v : root["downloadedModels"]) {
+        downloaded.push_back(v.asString());
+      }
+      std::vector<std::string> avails;
+      for (auto const& v : root["availableModels"]) {
+        avails.push_back(v.asString());
+      }
+      std::string download_url = root["downloadUrl"].asString();
+
+      if (downloaded.empty() && avails.empty()) {
+        model_id = id;
+        model = download_url;
+      } else {
+        if (is_cortexso) {
+          auto selection = cli_selection_utils::PrintModelSelection(
+              downloaded, avails,
+              default_branch.empty()
+                  ? std::nullopt
+                  : std::optional<std::string>(default_branch));
+
+          if (!selection.has_value()) {
+            CLI_LOG("Invalid selection");
+            return std::nullopt;
+          }
+          model_id = selection.value();
+          model = model_id;
+        } else {
+          auto selection = cli_selection_utils::PrintSelection(avails);
+          CLI_LOG("Selected: " << selection.value());
+          model_id = id + ":" + selection.value();
+          model = download_url + selection.value();
+        }
+      }
+    } else {
+      auto root = json_helper::ParseJsonString(res->body);
+      CLI_LOG(root["message"].asString());
+      return std::nullopt;
+    }
+  } else {
+    auto err = res.error();
+    CTL_ERR("HTTP error: " << httplib::to_string(err));
+    return std::nullopt;
+  }
+
+  // Send request download model to server
+  Json::Value json_data;
+  json_data["model"] = model;
+  auto data_str = json_data.toStyledString();
+  cli.set_read_timeout(std::chrono::seconds(60));
+  res = cli.Post("/v1/models/pull", httplib::Headers(), data_str.data(),
+                 data_str.size(), "application/json");
+
+  if (res) {
+    if (res->status != httplib::StatusCode::OK_200) {
+      auto root = json_helper::ParseJsonString(res->body);
+      CLI_LOG(root["message"].asString());
+      return std::nullopt;
+    }
+  } else {
+    auto err = res.error();
+    CTL_ERR("HTTP error: " << httplib::to_string(err));
+    return std::nullopt;
+  }
+
+  CLI_LOG("Start downloading ...")
+  DownloadProgress dp;
+  bool force_stop = false;
+
+  shutdown_handler = [this, &dp, &host, &port, &model_id, &force_stop](int) {
+    force_stop = true;
+    AbortModelPull(host, port, model_id);
+    dp.ForceStop();
+  };
+
+  utils::ScopeExit se([]() { shutdown_handler = {}; });
+#if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
+  struct sigaction sigint_action;
+  sigint_action.sa_handler = signal_handler;
+  sigemptyset(&sigint_action.sa_mask);
+  sigint_action.sa_flags = 0;
+  sigaction(SIGINT, &sigint_action, NULL);
+  sigaction(SIGTERM, &sigint_action, NULL);
+#elif defined(_WIN32)
+  auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
+    return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
+  };
+  SetConsoleCtrlHandler(
+      reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
+#endif
+  dp.Connect(host, port);
+  if (!dp.Handle(model_id))
+    return std::nullopt;
+  if (force_stop)
+    return std::nullopt;
+  CLI_LOG("Model " << model_id << " downloaded successfully!")
+  return model_id;
+}
+
+bool ModelPullCmd::AbortModelPull(const std::string& host, int port,
+                                  const std::string& task_id) {
+  Json::Value json_data;
+  json_data["taskId"] = task_id;
+  auto data_str = json_data.toStyledString();
+  httplib::Client cli(host + ":" + std::to_string(port));
+  cli.set_read_timeout(std::chrono::seconds(60));
+  auto res = cli.Delete("/v1/models/pull", httplib::Headers(), data_str.data(),
+                        data_str.size(), "application/json");
+  if (res) {
+    if (res->status == httplib::StatusCode::OK_200) {
+      CTL_INF("Abort model pull successfully: " << task_id);
+      return true;
+    } else {
+      auto root = json_helper::ParseJsonString(res->body);
+      CLI_LOG(root["message"].asString());
+      return false;
+    }
+  } else {
+    auto err = res.error();
+    CTL_ERR("HTTP error: " << httplib::to_string(err));
+    return false;
   }
 }
 };  // namespace commands

engine/cli/commands/model_pull_cmd.h

@@ -8,7 +8,14 @@ class ModelPullCmd {
  public:
   explicit ModelPullCmd(std::shared_ptr<DownloadService> download_service)
       : model_service_{ModelService(download_service)} {};
-  void Exec(const std::string& input);
+  explicit ModelPullCmd(const ModelService& model_service)
+      : model_service_{model_service} {};
+  std::optional<std::string> Exec(const std::string& host, int port,
+                                  const std::string& input);
+
+ private:
+  bool AbortModelPull(const std::string& host, int port,
+                      const std::string& task_id);
 
  private:
   ModelService model_service_;

engine/cli/commands/model_start_cmd.cc

@@ -14,7 +14,7 @@ bool ModelStartCmd::Exec(const std::string& host, int port,
                          const std::string& model_handle,
                          bool print_success_log) {
   std::optional<std::string> model_id =
-      SelectLocalModel(model_service_, model_handle);
+      SelectLocalModel(host, port, model_service_, model_handle);
 
   if (!model_id.has_value()) {
     return false;

engine/cli/commands/run_cmd.cc

@@ -3,6 +3,8 @@
 #include "config/yaml_config.h"
 #include "cortex_upd_cmd.h"
 #include "database/models.h"
+#include "engine_install_cmd.h"
+#include "model_pull_cmd.h"
 #include "model_start_cmd.h"
 #include "model_status_cmd.h"
 #include "server_start_cmd.h"
@@ -11,7 +13,8 @@
 
 namespace commands {
 
-std::optional<std::string> SelectLocalModel(ModelService& model_service,
+std::optional<std::string> SelectLocalModel(std::string host, int port,
+                                            ModelService& model_service,
                                             const std::string& model_handle) {
   std::optional<std::string> model_id = model_handle;
   cortex::db::Models modellist_handler;
@@ -42,8 +45,8 @@ std::optional<std::string> SelectLocalModel(ModelService& model_service,
   } else {
     auto related_models_ids = modellist_handler.FindRelatedModel(model_handle);
     if (related_models_ids.has_error() || related_models_ids.value().empty()) {
-      auto result = model_service.DownloadModel(model_handle);
-      if (result.has_error()) {
+      auto result = ModelPullCmd(model_service).Exec(host, port, model_handle);
+      if (!result) {
         CLI_LOG("Model " << model_handle << " not found!");
         return std::nullopt;
       }
@@ -79,7 +82,7 @@ std::string Repo2Engine(const std::string& r) {
 
 void RunCmd::Exec(bool run_detach) {
   std::optional<std::string> model_id =
-      SelectLocalModel(model_service_, model_handle_);
+      SelectLocalModel(host_, port_, model_service_, model_handle_);
   if (!model_id.has_value()) {
     return;
   }
@@ -114,9 +117,9 @@ void RunCmd::Exec(bool run_detach) {
         throw std::runtime_error("Engine " + mc.engine + " is incompatible");
       }
       if (required_engine.value().status == EngineService::kNotInstalled) {
-        auto install_engine_result = engine_service_.InstallEngine(mc.engine);
-        if (install_engine_result.has_error()) {
-          throw std::runtime_error(install_engine_result.error());
+        if (!EngineInstallCmd(download_service_, host_, port_)
+                 .Exec(mc.engine)) {
+          return;
         }
       }
     }