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 result, encoded = run_compression_benchmark(
155 data,
156 alg_name,
157 algorithm.encode,
158 algorithm.decode,
159 verbose,
160 )
162 # Add metadata to result
163 result.update(
164 {
165 "dataset": dataset.name,
166 "algorithm": alg_name,
167 "algorithm_version": algorithm.version,
168 "dataset_version": dataset.version,
169 "system_version": system_version,
170 }
171 )
172 results.append(result)
174 # Save result and compressed data
175 save_result_to_cache(
176 result,
177 encoded,
178 cache_dir,
179 dataset.name,
180 alg_name,
181 )
182 print(
183 f" Results saved to: {os.path.join(cache_dir, dataset.name, alg_name)}"
184 )
186 # Upload to memobin if enabled
187 if memobin_api_key and upload_enabled:
188 try:
189 memobin_url = construct_memobin_url(
190 alg_name,
191 dataset.name,
192 algorithm.version,
193 dataset.version,
194 system_version,
195 )
196 upload_to_memobin(
197 {"result": result},
198 memobin_url,
199 memobin_api_key,
200 )
201 if verbose:
202 print(" Successfully uploaded to memobin")
203 except Exception as e:
204 print(f" Warning: Failed to upload to memobin: {str(e)}")
206 print("\n=== Benchmark Run Complete ===\n")
208 # Collect algorithm and dataset information
209 algorithm_info = collect_algorithm_info(algorithms)
210 dataset_info = collect_dataset_info(datasets)
212 # Upload final benchmark status
213 if memobin_api_key and upload_enabled:
214 try:
215 upload_benchmark_status(
216 memobin_api_key,
217 "All datasets",
218 "All algorithms",
219 results,
220 total_benchmarks,
221 start_time,
222 )
223 except Exception as e:
224 print(f" Warning: Failed to upload final status to memobin: {str(e)}")
226 return {"results": results, "algorithms": algorithm_info, "datasets": dataset_info}