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