fix: cortex stop child process (menloresearch#1841)

nguyenhoangthuan99 · web-flow · commit c4b370f178d0 · 2025-01-06T17:17:31.000+07:00
* fix: download recursive

* fix: handle model not loaded

* fix: set permission for all venv folder

* format code

* fix: cortex stop will stop all running models

* Fix: comment
diff --git a/engine/controllers/process_manager.cc b/engine/controllers/process_manager.cc
@@ -1,13 +1,16 @@
 #include "process_manager.h"
-#include "utils/cortex_utils.h"
-
 #include <trantor/utils/Logger.h>
 #include <cstdlib>
+#include "json/json.h"
+#include "utils/cortex_utils.h"
 
 void ProcessManager::destroy(
     const HttpRequestPtr& req,
     std::function<void(const HttpResponsePtr&)>&& callback) {
-
+  auto loaded_engines = engine_service_->GetSupportedEngineNames();
+  for (const auto& engine : loaded_engines.value()) {
+    engine_service_->UnloadEngine(engine);
+  }
   app().quit();
   Json::Value ret;
   ret["message"] = "Program is exitting, goodbye!";
diff --git a/engine/controllers/process_manager.h b/engine/controllers/process_manager.h
@@ -2,6 +2,7 @@
 
 #include <drogon/HttpController.h>
 #include <drogon/HttpTypes.h>
+#include "services/engine_service.h"
 
 using namespace drogon;
 
@@ -13,4 +14,10 @@ class ProcessManager : public drogon::HttpController<ProcessManager, false> {
 
   void destroy(const HttpRequestPtr& req,
                std::function<void(const HttpResponsePtr&)>&& callback);
+
+  ProcessManager(std::shared_ptr<EngineService> engine_service)
+      : engine_service_(engine_service) {}
+
+ private:
+  std::shared_ptr<EngineService> engine_service_;
 };
diff --git a/engine/extensions/python-engine/python_engine.cc b/engine/extensions/python-engine/python_engine.cc
@@ -18,7 +18,6 @@ static size_t WriteCallback(char* ptr, size_t size, size_t nmemb,
 
 PythonEngine::PythonEngine() : q_(4 /*n_parallel*/, "python_engine") {}
 
-
 PythonEngine::~PythonEngine() {
   curl_global_cleanup();
 }
@@ -507,7 +506,6 @@ CurlResponse PythonEngine::MakeStreamPostRequest(
   return response;
 }
 
-
 void PythonEngine::HandleInference(
     std::shared_ptr<Json::Value> json_body,
     std::function<void(Json::Value&&, Json::Value&&)>&& callback) {
@@ -943,7 +941,11 @@ void PythonEngine::Load(EngineLoadOption opts) {
   // Develop register model here on loading engine
 };
 
-void PythonEngine::Unload(EngineUnloadOption opts) {};
+void PythonEngine::Unload(EngineUnloadOption opts) {
+  for (const auto& pair : models_) {
+    TerminateModelProcess(pair.first);
+  }
+};
 
 // extern "C" {
 // EngineI* get_engine() {
diff --git a/engine/main.cc b/engine/main.cc
@@ -182,7 +182,7 @@ void RunServer(std::optional<std::string> host, std::optional<int> port,
   auto model_ctl = std::make_shared<Models>(db_service, model_service,
                                             engine_service, model_src_svc);
   auto event_ctl = std::make_shared<Events>(event_queue_ptr);
-  auto pm_ctl = std::make_shared<ProcessManager>();
+  auto pm_ctl = std::make_shared<ProcessManager>(engine_service);
   auto hw_ctl = std::make_shared<Hardware>(engine_service, hw_service);
   auto server_ctl =
       std::make_shared<inferences::server>(inference_svc, engine_service);
diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc
@@ -567,7 +567,6 @@ ModelService::DownloadModelFromCortexsoAsync(
                          std::filesystem::path("python.exe"))
                             .string()
                      << std::endl;
-
 #else
           pyvenv_cfg << "home = "
                      << (venv_path / std::filesystem::path("bin/")).string()
@@ -577,14 +576,10 @@ ModelService::DownloadModelFromCortexsoAsync(
               << (venv_path / std::filesystem::path("bin/python")).string()
               << std::endl;
 #endif
-
           // Close the file
           pyvenv_cfg.close();
           // Add executable permission to python
-
-          set_permission_utils::SetExecutePermissionsRecursive(
-              venv_path );
-
+          set_permission_utils::SetExecutePermissionsRecursive(venv_path);
         } else {
           CTL_ERR("Failed to extract venv.zip");
         };
