Merge pull request #122 from janhq/unload-model-stop-background

Unload model stop background
janhq · Nov 13, 2023 · f49bd78 · f49bd78
2 parents a438284 + d39da8e
commit f49bd78
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/controllers/llamaCPP.cc b/controllers/llamaCPP.cc
@@ -196,6 +196,25 @@ void llamaCPP::embedding(
   return;
 }
 
+void llamaCPP::unloadModel(
+    const HttpRequestPtr &req,
+    std::function<void(const HttpResponsePtr &)> &&callback) {
+  Json::Value jsonResp;
+  jsonResp["message"] = "No model loaded";
+  if (model_loaded) {
+    stopBackgroundTask();
+
+    llama_free(llama.ctx);
+    llama_free_model(llama.model);
+    llama.ctx = nullptr;
+    llama.model = nullptr;
+    jsonResp["message"] = "Model unloaded successfully";
+  }
+  auto resp = nitro_utils::nitroHttpJsonResponse(jsonResp);
+  callback(resp);
+  return;
+}
+
 void llamaCPP::loadModel(
     const HttpRequestPtr &req,
     std::function<void(const HttpResponsePtr &)> &&callback) {
@@ -274,7 +293,19 @@ void llamaCPP::loadModel(
 
 void llamaCPP::backgroundTask() {
   while (model_loaded) {
-    model_loaded = llama.update_slots();
+    // model_loaded =
+    llama.update_slots();
   }
+  LOG_INFO << "Background task stopped!";
   return;
 }
+
+void llamaCPP::stopBackgroundTask() {
+  if (model_loaded) {
+    model_loaded = false;
+    LOG_INFO << "changed to false";
+    if (backgroundThread.joinable()) {
+      backgroundThread.join();
+    }
+  }
+}
diff --git a/controllers/llamaCPP.h b/controllers/llamaCPP.h
@@ -2124,6 +2124,8 @@ class llamaCPP : public drogon::HttpController<llamaCPP> {
   METHOD_ADD(llamaCPP::chatCompletion, "chat_completion", Post);
   METHOD_ADD(llamaCPP::embedding, "embedding", Post);
   METHOD_ADD(llamaCPP::loadModel, "loadmodel", Post);
+  METHOD_ADD(llamaCPP::unloadModel, "unloadmodel", Get);
+
   // PATH_ADD("/llama/chat_completion", Post);
   METHOD_LIST_END
   void chatCompletion(const HttpRequestPtr &req,
@@ -2132,13 +2134,17 @@ class llamaCPP : public drogon::HttpController<llamaCPP> {
                  std::function<void(const HttpResponsePtr &)> &&callback);
   void loadModel(const HttpRequestPtr &req,
                  std::function<void(const HttpResponsePtr &)> &&callback);
+  void unloadModel(const HttpRequestPtr &req,
+                   std::function<void(const HttpResponsePtr &)> &&callback);
   void warmupModel();
 
   void backgroundTask();
 
+  void stopBackgroundTask();
+
 private:
   llama_server_context llama;
-  bool model_loaded = false;
+  std::atomic<bool> model_loaded = false;
   size_t sent_count = 0;
   size_t sent_token_probs_index = 0;
   std::thread backgroundThread;