From 6a6872d32d08af381db50535ba58ad3d42c5dac3 Mon Sep 17 00:00:00 2001
From: clore <support@clore.ai>
Date: Sat, 11 May 2024 01:18:09 +0000
Subject: [PATCH] disk benchmarking

---
 lib/get_specs.py | 88 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 86 insertions(+), 2 deletions(-)

diff --git a/lib/get_specs.py b/lib/get_specs.py
index 5add793..f0fddaf 100644
--- a/lib/get_specs.py
+++ b/lib/get_specs.py
@@ -12,6 +12,7 @@ import aiohttp
 import asyncio
 import shutil
 import psutil
+import time
 import sys
 import os
 import re
@@ -42,6 +43,77 @@ def get_kernel():
 def is_hive():
     return "hive" in get_kernel()
 
+def drop_caches():
+    try:
+        with open('/proc/sys/vm/drop_caches', 'w') as f:
+            f.write('3\n')
+    except Exception as e:
+        pass
+
+def write_test(file_path, block_size, num_blocks):
+    data = os.urandom(block_size)
+    total_bytes = block_size * num_blocks
+
+    start_time = time.time()
+    
+    with open(file_path, 'wb') as f:
+        for _ in range(num_blocks):
+            f.write(data)
+            f.flush()
+            os.fsync(f.fileno())
+    
+    elapsed_time = time.time() - start_time
+    write_speed = total_bytes / elapsed_time / (1024 * 1024)
+    
+    return write_speed, elapsed_time
+
+def read_test(file_path, block_size, num_blocks):
+    total_bytes = block_size * num_blocks
+
+    # Drop caches to avoid OS-level caching effects
+    drop_caches()
+    
+    start_time = time.time()
+    
+    with open(file_path, 'rb') as f:
+        for _ in range(num_blocks):
+            data = f.read(block_size)
+            if not data:
+                break
+    
+    elapsed_time = time.time() - start_time
+    read_speed = total_bytes / elapsed_time / (1024 * 1024)
+    
+    return read_speed, elapsed_time
+
+def disk_benchmark():
+    total, used, free = shutil.disk_usage("/")
+
+    free_gb = free/1024/1024/1024
+
+    if free_gb<1:
+        return 0,0
+    
+    block_size = 1024*1024
+    num_blocks = 250 if free_gb < 3 else 1500
+    
+    file_path="/tmp/output"
+    
+    print("Running disk benchmark...")
+    print(f"Block Size: {block_size} bytes, Number of Blocks: {num_blocks}")
+
+    # Run write test
+    write_speed, write_time = write_test(file_path, block_size, num_blocks)
+    print(f"Write Speed: {write_speed:.2f} MB/s, Time: {write_time:.2f} seconds")
+
+    # Run read test
+    read_speed, read_time = read_test(file_path, block_size, num_blocks)
+    print(f"Read Speed: {read_speed:.2f} MB/s, Time: {read_time:.2f} seconds")
+
+    # Cleanup
+    os.remove(file_path)
+    return float(round(write_speed,2)), float(round(read_speed,2))
+
 def get_nvidia_version():
     try:
         output = subprocess.check_output(['nvidia-smi', '-x', '-q'], encoding='utf-8')
@@ -69,6 +141,15 @@ async def measure_internet_speed():
     except Exception as e:
         return '',0, 0
     
+async def disk_speed():
+    try:
+        loop = asyncio.get_event_loop()
+        write_speed, read_speed = await loop.run_in_executor(None, disk_benchmark)
+        return write_speed, read_speed
+    except Exception as e:
+        print("disk benchmark exception",e)
+        return 0, 0
+    
 async def get_country_code():
     async with aiohttp.ClientSession() as session:
         try:
@@ -210,7 +291,7 @@ class Specs:
     def __init__(self):
         self.motherboard_name_file = "/sys/devices/virtual/dmi/id/board_name"
 
-    async def get(self, benchmark_internet=False):
+    async def get(self, benchmark_internet=False, benchmark_disk=False):
         total_threads, total_cores, model_name = self.get_cpu_info()
         gpu_str, gpu_mem, gpus, nvml_err = get_gpu_info()
         docker_daemon_config = docker_interface.get_daemon_config()
@@ -255,6 +336,9 @@ class Specs:
             else: # Disk is overlay
                 data_root_location="separate"
                 disk_str = f"{disk_type} {overlay_total_size}GB"
+
+        disk_speeds = await disk_speed()
+
         response = {
             "mb": await self.motherboard_type(),
             "cpu":model_name,
@@ -263,7 +347,7 @@ class Specs:
             "swap": self.get_swap_size(),
             "data_root_location":data_root_location,
             "disk": disk_str,
-            "disk_speed":0,
+            "disk_speed": disk_speeds[1],
             "gpu":gpu_str,
             "gpuram": gpu_mem,
             "gpus": gpus,