Global context manager for all benchmarks with lazy loading support.
| 19 | from src.registry import BENCHMARK |
| 20 | |
| 21 | class BenchmarkContextManager(BaseModel): |
| 22 | """Global context manager for all benchmarks with lazy loading support.""" |
| 23 | model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow") |
| 24 | |
| 25 | base_dir: str = Field(default=None, description="The base directory to use for the benchmarks") |
| 26 | save_path: str = Field(default=None, description="The path to save the benchmarks") |
| 27 | contract_path: str = Field(default=None, description="The path to save the benchmark contract") |
| 28 | |
| 29 | def __init__(self, |
| 30 | base_dir: Optional[str] = None, |
| 31 | save_path: Optional[str] = None, |
| 32 | contract_path: Optional[str] = None, |
| 33 | **kwargs): |
| 34 | """Initialize the benchmark context manager. |
| 35 | |
| 36 | Args: |
| 37 | base_dir: Base directory for storing benchmark data |
| 38 | save_path: Path to save benchmark configurations |
| 39 | contract_path: Path to save benchmark contract |
| 40 | """ |
| 41 | super().__init__(**kwargs) |
| 42 | |
| 43 | if base_dir is not None: |
| 44 | self.base_dir = assemble_project_path(base_dir) |
| 45 | else: |
| 46 | self.base_dir = assemble_project_path(os.path.join(config.workdir, "benchmark")) |
| 47 | logger.info(f"| 📁 Benchmark context manager base directory: {self.base_dir}.") |
| 48 | os.makedirs(self.base_dir, exist_ok=True) |
| 49 | |
| 50 | if save_path is not None: |
| 51 | self.save_path = assemble_project_path(save_path) |
| 52 | else: |
| 53 | self.save_path = os.path.join(self.base_dir, "benchmark.json") |
| 54 | logger.info(f"| 📁 Benchmark context manager save path: {self.save_path}.") |
| 55 | |
| 56 | if contract_path is not None: |
| 57 | self.contract_path = assemble_project_path(contract_path) |
| 58 | else: |
| 59 | self.contract_path = os.path.join(self.base_dir, "contract.md") |
| 60 | logger.info(f"| 📁 Benchmark context manager contract path: {self.contract_path}.") |
| 61 | |
| 62 | self._benchmark_configs: Dict[str, BenchmarkConfig] = {} # Current active configs (latest version) |
| 63 | self._benchmark_history_versions: Dict[str, Dict[str, BenchmarkConfig]] = {} |
| 64 | |
| 65 | self._cleanup_registered = False |
| 66 | |
| 67 | async def initialize(self, benchmark_names: Optional[List[str]] = None): |
| 68 | """Initialize the benchmark context manager.""" |
| 69 | # Register benchmark-related symbols for auto-injection in dynamic code |
| 70 | dynamic_manager.register_symbol("BENCHMARK", BENCHMARK) |
| 71 | dynamic_manager.register_symbol("Benchmark", Benchmark) |
| 72 | |
| 73 | # Register benchmark context provider for automatic import injection |
| 74 | def benchmark_context_provider(): |
| 75 | """Provide benchmark-related imports for dynamic benchmark classes.""" |
| 76 | return { |
| 77 | "BENCHMARK": BENCHMARK, |
| 78 | "Benchmark": Benchmark, |