/ concept-collection / ephys_compression_tests
Sign in
concept-collection / ephys_compression_tests
118 lines · 3.9 KBBlameHistoryRaw
1import os
2import json
3from typing import Optional, Dict, Any
4from ._memobin import (
5 construct_memobin_url,
6 download_from_memobin,
7)
10def check_cached_result(
11 cache_dir: str,
12 dataset_name: str,
13 algorithm_name: str,
14 algorithm_version: str,
15 dataset_version: str,
16 system_version: str,
17 force: bool = False,
18 verbose: bool = True,
19) -> Optional[Dict[str, Any]]:
20 """Check for cached benchmark results locally and in memobin.
22 Args:
23 cache_dir: Directory containing cached results
24 dataset_name: Name of the dataset
25 algorithm_name: Name of the algorithm
26 algorithm_version: Version of the algorithm
27 dataset_version: Version of the dataset
28 system_version: Version of the system
29 force: If True, ignore cached results
30 verbose: Whether to print progress messages
32 Returns:
33 Cached result dictionary if found and valid, None otherwise
34 """
35 test_dir = os.path.join(cache_dir, dataset_name, algorithm_name)
36 metadata_file = os.path.join(test_dir, "metadata.json")
38 # First try local cache (unless force flag is set)
39 cached_data = None
40 if not force and os.path.exists(metadata_file):
41 with open(metadata_file, "r") as f:
42 cached_data = json.load(f)
43 # if versions do not match, then set to None
44 if isinstance(cached_data, dict) and "result" in cached_data:
45 result = cached_data["result"]
46 if (
47 result["algorithm_version"] != algorithm_version
48 or result["dataset_version"] != dataset_version
49 or result.get("system_version", "") != system_version
50 ):
51 cached_data = None
53 # If not in local cache, try memobin (unless force flag is set)
54 if cached_data is None and not force:
55 memobin_url = construct_memobin_url(
56 algorithm_name,
57 dataset_name,
58 algorithm_version,
59 dataset_version,
60 system_version,
61 "metadata.json",
62 )
63 if verbose:
64 print(" Looking for cached result in memobin...")
65 cached_data = download_from_memobin(memobin_url)
66 if cached_data is not None:
67 if verbose:
68 print(" Found result in memobin, saving locally...")
69 # Save to local cache
70 os.makedirs(test_dir, exist_ok=True)
71 with open(metadata_file, "w") as f:
72 json.dump(cached_data, f, indent=2)
74 if (
75 cached_data is not None
76 and isinstance(cached_data, dict)
77 and "result" in cached_data
78 ):
79 result = cached_data["result"]
80 if (
81 isinstance(result, dict)
82 and result.get("algorithm_version") == algorithm_version
83 and result.get("dataset_version") == dataset_version
84 and result.get("system_version", "") == system_version
85 ):
86 result["cache_status"] = "cached"
87 return result
89 return None
92def save_result_to_cache(
93 result: Dict[str, Any],
94 encoded_data: bytes,
95 cache_dir: str,
96 dataset_name: str,
97 algorithm_name: str,
98) -> None:
99 """Save benchmark result and compressed data to cache.
101 Args:
102 result: Benchmark result dictionary
103 encoded_data: Compressed data bytes
104 cache_dir: Directory to store cached results
105 dataset_name: Name of the dataset
106 algorithm_name: Name of the algorithm
107 """
108 test_dir = os.path.join(cache_dir, dataset_name, algorithm_name)
109 metadata_file = os.path.join(test_dir, "metadata.json")
110 compressed_file = os.path.join(test_dir, "compressed.dat")
112 os.makedirs(test_dir, exist_ok=True)
113 cache_data = {"result": result}
115 with open(metadata_file, "w") as f:
116 json.dump(cache_data, f, indent=2)
117 with open(compressed_file, "wb") as f:
118 f.write(encoded_data)
moveopenescclose