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