|
| 1 | +import json |
| 2 | +import logging |
| 3 | +from dataclasses import asdict |
| 4 | +from functools import cached_property |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +from ramalama.benchmarks.errors import MissingStorageFolderError |
| 8 | +from ramalama.benchmarks.schemas import BenchmarkRecord, DeviceInfoV1, get_benchmark_record |
| 9 | +from ramalama.benchmarks.utilities import parse_jsonl |
| 10 | +from ramalama.config import CONFIG |
| 11 | +from ramalama.log_levels import LogLevel |
| 12 | + |
| 13 | +logger = logging.getLogger("ramalama.benchmarks") |
| 14 | +logger.setLevel(CONFIG.log_level or LogLevel.WARNING) |
| 15 | + |
| 16 | +SCHEMA_VERSION = 1 |
| 17 | +BENCHMARKS_FILENAME = "benchmarks.jsonl" |
| 18 | + |
| 19 | + |
| 20 | +class BenchmarksManager: |
| 21 | + def __init__(self, storage_folder: str | Path | None): |
| 22 | + if storage_folder is None: |
| 23 | + raise MissingStorageFolderError |
| 24 | + |
| 25 | + self.storage_folder = Path(storage_folder) |
| 26 | + self.storage_file = self.storage_folder / BENCHMARKS_FILENAME |
| 27 | + self.storage_file.parent.mkdir(parents=True, exist_ok=True) |
| 28 | + |
| 29 | + @cached_property |
| 30 | + def device_info(self) -> DeviceInfoV1: |
| 31 | + return DeviceInfoV1.current_device_info() |
| 32 | + |
| 33 | + def save(self, results: list[BenchmarkRecord] | BenchmarkRecord): |
| 34 | + if not isinstance(results, list): |
| 35 | + results = [results] |
| 36 | + |
| 37 | + if len(results) == 0: |
| 38 | + return |
| 39 | + |
| 40 | + with self.storage_file.open("a", encoding="utf-8") as handle: |
| 41 | + for record in results: |
| 42 | + handle.write(json.dumps(asdict(record), ensure_ascii=True)) |
| 43 | + handle.write("\n") |
| 44 | + |
| 45 | + def list(self) -> list[BenchmarkRecord]: |
| 46 | + """List benchmark results from JSONL storage.""" |
| 47 | + if not self.storage_file.exists(): |
| 48 | + return [] |
| 49 | + content = self.storage_file.read_text(encoding="utf-8") |
| 50 | + return [get_benchmark_record(result) for result in parse_jsonl(content)] |
0 commit comments