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,