File size: 5,948 Bytes
f961ed2 95b7f68 f961ed2 aa5f00f f961ed2 95b7f68 07ef5dd 95b7f68 f961ed2 95b7f68 07ef5dd 3f732a9 95b7f68 07ef5dd 95b7f68 f961ed2 95b7f68 f961ed2 0144468 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import os
import time
import math
import shutil
import platform
import pandas as pd
import gradio as gr
# DEFAULT_DIR = "/bucket_data"
def run_benchmark(storage_choice, file_size_mb, block_size_kb):
# Ensure directory exists
if storage_choice == "Local Data Space (/local_data)":
target_dir = "/local_data"
elif storage_choice == "Same Acc Bucket Data Space (/sameAcc_bucket_data)":
target_dir = "/sameAcc_bucket_data"
elif storage_choice == "Different Acc Bucket Data Space (/diffAcc_bucket_data)":
target_dir = "/diffAcc_bucket_data"
else:
target_dir = "/local_data"
os.makedirs(target_dir, exist_ok=True)
test_file_path = os.path.join(target_dir, "__io_perf_test__.bin")
# Conversions
file_size_bytes = int(file_size_mb * 1024 * 1024)
block_size_bytes = int(block_size_kb * 1024)
num_blocks = math.ceil(file_size_bytes / block_size_bytes)
# Pre-generate a data block to avoid processing overhead during I/O loop
dummy_block = os.urandom(block_size_bytes)
results = {}
try:
# --- WRITE TEST ---
# Attempt O_DIRECT to bypass OS caching if supported by the OS/FileSystem
flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
if hasattr(os, 'O_DIRECT') and platform.system() == 'Linux':
try:
# O_DIRECT requires block sizes to be aligned (usually 4096 bytes)
fd = os.open(test_file_path, flags | os.O_DIRECT)
except OSError:
fd = os.open(test_file_path, flags) # Fallback
else:
fd = os.open(test_file_path, flags)
start_time = time.perf_counter()
for _ in range(num_blocks):
os.write(fd, dummy_block)
# Force write out to physical disk before stopping timer
try:
os.fdatasync(fd)
except AttributeError:
os.fsync(fd)
end_time = time.perf_counter()
os.close(fd)
write_duration = end_time - start_time
write_speed = file_size_mb / write_duration
results['Write Speed'] = f"{write_speed:.2f} MB/s"
results['Write Time'] = f"{write_duration:.4f} seconds"
# --- READ TEST ---
flags = os.O_RDONLY
if hasattr(os, 'O_DIRECT') and platform.system() == 'Linux':
try:
fd = os.open(test_file_path, flags | os.O_DIRECT)
except OSError:
fd = os.open(test_file_path, flags)
else:
fd = os.open(test_file_path, flags)
start_time = time.perf_counter()
while True:
read_data = os.read(fd, block_size_bytes)
if not read_data:
break
end_time = time.perf_counter()
os.close(fd)
read_duration = end_time - start_time
read_speed = file_size_mb / read_duration
results['Read Speed'] = f"{read_speed:.2f} MB/s"
results['Read Time'] = f"{read_duration:.4f} seconds"
# --- SYSTEM INFO ---
total, used, free = shutil.disk_usage(target_dir)
results['Storage Total'] = f"{total / (1024**3):.2f} GB"
results['Storage Free'] = f"{free / (1024**3):.2f} GB"
results['Status'] = "Success"
except Exception as e:
results['Status'] = f"Error: {str(e)}"
results['Write Speed'], results['Read Speed'] = "N/A", "N/A"
results['Write Time'], results['Read Time'] = "N/A", "N/A"
results['Storage Total'], results['Storage Free'] = "N/A", "N/A"
finally:
# Cleanup test file safely
if os.path.exists(test_file_path):
os.remove(test_file_path)
# Format to Pandas DataFrame for a pretty UI presentation
df_data = {
"Metric": list(results.keys()),
"Value": list(results.values())
}
return pd.DataFrame(df_data)
# --- GRADIO INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# 💾 Hugging Face Space local /data Speed Test
Use this utility to benchmark the real-world read/write performance of ephemeral paths or attached persistent storage (`/data`).
"""
)
with gr.Row():
with gr.Column():
storage_choice = gr.Radio(
label="Select Target Storage Volume",
choices=["Local Data Space (/local_data)", "Same Acc Bucket Data Space (/sameAcc_bucket_data)", "Different Acc Bucket Data Space (/diffAcc_bucket_data)"],
value="Local Data Space (/local_data)",
info="Switches underlying workspace mounts directly."
)
# test_path = gr.Textbox(
# label="Target Directory",
# value=DEFAULT_DIR,
# placeholder="e.g. /data or ./test_dir"
# )
size_slider = gr.Slider(
label="Test File Size (MB)",
minimum=10,
maximum=10000,
value=100,
step=10
)
block_slider = gr.Slider(
label="Block Size (KB) - UI Buffer sizing",
minimum=4,
maximum=4096,
value=1024,
step=4
)
btn_run = gr.Button("🚀 Run IO Benchmark", variant="primary")
with gr.Column():
gr.Markdown("### 📊 Benchmark Results")
output_table = gr.DataFrame(
headers=["Metric", "Value"],
datatype=["str", "str"],
interactive=False
)
btn_run.click(
fn=run_benchmark,
inputs=[storage_choice, size_slider, block_slider],
outputs=output_table
)
if __name__ == "__main__":
demo.launch(
ssr_mode=False
) |