1import time
2import json
3import os
4from typing import Dict, Any, Tuple
5import numpy as np
6from statistics import median
7from .algorithms import algorithms
8from .datasets import datasets
9from ._memobin import construct_memobin_url, upload_to_memobin, download_from_memobin
12system_version = 'v3'
14def run_benchmarks(cache_dir: str = '.benchmark_cache', verbose: bool = True) -> Dict[str, Any]:
15 """Run all benchmarks, with caching based on algorithm and dataset versions.
17 Results are stored in separate directories for each dataset/algorithm combination:
18 cache_dir/
19 dataset_name/
20 algorithm_name/
21 metadata.json # Contains algorithm version, dataset version, and results
22 compressed.dat # The actual compressed data
24 Args:
25 cache_dir: Directory to store cached results
27 Returns:
28 Dictionary containing benchmark results and metadata
29 """
30 print("\n=== Starting Benchmark Run ===")
31 print(f"Cache directory: {cache_dir}")
33 os.makedirs(cache_dir, exist_ok=True)
35 results = []
36 print("\nRunning benchmarks for all dataset-algorithm combinations...")
38 # Run benchmarks for each dataset and algorithm combination
39 for dataset in datasets:
40 print(f"\n--- Dataset: {dataset['name']} ---")
41 # Create dataset once for all algorithms
42 data = dataset['create']()
43 dtype = str(data.dtype)
44 original_size = len(data.tobytes())
45 print(f"Created dataset: shape={data.shape}, dtype={dtype}")
46 print(f"Original size: {original_size:,} bytes")
48 for algorithm in algorithms:
49 alg_name = algorithm['name']
50 print(f"\nTesting algorithm: {alg_name}")
52 # Check if we can use cached result
53 test_dir = os.path.join(cache_dir, dataset['name'], alg_name)
54 metadata_file = os.path.join(test_dir, 'metadata.json')
55 compressed_file = os.path.join(test_dir, 'compressed.dat')
57 # First try local cache
58 cached_data = None
59 if os.path.exists(metadata_file) and os.path.exists(compressed_file):
60 with open(metadata_file, 'r') as f:
61 cached_data = json.load(f)
63 # If not in local cache, try memobin
64 if cached_data is None:
65 memobin_url = construct_memobin_url(
66 alg_name, dataset['name'],
67 algorithm['version'], dataset['version'],
68 system_version
69 )
70 if verbose:
71 print(" Looking for cached result in memobin...")
72 cached_data = download_from_memobin(memobin_url)
73 if cached_data is not None:
74 if verbose:
75 print(" Found result in memobin, saving locally...")
76 # Save to local cache
77 os.makedirs(test_dir, exist_ok=True)
78 with open(metadata_file, 'w') as f:
79 json.dump(cached_data, f, indent=2)
81 if cached_data is not None and (
82 cached_data['result']['algorithm_version'] == algorithm['version'] and
83 cached_data['result']['dataset_version'] == dataset['version'] and
84 cached_data['result'].get('system_version', '') == system_version
85 ):
86 print(" Using cached result:")
87 results.append(cached_data['result'])
88 continue
90 print(" Running new benchmark...")
92 def run_timed_trials(operation, *args) -> Tuple[float, float]:
93 """Run multiple trials of an operation until total time exceeds 1 second.
94 Returns (median_time, mb_per_sec)"""
95 times = []
96 total_time = 0
97 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
99 while total_time < 1.0:
100 start_time = time.perf_counter()
101 _ = operation(*args) # Execute operation but discard result
102 trial_time = time.perf_counter() - start_time
103 times.append(trial_time)
104 total_time += trial_time
106 median_time = median(times)
107 mb_per_sec = array_size_mb / median_time
108 return median_time, mb_per_sec
110 # Measure encoding with multiple trials
111 encode_time, encode_mb_per_sec = run_timed_trials(algorithm['encode'], data)
112 encoded = algorithm['encode'](data) # One final encode to get the result
113 compressed_size = len(encoded)
114 compression_ratio = original_size / compressed_size
115 print(" Compression complete:")
116 print(f" Compressed size: {compressed_size:,} bytes")
117 print(f" Compression ratio: {compression_ratio:.2f}x")
118 print(f" Encode time: {encode_time*1000:.2f}ms")
119 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
121 print(" Verifying decompression...")
122 # Measure decoding with multiple trials
123 decode_time, decode_mb_per_sec = run_timed_trials(algorithm['decode'], encoded, dtype)
124 decoded = algorithm['decode'](encoded, dtype) # One final decode to verify
125 print(f" Decode time: {decode_time*1000:.2f}ms")
126 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
128 # Verify correctness
129 if not np.array_equal(data, decoded):
130 raise ValueError(
131 f"Decompression verification failed for {alg_name} on {dataset['name']}"
132 )
133 print(" Verification successful!")
135 # Store result
136 result = {
137 'dataset': dataset['name'],
138 'algorithm': alg_name,
139 'algorithm_version': algorithm['version'],
140 'dataset_version': dataset['version'],
141 'system_version': system_version,
142 'compression_ratio': compression_ratio,
143 'encode_time': encode_time,
144 'decode_time': decode_time,
145 'encode_mb_per_sec': encode_mb_per_sec,
146 'decode_mb_per_sec': decode_mb_per_sec,
147 'original_size': original_size,
148 'compressed_size': compressed_size,
149 'array_shape': data.shape,
150 'array_dtype': dtype,
151 'timestamp': time.time()
152 }
153 results.append(result)
155 # Save result and compressed data
156 os.makedirs(test_dir, exist_ok=True)
157 cache_data = {
158 'result': result
159 }
160 with open(metadata_file, 'w') as f:
161 json.dump(cache_data, f, indent=2)
162 with open(compressed_file, 'wb') as f:
163 f.write(encoded)
164 print(f" Results saved to: {test_dir}")
166 # Upload to memobin if API key is set
167 memobin_api_key = os.environ.get('MEMOBIN_API_KEY')
168 if memobin_api_key:
169 if verbose:
170 print(" Uploading results to memobin...")
171 try:
172 memobin_url = construct_memobin_url(
173 alg_name, dataset['name'],
174 algorithm['version'], dataset['version'],
175 system_version
176 )
177 upload_to_memobin(
178 cache_data,
179 memobin_url,
180 os.environ.get('MEMOBIN_USER_ID', 'default'),
181 memobin_api_key
182 )
183 if verbose:
184 print(" Successfully uploaded to memobin")
185 except Exception as e:
186 print(f" Warning: Failed to upload to memobin: {str(e)}")
188 print("\n=== Benchmark Run Complete ===\n")
189 return {'results': results}