fix: update download api

Author: namchuai

docs/static/openapi/jan.json

@@ -1651,6 +1651,15 @@
                 "value": "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/blob/main/mistral-7b-instruct-v0.1.Q2_K.gguf"
               }
             ]
+          },
+          "id": {
+            "type": "string",
+            "description": "The id which will be used to register the model.",
+            "examples": [
+              {
+                "value": "my-custom-model-id"
+              }
+            ]
           }
         }
       },

engine/cli/commands/model_list_cmd.cc

@@ -17,8 +17,9 @@ using namespace tabulate;
 using Row_t =
     std::vector<variant<std::string, const char*, string_view, Table>>;
 
-void ModelListCmd::Exec(const std::string& host, int port, std::string filter,
-                        bool display_engine, bool display_version) {
+void ModelListCmd::Exec(const std::string& host, int port,
+                        const std::string& filter, bool display_engine,
+                        bool display_version) {
   // Start server if server is not started yet
   if (!commands::IsServerAlive(host, port)) {
     CLI_LOG("Starting server ...");

engine/cli/commands/model_list_cmd.h

@@ -6,7 +6,7 @@ namespace commands {
 
 class ModelListCmd {
  public:
-  void Exec(const std::string& host, int port, std::string filter,
+  void Exec(const std::string& host, int port, const std::string& filter,
             bool display_engine = false, bool display_version = false);
 };
 }  // namespace commands

engine/common/download_task.h

@@ -7,9 +7,11 @@
 #include <string>
 
 enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };
+
 using namespace nlohmann;
 
 struct DownloadItem {
+
   std::string id;
 
   std::string downloadUrl;
@@ -54,8 +56,12 @@ inline std::string DownloadTypeToString(DownloadType type) {
 }
 
 struct DownloadTask {
+  enum class Status { Pending, InProgress, Completed, Cancelled, Error };
+
   std::string id;
 
+  Status status;
+
   DownloadType type;
 
   std::vector<DownloadItem> items;

engine/common/download_task_queue.h

@@ -0,0 +1,77 @@
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <optional>
+#include <shared_mutex>
+#include <string>
+#include <unordered_map>
+#include "common/download_task.h"
+
+class DownloadTaskQueue {
+ private:
+  std::deque<DownloadTask> taskQueue;
+  std::unordered_map<std::string, typename std::deque<DownloadTask>::iterator>
+      taskMap;
+  mutable std::shared_mutex mutex;
+  std::condition_variable_any cv;
+
+ public:
+  void push(DownloadTask task) {
+    std::unique_lock lock(mutex);
+    taskQueue.push_back(std::move(task));
+    taskMap[taskQueue.back().id] = std::prev(taskQueue.end());
+    cv.notify_one();
+  }
+
+  std::optional<DownloadTask> pop() {
+    std::unique_lock lock(mutex);
+    if (taskQueue.empty()) {
+      return std::nullopt;
+    }
+    DownloadTask task = std::move(taskQueue.front());
+    taskQueue.pop_front();
+    taskMap.erase(task.id);
+    return task;
+  }
+
+  bool cancelTask(const std::string& taskId) {
+    std::unique_lock lock(mutex);
+    auto it = taskMap.find(taskId);
+    if (it != taskMap.end()) {
+      it->second->status = DownloadTask::Status::Cancelled;
+      taskQueue.erase(it->second);
+      taskMap.erase(it);
+      return true;
+    }
+    return false;
+  }
+
+  bool updateTaskStatus(const std::string& taskId,
+                        DownloadTask::Status newStatus) {
+    std::unique_lock lock(mutex);
+    auto it = taskMap.find(taskId);
+    if (it != taskMap.end()) {
+      it->second->status = newStatus;
+      if (newStatus == DownloadTask::Status::Cancelled ||
+          newStatus == DownloadTask::Status::Error) {
+        taskQueue.erase(it->second);
+        taskMap.erase(it);
+      }
+      return true;
+    }
+    return false;
+  }
+
+  std::optional<DownloadTask> getNextPendingTask() {
+    std::shared_lock lock(mutex);
+    auto it = std::find_if(
+        taskQueue.begin(), taskQueue.end(), [](const DownloadTask& task) {
+          return task.status == DownloadTask::Status::Pending;
+        });
+
+    if (it != taskQueue.end()) {
+      return *it;
+    }
+    return std::nullopt;
+  }
+};

engine/controllers/models.cc

@@ -1,5 +1,6 @@
 #include "database/models.h"
 #include <drogon/HttpTypes.h>
+#include <optional>
 #include "config/gguf_parser.h"
 #include "config/yaml_config.h"
 #include "models.h"
@@ -26,15 +27,22 @@ void Models::PullModel(const HttpRequestPtr& req,
     return;
   }
 
+  std::optional<std::string> desired_model_id = std::nullopt;
+  auto id = (*(req->getJsonObject())).get("id", "").asString();
+  if (!id.empty()) {
+    desired_model_id = id;
+  }
+
   auto handle_model_input =
       [&, model_handle]() -> cpp::result<DownloadTask, std::string> {
     CTL_INF("Handle model input, model handle: " + model_handle);
     if (string_utils::StartsWith(model_handle, "https")) {
-      return model_service_->HandleDownloadUrlAsync(model_handle);
+      return model_service_->HandleDownloadUrlAsync(model_handle,
+                                                    desired_model_id);
     } else if (model_handle.find(":") != std::string::npos) {
       auto model_and_branch = string_utils::SplitBy(model_handle, ":");
       return model_service_->DownloadModelFromCortexsoAsync(
-          model_and_branch[0], model_and_branch[1]);
+          model_and_branch[0], model_and_branch[1], desired_model_id);
     }
 
     return cpp::fail("Invalid model handle or not supported!");
@@ -107,7 +115,6 @@ void Models::ListModel(
   auto list_entry = modellist_handler.LoadModelList();
   if (list_entry) {
     for (const auto& model_entry : list_entry.value()) {
-      // auto model_entry = modellist_handler.GetModelInfo(model_handle);
       try {
         yaml_handler.ModelConfigFromFile(
             fmu::ToAbsoluteCortexDataPath(
@@ -116,7 +123,6 @@ void Models::ListModel(
         auto model_config = yaml_handler.GetModelConfig();
         Json::Value obj = model_config.ToJson();
         obj["id"] = model_entry.model;
-        obj["model_alias"] = model_entry.model_alias;
         obj["model"] = model_entry.model;
         data.append(std::move(obj));
         yaml_handler.Reset();
@@ -333,71 +339,6 @@ void Models::ImportModel(
   }
 }
 
-void Models::SetModelAlias(
-    const HttpRequestPtr& req,
-    std::function<void(const HttpResponsePtr&)>&& callback) const {
-  if (!http_util::HasFieldInReq(req, callback, "model") ||
-      !http_util::HasFieldInReq(req, callback, "modelAlias")) {
-    return;
-  }
-  auto model_handle = (*(req->getJsonObject())).get("model", "").asString();
-  auto model_alias = (*(req->getJsonObject())).get("modelAlias", "").asString();
-  LOG_DEBUG << "GetModel, Model handle: " << model_handle
-            << ", Model alias: " << model_alias;
-
-  cortex::db::Models modellist_handler;
-  try {
-    auto result = modellist_handler.UpdateModelAlias(model_handle, model_alias);
-    if (result.has_error()) {
-      std::string message = result.error();
-      LOG_ERROR << message;
-      Json::Value ret;
-      ret["result"] = "Set alias failed!";
-      ret["modelHandle"] = model_handle;
-      ret["message"] = message;
-      auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret);
-      resp->setStatusCode(k400BadRequest);
-      callback(resp);
-    } else {
-      if (result.value()) {
-        std::string message = "Successfully set model alias '" + model_alias +
-                              "' for modeID '" + model_handle + "'.";
-        LOG_INFO << message;
-        Json::Value ret;
-        ret["result"] = "OK";
-        ret["modelHandle"] = model_handle;
-        ret["message"] = message;
-        auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret);
-        resp->setStatusCode(k200OK);
-        callback(resp);
-      } else {
-        std::string message = "Unable to set model alias for modelID '" +
-                              model_handle + "': model alias '" + model_alias +
-                              "' is not unique!";
-        LOG_ERROR << message;
-        Json::Value ret;
-        ret["result"] = "Set alias failed!";
-        ret["modelHandle"] = model_handle;
-        ret["message"] = message;
-        auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret);
-        resp->setStatusCode(k400BadRequest);
-        callback(resp);
-      }
-    }
-  } catch (const std::exception& e) {
-    std::string message = "Error when setting model alias ('" + model_alias +
-                          "') for modelID '" + model_handle + "':" + e.what();
-    LOG_ERROR << message;
-    Json::Value ret;
-    ret["result"] = "Set alias failed!";
-    ret["modelHandle"] = model_handle;
-    ret["message"] = message;
-    auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret);
-    resp->setStatusCode(k400BadRequest);
-    callback(resp);
-  }
-}
-
 void Models::StartModel(
     const HttpRequestPtr& req,
     std::function<void(const HttpResponsePtr&)>&& callback) {

engine/controllers/models.h

@@ -16,7 +16,6 @@ class Models : public drogon::HttpController<Models, false> {
   METHOD_ADD(Models::UpdateModel, "/{1}", Patch);
   METHOD_ADD(Models::ImportModel, "/import", Post);
   METHOD_ADD(Models::DeleteModel, "/{1}", Delete);
-  METHOD_ADD(Models::SetModelAlias, "/alias", Post);
   METHOD_ADD(Models::StartModel, "/start", Post);
   METHOD_ADD(Models::StopModel, "/stop", Post);
   METHOD_ADD(Models::GetModelStatus, "/status/{1}", Get);
@@ -28,7 +27,6 @@ class Models : public drogon::HttpController<Models, false> {
   ADD_METHOD_TO(Models::UpdateModel, "/v1/models/{1}", Patch);
   ADD_METHOD_TO(Models::ImportModel, "/v1/models/import", Post);
   ADD_METHOD_TO(Models::DeleteModel, "/v1/models/{1}", Delete);
-  ADD_METHOD_TO(Models::SetModelAlias, "/v1/models/alias", Post);
   ADD_METHOD_TO(Models::StartModel, "/v1/models/start", Post);
   ADD_METHOD_TO(Models::StopModel, "/v1/models/stop", Post);
   ADD_METHOD_TO(Models::GetModelStatus, "/v1/models/status/{1}", Get);

engine/e2e-test/test_api_model_alias.py

@@ -12,8 +12,7 @@ def setup_and_teardown(self):
         success = start_server()
         if not success:
             raise Exception("Failed to start server")
-
-        # TODO: using pull with branch for easy testing tinyllama:gguf for example
+        run("Install Engine", ["engines", "install", "llama-cpp"], timeout=None)
         run("Delete model", ["models", "delete", "tinyllama:gguf"])
         run("Pull model", ["pull", "tinyllama:gguf"], timeout=None,)