fix: (#4091) model size not displaying

janhq · Nov 26, 2024 · 477f6d5 · 477f6d5
1 parent ad5ecc1
commit 477f6d5
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 15 deletions.
diff --git a/engine/cli/commands/model_import_cmd.cc b/engine/cli/commands/model_import_cmd.cc
@@ -1,13 +1,7 @@
 #include "model_import_cmd.h"
-#include <filesystem>
-#include <vector>
-#include "config/gguf_parser.h"
-#include "config/yaml_config.h"
-#include "database/models.h"
+#include <json/value.h>
 #include "httplib.h"
-#include "json/json.h"
 #include "server_start_cmd.h"
-#include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
 
 namespace commands {

diff --git a/engine/common/download_task.h b/engine/common/download_task.h
@@ -2,9 +2,9 @@
 
 #include <json/json.h>
 #include <filesystem>
+#include <optional>
 #include <sstream>
 #include <string>
-#include <optional>
 
 enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };
 
@@ -161,4 +161,4 @@ inline DownloadTask GetDownloadTaskFromJson(const Json::Value item_json) {
   }
   return task;
 }
-}  // namespace common
+}  // namespace common
diff --git a/engine/services/download_service.cc b/engine/services/download_service.cc
@@ -544,10 +544,18 @@ void DownloadService::EmitTaskCompleted(const std::string& task_id) {
 }
 
 void DownloadService::ExecuteCallback(const DownloadTask& task) {
-  std::lock_guard<std::mutex> lock(callbacks_mutex_);
-  auto it = callbacks_.find(task.id);
-  if (it != callbacks_.end()) {
-    it->second(task);
-    callbacks_.erase(it);
+  std::lock_guard<std::mutex> active_task_lock(active_tasks_mutex_);
+  if (auto it = active_tasks_.find(task.id); it != active_tasks_.end()) {
+    CTL_ERR("Active task found for task: " << task.id);
+    for (auto& item : it->second->items) {
+      item.downloadedBytes = item.bytes;
+    }
+    CTL_ERR("Task: " << it->second->ToString());
+    std::lock_guard<std::mutex> lock(callbacks_mutex_);
+    auto callback = callbacks_.find(task.id);
+    if (callback != callbacks_.end()) {
+      callback->second(*it->second);
+      callbacks_.erase(callback);
+    }
   }
 }
diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc
@@ -463,6 +463,12 @@ ModelService::DownloadModelFromCortexsoAsync(
     yaml_handler.ModelConfigFromFile(model_yml_item->localPath.string());
     auto mc = yaml_handler.GetModelConfig();
     mc.model = unique_model_id;
+
+    uint64_t model_size = 0;
+    for (const auto& item : finishedTask.items) {
+      model_size = model_size + item.bytes.value_or(0);
+    }
+    mc.size = model_size;
     yaml_handler.UpdateModelConfig(mc);
     yaml_handler.WriteYamlFile(model_yml_item->localPath.string());
 

diff --git a/engine/services/model_service.h b/engine/services/model_service.h
@@ -5,7 +5,6 @@
 #include <string>
 #include "common/engine_servicei.h"
 #include "config/model_config.h"
-#include "database/models.h"
 #include "services/download_service.h"
 #include "services/inference_service.h"