benchmark.py

import dask.array as da
from dask.distributed import Client, LocalCluster
import psutil
import time
import socket
import platform
import multiprocessing

def main(n, n_workers, threads_per_worker, memory_per_worker, local_directory):
    cluster = LocalCluster(n_workers=n_workers, 
                           threads_per_worker=threads_per_worker, 
                           memory_limit=memory_per_worker,
                           local_directory=local_directory)  
    client = Client(cluster)

    start_time = time.time()

    x = da.random.random((n, n), chunks=(n // n_workers, n // n_workers))
    y = da.random.random((n, n), chunks=(n // n_workers, n // n_workers))
    z = x @ y
    result = z.compute()

    end_time = time.time()

    cpu_cores = multiprocessing.cpu_count()
    cpu_threads = psutil.cpu_count(logical=True)
    cpu_model = platform.processor()
    ram_total = psutil.virtual_memory().total / (1024 ** 3)  # Convert to GB
    ram_used = psutil.virtual_memory().used / (1024 ** 3)    # Convert to GB

    print(f"=== Machine Benchmark Report ===")
    print(f"Machine: {socket.gethostname()}")
    print(f"CPU Model: {cpu_model}")
    print(f"CPU Cores: {cpu_cores}")
    print(f"CPU Threads: {cpu_threads}")
    print(f"RAM Total: {ram_total:.2f} GB")
    print(f"RAM Used: {ram_used:.2f} GB")
    print(f"Computation time: {end_time - start_time:.2f} seconds")

    client.close()

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int, default=10000, help="Size of the matrix")
    parser.add_argument("--n_workers", type=int, default=8, help="Number of workers")
    parser.add_argument("--threads_per_worker", type=int, default=1, help="Number of threads per worker")
    parser.add_argument("--memory_per_worker", type=str, default='2GB', help="Memory limit per worker")
    parser.add_argument("--local_directory", type=str, default='.', help="Local directory for scratch data")
    args = parser.parse_args()

    main(args.n, args.n_workers, args.threads_per_worker, args.memory_per_worker, args.local_directory)