1import os
2import time
3from typing import Dict, Any, List, Optional
4import numpy as np
6from ..algorithms import algorithms
7from ..datasets import datasets
8from ._memobin import construct_memobin_url, upload_to_memobin
9from .upload_dataset import upload_dataset_to_memobin
10from .upload_reconstructed import upload_reconstructed_to_memobin
11from .cache_management import check_cached_result, save_result_to_cache
12from .benchmark_timing import run_compression_benchmark
13from .collect_info import collect_algorithm_info, collect_dataset_info
14from .is_compatible import is_compatible
15from .upload_benchmark_status import upload_benchmark_status
16from ..types import Algorithm, Dataset
18system_version = "v6"
21def run_benchmarks(
22 cache_dir: str = ".benchmark_cache",
23 verbose: bool = True,
24 selected_algorithms: Optional[List[Algorithm]] = None,
25 selected_datasets: Optional[List[Dataset]] = None,
26 force: bool = False,
27) -> Dict[str, Any]:
28 """Run all benchmarks, with caching based on algorithm and dataset versions.
30 Results are stored in separate directories for each dataset/algorithm combination:
31 cache_dir/
32 dataset_name/
33 algorithm_name/
34 metadata.json # Contains algorithm version, dataset version, and results
35 compressed.dat # The actual compressed data
37 Args:
38 cache_dir: Directory to store cached results
39 verbose: Whether to print progress messages
40 selected_algorithms: Optional list of specific algorithms to run
41 selected_datasets: Optional list of specific datasets to run
42 force: If True, ignore cached results
44 Returns:
45 Dictionary containing benchmark results and metadata
46 """
47 print("\n=== Starting Benchmark Run ===")
48 print(f"Cache directory: {cache_dir}")
50 os.makedirs(cache_dir, exist_ok=True)
52 start_time = time.time()
53 last_status_upload = 0 # Track last status upload time
54 results = []
55 print("\nRunning benchmarks for all dataset-algorithm combinations...")
57 # Use selected datasets/algorithms or fall back to all
58 datasets_to_run = selected_datasets if selected_datasets is not None else datasets
59 algorithms_to_run = (
60 selected_algorithms if selected_algorithms is not None else algorithms
61 )
63 # Calculate total number of benchmarks
64 total_benchmarks = sum(
65 1
66 for dataset in datasets_to_run
67 for algorithm in algorithms_to_run
68 if is_compatible(algorithm.tags, dataset.tags)
69 )
71 # Run benchmarks for each dataset and algorithm combination
72 memobin_api_key = os.environ.get("MEMOBIN_API_KEY")
73 upload_enabled = os.environ.get("UPLOAD_TO_MEMOBIN") == "1"
75 for dataset in datasets_to_run:
76 dataset_tags = dataset.tags
77 print(f"\n*** Dataset: {dataset.name} (tags: {dataset_tags}) ***")
79 # only create the dataset if it is needed
80 data = None
82 for algorithm in algorithms_to_run:
83 alg_name = algorithm.name
84 alg_tags = algorithm.tags
86 # Skip if algorithm and dataset are not compatible based on tags
87 if not is_compatible(alg_tags, dataset_tags):
88 if verbose:
89 print(
90 f"\nSkipping algorithm {alg_name} (tags: {alg_tags}) - incompatible with dataset tags"
91 )
92 continue
94 print(f"\nTesting algorithm: {alg_name} on dataset: {dataset.name}")
96 # Upload current status to memobin if enabled (once per minute)
97 current_time = time.time()
98 if (
99 memobin_api_key
100 and upload_enabled
101 and (current_time - last_status_upload >= 60)
102 ): # Check if 60 seconds have passed
103 try:
104 upload_benchmark_status(
105 memobin_api_key,
106 dataset.name,
107 alg_name,
108 results,
109 total_benchmarks,
110 start_time,
111 )
112 last_status_upload = current_time # Update last upload time
113 except Exception as e:
114 print(f" Warning: Failed to upload status to memobin: {str(e)}")
116 # Check if we can use cached result
117 cached_result = check_cached_result(
118 cache_dir,
119 dataset.name,
120 alg_name,
121 algorithm.version,
122 dataset.version,
123 system_version,
124 force,
125 verbose,
126 )
128 if cached_result is not None:
129 print(" Using cached result")
130 results.append(cached_result)
131 continue
133 print(f" Running benchmark for {alg_name} on {dataset.name}...")
134 if data is None:
135 data = dataset.create()
136 print(f"Created dataset: shape={data.shape}, dtype={data.dtype}")
137 else:
138 print("Dataset already created")
140 # Upload dataset to memobin if enabled
141 if memobin_api_key and upload_enabled:
142 try:
143 upload_dataset_to_memobin(
144 data,
145 dataset.name,
146 dataset.version,
147 memobin_api_key,
148 cache_dir,
149 verbose,
150 )
151 except Exception as e:
152 print(f" Warning: Failed to upload dataset to memobin: {str(e)}")
154 # Run the benchmark
155 lossy = "lossy" in alg_tags
156 result, encoded, decoded = run_compression_benchmark(
157 data,
158 alg_name,
159 algorithm.encode,
160 algorithm.decode,
161 verbose,
162 lossy=lossy
163 )
165 # Add metadata to result
166 result.update(
167 {
168 "dataset": dataset.name,
169 "algorithm": alg_name,
170 "algorithm_version": algorithm.version,
171 "dataset_version": dataset.version,
172 "system_version": system_version,
173 }
174 )
175 results.append(result)
177 # Save result, compressed data, and reconstructed data (for lossy algorithms)
178 save_result_to_cache(
179 result,
180 encoded,
181 cache_dir,
182 dataset.name,
183 alg_name,
184 reconstructed_data=decoded if lossy else None,
185 )
186 print(
187 f" Results saved to: {os.path.join(cache_dir, dataset.name, alg_name)}"
188 )
190 # Upload to memobin if enabled
191 if memobin_api_key and upload_enabled:
192 try:
193 memobin_url = construct_memobin_url(
194 alg_name,
195 dataset.name,
196 algorithm.version,
197 dataset.version,
198 system_version,
199 )
200 upload_to_memobin(
201 {"result": result},
202 memobin_url,
203 memobin_api_key,
204 )
205 if verbose:
206 print(" Successfully uploaded benchmark result to memobin")
207 except Exception as e:
208 print(f" Warning: Failed to upload to memobin: {str(e)}")
210 # Upload reconstructed array for lossy algorithms
211 if lossy:
212 try:
213 upload_reconstructed_to_memobin(
214 decoded,
215 alg_name,
216 dataset.name,
217 algorithm.version,
218 dataset.version,
219 system_version,
220 memobin_api_key,
221 verbose,
222 )
223 except Exception as e:
224 print(f" Warning: Failed to upload reconstructed data to memobin: {str(e)}")
226 print("\n=== Benchmark Run Complete ===\n")
228 # Collect algorithm and dataset information
229 algorithm_info = collect_algorithm_info(algorithms)
230 dataset_info = collect_dataset_info(datasets)
232 # Add reconstructed URLs to results for lossy algorithms
233 from .collect_info import add_reconstructed_urls_to_results
234 add_reconstructed_urls_to_results(results, algorithms_to_run)
236 # Upload final benchmark status
237 if memobin_api_key and upload_enabled:
238 try:
239 upload_benchmark_status(
240 memobin_api_key,
241 "All datasets",
242 "All algorithms",
243 results,
244 total_benchmarks,
245 start_time,
246 )
247 except Exception as e:
248 print(f" Warning: Failed to upload final status to memobin: {str(e)}")
250 return {"results": results, "algorithms": algorithm_info, "datasets": dataset_info}