Spaces:

kphfrepo
/

hf_readwrite_speedtest

Paused

Update app.py

0144468 verified 4 months ago

5.95 kB

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
    )

	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
	)