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