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 = 'v4'
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):
60 with open(metadata_file, 'r') as f:
61 cached_data = json.load(f)
62 # if versions do not match, then set to None
63 if (
64 cached_data['result']['algorithm_version'] != algorithm['version'] or
65 cached_data['result']['dataset_version'] != dataset['version'] or
66 cached_data['result'].get('system_version', '') != system_version
67 ):
68 cached_data = None
70 # If not in local cache, try memobin
71 if cached_data is None:
72 memobin_url = construct_memobin_url(
73 alg_name, dataset['name'],
74 algorithm['version'], dataset['version'],
75 system_version
76 )
77 if verbose:
78 print(" Looking for cached result in memobin...")
79 cached_data = download_from_memobin(memobin_url)
80 if cached_data is not None:
81 if verbose:
82 print(" Found result in memobin, saving locally...")
83 # Save to local cache
84 os.makedirs(test_dir, exist_ok=True)
85 with open(metadata_file, 'w') as f:
86 json.dump(cached_data, f, indent=2)
88 if cached_data is not None and (
89 cached_data['result']['algorithm_version'] == algorithm['version'] and
90 cached_data['result']['dataset_version'] == dataset['version'] and
91 cached_data['result'].get('system_version', '') == system_version
92 ):
93 print(" Using cached result:")
94 results.append(cached_data['result'])
95 continue
97 print(" Running new benchmark...")
99 def run_timed_trials(operation, *args) -> Tuple[float, float]:
100 """Run multiple trials of an operation until total time exceeds 1 second.
101 Returns (median_time, mb_per_sec)"""
102 times = []
103 total_time = 0
104 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
106 while total_time < 1.0:
107 start_time = time.perf_counter()
108 _ = operation(*args) # Execute operation but discard result
109 trial_time = time.perf_counter() - start_time
110 times.append(trial_time)
111 total_time += trial_time
113 median_time = median(times)
114 mb_per_sec = array_size_mb / median_time
115 return median_time, mb_per_sec
117 # Measure encoding with multiple trials
118 encode_time, encode_mb_per_sec = run_timed_trials(algorithm['encode'], data)
119 encoded = algorithm['encode'](data) # One final encode to get the result
120 compressed_size = len(encoded)
121 compression_ratio = original_size / compressed_size
122 print(" Compression complete:")
123 print(f" Compressed size: {compressed_size:,} bytes")
124 print(f" Compression ratio: {compression_ratio:.2f}x")
125 print(f" Encode time: {encode_time*1000:.2f}ms")
126 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
128 print(" Verifying decompression...")
129 # Measure decoding with multiple trials
130 decode_time, decode_mb_per_sec = run_timed_trials(algorithm['decode'], encoded, dtype)
131 decoded = algorithm['decode'](encoded, dtype) # One final decode to verify
132 print(f" Decode time: {decode_time*1000:.2f}ms")
133 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
135 # Verify correctness
136 if not np.array_equal(data, decoded):
137 raise ValueError(
138 f"Decompression verification failed for {alg_name} on {dataset['name']}"
139 )
140 print(" Verification successful!")
142 # Store result
143 result = {
144 'dataset': dataset['name'],
145 'algorithm': alg_name,
146 'algorithm_version': algorithm['version'],
147 'dataset_version': dataset['version'],
148 'system_version': system_version,
149 'compression_ratio': compression_ratio,
150 'encode_time': encode_time,
151 'decode_time': decode_time,
152 'encode_mb_per_sec': encode_mb_per_sec,
153 'decode_mb_per_sec': decode_mb_per_sec,
154 'original_size': original_size,
155 'compressed_size': compressed_size,
156 'array_shape': data.shape,
157 'array_dtype': dtype,
158 'timestamp': time.time()
159 }
160 results.append(result)
162 # Save result and compressed data
163 os.makedirs(test_dir, exist_ok=True)
164 cache_data = {
165 'result': result
166 }
167 with open(metadata_file, 'w') as f:
168 json.dump(cache_data, f, indent=2)
169 with open(compressed_file, 'wb') as f:
170 f.write(encoded)
171 print(f" Results saved to: {test_dir}")
173 # Upload to memobin if API key is set and upload is enabled
174 memobin_api_key = os.environ.get('MEMOBIN_API_KEY')
175 upload_enabled = os.environ.get('UPLOAD_TO_MEMOBIN') == '1'
176 if memobin_api_key and upload_enabled:
177 if verbose:
178 print(" Uploading results to memobin...")
179 try:
180 memobin_url = construct_memobin_url(
181 alg_name, dataset['name'],
182 algorithm['version'], dataset['version'],
183 system_version
184 )
185 upload_to_memobin(
186 cache_data,
187 memobin_url,
188 os.environ.get('MEMOBIN_USER_ID', 'default'),
189 memobin_api_key
190 )
191 if verbose:
192 print(" Successfully uploaded to memobin")
193 except Exception as e:
194 print(f" Warning: Failed to upload to memobin: {str(e)}")
196 print("\n=== Benchmark Run Complete ===\n")
197 return {'results': results}