1import time
2import json
3import os
4from typing import Dict, Any, Tuple, List, Optional
5import numpy as np
6from statistics import median
7from .algorithms import algorithms
8from .datasets import datasets
9from ._memobin import (
10 construct_memobin_url,
11 construct_dataset_url,
12 upload_to_memobin,
13 download_from_memobin,
14 exists_in_memobin,
15)
18system_version = "v5"
19GITHUB_ALGORITHMS_PREFIX = "https://github.com/magland/zia/blob/main/zia_benchmark/src/zia_benchmark/algorithms/"
20GITHUB_DATASETS_PREFIX = (
21 "https://github.com/magland/zia/blob/main/zia_benchmark/src/zia_benchmark/datasets/"
22)
25def is_compatible(algorithm_tags: List[str], dataset_tags: List[str]) -> bool:
26 """Check if an algorithm is compatible with a dataset based on their tags.
28 Args:
29 algorithm_tags: List of tags for the algorithm
30 dataset_tags: List of tags for the dataset
32 Returns:
33 True if the algorithm should be applied to the dataset
34 """
35 # If algorithm has delta_encoding tag, dataset must have continuous tag
36 if "delta_encoding" in algorithm_tags and "continuous" not in dataset_tags:
37 return False
38 if "markov_prediction" in algorithm_tags and "continuous" not in dataset_tags:
39 return False
40 return True
43def run_benchmarks(
44 cache_dir: str = ".benchmark_cache",
45 verbose: bool = True,
46 selected_algorithms: Optional[List[dict]] = None,
47 selected_datasets: Optional[List[dict]] = None,
48) -> Dict[str, Any]:
49 """Run all benchmarks, with caching based on algorithm and dataset versions.
51 Results are stored in separate directories for each dataset/algorithm combination:
52 cache_dir/
53 dataset_name/
54 algorithm_name/
55 metadata.json # Contains algorithm version, dataset version, and results
56 compressed.dat # The actual compressed data
58 Args:
59 cache_dir: Directory to store cached results
61 Returns:
62 Dictionary containing benchmark results and metadata
63 """
64 print("\n=== Starting Benchmark Run ===")
65 print(f"Cache directory: {cache_dir}")
67 os.makedirs(cache_dir, exist_ok=True)
69 results = []
70 print("\nRunning benchmarks for all dataset-algorithm combinations...")
72 # Use selected datasets/algorithms or fall back to all
73 datasets_to_run = selected_datasets if selected_datasets is not None else datasets
74 algorithms_to_run = (
75 selected_algorithms if selected_algorithms is not None else algorithms
76 )
78 # Run benchmarks for each dataset and algorithm combination
79 for dataset in datasets_to_run:
80 dataset_tags = dataset.get("tags", [])
81 print(f"\n*** Dataset: {dataset['name']} (tags: {dataset_tags}) ***")
83 # data will only be created if needed
84 data = None
85 original_size = None
86 dtype = None
88 for algorithm in algorithms_to_run:
89 alg_name = algorithm["name"]
90 alg_tags = algorithm.get("tags", [])
92 # Skip if algorithm and dataset are not compatible based on tags
93 if not is_compatible(alg_tags, dataset_tags):
94 if verbose:
95 print(
96 f"\nSkipping algorithm {alg_name} (tags: {alg_tags}) - incompatible with dataset tags"
97 )
98 continue
100 print(f"\nTesting algorithm: {alg_name} (tags: {alg_tags})")
102 # Check if we can use cached result
103 test_dir = os.path.join(cache_dir, dataset["name"], alg_name)
104 metadata_file = os.path.join(test_dir, "metadata.json")
105 compressed_file = os.path.join(test_dir, "compressed.dat")
107 # First try local cache
108 cached_data = None
109 if os.path.exists(metadata_file):
110 with open(metadata_file, "r") as f:
111 cached_data = json.load(f)
112 # if versions do not match, then set to None
113 if isinstance(cached_data, dict) and "result" in cached_data:
114 result = cached_data["result"]
115 if (
116 result["algorithm_version"] != algorithm["version"]
117 or result["dataset_version"] != dataset["version"]
118 or result.get("system_version", "") != system_version
119 ):
120 cached_data = None
122 # If not in local cache, try memobin
123 if cached_data is None:
124 memobin_url = construct_memobin_url(
125 alg_name,
126 dataset["name"],
127 algorithm["version"],
128 dataset["version"],
129 system_version,
130 "metadata.json",
131 )
132 if verbose:
133 print(" Looking for cached result in memobin...")
134 cached_data = download_from_memobin(memobin_url)
135 if cached_data is not None:
136 if verbose:
137 print(" Found result in memobin, saving locally...")
138 # Save to local cache
139 os.makedirs(test_dir, exist_ok=True)
140 with open(metadata_file, "w") as f:
141 json.dump(cached_data, f, indent=2)
143 if (
144 cached_data is not None
145 and isinstance(cached_data, dict)
146 and "result" in cached_data
147 ):
148 result = cached_data["result"]
149 if (
150 isinstance(result, dict)
151 and result.get("algorithm_version") == algorithm["version"]
152 and result.get("dataset_version") == dataset["version"]
153 and result.get("system_version", "") == system_version
154 ):
155 print(" Using cached result:")
156 results.append(result)
157 continue
159 print(" Running new benchmark...")
160 if data is None:
161 # only create data if needed
162 data = dataset["create"]()
163 dtype = str(data.dtype)
164 original_size = len(data.tobytes())
165 print(f"Created dataset: shape={data.shape}, dtype={dtype}")
166 print(f"Original size: {original_size:,} bytes")
168 # Upload dataset to memobin if enabled
169 memobin_api_key = os.environ.get("MEMOBIN_API_KEY")
170 upload_enabled = os.environ.get("UPLOAD_TO_MEMOBIN") == "1"
171 if memobin_api_key and upload_enabled:
172 try:
173 # Upload array metadata as JSON
174 dataset_url_json = construct_dataset_url(
175 dataset["name"], dataset["version"], "json"
176 )
177 if not exists_in_memobin(dataset_url_json):
178 if verbose:
179 print(" Uploading dataset metadata to memobin...")
180 metadata = {"dtype": str(data.dtype), "shape": data.shape}
181 upload_to_memobin(
182 metadata,
183 dataset_url_json,
184 memobin_api_key,
185 content_type="application/json",
186 )
187 if verbose:
188 print(" Successfully uploaded metadata")
190 # Upload raw .dat format
191 dataset_url_raw = construct_dataset_url(
192 dataset["name"], dataset["version"], "dat"
193 )
194 if not exists_in_memobin(dataset_url_raw):
195 if verbose:
196 print(" Uploading dataset (raw) to memobin...")
197 upload_to_memobin(
198 data.tobytes(),
199 dataset_url_raw,
200 memobin_api_key,
201 content_type="application/octet-stream",
202 )
203 if verbose:
204 print(" Successfully uploaded raw dataset")
206 # Upload .npy format
207 dataset_url_npy = construct_dataset_url(
208 dataset["name"], dataset["version"], "npy"
209 )
210 if not exists_in_memobin(dataset_url_npy):
211 if verbose:
212 print(" Uploading dataset (npy) to memobin...")
213 # Save array to a temporary .npy file
214 temp_npy = os.path.join(cache_dir, "temp.npy")
215 np.save(temp_npy, data)
216 with open(temp_npy, "rb") as f:
217 npy_bytes = f.read()
218 os.remove(temp_npy) # Clean up temp file
220 upload_to_memobin(
221 npy_bytes,
222 dataset_url_npy,
223 memobin_api_key,
224 content_type="application/octet-stream",
225 )
226 if verbose:
227 print(" Successfully uploaded npy dataset")
228 except Exception as e:
229 print(
230 f" Warning: Failed to upload dataset to memobin: {str(e)}"
231 )
233 assert data is not None
234 assert isinstance(data, np.ndarray)
235 assert isinstance(original_size, int)
236 assert isinstance(dtype, str)
238 def run_timed_trials(operation, *args) -> Tuple[float, float]:
239 """Run multiple trials of an operation until total time exceeds 1 second.
240 Returns (median_time, mb_per_sec)"""
241 assert data is not None
242 assert isinstance(data, np.ndarray)
243 times = []
244 total_time = 0
245 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
247 while total_time < 1.0:
248 start_time = time.perf_counter()
249 _ = operation(*args) # Execute operation but discard result
250 trial_time = time.perf_counter() - start_time
251 times.append(trial_time)
252 total_time += trial_time
254 median_time = median(times)
255 mb_per_sec = array_size_mb / median_time
256 return median_time, mb_per_sec
258 # Measure encoding with multiple trials
259 encode_time, encode_mb_per_sec = run_timed_trials(algorithm["encode"], data)
260 encoded = algorithm["encode"](data) # One final encode to get the result
261 compressed_size = len(encoded)
262 compression_ratio = original_size / compressed_size
263 print(" Compression complete:")
264 print(f" Compressed size: {compressed_size:,} bytes")
265 print(f" Compression ratio: {compression_ratio:.2f}x")
266 print(f" Encode time: {encode_time*1000:.2f}ms")
267 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
269 print(" Verifying decompression...")
270 # Measure decoding with multiple trials
271 decode_time, decode_mb_per_sec = run_timed_trials(
272 algorithm["decode"], encoded, dtype
273 )
274 decoded = algorithm["decode"](encoded, dtype) # One final decode to verify
275 print(f" Decode time: {decode_time*1000:.2f}ms")
276 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
278 if len(data) != len(decoded):
279 raise ValueError(
280 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
281 )
283 # Verify correctness
284 if not np.array_equal(data, decoded):
285 print(data[:100])
286 print(decoded[:100])
287 for j in range(len(data)):
288 if data[j] != decoded[j]:
289 print(f"Error at index {j}: {data[j]} != {decoded[j]}")
290 raise ValueError(
291 f"Decompression verification failed for {alg_name} on {dataset['name']}"
292 )
293 print(" Verification successful!")
295 # Store result
296 result = {
297 "dataset": dataset["name"],
298 "algorithm": alg_name,
299 "algorithm_version": algorithm["version"],
300 "dataset_version": dataset["version"],
301 "system_version": system_version,
302 "compression_ratio": compression_ratio,
303 "encode_time": encode_time,
304 "decode_time": decode_time,
305 "encode_mb_per_sec": encode_mb_per_sec,
306 "decode_mb_per_sec": decode_mb_per_sec,
307 "original_size": original_size,
308 "compressed_size": compressed_size,
309 "array_shape": data.shape,
310 "array_dtype": dtype,
311 "timestamp": time.time(),
312 }
313 results.append(result)
315 # Save result and compressed data
316 os.makedirs(test_dir, exist_ok=True)
317 cache_data = {"result": result}
318 with open(metadata_file, "w") as f:
319 json.dump(cache_data, f, indent=2)
320 with open(compressed_file, "wb") as f:
321 f.write(encoded)
322 print(f" Results saved to: {test_dir}")
324 # Upload to memobin if API key is set and upload is enabled
325 memobin_api_key = os.environ.get("MEMOBIN_API_KEY")
326 upload_enabled = os.environ.get("UPLOAD_TO_MEMOBIN") == "1"
327 if memobin_api_key and upload_enabled:
328 if verbose:
329 print(" Uploading results to memobin...")
330 try:
331 memobin_url = construct_memobin_url(
332 alg_name,
333 dataset["name"],
334 algorithm["version"],
335 dataset["version"],
336 system_version,
337 )
338 upload_to_memobin(
339 cache_data,
340 memobin_url,
341 memobin_api_key,
342 )
343 if verbose:
344 print(" Successfully uploaded to memobin")
345 except Exception as e:
346 print(f" Warning: Failed to upload to memobin: {str(e)}")
348 print("\n=== Benchmark Run Complete ===\n")
350 # Collect algorithm and dataset information as lists
351 algorithm_info = []
352 for algorithm in algorithms:
353 info = {
354 "name": algorithm["name"],
355 "description": algorithm.get("description", ""),
356 "version": algorithm["version"],
357 "tags": algorithm.get("tags", []),
358 }
359 if "source_file" in algorithm:
360 info["source_file"] = GITHUB_ALGORITHMS_PREFIX + algorithm["source_file"]
361 algorithm_info.append(info)
363 dataset_info = []
364 for dataset in datasets:
365 info = {
366 "name": dataset["name"],
367 "description": dataset.get("description", ""),
368 "version": dataset["version"],
369 "tags": dataset.get("tags", []),
370 "data_url_raw": construct_dataset_url(
371 dataset["name"], dataset["version"], "dat"
372 ),
373 "data_url_npy": construct_dataset_url(
374 dataset["name"], dataset["version"], "npy"
375 ),
376 "data_url_json": construct_dataset_url(
377 dataset["name"], dataset["version"], "json"
378 ),
379 }
380 if "source_file" in dataset:
381 info["source_file"] = GITHUB_DATASETS_PREFIX + dataset["source_file"]
382 dataset_info.append(info)
384 return {"results": results, "algorithms": algorithm_info, "datasets": dataset_info}