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