Merge branch 'dev' into chore/update-star-history

Gabrielle Ong · web-flow · commit e7ad487c97d2 · 2024-11-27T10:23:56.000+08:00
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -56,7 +56,7 @@ COPY ./docs/static/openapi/cortex.json /app/docs/static/openapi/cortex.json
 RUN cd engine && make configure-vcpkg && make build CMAKE_EXTRA_FLAGS="-DCORTEX_CPP_VERSION=${CORTEX_CPP_VERSION} -DCMAKE_BUILD_TEST=OFF -DCMAKE_TOOLCHAIN_FILE=vcpkg/scripts/buildsystems/vcpkg.cmake ${CMAKE_EXTRA_FLAGS}"
 
 # Stage 3: Runtime
-
+FROM common as runtime
 WORKDIR /app
 COPY --from=build /app/engine/build/cortex /usr/local/bin/cortex
 COPY --from=build /app/engine/build/cortex-server /usr/local/bin/cortex-server
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,16 @@ 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()) {
+    for (auto& item : it->second->items) {
+      item.downloadedBytes = item.bytes;
+    }
+    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"
 

Original file line number	Diff line number	Diff line change
`@@ -2,9 +2,9 @@`
`2`	`2`
`3`	`3`	`#include <json/json.h>`
`4`	`4`	`#include <filesystem>`
	`5`	`+#include <optional>`
`5`	`6`	`#include <sstream>`
`6`	`7`	`#include <string>`
`7`		`-#include <optional>`
`8`	`8`
`9`	`9`	`enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };`
`10`	`10`
`@@ -161,4 +161,4 @@ inline DownloadTask GetDownloadTaskFromJson(const Json::Value item_json) {`
`161`	`161`	`}`
`162`	`162`	`return task;`
`163`	`163`	`}`
`164`		`-} // namespace common`
	`164`	`+} // namespace common`