1import time
2import json
3import os
4from typing import Dict, Any, Tuple, List
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"
15def is_compatible(algorithm_tags: List[str], dataset_tags: List[str]) -> bool:
16 """Check if an algorithm is compatible with a dataset based on their tags.
18 Args:
19 algorithm_tags: List of tags for the algorithm
20 dataset_tags: List of tags for the dataset
22 Returns:
23 True if the algorithm should be applied to the dataset
24 """
25 # If algorithm has delta_encoding tag, dataset must have continuous tag
26 if "delta_encoding" in algorithm_tags and "continuous" not in dataset_tags:
27 return False
28 if "markov_prediction" in algorithm_tags and "continuous" not in dataset_tags:
29 return False
30 return True
33def run_benchmarks(
34 cache_dir: str = ".benchmark_cache", verbose: bool = True
35) -> Dict[str, Any]:
36 """Run all benchmarks, with caching based on algorithm and dataset versions.
38 Results are stored in separate directories for each dataset/algorithm combination:
39 cache_dir/
40 dataset_name/
41 algorithm_name/
42 metadata.json # Contains algorithm version, dataset version, and results
43 compressed.dat # The actual compressed data
45 Args:
46 cache_dir: Directory to store cached results
48 Returns:
49 Dictionary containing benchmark results and metadata
50 """
51 print("\n=== Starting Benchmark Run ===")
52 print(f"Cache directory: {cache_dir}")
54 os.makedirs(cache_dir, exist_ok=True)
56 results = []
57 print("\nRunning benchmarks for all dataset-algorithm combinations...")
59 # Run benchmarks for each dataset and algorithm combination
60 for dataset in datasets:
61 dataset_tags = dataset.get("tags", [])
62 print(f"\n--- Dataset: {dataset['name']} (tags: {dataset_tags}) ---")
63 # Create dataset once for all algorithms
64 data = dataset["create"]()
65 dtype = str(data.dtype)
66 original_size = len(data.tobytes())
67 print(f"Created dataset: shape={data.shape}, dtype={dtype}")
68 print(f"Original size: {original_size:,} bytes")
70 for algorithm in algorithms:
71 alg_name = algorithm["name"]
72 alg_tags = algorithm.get("tags", [])
74 # Skip if algorithm and dataset are not compatible based on tags
75 if not is_compatible(alg_tags, dataset_tags):
76 if verbose:
77 print(
78 f"\nSkipping algorithm {alg_name} (tags: {alg_tags}) - incompatible with dataset tags"
79 )
80 continue
82 print(f"\nTesting algorithm: {alg_name} (tags: {alg_tags})")
84 # Check if we can use cached result
85 test_dir = os.path.join(cache_dir, dataset["name"], alg_name)
86 metadata_file = os.path.join(test_dir, "metadata.json")
87 compressed_file = os.path.join(test_dir, "compressed.dat")
89 # First try local cache
90 cached_data = None
91 if os.path.exists(metadata_file):
92 with open(metadata_file, "r") as f:
93 cached_data = json.load(f)
94 # if versions do not match, then set to None
95 if (
96 cached_data["result"]["algorithm_version"]
97 != algorithm["version"]
98 or cached_data["result"]["dataset_version"]
99 != dataset["version"]
100 or cached_data["result"].get("system_version", "")
101 != system_version
102 ):
103 cached_data = None
105 # If not in local cache, try memobin
106 if cached_data is None:
107 memobin_url = construct_memobin_url(
108 alg_name,
109 dataset["name"],
110 algorithm["version"],
111 dataset["version"],
112 system_version,
113 )
114 if verbose:
115 print(" Looking for cached result in memobin...")
116 cached_data = download_from_memobin(memobin_url)
117 if cached_data is not None:
118 if verbose:
119 print(" Found result in memobin, saving locally...")
120 # Save to local cache
121 os.makedirs(test_dir, exist_ok=True)
122 with open(metadata_file, "w") as f:
123 json.dump(cached_data, f, indent=2)
125 if cached_data is not None and (
126 cached_data["result"]["algorithm_version"] == algorithm["version"]
127 and cached_data["result"]["dataset_version"] == dataset["version"]
128 and cached_data["result"].get("system_version", "") == system_version
129 ):
130 print(" Using cached result:")
131 results.append(cached_data["result"])
132 continue
134 print(" Running new benchmark...")
136 def run_timed_trials(operation, *args) -> Tuple[float, float]:
137 """Run multiple trials of an operation until total time exceeds 1 second.
138 Returns (median_time, mb_per_sec)"""
139 times = []
140 total_time = 0
141 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
143 while total_time < 1.0:
144 start_time = time.perf_counter()
145 _ = operation(*args) # Execute operation but discard result
146 trial_time = time.perf_counter() - start_time
147 times.append(trial_time)
148 total_time += trial_time
150 median_time = median(times)
151 mb_per_sec = array_size_mb / median_time
152 return median_time, mb_per_sec
154 # Measure encoding with multiple trials
155 encode_time, encode_mb_per_sec = run_timed_trials(algorithm["encode"], data)
156 encoded = algorithm["encode"](data) # One final encode to get the result
157 compressed_size = len(encoded)
158 compression_ratio = original_size / compressed_size
159 print(" Compression complete:")
160 print(f" Compressed size: {compressed_size:,} bytes")
161 print(f" Compression ratio: {compression_ratio:.2f}x")
162 print(f" Encode time: {encode_time*1000:.2f}ms")
163 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
165 print(" Verifying decompression...")
166 # Measure decoding with multiple trials
167 decode_time, decode_mb_per_sec = run_timed_trials(
168 algorithm["decode"], encoded, dtype
169 )
170 decoded = algorithm["decode"](encoded, dtype) # One final decode to verify
171 print(f" Decode time: {decode_time*1000:.2f}ms")
172 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
174 if len(data) != len(decoded):
175 raise ValueError(
176 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
177 )
179 # Verify correctness
180 if not np.array_equal(data, decoded):
181 print(data[:100])
182 print(decoded[:100])
183 for j in range(len(data)):
184 if data[j] != decoded[j]:
185 print(f"Error at index {j}: {data[j]} != {decoded[j]}")
186 raise ValueError(
187 f"Decompression verification failed for {alg_name} on {dataset['name']}"
188 )
189 print(" Verification successful!")
191 # Store result
192 result = {
193 "dataset": dataset["name"],
194 "algorithm": alg_name,
195 "algorithm_version": algorithm["version"],
196 "dataset_version": dataset["version"],
197 "system_version": system_version,
198 "compression_ratio": compression_ratio,
199 "encode_time": encode_time,
200 "decode_time": decode_time,
201 "encode_mb_per_sec": encode_mb_per_sec,
202 "decode_mb_per_sec": decode_mb_per_sec,
203 "original_size": original_size,
204 "compressed_size": compressed_size,
205 "array_shape": data.shape,
206 "array_dtype": dtype,
207 "timestamp": time.time(),
208 }
209 results.append(result)
211 # Save result and compressed data
212 os.makedirs(test_dir, exist_ok=True)
213 cache_data = {"result": result}
214 with open(metadata_file, "w") as f:
215 json.dump(cache_data, f, indent=2)
216 with open(compressed_file, "wb") as f:
217 f.write(encoded)
218 print(f" Results saved to: {test_dir}")
220 # Upload to memobin if API key is set and upload is enabled
221 memobin_api_key = os.environ.get("MEMOBIN_API_KEY")
222 upload_enabled = os.environ.get("UPLOAD_TO_MEMOBIN") == "1"
223 if memobin_api_key and upload_enabled:
224 if verbose:
225 print(" Uploading results to memobin...")
226 try:
227 memobin_url = construct_memobin_url(
228 alg_name,
229 dataset["name"],
230 algorithm["version"],
231 dataset["version"],
232 system_version,
233 )
234 upload_to_memobin(
235 cache_data,
236 memobin_url,
237 os.environ.get("MEMOBIN_USER_ID", "default"),
238 memobin_api_key,
239 )
240 if verbose:
241 print(" Successfully uploaded to memobin")
242 except Exception as e:
243 print(f" Warning: Failed to upload to memobin: {str(e)}")
245 print("\n=== Benchmark Run Complete ===\n")
247 # Collect algorithm and dataset information as lists
248 algorithm_info = []
249 for algorithm in algorithms:
250 algorithm_info.append(
251 {
252 "name": algorithm["name"],
253 "description": algorithm.get("description", ""),
254 "version": algorithm["version"],
255 "tags": algorithm.get("tags", []),
256 }
257 )
259 dataset_info = []
260 for dataset in datasets:
261 dataset_info.append(
262 {
263 "name": dataset["name"],
264 "description": dataset.get("description", ""),
265 "version": dataset["version"],
266 "tags": dataset.get("tags", []),
267 }
268 )
270 return {"results": results, "algorithms": algorithm_info, "datasets": dataset_info}