From 9eed6b14bf7e54c154412f6c9bd62e71404768cf Mon Sep 17 00:00:00 2001
From: Lizonghang <870644199@qq.com>
Date: Tue, 5 Nov 2024 21:12:02 +0400
Subject: [PATCH] add disk read speed test

---
 common/common.cpp   |  2 ++
 common/profiler.cpp | 72 +++++++++++++++++++++++++++++++++++++++++----
 common/profiler.h   |  3 ++
 3 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 72a968ab..dd09cc05 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -838,9 +838,11 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
     uint32_t n_cpu_cores      = profiler::device_cpu_cores();
     uint64_t total_memory     = profiler::device_physical_memory(false);
     uint64_t available_memory = profiler::device_physical_memory(true);
+    uint64_t disk_read_bw     = profiler::get_disk_read_speed(params.model.c_str(), 500);
     LOG_INF("Number of CPU cores:       %u\n", n_cpu_cores);
     LOG_INF("Total Physical Memory:     %.2f GB\n", total_memory / (double)(1 << 30));
     LOG_INF("Available Physical Memory: %.2f GB\n", available_memory / (double)(1 << 30));
+    LOG_INF("Disk Read Bandwidth:       %.2f GB/s\n", disk_read_bw / (double)(1 << 30));
 
     if (model == NULL) {
         LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.c_str());
diff --git a/common/profiler.cpp b/common/profiler.cpp
index f6311f41..2c23b2c4 100644
--- a/common/profiler.cpp
+++ b/common/profiler.cpp
@@ -12,7 +12,12 @@
     #include <unistd.h>
 #endif
 
-#include <sys/types.h>  
+#include <chrono>
+#include <fstream>
+#include <string>
+#include <sstream>
+#include <sys/types.h>
+#include <vector>
 
 namespace profiler {
 
@@ -62,10 +67,29 @@ uint64_t device_physical_memory(bool available) {
     }
 
 #elif defined(__linux__)
-    struct sysinfo info;
-    if (sysinfo(&info) == 0) {
-        memory = available ? info.freeram : info.totalram;
-        memory *= info.mem_unit;
+    if (available) {
+        // read available memory from /proc/meminfo
+        std::ifstream meminfo("/proc/meminfo");
+        std::string line;
+        if (meminfo.is_open()) {
+            while (std::getline(meminfo, line)) {
+                if (line.find("MemAvailable:") == 0) {
+                    std::istringstream iss(line);
+                    std::string key;
+                    uint64_t kb;
+                    iss >> key >> kb;
+                    memory = kb * 1024;
+                    break;
+                }
+            }
+            meminfo.close();
+        }
+    } else {
+        // get total memory using sysinfo
+        struct sysinfo info;
+        if (sysinfo(&info) == 0) {
+            memory = info.totalram * info.mem_unit;
+        }
     }
 
 #elif defined(__APPLE__) && defined(__MACH__)
@@ -88,4 +112,42 @@ uint64_t device_physical_memory(bool available) {
 
     return memory;
 }
+
+uint64_t get_disk_read_speed(const char * test_file, size_t buffer_size_mb) {
+    uint64_t speed = 0;
+    size_t buffer_size = buffer_size_mb * 1024 * 1024; // buffer size in bytes
+
+    try {
+        // open a file for reading
+        std::ifstream file(test_file, std::ios::binary | std::ios::in);
+        if (!file) {
+            LOG_ERR("Unable to open the file at path: %s\n", test_file);
+            return speed;
+        }
+
+        // prepare buffer for reading
+        std::vector<char> buffer(buffer_size);
+
+        auto start_time = std::chrono::high_resolution_clock::now();
+
+        // read file into buffer
+        file.read(buffer.data(), buffer.size());
+        if (!file) {
+            LOG_ERR("Failed to read enough data from the test file\n");
+            return speed;
+        }
+
+        auto end_time = std::chrono::high_resolution_clock::now();
+        std::chrono::duration<double> elapsed_time = end_time - start_time;
+
+        // Calculate speed in bytes per second
+        if (elapsed_time.count() > 0) {
+            speed = static_cast<uint64_t>(buffer.size() / elapsed_time.count());
+        }
+    } catch (const std::exception &e) {
+        LOG_ERR("Exception while calculating disk read speed: %s\n", e.what());
+    }
+
+    return speed;
+}
 } // namespace profiler
\ No newline at end of file
diff --git a/common/profiler.h b/common/profiler.h
index c21615a8..1f3374f1 100644
--- a/common/profiler.h
+++ b/common/profiler.h
@@ -1,9 +1,12 @@
 #ifndef PROFILER_H
 #define PROFILER_H
 
+#include <string>
+
 namespace profiler {
     uint32_t device_cpu_cores();
     uint64_t device_physical_memory(bool available = true);
+    uint64_t get_disk_read_speed(const char * test_file, size_t buffer_size_mb = 500);
 } // namespace profiler
 
 #endif // PROFILER_H
\ No newline at end of file