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