chore: update model's size on import (#1630)

* chore: update model's size on import * refactor: consolidate get file size using filesystem API * chore: remove auto declare * chore: correct file size retrieval error handling * fix: missing import gguf_parser * chore: change back to auto type declare
janhq · Nov 4, 2024 · 206650f · 206650f
1 parent 55bbe0d
commit 206650f
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 48 deletions.
diff --git a/engine/config/gguf_parser.cc b/engine/config/gguf_parser.cc
@@ -9,6 +9,7 @@
 #include <stdexcept>
 #include <string>
 #include <vector>
+#include <filesystem>
 
 #ifdef _WIN32
 #include <io.h>
@@ -71,24 +72,9 @@ void GGUFHandler::OpenFile(const std::string& file_path) {
   CloseHandle(file_handle_);
 
 #else
-  FILE* fd = fopen(file_path.c_str(), "rb");
-  if (!fd) {
-    perror("Error opening file");
-    throw std::runtime_error("Failed to open file");
-  }
-
-  // Get file size
-  // file_size_ = lseek(fd, 0, SEEK_END);
-  fseek(fd, 0, SEEK_END);  // move file pointer to end of file
-  file_size_ = ftell(fd);  // get the file size, in bytes
-  fclose(fd);
-  if (file_size_ == -1) {
-    perror("Error getting file size");
-    // close(fd);
-    throw std::runtime_error("Failed to get file size");
-  }
+  file_size_ = std::filesystem::file_size(file_path);
+
   int file_descriptor = open(file_path.c_str(), O_RDONLY);
-  ;
   // Memory-map the file
   data_ = static_cast<uint8_t*>(
       mmap(nullptr, file_size_, PROT_READ, MAP_PRIVATE, file_descriptor, 0));

diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc
@@ -351,8 +351,12 @@ void Models::ImportModel(
           modelPath, file_path,
           std::filesystem::copy_options::update_existing);
       model_config.files.push_back(file_path.string());
+      auto size = std::filesystem::file_size(file_path);
+      model_config.size = size;
     } else {
       model_config.files.push_back(modelPath);
+      auto size = std::filesystem::file_size(modelPath);
+      model_config.size = size;
     }
     model_config.model = modelHandle;
     model_config.name = modelName.empty() ? model_config.name : modelName;

diff --git a/engine/services/download_service.cc b/engine/services/download_service.cc
@@ -7,7 +7,6 @@
 #include <optional>
 #include <ostream>
 #include <utility>
-#include "download_service.h"
 #include "utils/format_utils.h"
 #include "utils/huggingface_utils.h"
 #include "utils/logging_utils.h"
@@ -148,11 +147,10 @@ cpp::result<bool, std::string> DownloadService::Download(
   std::string mode = "wb";
   if (std::filesystem::exists(download_item.localPath) &&
       download_item.bytes.has_value()) {
-    curl_off_t existing_file_size = GetLocalFileSize(download_item.localPath);
-    if (existing_file_size == -1) {
-      CLI_LOG("Cannot get file size: " << download_item.localPath.string()
-                                       << " . Start download over!");
-    } else {
+    try {
+      curl_off_t existing_file_size =
+          std::filesystem::file_size(download_item.localPath);
+
       CTL_INF("Existing file size: " << download_item.downloadUrl << " - "
                                      << download_item.localPath.string()
                                      << " - " << existing_file_size);
@@ -186,6 +184,9 @@ cpp::result<bool, std::string> DownloadService::Download(
           return false;
         }
       }
+    } catch (const std::filesystem::filesystem_error& e) {
+      CLI_LOG("Cannot get file size: "
+              << e.what() << download_item.localPath.string() << "\n");
     }
   }
 
@@ -205,12 +206,12 @@ cpp::result<bool, std::string> DownloadService::Download(
   curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
 
   if (mode == "ab") {
-    auto local_file_size = GetLocalFileSize(download_item.localPath);
-    if (local_file_size != -1) {
-      curl_easy_setopt(curl, CURLOPT_RESUME_FROM_LARGE,
-                       GetLocalFileSize(download_item.localPath));
-    } else {
-      CTL_ERR("Cannot get file size: " << download_item.localPath.string());
+    try {
+      curl_off_t local_file_size =
+          std::filesystem::file_size(download_item.localPath);
+      curl_easy_setopt(curl, CURLOPT_RESUME_FROM_LARGE, local_file_size);
+    } catch (const std::filesystem::filesystem_error& e) {
+      CTL_ERR("Cannot get file size: " << e.what() << '\n');
     }
   }
 
@@ -225,23 +226,6 @@ cpp::result<bool, std::string> DownloadService::Download(
   curl_easy_cleanup(curl);
   return true;
 }
-
-curl_off_t DownloadService::GetLocalFileSize(
-    const std::filesystem::path& path) const {
-  auto file = fopen(path.string().c_str(), "r");
-  if (!file) {
-    return -1;
-  }
-
-  if (fseek64(file, 0, SEEK_END) != 0) {
-    return -1;
-  }
-
-  auto file_size = ftell64(file);
-  fclose(file);
-  return file_size;
-}
-
 void DownloadService::WorkerThread() {
   while (!stop_flag_) {
     DownloadTask task;

diff --git a/engine/services/download_service.h b/engine/services/download_service.h
@@ -90,8 +90,6 @@ class DownloadService {
       const std::string& download_id,
       const DownloadItem& download_item) noexcept;
 
-  curl_off_t GetLocalFileSize(const std::filesystem::path& path) const;
-
   std::shared_ptr<EventQueue> event_queue_;
 
   CURLM* multi_handle_;