From 5766b83f249f7c21b260966630fb7407d8af6d07 Mon Sep 17 00:00:00 2001
From: ozanarmagan <o.armagan2020@gtu.edu.tr>
Date: Tue, 17 Oct 2023 00:26:36 +0300
Subject: [PATCH] Avoid loading models that larger than free memory

---
 include/system_metrics.h      |  4 ++++
 src/text_embedder_manager.cpp | 15 +++++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/include/system_metrics.h b/include/system_metrics.h
index 678ae4f9..94363afa 100644
--- a/include/system_metrics.h
+++ b/include/system_metrics.h
@@ -173,4 +173,8 @@ public:
         // compute
         return compute_cpu_stats(cpu_data_prev, cpu_data_now);
     }
+
+    static uint64_t get_memory_free_bytes() {
+        return get_memory_total_bytes() - get_memory_used_bytes();
+    }
 };
\ No newline at end of file
diff --git a/src/text_embedder_manager.cpp b/src/text_embedder_manager.cpp
index 89400a79..912336dc 100644
--- a/src/text_embedder_manager.cpp
+++ b/src/text_embedder_manager.cpp
@@ -1,4 +1,5 @@
 #include "text_embedder_manager.h"
+#include "system_metrics.h"
 
 
 TextEmbedderManager& TextEmbedderManager::get_instance() {
@@ -107,8 +108,18 @@ Option<bool> TextEmbedderManager::validate_and_init_local_model(const nlohmann::
         return Option<bool>(true);
     }
 
-    const std::shared_ptr<TextEmbedder>& embedder = std::make_shared<TextEmbedder>(
-            get_model_name_without_namespace(model_name));
+    const auto& model_name_without_namespace = get_model_name_without_namespace(model_name);
+    const auto& free_memory = SystemMetrics::get_memory_free_bytes();
+    const auto& model_file_size = std::filesystem::file_size(abs_path);
+    
+    // return error if model file size is greater than free memory
+    if(model_file_size > free_memory) {
+        LOG(ERROR) << "Model file size is greater than free memory: " << model_file_size << " > " << free_memory;
+        return Option<bool>(400, "Model file size is greater than free memory");
+    }
+
+
+    const std::shared_ptr<TextEmbedder>& embedder = std::make_shared<TextEmbedder>(model_name_without_namespace);
 
     auto validate_op = embedder->validate();
     if(!validate_op.ok()) {