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 return True
31def run_benchmarks(
32 cache_dir: str = ".benchmark_cache", verbose: bool = True
33) -> Dict[str, Any]:
34 """Run all benchmarks, with caching based on algorithm and dataset versions.
36 Results are stored in separate directories for each dataset/algorithm combination:
37 cache_dir/
38 dataset_name/
39 algorithm_name/
40 metadata.json # Contains algorithm version, dataset version, and results
41 compressed.dat # The actual compressed data
43 Args:
44 cache_dir: Directory to store cached results
46 Returns:
47 Dictionary containing benchmark results and metadata
48 """
49 print("\n=== Starting Benchmark Run ===")
50 print(f"Cache directory: {cache_dir}")
52 os.makedirs(cache_dir, exist_ok=True)
54 results = []
55 print("\nRunning benchmarks for all dataset-algorithm combinations...")
57 # Run benchmarks for each dataset and algorithm combination
58 for dataset in datasets:
59 dataset_tags = dataset.get("tags", [])
60 print(f"\n--- Dataset: {dataset['name']} (tags: {dataset_tags}) ---")
61 # Create dataset once for all algorithms
62 data = dataset["create"]()
63 dtype = str(data.dtype)
64 original_size = len(data.tobytes())
65 print(f"Created dataset: shape={data.shape}, dtype={dtype}")
66 print(f"Original size: {original_size:,} bytes")
68 for algorithm in algorithms:
69 alg_name = algorithm["name"]
70 alg_tags = algorithm.get("tags", [])
72 # Skip if algorithm and dataset are not compatible based on tags
73 if not is_compatible(alg_tags, dataset_tags):
74 if verbose:
75 print(
76 f"\nSkipping algorithm {alg_name} (tags: {alg_tags}) - incompatible with dataset tags"
77 )
78 continue
80 print(f"\nTesting algorithm: {alg_name} (tags: {alg_tags})")
82 # Check if we can use cached result
83 test_dir = os.path.join(cache_dir, dataset["name"], alg_name)
84 metadata_file = os.path.join(test_dir, "metadata.json")
85 compressed_file = os.path.join(test_dir, "compressed.dat")
87 # First try local cache
88 cached_data = None
89 if os.path.exists(metadata_file):
90 with open(metadata_file, "r") as f:
91 cached_data = json.load(f)
92 # if versions do not match, then set to None
93 if (
94 cached_data["result"]["algorithm_version"]
95 != algorithm["version"]
96 or cached_data["result"]["dataset_version"]
97 != dataset["version"]
98 or cached_data["result"].get("system_version", "")
99 != system_version
100 ):
101 cached_data = None
103 # If not in local cache, try memobin
104 if cached_data is None:
105 memobin_url = construct_memobin_url(
106 alg_name,
107 dataset["name"],
108 algorithm["version"],
109 dataset["version"],
110 system_version,
111 )
112 if verbose:
113 print(" Looking for cached result in memobin...")
114 cached_data = download_from_memobin(memobin_url)
115 if cached_data is not None:
116 if verbose:
117 print(" Found result in memobin, saving locally...")
118 # Save to local cache
119 os.makedirs(test_dir, exist_ok=True)
120 with open(metadata_file, "w") as f:
121 json.dump(cached_data, f, indent=2)
123 if cached_data is not None and (
124 cached_data["result"]["algorithm_version"] == algorithm["version"]
125 and cached_data["result"]["dataset_version"] == dataset["version"]
126 and cached_data["result"].get("system_version", "") == system_version
127 ):
128 print(" Using cached result:")
129 results.append(cached_data["result"])
130 continue
132 print(" Running new benchmark...")
134 def run_timed_trials(operation, *args) -> Tuple[float, float]:
135 """Run multiple trials of an operation until total time exceeds 1 second.
136 Returns (median_time, mb_per_sec)"""
137 times = []
138 total_time = 0
139 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
141 while total_time < 1.0:
142 start_time = time.perf_counter()
143 _ = operation(*args) # Execute operation but discard result
144 trial_time = time.perf_counter() - start_time
145 times.append(trial_time)
146 total_time += trial_time
148 median_time = median(times)
149 mb_per_sec = array_size_mb / median_time
150 return median_time, mb_per_sec
152 # Measure encoding with multiple trials
153 encode_time, encode_mb_per_sec = run_timed_trials(algorithm["encode"], data)
154 encoded = algorithm["encode"](data) # One final encode to get the result
155 compressed_size = len(encoded)
156 compression_ratio = original_size / compressed_size
157 print(" Compression complete:")
158 print(f" Compressed size: {compressed_size:,} bytes")
159 print(f" Compression ratio: {compression_ratio:.2f}x")
160 print(f" Encode time: {encode_time*1000:.2f}ms")
161 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
163 print(" Verifying decompression...")
164 # Measure decoding with multiple trials
165 decode_time, decode_mb_per_sec = run_timed_trials(
166 algorithm["decode"], encoded, dtype
167 )
168 decoded = algorithm["decode"](encoded, dtype) # One final decode to verify
169 print(f" Decode time: {decode_time*1000:.2f}ms")
170 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
172 # Verify correctness
173 if not np.array_equal(data, decoded):
174 raise ValueError(
175 f"Decompression verification failed for {alg_name} on {dataset['name']}"
176 )
177 print(" Verification successful!")
179 # Store result
180 result = {
181 "dataset": dataset["name"],
182 "algorithm": alg_name,
183 "algorithm_version": algorithm["version"],
184 "dataset_version": dataset["version"],
185 "system_version": system_version,
186 "compression_ratio": compression_ratio,
187 "encode_time": encode_time,
188 "decode_time": decode_time,
189 "encode_mb_per_sec": encode_mb_per_sec,
190 "decode_mb_per_sec": decode_mb_per_sec,
191 "original_size": original_size,
192 "compressed_size": compressed_size,
193 "array_shape": data.shape,
194 "array_dtype": dtype,
195 "timestamp": time.time(),
196 }
197 results.append(result)
199 # Save result and compressed data
200 os.makedirs(test_dir, exist_ok=True)
201 cache_data = {"result": result}
202 with open(metadata_file, "w") as f:
203 json.dump(cache_data, f, indent=2)
204 with open(compressed_file, "wb") as f:
205 f.write(encoded)
206 print(f" Results saved to: {test_dir}")
208 # Upload to memobin if API key is set and upload is enabled
209 memobin_api_key = os.environ.get("MEMOBIN_API_KEY")
210 upload_enabled = os.environ.get("UPLOAD_TO_MEMOBIN") == "1"
211 if memobin_api_key and upload_enabled:
212 if verbose:
213 print(" Uploading results to memobin...")
214 try:
215 memobin_url = construct_memobin_url(
216 alg_name,
217 dataset["name"],
218 algorithm["version"],
219 dataset["version"],
220 system_version,
221 )
222 upload_to_memobin(
223 cache_data,
224 memobin_url,
225 os.environ.get("MEMOBIN_USER_ID", "default"),
226 memobin_api_key,
227 )
228 if verbose:
229 print(" Successfully uploaded to memobin")
230 except Exception as e:
231 print(f" Warning: Failed to upload to memobin: {str(e)}")
233 print("\n=== Benchmark Run Complete ===\n")
234 return {"results": results}