diff --git a/.semversioner/next-release/patch-20260911141730000000.json b/.semversioner/next-release/patch-20260911141730000000.json new file mode 100644 index 0000000..a367ccd --- /dev/null +++ b/.semversioner/next-release/patch-20260911141730000000.json @@ -0,0 +1,4 @@ +{ + "type": "patch", + "description": "Make evaluation caches concurrency-safe and record reproducibility metadata" +} diff --git a/benchmark_qed/__main__.py b/benchmark_qed/__main__.py index 3f5f9c5..979fddd 100644 --- a/benchmark_qed/__main__.py +++ b/benchmark_qed/__main__.py @@ -8,6 +8,7 @@ from benchmark_qed.autoe.cli import app as autoe_cli from benchmark_qed.autoq.cli import app as autoq_cli +from benchmark_qed.cli.cache import app as cache_cli from benchmark_qed.cli.init_config import app as init_cli from benchmark_qed.cli.interactive import interactive_init from benchmark_qed.data.cli import app as data_cli @@ -16,6 +17,7 @@ app.add_typer(autoe_cli, name="autoe", help="Relative scores CLI.") app.add_typer(autoq_cli, help="Question generation CLI.") +app.add_typer(cache_cli, name="cache", help="Local cache inspection CLI.") app.add_typer(init_cli, name="config", help="Configuration initialization CLI.") app.command(name="init", help="Interactively create a benchmark-qed configuration.")( interactive_init diff --git a/benchmark_qed/autoe/chunk_assertion/__init__.py b/benchmark_qed/autoe/chunk_assertion/__init__.py index d56d47e..d1c0f12 100644 --- a/benchmark_qed/autoe/chunk_assertion/__init__.py +++ b/benchmark_qed/autoe/chunk_assertion/__init__.py @@ -4,7 +4,10 @@ from benchmark_qed.autoe.chunk_assertion.aggregation import summarize_at_k from benchmark_qed.autoe.chunk_assertion.cache import ( ContentAddressedCache, + build_cache_metadata, compute_cache_key, + compute_config_fingerprint, + compute_logical_key, ) from benchmark_qed.autoe.chunk_assertion.scoring import ( run_assertion_eval_chunk_mode, @@ -12,7 +15,10 @@ __all__ = [ "ContentAddressedCache", + "build_cache_metadata", "compute_cache_key", + "compute_config_fingerprint", + "compute_logical_key", "run_assertion_eval_chunk_mode", "summarize_at_k", ] diff --git a/benchmark_qed/autoe/chunk_assertion/cache.py b/benchmark_qed/autoe/chunk_assertion/cache.py index 2cbd790..151308b 100644 --- a/benchmark_qed/autoe/chunk_assertion/cache.py +++ b/benchmark_qed/autoe/chunk_assertion/cache.py @@ -9,90 +9,224 @@ from pathlib import Path from typing import Any +from benchmark_qed.cache import ( + SQLiteCache, + changed_configuration_fields, + redact_sensitive_values, + stable_fingerprint, +) + log: logging.Logger = logging.getLogger(__name__) +_CACHE_NAMESPACE = "chunk_assertion" +_CACHE_SCHEMA_VERSION = 1 +_DEFAULT_LEASE_TTL_SECONDS = 120.0 + + +def build_cache_metadata( + *, + model: str = "", + call_args: dict[str, Any] | None = None, + system_prompt: str = "", + user_prompt: str = "", +) -> dict[str, Any]: + """Build inspectable, credential-safe metadata for a cached judgement.""" + return { + "schema_version": _CACHE_SCHEMA_VERSION, + "evaluator": _CACHE_NAMESPACE, + "model": model, + "call_args": redact_sensitive_values(call_args or {}), + "system_prompt": system_prompt, + "user_prompt": user_prompt, + } + + +def compute_logical_key(assertion_text: str, chunk_content: str) -> str: + """Fingerprint the inputs independently from judge configuration.""" + return stable_fingerprint({ + "assertion": assertion_text, + "chunk": chunk_content, + }) + + +def compute_config_fingerprint( + *, + model: str = "", + call_args: dict[str, Any] | None = None, + system_prompt: str = "", + user_prompt: str = "", +) -> str: + """Fingerprint the judge configuration independently from inputs.""" + return stable_fingerprint( + build_cache_metadata( + model=model, + call_args=call_args, + system_prompt=system_prompt, + user_prompt=user_prompt, + ) + ) + class ContentAddressedCache: - r"""Persistent SHA256-based cache for (assertion, chunk) -> grade. + r"""Persistent SQLite cache for (assertion, chunk) -> grade. - Key is computed by :func:`compute_cache_key`, which hashes the assertion, - chunk, judge model configuration, and prompt templates so that changes to - the model or prompts invalidate stale entries while identical configurations - reuse cached grades across runs. + SQLite WAL mode allows readers and writers in separate processes to share + the cache safely. Writes use first-writer-wins semantics for a cache key. + Existing JSONL caches are imported once when the database is initialized. """ - def __init__(self, cache_path: Path | str) -> None: - """Initialize cache at given path. - - Args: - cache_path: Path to JSONL cache file. Parent directory is created if needed. - """ - self.cache_path: Path = Path(cache_path) - self.cache_path.parent.mkdir(parents=True, exist_ok=True) - self._data: dict[str, str] = {} + def __init__( + self, + cache_path: Path | str, + *, + lease_ttl_seconds: float = _DEFAULT_LEASE_TTL_SECONDS, + ) -> None: + """Initialize the cache and import a legacy JSONL cache when present.""" + requested_path = Path(cache_path) + self.legacy_cache_path: Path | None = None + if requested_path.suffix.lower() == ".jsonl": + self.legacy_cache_path = requested_path + self.cache_path = requested_path.with_suffix(".sqlite3") + else: + self.cache_path = requested_path + legacy_path = requested_path.with_suffix(".jsonl") + if legacy_path.exists(): + self.legacy_cache_path = legacy_path + + self._store = SQLiteCache(self.cache_path, _CACHE_NAMESPACE) + self.lease_ttl_seconds = lease_ttl_seconds + self._pending: dict[str, tuple[str, dict[str, Any]]] = {} self.new_count: int = 0 - self._load() + self._migrate_legacy_cache() - def _load(self) -> None: - """Load existing cache from disk.""" - if not self.cache_path.exists(): + def _migrate_legacy_cache(self) -> None: + """Import an existing JSONL cache once without modifying the source.""" + if self.legacy_cache_path is None or not self.legacy_cache_path.exists(): return + + migration_name = f"legacy_import:{self.legacy_cache_path.resolve()}" + if self._store.get_property(migration_name) is not None: + return + + records: list[tuple[str, str, dict[str, Any]]] = [] try: - with self.cache_path.open(encoding="utf-8") as f: - for line in f: - stripped = line.strip() - if not stripped: - continue + with self.legacy_cache_path.open(encoding="utf-8") as legacy_file: + for line in legacy_file: try: - record = json.loads(stripped) - cache_key = record.get("key") - grade = record.get("grade") - if cache_key and grade: - self._data[cache_key] = grade + record = json.loads(line) except json.JSONDecodeError: continue + cache_key = record.get("key") + grade = record.get("grade") + if cache_key and grade: + records.append(( + str(cache_key), + str(grade), + { + "schema_version": _CACHE_SCHEMA_VERSION, + "source": "legacy_jsonl", + }, + )) except OSError as exc: - log.warning("Failed to load cache from %s: %s", self.cache_path, exc) + log.warning( + "Failed to import legacy cache from %s: %s", self.legacy_cache_path, exc + ) + return + + self._store.put_many(records) + self._store.set_property_if_absent(migration_name, "complete") def get(self, cache_key: str) -> str | None: - """Retrieve grade for cache key. - - Args: - cache_key: SHA256 hash of (assertion + chunk) - - Returns - ------- - Grade string ('full_support', 'partial_support', 'no_support') or None - """ - return self._data.get(cache_key) - - def put(self, cache_key: str, grade: str) -> None: - """Store grade for cache key. - - Args: - cache_key: SHA256 hash of (assertion + chunk) - grade: Grade string ('full_support', 'partial_support', 'no_support') - """ - if cache_key not in self._data: - self._data[cache_key] = grade - self.new_count += 1 - - def flush(self) -> None: - """Write cache contents to disk. - - Rewrites the full JSONL file atomically to avoid duplicating existing - entries when flushing incremental updates. - """ - if self.new_count == 0: + """Retrieve a grade for a cache key.""" + pending = self._pending.get(cache_key) + if pending is not None: + return pending[0] + entry = self._store.get(cache_key) + return str(entry[0]) if entry is not None else None + + def get_metadata(self, cache_key: str) -> dict[str, Any] | None: + """Retrieve the inspectable metadata stored with a cache entry.""" + pending = self._pending.get(cache_key) + if pending is not None: + return pending[1] + entry = self._store.get(cache_key) + return entry[1] if entry is not None else None + + def put( + self, + cache_key: str, + grade: str, + metadata: dict[str, Any] | None = None, + ) -> None: + """Stage a grade for insertion on the next flush.""" + if cache_key in self._pending or self.get(cache_key) is not None: return - tmp_path = self.cache_path.with_suffix(self.cache_path.suffix + ".tmp") - with tmp_path.open("w", encoding="utf-8") as f: - f.writelines( - json.dumps({"key": cache_key, "grade": grade}) + "\n" - for cache_key, grade in self._data.items() - ) - tmp_path.replace(self.cache_path) + self._pending[cache_key] = (grade, metadata or {}) + self.new_count = len(self._pending) + + def claim(self, cache_key: str, owner_id: str) -> bool: + """Claim an uncached judgement for this worker.""" + return self._store.try_acquire( + cache_key, + owner_id, + ttl_seconds=self.lease_ttl_seconds, + ) + + def renew(self, cache_keys: list[str], owner_id: str) -> int: + """Renew active judgement leases.""" + return self._store.renew_leases( + cache_keys, + owner_id, + ttl_seconds=self.lease_ttl_seconds, + ) + + def release(self, cache_keys: list[str], owner_id: str) -> None: + """Release judgement leases without publishing.""" + self._store.release_leases(cache_keys, owner_id) + + def publish( + self, + cache_key: str, + grade: str, + metadata: dict[str, Any], + *, + owner_id: str, + logical_key: str, + config_fingerprint: str, + ) -> bool: + """Publish a grade and release its lease atomically.""" + return self._store.publish( + cache_key, + grade, + metadata, + owner_id=owner_id, + logical_key=logical_key, + config_fingerprint=config_fingerprint, + ) + + def find_configuration_mismatches( + self, + identities: list[tuple[str, str]], + current_metadata: dict[str, Any], + ) -> tuple[int, list[str]]: + """Return the mismatch count and changed configuration fields.""" + alternatives = self._store.find_alternate_configurations(identities) + changed_fields = { + field + for metadata in alternatives + for field in changed_configuration_fields(current_metadata, metadata) + } + return len(alternatives), sorted(changed_fields) + + def flush(self) -> int: + """Atomically insert staged entries and return the number persisted.""" + inserted_count = self._store.put_many([ + (cache_key, grade, metadata) + for cache_key, (grade, metadata) in self._pending.items() + ]) + self._pending.clear() self.new_count = 0 + return inserted_count def compute_cache_key( @@ -104,30 +238,12 @@ def compute_cache_key( system_prompt: str = "", user_prompt: str = "", ) -> str: - """Compute stable SHA256 cache key for an (assertion, chunk) judgement. - - The judge model identifier, call arguments (for example temperature and - seed), and the prompt templates are folded into the key so that changing - the model or prompts invalidates stale cache entries instead of returning - grades produced under a different configuration. - - Args: - assertion_text: Assertion statement - chunk_content: Chunk/passage text - model: Judge model identifier (for example ``gpt-4.1``) - call_args: Additional LLM call arguments (for example temperature, seed) - system_prompt: System prompt template used for judging - user_prompt: User prompt template used for judging - - Returns - ------- - SHA256 hexdigest - """ + """Compute a stable SHA256 key for an (assertion, chunk) judgement.""" payload = { "assertion": assertion_text, "chunk": chunk_content, "model": model, - "call_args": call_args or {}, + "call_args": redact_sensitive_values(call_args or {}), "system_prompt": system_prompt, "user_prompt": user_prompt, } diff --git a/benchmark_qed/autoe/chunk_assertion/scoring.py b/benchmark_qed/autoe/chunk_assertion/scoring.py index 6109a08..5d788fa 100644 --- a/benchmark_qed/autoe/chunk_assertion/scoring.py +++ b/benchmark_qed/autoe/chunk_assertion/scoring.py @@ -4,10 +4,12 @@ from __future__ import annotations import asyncio +import contextlib import json import logging import operator import re +import uuid from pathlib import Path from typing import TYPE_CHECKING, Any @@ -16,7 +18,10 @@ from benchmark_qed.autoe.chunk_assertion.aggregation import summarize_at_k from benchmark_qed.autoe.chunk_assertion.cache import ( ContentAddressedCache, + build_cache_metadata, compute_cache_key, + compute_config_fingerprint, + compute_logical_key, ) from benchmark_qed.autoe.data_model.chunk_assertion import ( ChunkAssertionGrade, @@ -33,6 +38,9 @@ log: logging.Logger = logging.getLogger(__name__) +ChunkTarget = tuple[int, int, int] +ChunkWork = tuple[str, str, str, str, str, list[ChunkTarget]] + def _collect_chunks(eval_result: RetrievalResult) -> list[dict[str, Any]]: """Flatten one question's retrieval into an ordered chunk list. @@ -190,16 +198,18 @@ def _extract_grade(response_text: str) -> str: async def _evaluate_uncached_pairs( - uncached_work: list[tuple[str, str, str, int, int, int]], + uncached_work: list[ChunkWork], *, llm_client: LLMCompletion, llm_config: LLMConfig, system_prompt: str, user_prompt: str, + cache_metadata: dict[str, Any], cache: ContentAddressedCache, + lease_owner: str, per_chunk_grades: dict[tuple[int, int], list[tuple[int, float, str]]], assertion_call_stats: dict[tuple[int, int], dict[str, int]], -) -> int: +) -> tuple[int, int]: """Judge uncached (assertion, chunk) pairs via a bounded worker pool. Uses a shared queue with at most ``concurrent_requests`` worker tasks so @@ -208,40 +218,60 @@ async def _evaluate_uncached_pairs( shared counters and grade maps are mutated safely between awaits. Args: - uncached_work: Work items of (assertion, chunk, cache_key, q_idx, a_idx, rank). + uncached_work: Unique work items with all result targets for each cache key. llm_client: LLM client for judging. llm_config: LLM configuration. system_prompt: System prompt template. user_prompt: User prompt template. + cache_metadata: Reproducibility metadata stored with successful grades. cache: Cache to store new grades in. + lease_owner: Unique owner of the work-item leases. per_chunk_grades: Mutated in place with (rank, score, grade) per (q, a). assertion_call_stats: Mutated in place with successful/failed call counts. Returns ------- - Number of LLM calls that failed. + Number of logical pairs that failed and number of cache entries inserted. """ - n_total = len(uncached_work) + n_total = sum(len(targets) for *_inputs, targets in uncached_work) next_progress_threshold = max(1, n_total // 10) call_errors = 0 + inserted_count = 0 n_done = 0 - work_queue: asyncio.Queue[tuple[str, str, str, int, int, int]] = asyncio.Queue() + work_queue: asyncio.Queue[ChunkWork] = asyncio.Queue() for work_item in uncached_work: work_queue.put_nowait(work_item) def _record_result( - work_item: tuple[str, str, str, int, int, int], grade: str + work_item: ChunkWork, + grade: str, ) -> None: - nonlocal call_errors, n_done, next_progress_threshold - _assertion_text, _chunk_content, cache_key, q_idx, a_idx, rank = work_item - - q_key = (q_idx, a_idx) + nonlocal call_errors, inserted_count, n_done, next_progress_threshold + ( + _assertion_text, + _chunk_content, + cache_key, + logical_key, + config_fingerprint, + targets, + ) = work_item if grade == "__error__": - call_errors += 1 - assertion_call_stats[q_key]["failed"] += 1 + call_errors += len(targets) + cache.release([cache_key], lease_owner) + for q_idx, a_idx, _rank in targets: + assertion_call_stats[(q_idx, a_idx)]["failed"] += 1 else: - cache.put(cache_key, grade) + inserted_count += int( + cache.publish( + cache_key, + grade, + cache_metadata, + owner_id=lease_owner, + logical_key=logical_key, + config_fingerprint=config_fingerprint, + ) + ) mapped_grade = ( grade if grade @@ -252,14 +282,16 @@ def _record_result( } else ChunkAssertionGrade.NO_SUPPORT ) - per_chunk_grades[q_key].append(( - rank, - grade_to_score(grade), - mapped_grade, - )) - assertion_call_stats[q_key]["successful"] += 1 - - n_done += 1 + for q_idx, a_idx, rank in targets: + q_key = (q_idx, a_idx) + per_chunk_grades[q_key].append(( + rank, + grade_to_score(grade), + mapped_grade, + )) + assertion_call_stats[q_key]["successful"] += 1 + + n_done += len(targets) if n_done >= next_progress_threshold: rich_print( f" LLM judging progress: {n_done}/{n_total} pairs ({(100.0 * n_done / n_total):.1f}%)", @@ -272,7 +304,14 @@ async def _worker() -> None: work_item = work_queue.get_nowait() except asyncio.QueueEmpty: return - assertion_text, chunk_content, _cache_key, _q_idx, _a_idx, _rank = work_item + ( + assertion_text, + chunk_content, + _cache_key, + _logical_key, + _config_fingerprint, + _targets, + ) = work_item grade, _reasoning = await _label_chunk_assertion( llm_client, assertion_text, @@ -283,10 +322,96 @@ async def _worker() -> None: ) _record_result(work_item, grade) - n_workers = max(1, min(llm_config.concurrent_requests, n_total)) - await asyncio.gather(*[asyncio.ensure_future(_worker()) for _ in range(n_workers)]) + async def _heartbeat() -> None: + cache_keys = [work_item[2] for work_item in uncached_work] + while True: + await asyncio.sleep(max(0.1, cache.lease_ttl_seconds / 3)) + cache.renew(cache_keys, lease_owner) - return call_errors + n_workers = max(1, min(llm_config.concurrent_requests, len(uncached_work))) + heartbeat = asyncio.create_task(_heartbeat()) + try: + await asyncio.gather(*[ + asyncio.ensure_future(_worker()) for _ in range(n_workers) + ]) + except (Exception, asyncio.CancelledError): + cache.release([work_item[2] for work_item in uncached_work], lease_owner) + raise + finally: + heartbeat.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat + + return call_errors, inserted_count + + +async def _evaluate_with_leases( + uncached_work: list[ChunkWork], + *, + llm_client: LLMCompletion, + llm_config: LLMConfig, + system_prompt: str, + user_prompt: str, + cache_metadata: dict[str, Any], + cache: ContentAddressedCache, + per_chunk_grades: dict[tuple[int, int], list[tuple[int, float, str]]], + assertion_call_stats: dict[tuple[int, int], dict[str, int]], +) -> tuple[int, int, int]: + """Evaluate owned work and reuse results published by other processes.""" + lease_owner = uuid.uuid4().hex + owned_work: list[ChunkWork] = [] + waiting_work: dict[str, ChunkWork] = {} + for work_item in uncached_work: + if cache.claim(work_item[2], lease_owner): + owned_work.append(work_item) + else: + waiting_work[work_item[2]] = work_item + + call_errors = 0 + inserted_count = 0 + concurrent_hits = 0 + while owned_work or waiting_work: + if owned_work: + failed, inserted = await _evaluate_uncached_pairs( + owned_work, + llm_client=llm_client, + llm_config=llm_config, + system_prompt=system_prompt, + user_prompt=user_prompt, + cache_metadata=cache_metadata, + cache=cache, + lease_owner=lease_owner, + per_chunk_grades=per_chunk_grades, + assertion_call_stats=assertion_call_stats, + ) + call_errors += failed + inserted_count += inserted + owned_work = [] + + completed_keys: list[str] = [] + for cache_key, work_item in waiting_work.items(): + cached_grade = cache.get(cache_key) + if cached_grade is not None: + targets = work_item[5] + for q_idx, a_idx, rank in targets: + q_key = (q_idx, a_idx) + per_chunk_grades[q_key].append(( + rank, + grade_to_score(cached_grade), + cached_grade, + )) + assertion_call_stats[q_key]["successful"] += 1 + concurrent_hits += len(targets) + completed_keys.append(cache_key) + elif cache.claim(cache_key, lease_owner): + owned_work.append(work_item) + completed_keys.append(cache_key) + for cache_key in completed_keys: + waiting_work.pop(cache_key) + if waiting_work and not owned_work: + await asyncio.sleep(0.1) + + return call_errors, inserted_count, concurrent_hits async def run_assertion_eval_chunk_mode( @@ -327,13 +452,25 @@ async def run_assertion_eval_chunk_mode( Dict of {label: EvalSummary} keyed by f"k{k}" plus "all" entry """ if cache_path is None: - cache_path = Path.cwd() / ".benchmark_qed_cache" / "chunk_assertions.jsonl" + cache_path = Path.cwd() / ".benchmark_qed_cache" / "chunk_assertions.sqlite3" cache = ContentAddressedCache(cache_path) + cache_metadata = build_cache_metadata( + model=llm_config.model, + call_args=llm_config.call_args, + system_prompt=system_prompt, + user_prompt=user_prompt, + ) + config_fingerprint = compute_config_fingerprint( + model=llm_config.model, + call_args=llm_config.call_args, + system_prompt=system_prompt, + user_prompt=user_prompt, + ) # Track per-(q, a) chunk grades: (question_idx, assertion_idx) -> [(rank, score, grade), ...] per_chunk_grades: dict[tuple[int, int], list[tuple[int, float, str]]] = {} assertion_call_stats: dict[tuple[int, int], dict[str, int]] = {} - uncached_work: list[tuple[str, str, str, int, int, int]] = [] + uncached_by_key: dict[str, tuple[str, str, str, str, list[ChunkTarget]]] = {} retrieved_chunk_counts: list[int] = [] total_checks = 0 @@ -389,6 +526,7 @@ async def run_assertion_eval_chunk_mode( system_prompt=system_prompt, user_prompt=user_prompt, ) + logical_key = compute_logical_key(assertion_text, chunk_content) cached_grade = cache.get(cache_key) if cached_grade is not None: cache_hits += 1 @@ -399,14 +537,17 @@ async def run_assertion_eval_chunk_mode( )) assertion_call_stats[q_key]["successful"] += 1 else: - uncached_work.append(( - assertion_text, - chunk_content, - cache_key, - q_idx, - a_idx, - rank, - )) + existing_work = uncached_by_key.get(cache_key) + if existing_work is None: + uncached_by_key[cache_key] = ( + assertion_text, + chunk_content, + logical_key, + config_fingerprint, + [(q_idx, a_idx, rank)], + ) + else: + existing_work[4].append((q_idx, a_idx, rank)) if positional_fallbacks: rich_print( @@ -425,24 +566,62 @@ async def run_assertion_eval_chunk_mode( rich_print(" No assertions or chunks to evaluate.") # Second pass: evaluate uncached pairs - if uncached_work: + if uncached_by_key: + uncached_work = [ + ( + assertion_text, + chunk_content, + cache_key, + logical_key, + work_config_fingerprint, + targets, + ) + for cache_key, ( + assertion_text, + chunk_content, + logical_key, + work_config_fingerprint, + targets, + ) in uncached_by_key.items() + ] + mismatch_count, changed_fields = cache.find_configuration_mismatches( + [ + (logical_key, work_config_fingerprint) + for ( + _assertion, + _chunk, + _key, + logical_key, + work_config_fingerprint, + _targets, + ) in uncached_work + ], + cache_metadata, + ) + if mismatch_count: + changed_text = ", ".join(changed_fields) or "configuration metadata" + rich_print( + f" [yellow]WARNING: Found {mismatch_count} cached result(s) for " + f"the same inputs with different {changed_text}; they will not " + "be reused.[/yellow]" + ) rich_print( f" Running {len(uncached_work)} LLM judgements for uncached (assertion, chunk) pairs...", ) - call_errors += await _evaluate_uncached_pairs( + failed, inserted_count, concurrent_hits = await _evaluate_with_leases( uncached_work, llm_client=llm_client, llm_config=llm_config, system_prompt=system_prompt, user_prompt=user_prompt, + cache_metadata=cache_metadata, cache=cache, per_chunk_grades=per_chunk_grades, assertion_call_stats=assertion_call_stats, ) - - # Flush cache - cache.flush() - rich_print(f" Cached {len(uncached_work)} new entries to {cache_path.name}") + call_errors += failed + cache_hits += concurrent_hits + rich_print(f" Cached {inserted_count} new entries to {cache.cache_path.name}") if call_errors: rich_print(f" WARNING: {call_errors} LLM calls failed") diff --git a/benchmark_qed/autoe/cli.py b/benchmark_qed/autoe/cli.py index 6354f62..03d23c3 100644 --- a/benchmark_qed/autoe/cli.py +++ b/benchmark_qed/autoe/cli.py @@ -2079,9 +2079,9 @@ def chunk_assertion_scores( # Set cache directory cache_path = None if config.cache_dir: - cache_path = Path(config.cache_dir) / "chunk_assertions.jsonl" + cache_path = Path(config.cache_dir) / "chunk_assertions.sqlite3" else: - cache_path = Path.cwd() / ".benchmark_qed_cache" / "chunk_assertions.jsonl" + cache_path = Path.cwd() / ".benchmark_qed_cache" / "chunk_assertions.sqlite3" # Run chunk-level evaluation rich_print("\n[bold]Running chunk-level assertion evaluation...[/bold]") diff --git a/benchmark_qed/autoe/retrieval_metrics/relevance_assessment/base.py b/benchmark_qed/autoe/retrieval_metrics/relevance_assessment/base.py index 65679fe..3315f56 100644 --- a/benchmark_qed/autoe/retrieval_metrics/relevance_assessment/base.py +++ b/benchmark_qed/autoe/retrieval_metrics/relevance_assessment/base.py @@ -1,10 +1,12 @@ # Copyright (c) 2025 Microsoft Corporation. """Base classes for relevance assessment.""" +import asyncio import contextlib import hashlib import json import logging +import uuid from abc import ABC, abstractmethod from datetime import UTC, datetime from pathlib import Path @@ -15,8 +17,15 @@ RelevanceAssessmentItem, RelevanceAssessmentResponse, ) +from benchmark_qed.cache import ( + SQLiteCache, + changed_configuration_fields, + redact_sensitive_values, + stable_fingerprint, +) log: logging.Logger = logging.getLogger(__name__) +_LEASE_TTL_SECONDS = 120.0 class RelevanceRater(ABC): @@ -36,9 +45,14 @@ def __init__( self.cache_enabled: bool = cache_enabled and cache_dir is not None self.cache_hits: int = 0 self.cache_misses: int = 0 + self._cache_store: SQLiteCache | None = None if self.cache_enabled and self.cache_dir: - self.cache_dir.mkdir(parents=True, exist_ok=True) + self._cache_store = SQLiteCache( + self.cache_dir / "relevance_cache.sqlite3", + namespace=self.__class__.__name__, + ) + self._migrate_legacy_cache_files() async def rate_relevance( self, query: str, text_units: list[TextUnit] @@ -59,9 +73,8 @@ async def rate_relevance( return await self._rate_relevance_impl(query, text_units) # With caching enabled, check each text unit individually - cached_assessments = [] - uncached_text_units = [] - uncached_indices = [] + cached_assessments: list[tuple[int, RelevanceAssessmentItem]] = [] + uncached_by_key: dict[str, tuple[TextUnit, list[int]]] = {} rater_params = self._get_cache_relevant_params() @@ -73,36 +86,171 @@ async def rate_relevance( cached_assessments.append((i, cached_assessment)) self.cache_hits += 1 else: - uncached_text_units.append(text_unit) - uncached_indices.append(i) self.cache_misses += 1 + pending = uncached_by_key.get(cache_key) + if pending is None: + uncached_by_key[cache_key] = (text_unit, [i]) + else: + pending[1].append(i) # Process uncached text units if any - uncached_results = [] - if uncached_text_units: - uncached_response = await self._rate_relevance_impl( - query, uncached_text_units - ) - uncached_results = uncached_response.assessment - - # Cache individual results - for j, text_unit in enumerate(uncached_text_units): - cache_key = self._generate_cache_key(query, text_unit, rater_params) - self._save_to_cache(cache_key, query, uncached_results[j]) + if uncached_by_key: + if self._cache_store is None: + msg = "Caching is enabled but the cache store is unavailable" + raise RuntimeError(msg) + + config_fingerprint = self._generate_config_fingerprint(rater_params) + current_metadata = self._build_cache_metadata(query, rater_params) + alternatives = self._cache_store.find_alternate_configurations([ + ( + self._generate_logical_key(query, text_unit), + config_fingerprint, + ) + for text_unit, _indices in uncached_by_key.values() + ]) + if alternatives: + changed_fields = sorted({ + field + for metadata in alternatives + for field in changed_configuration_fields( + current_metadata, metadata + ) + }) + log.warning( + "Found %d cached relevance result(s) for the same inputs " + "with different %s; they will not be reused", + len(alternatives), + ", ".join(changed_fields) or "configuration metadata", + ) + + lease_owner = uuid.uuid4().hex + owned = { + cache_key: item + for cache_key, item in uncached_by_key.items() + if self._cache_store.try_acquire( + cache_key, + lease_owner, + ttl_seconds=_LEASE_TTL_SECONDS, + ) + } + waiting = { + cache_key: item + for cache_key, item in uncached_by_key.items() + if cache_key not in owned + } + + while owned or waiting: + if owned: + await self._evaluate_and_publish_claimed( + query=query, + claimed=owned, + rater_params=rater_params, + config_fingerprint=config_fingerprint, + lease_owner=lease_owner, + cached_assessments=cached_assessments, + ) + owned = {} + + completed_keys: list[str] = [] + recovered: dict[str, tuple[TextUnit, list[int]]] = {} + for cache_key, item in waiting.items(): + cached_assessment = self._load_from_cache(cache_key) + if cached_assessment is not None: + cached_assessments.extend( + (index, cached_assessment) for index in item[1] + ) + completed_keys.append(cache_key) + elif self._cache_store.try_acquire( + cache_key, + lease_owner, + ttl_seconds=_LEASE_TTL_SECONDS, + ): + recovered[cache_key] = item + completed_keys.append(cache_key) + for cache_key in completed_keys: + waiting.pop(cache_key) + owned = recovered + if waiting and not owned: + await asyncio.sleep(0.1) # Combine cached and uncached results in original order - all_assessments: list[RelevanceAssessmentItem] = [None] * len(text_units) # type: ignore - - # Place cached results - for original_idx, cached_assessment in cached_assessments: - all_assessments[original_idx] = cached_assessment - - # Place uncached results - for j, original_idx in enumerate(uncached_indices): - all_assessments[original_idx] = uncached_results[j] + assessments_by_index = dict(cached_assessments) + all_assessments = [ + assessments_by_index[index] for index in range(len(text_units)) + ] return RelevanceAssessmentResponse(assessment=all_assessments) + async def _evaluate_and_publish_claimed( + self, + *, + query: str, + claimed: dict[str, tuple[TextUnit, list[int]]], + rater_params: dict[str, Any], + config_fingerprint: str, + lease_owner: str, + cached_assessments: list[tuple[int, RelevanceAssessmentItem]], + ) -> None: + """Evaluate claimed units while renewing and atomically releasing leases.""" + if self._cache_store is None: + msg = "Cache store is unavailable" + raise RuntimeError(msg) + cache_store = self._cache_store + cache_keys = list(claimed) + + async def _heartbeat() -> None: + while True: + await asyncio.sleep(max(0.1, _LEASE_TTL_SECONDS / 3)) + cache_store.renew_leases( + cache_keys, + lease_owner, + ttl_seconds=_LEASE_TTL_SECONDS, + ) + + heartbeat = asyncio.create_task(_heartbeat()) + try: + response = await self._rate_relevance_impl( + query, + [text_unit for text_unit, _indices in claimed.values()], + ) + self._validate_result_count(response, len(claimed)) + + for ( + cache_key, + (text_unit, indices), + ), assessment in zip(claimed.items(), response.assessment, strict=True): + cache_store.publish( + cache_key, + assessment.model_dump(exclude={"text_unit": {"text_embedding"}}), + self._build_cache_metadata(query, rater_params), + owner_id=lease_owner, + logical_key=self._generate_logical_key(query, text_unit), + config_fingerprint=config_fingerprint, + ) + cached_assessments.extend( + (original_idx, assessment) for original_idx in indices + ) + except (Exception, asyncio.CancelledError): + cache_store.release_leases(cache_keys, lease_owner) + raise + finally: + heartbeat.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat + + @staticmethod + def _validate_result_count( + response: RelevanceAssessmentResponse, expected_count: int + ) -> None: + """Raise when a rater violates its one-result-per-input contract.""" + if len(response.assessment) != expected_count: + msg = ( + "Relevance rater returned " + f"{len(response.assessment)} results for " + f"{expected_count} uncached text units" + ) + raise RuntimeError(msg) + @abstractmethod async def _rate_relevance_impl( self, query: str, text_units: list[TextUnit] @@ -127,56 +275,89 @@ def _generate_cache_key( "query": query.strip().lower(), "text_content": text_unit.text.strip().lower(), "rater_type": self.__class__.__name__, - "rater_params": rater_params, + "rater_params": redact_sensitive_values(rater_params), } - content_str = json.dumps(cache_data, sort_keys=True) + content_str = json.dumps(cache_data, sort_keys=True, default=str) return hashlib.sha256(content_str.encode()).hexdigest() + def _generate_logical_key(self, query: str, text_unit: TextUnit) -> str: + """Fingerprint query and text independently from rater configuration.""" + return stable_fingerprint({ + "query": query.strip().lower(), + "text_content": text_unit.text.strip().lower(), + "rater_type": self.__class__.__name__, + }) + + @staticmethod + def _generate_config_fingerprint(rater_params: dict[str, Any]) -> str: + """Fingerprint relevance settings independently from query and text.""" + return stable_fingerprint(rater_params) + + def _build_cache_metadata( + self, query: str, rater_params: dict[str, Any] + ) -> dict[str, Any]: + """Build inspectable metadata for a relevance assessment.""" + return { + "timestamp": datetime.now(tz=UTC).isoformat(), + "query": query.strip().lower(), + "rater_type": self.__class__.__name__, + "rater_params": redact_sensitive_values(rater_params), + } + def _load_from_cache(self, cache_key: str) -> RelevanceAssessmentItem | None: """Load cached result for a single text unit if available.""" - if not self.cache_enabled or not self.cache_dir: - return None - - cache_file = self.cache_dir / f"{cache_key}.json" - if not cache_file.exists(): + if self._cache_store is None: return None try: - with cache_file.open(encoding="utf-8") as f: - data = json.load(f) - return RelevanceAssessmentItem(**data["assessment_item"]) - except (OSError, json.JSONDecodeError, KeyError, TypeError): - # If cache file is corrupted, ignore and continue + entry = self._cache_store.get(cache_key) + if entry is None: + return None + assessment_data, _metadata = entry + return RelevanceAssessmentItem(**assessment_data) + except (json.JSONDecodeError, TypeError, ValueError) as exc: + log.warning("Ignoring invalid relevance cache entry %s: %s", cache_key, exc) return None - def _save_to_cache( - self, cache_key: str, query: str, assessment_item: RelevanceAssessmentItem - ) -> None: - """Save single text unit assessment result to cache.""" - if not self.cache_enabled or not self.cache_dir: + def _migrate_legacy_cache_files(self) -> None: + """Import legacy per-key JSON files once without deleting them.""" + if self._cache_store is None or self.cache_dir is None: + return + migration_name = f"legacy_relevance_import:{self.__class__.__name__}" + if self._cache_store.get_property(migration_name) is not None: return - cache_file = self.cache_dir / f"{cache_key}.json" - - # Create a copy for serialization to avoid mutating the original - # We exclude the embedding to avoid Pydantic warnings about numpy arrays - assessment_data = assessment_item.model_dump( - exclude={"text_unit": {"text_embedding"}} - ) - - cache_data = { - "timestamp": datetime.now(tz=UTC).isoformat(), - "query": query.strip().lower(), - "assessment_item": assessment_data, - } - - try: - with cache_file.open("w", encoding="utf-8") as f: - json.dump(cache_data, f, indent=2) - except OSError: - # If we can't write cache, continue without it - log.debug("Failed to write cache file: %s", cache_file) + entries: list[tuple[str, Any, dict[str, Any]]] = [] + for cache_file in self.cache_dir.glob("*.json"): + try: + with cache_file.open(encoding="utf-8") as file: + data = json.load(file) + assessment_data = data["assessment_item"] + RelevanceAssessmentItem(**assessment_data) + except ( + OSError, + json.JSONDecodeError, + KeyError, + TypeError, + ValueError, + ) as exc: + log.warning( + "Could not import legacy relevance cache file %s: %s", + cache_file, + exc, + ) + continue + metadata = { + "source": "legacy_json", + "timestamp": data.get("timestamp"), + "query": data.get("query"), + "rater_type": self.__class__.__name__, + } + entries.append((cache_file.stem, assessment_data, metadata)) + + self._cache_store.put_many(entries) + self._cache_store.set_property_if_absent(migration_name, "complete") def _get_cache_relevant_params(self) -> dict[str, Any]: """ @@ -233,16 +414,12 @@ def get_cache_stats(self) -> dict[str, Any]: total_requests = self.cache_hits + self.cache_misses hit_rate = (self.cache_hits / total_requests * 100) if total_requests > 0 else 0 - # Count cache files cache_files = 0 cache_size_mb = 0 - if self.cache_dir and self.cache_dir.exists(): - all_cache_files = list(self.cache_dir.glob("*.json")) - cache_files = len(all_cache_files) - cache_size_mb = sum(f.stat().st_size for f in all_cache_files) / ( - 1024 * 1024 - ) + if self._cache_store is not None: + cache_files = self._cache_store.count() + cache_size_mb = self._cache_store.size_bytes() / (1024 * 1024) return { "caching_enabled": True, @@ -256,12 +433,19 @@ def get_cache_stats(self) -> dict[str, Any]: def clear_cache(self) -> None: """Clear all cached results.""" - if not self.cache_enabled or not self.cache_dir: + if self._cache_store is None or self.cache_dir is None: return + self._cache_store.clear() for cache_file in self.cache_dir.glob("*.json"): - with contextlib.suppress(Exception): + try: cache_file.unlink() + except OSError as exc: + log.warning( + "Failed to remove legacy relevance cache file %s: %s", + cache_file, + exc, + ) # Reset statistics self.cache_hits = 0 diff --git a/benchmark_qed/cache.py b/benchmark_qed/cache.py new file mode 100644 index 0000000..f50fe3b --- /dev/null +++ b/benchmark_qed/cache.py @@ -0,0 +1,552 @@ +# Copyright (c) 2025 Microsoft Corporation. +"""Shared concurrency-safe local cache storage.""" + +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import time +from contextlib import contextmanager +from itertools import starmap +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterator + +_CACHE_SCHEMA_VERSION = 2 +_INITIALIZATION_TIMEOUT_SECONDS = 30 +_SENSITIVE_KEYS = { + "api_key", + "apikey", + "authorization", + "connection_string", + "credential", + "credentials", + "password", + "secret", + "token", +} +_SENSITIVE_KEY_SUFFIXES = ( + "_access_token", + "_api_key", + "_auth_token", + "_connection_string", + "_credential", + "_password", + "_secret", +) + + +def _is_sensitive_key(key: object) -> bool: + """Return whether a configuration key conventionally contains a secret.""" + normalized = str(key).lower() + return normalized in _SENSITIVE_KEYS or normalized.endswith(_SENSITIVE_KEY_SUFFIXES) + + +def redact_sensitive_values(value: Any) -> Any: + """Remove credentials from configuration persisted in local caches.""" + if isinstance(value, dict): + return { + str(key): ( + "" + if _is_sensitive_key(key) + else redact_sensitive_values(item) + ) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [redact_sensitive_values(item) for item in value] + return value + + +def stable_fingerprint(value: Any) -> str: + """Return a deterministic SHA256 fingerprint for JSON-like data.""" + content = json.dumps( + redact_sensitive_values(value), sort_keys=True, default=str + ).encode() + return hashlib.sha256(content).hexdigest() + + +def changed_configuration_fields( + current: dict[str, Any], cached: dict[str, Any] +) -> list[str]: + """Return top-level configuration fields whose values differ.""" + ignored = {"created_at", "schema_version", "source", "timestamp"} + return sorted( + key + for key in current.keys() | cached.keys() + if key not in ignored and current.get(key) != cached.get(key) + ) + + +class SQLiteCache: + """Namespaced JSON cache backed by SQLite in WAL mode.""" + + def __init__(self, database_path: Path | str, namespace: str) -> None: + """Initialize a cache namespace in a shared SQLite database.""" + self.database_path = Path(database_path) + self.database_path.parent.mkdir(parents=True, exist_ok=True) + self.namespace = namespace + self._initialize() + + @contextmanager + def _connect(self) -> Iterator[sqlite3.Connection]: + """Open a configured connection for one short cache operation.""" + connection = sqlite3.connect(self.database_path, timeout=30) + connection.execute("PRAGMA busy_timeout = 30000") + try: + with connection: + yield connection + finally: + connection.close() + + def _initialize(self) -> None: + """Initialize the schema, retrying only first-start lock contention.""" + deadline = time.monotonic() + _INITIALIZATION_TIMEOUT_SECONDS + while True: + try: + self._initialize_once() + except sqlite3.OperationalError as exc: + if "locked" not in str(exc).lower() or time.monotonic() >= deadline: + raise + time.sleep(0.05) + else: + return + + def _initialize_once(self) -> None: + """Create or migrate the cache schema.""" + with self._connect() as connection: + connection.execute("PRAGMA journal_mode = WAL") + connection.execute("PRAGMA synchronous = NORMAL") + connection.execute( + """ + CREATE TABLE IF NOT EXISTS cache_properties ( + name TEXT PRIMARY KEY, + value TEXT NOT NULL + ) + """ + ) + version_row = connection.execute( + "SELECT value FROM cache_properties WHERE name = 'schema_version'" + ).fetchone() + if version_row is None: + self._create_latest_schema(connection) + connection.execute( + "INSERT INTO cache_properties(name, value) VALUES (?, ?)", + ("schema_version", str(_CACHE_SCHEMA_VERSION)), + ) + return + + version = int(version_row[0]) + if version > _CACHE_SCHEMA_VERSION: + msg = ( + f"Unsupported cache schema version {version} " + f"in {self.database_path}" + ) + raise RuntimeError(msg) + if version == 1: + self._migrate_v1_to_v2(connection) + version = 2 + if version != _CACHE_SCHEMA_VERSION: + msg = f"No migration available for cache schema version {version}" + raise RuntimeError(msg) + self._create_latest_schema(connection) + + @staticmethod + def _create_latest_schema(connection: sqlite3.Connection) -> None: + """Create all tables and indexes for the latest schema.""" + connection.execute( + """ + CREATE TABLE IF NOT EXISTS cache_entries ( + namespace TEXT NOT NULL, + key TEXT NOT NULL, + value_json TEXT NOT NULL, + metadata_json TEXT NOT NULL, + logical_key TEXT, + config_fingerprint TEXT, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (namespace, key) + ) + """ + ) + connection.execute( + """ + CREATE INDEX IF NOT EXISTS idx_cache_entries_logical + ON cache_entries(namespace, logical_key) + """ + ) + connection.execute( + """ + CREATE TABLE IF NOT EXISTS cache_leases ( + namespace TEXT NOT NULL, + key TEXT NOT NULL, + owner_id TEXT NOT NULL, + expires_at REAL NOT NULL, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (namespace, key) + ) + """ + ) + + @staticmethod + def _migrate_v1_to_v2(connection: sqlite3.Connection) -> None: + """Add diagnostics and lease support to a v1 cache.""" + columns = { + str(row[1]) + for row in connection.execute("PRAGMA table_info(cache_entries)") + } + if "logical_key" not in columns: + connection.execute("ALTER TABLE cache_entries ADD COLUMN logical_key TEXT") + if "config_fingerprint" not in columns: + connection.execute( + "ALTER TABLE cache_entries ADD COLUMN config_fingerprint TEXT" + ) + SQLiteCache._create_latest_schema(connection) + connection.execute( + "UPDATE cache_properties SET value = '2' WHERE name = 'schema_version'" + ) + + def get(self, key: str) -> tuple[Any, dict[str, Any]] | None: + """Return a decoded value and metadata for a key.""" + with self._connect() as connection: + row = connection.execute( + """ + SELECT value_json, metadata_json + FROM cache_entries + WHERE namespace = ? AND key = ? + """, + (self.namespace, key), + ).fetchone() + if row is None: + return None + return json.loads(row[0]), json.loads(row[1]) + + def put_many( + self, + entries: list[tuple[str, Any, dict[str, Any]]], + *, + identities: dict[str, tuple[str, str]] | None = None, + ) -> int: + """Insert entries atomically with first-writer-wins semantics.""" + if not entries: + return 0 + records = [ + ( + self.namespace, + key, + json.dumps(value, sort_keys=True, default=str), + json.dumps(metadata, sort_keys=True, default=str), + (identities or {}).get(key, (None, None))[0], + (identities or {}).get(key, (None, None))[1], + ) + for key, value, metadata in entries + ] + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + changes_before = connection.total_changes + connection.executemany( + """ + INSERT INTO cache_entries( + namespace, key, value_json, metadata_json, + logical_key, config_fingerprint + ) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(namespace, key) DO NOTHING + """, + records, + ) + return connection.total_changes - changes_before + + def publish( + self, + key: str, + value: Any, + metadata: dict[str, Any], + *, + owner_id: str, + logical_key: str, + config_fingerprint: str, + ) -> bool: + """Publish a claimed result and release its lease atomically.""" + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + changes_before = connection.total_changes + connection.execute( + """ + INSERT INTO cache_entries( + namespace, key, value_json, metadata_json, + logical_key, config_fingerprint + ) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(namespace, key) DO NOTHING + """, + ( + self.namespace, + key, + json.dumps(value, sort_keys=True, default=str), + json.dumps(metadata, sort_keys=True, default=str), + logical_key, + config_fingerprint, + ), + ) + inserted = connection.total_changes > changes_before + connection.execute( + """ + DELETE FROM cache_leases + WHERE namespace = ? AND key = ? AND owner_id = ? + """, + (self.namespace, key, owner_id), + ) + return inserted + + def try_acquire( + self, + key: str, + owner_id: str, + *, + ttl_seconds: float, + now: float | None = None, + ) -> bool: + """Claim missing work, replacing an expired lease if necessary.""" + current_time = time.time() if now is None else now + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + if ( + connection.execute( + """ + SELECT 1 FROM cache_entries + WHERE namespace = ? AND key = ? + """, + (self.namespace, key), + ).fetchone() + is not None + ): + return False + connection.execute( + """ + DELETE FROM cache_leases + WHERE namespace = ? AND key = ? AND expires_at <= ? + """, + (self.namespace, key, current_time), + ) + changes_before = connection.total_changes + connection.execute( + """ + INSERT INTO cache_leases(namespace, key, owner_id, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(namespace, key) DO NOTHING + """, + (self.namespace, key, owner_id, current_time + ttl_seconds), + ) + return connection.total_changes > changes_before + + def renew_leases( + self, keys: list[str], owner_id: str, *, ttl_seconds: float + ) -> int: + """Extend leases still owned by a worker.""" + if not keys: + return 0 + with self._connect() as connection: + changes_before = connection.total_changes + connection.executemany( + """ + UPDATE cache_leases + SET expires_at = ? + WHERE namespace = ? AND owner_id = ? + AND key = ? + """, + [ + ( + time.time() + ttl_seconds, + self.namespace, + owner_id, + key, + ) + for key in keys + ], + ) + return connection.total_changes - changes_before + + def release_leases(self, keys: list[str], owner_id: str) -> None: + """Release leases owned by a worker without publishing results.""" + if not keys: + return + with self._connect() as connection: + connection.executemany( + """ + DELETE FROM cache_leases + WHERE namespace = ? AND owner_id = ? + AND key = ? + """, + [(self.namespace, owner_id, key) for key in keys], + ) + + def find_alternate_configurations( + self, requests: list[tuple[str, str]] + ) -> list[dict[str, Any]]: + """Return metadata for matching inputs produced by other configurations.""" + requested = dict(requests) + if not requested: + return [] + with self._connect() as connection: + connection.execute( + """ + CREATE TEMP TABLE requested_cache_configs ( + logical_key TEXT PRIMARY KEY, + config_fingerprint TEXT NOT NULL + ) + """ + ) + connection.executemany( + "INSERT INTO requested_cache_configs VALUES (?, ?)", + requested.items(), + ) + rows = connection.execute( + """ + SELECT entries.metadata_json + FROM cache_entries AS entries + INNER JOIN requested_cache_configs AS requested + ON requested.logical_key = entries.logical_key + WHERE entries.namespace = ? + AND entries.config_fingerprint IS NOT requested.config_fingerprint + """, + (self.namespace,), + ).fetchall() + return list(starmap(json.loads, rows)) + + def get_property(self, name: str) -> str | None: + """Return a database-level property.""" + with self._connect() as connection: + row = connection.execute( + "SELECT value FROM cache_properties WHERE name = ?", (name,) + ).fetchone() + return str(row[0]) if row is not None else None + + def set_property_if_absent(self, name: str, value: str) -> None: + """Set a database-level property without replacing an existing value.""" + with self._connect() as connection: + connection.execute( + """ + INSERT INTO cache_properties(name, value) + VALUES (?, ?) + ON CONFLICT(name) DO NOTHING + """, + (name, value), + ) + + def count(self) -> int: + """Return the number of entries in this namespace.""" + with self._connect() as connection: + row = connection.execute( + "SELECT COUNT(*) FROM cache_entries WHERE namespace = ?", + (self.namespace,), + ).fetchone() + return int(row[0]) if row is not None else 0 + + def clear(self) -> None: + """Delete all entries and leases in this namespace.""" + with self._connect() as connection: + connection.execute( + "DELETE FROM cache_entries WHERE namespace = ?", (self.namespace,) + ) + connection.execute( + "DELETE FROM cache_leases WHERE namespace = ?", (self.namespace,) + ) + + def size_bytes(self) -> int: + """Return total on-disk size of the database and WAL sidecars.""" + return sum( + path.stat().st_size + for path in ( + self.database_path, + Path(f"{self.database_path}-wal"), + Path(f"{self.database_path}-shm"), + ) + if path.exists() + ) + + +def inspect_cache(database_path: Path | str) -> dict[str, Any]: + """Return schema, namespace, provenance, and lease details without mutation.""" + path = Path(database_path) + if not path.exists(): + msg = f"Cache database does not exist: {path}" + raise FileNotFoundError(msg) + connection = sqlite3.connect(f"{path.resolve().as_uri()}?mode=ro", uri=True) + try: + tables = { + str(row[0]) + for row in connection.execute( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ) + } + if "cache_entries" not in tables or "cache_properties" not in tables: + msg = f"Not a benchmark-qed cache database: {path}" + raise ValueError(msg) + version_row = connection.execute( + "SELECT value FROM cache_properties WHERE name = 'schema_version'" + ).fetchone() + entry_columns = { + str(row[1]) + for row in connection.execute("PRAGMA table_info(cache_entries)") + } + if "config_fingerprint" in entry_columns: + namespace_rows = connection.execute( + """ + SELECT namespace, COUNT(*), COUNT(DISTINCT config_fingerprint) + FROM cache_entries + GROUP BY namespace + ORDER BY namespace + """ + ) + else: + namespace_rows = connection.execute( + """ + SELECT namespace, COUNT(*), 0 + FROM cache_entries + GROUP BY namespace + ORDER BY namespace + """ + ) + namespaces = [ + { + "namespace": namespace, + "entries": entry_count, + "configurations": config_count, + } + for namespace, entry_count, config_count in namespace_rows + ] + provenance = [ + { + "namespace": namespace, + "created_at": created_at, + "metadata": json.loads(metadata_json), + } + for namespace, created_at, metadata_json in connection.execute( + """ + SELECT namespace, created_at, metadata_json + FROM cache_entries + ORDER BY created_at DESC + LIMIT 10 + """ + ) + ] + lease_count = ( + connection.execute( + "SELECT COUNT(*) FROM cache_leases WHERE expires_at > ?", + (time.time(),), + ).fetchone()[0] + if "cache_leases" in tables + else 0 + ) + finally: + connection.close() + return { + "path": str(path), + "schema_version": int(version_row[0]) if version_row else None, + "namespaces": namespaces, + "active_leases": lease_count, + "recent_provenance": provenance, + } diff --git a/benchmark_qed/cli/cache.py b/benchmark_qed/cli/cache.py new file mode 100644 index 0000000..0d8d790 --- /dev/null +++ b/benchmark_qed/cli/cache.py @@ -0,0 +1,54 @@ +# Copyright (c) 2025 Microsoft Corporation. +"""Commands for inspecting local evaluation caches.""" + +import json +from pathlib import Path +from typing import Annotated + +import typer + +from benchmark_qed.cache import inspect_cache + +app: typer.Typer = typer.Typer( + pretty_exceptions_show_locals=False, + help="Inspect local evaluation caches.", +) + + +@app.command(name="inspect") +def inspect_command( + cache_path: Annotated[ + Path, + typer.Argument( + exists=True, + dir_okay=False, + readable=True, + help="Path to a benchmark-qed SQLite cache.", + ), + ], + json_output: Annotated[ + bool, + typer.Option("--json", help="Emit machine-readable JSON."), + ] = False, +) -> None: + """Show cache schema, namespaces, leases, and recent provenance.""" + details = inspect_cache(cache_path) + if json_output: + typer.echo(json.dumps(details, indent=2, sort_keys=True)) + return + + typer.echo(f"Cache: {details['path']}") + typer.echo(f"Schema version: {details['schema_version']}") + typer.echo(f"Active leases: {details['active_leases']}") + typer.echo("Namespaces:") + for namespace in details["namespaces"]: + typer.echo( + f" {namespace['namespace']}: {namespace['entries']} entries, " + f"{namespace['configurations']} configurations" + ) + typer.echo("Recent provenance:") + for record in details["recent_provenance"]: + typer.echo( + f" {record['created_at']} {record['namespace']}: " + f"{json.dumps(record['metadata'], sort_keys=True)}" + ) diff --git a/docs/cli/retrieval_scores.md b/docs/cli/retrieval_scores.md index 99f755c..97cb806 100644 --- a/docs/cli/retrieval_scores.md +++ b/docs/cli/retrieval_scores.md @@ -349,6 +349,18 @@ Always enable `cache_dir` to: - Avoid re-assessing the same query-chunk pairs - Share cache across multiple runs +Relevance assessments are stored in `relevance_cache.sqlite3` under this +directory. SQLite WAL mode supports concurrent readers and writers. Existing +per-assessment JSON cache files are imported automatically and retained for +backward compatibility. Cross-process work leases prevent duplicate LLM calls +and automatically expire after an interrupted worker exits. + +Use `benchmark-qed cache inspect /relevance_cache.sqlite3` to view +schema information, entry and configuration counts, active leases, and recent +provenance. If the same query and text were previously evaluated with another +model or prompt configuration, the evaluator logs a warning and computes a new +result instead of reusing incompatible data. + ### Concurrency - Start with `concurrent_requests: 16-32` for OpenAI diff --git a/docs/developing.md b/docs/developing.md index 9b6c3d8..fea1e1d 100644 --- a/docs/developing.md +++ b/docs/developing.md @@ -410,7 +410,20 @@ You'll provide this data to benchmark-qed in one of two formats (see Step 2 belo - `settings.yaml`: Contains pipeline settings, including: - `retrieval_path`: Path to your retrieved chunks JSON (RetrievalResult format above) - `k_list`: K values to report coverage metrics (e.g., [5, 10, 20, 50]) - - `cache_dir`: Directory for persistent (assertion, chunk) cache + - `cache_dir`: Directory for the persistent SQLite (assertion, chunk) + cache. The cache supports concurrent processes through WAL mode and + records the model, call arguments, and prompts used for each result. + Existing `chunk_assertions.jsonl` caches are imported automatically. + Concurrent processes claim uncached work so only one performs each LLM + request; leases are renewed during long calls and expire after an + interrupted process exits. + + Inspect cache provenance, configuration counts, and active work leases with: + + ```sh + uv run benchmark-qed cache inspect \ + .benchmark_qed_cache/chunk_assertions.sqlite3 + ``` The generated `settings.yaml` includes commented-out `input_storage` and `output_storage` sections for configuring Azure Blob Storage backends. @@ -459,5 +472,3 @@ To learn how to use AutoE programmatically, please see the [AutoE Notebook Examp To explore the query synthesis workflow in detail, please see the [AutoQ CLI Documentation](cli/autoq.md) for command-line usage and the [AutoQ Notebook Example](notebooks/autoq.ipynb) for a step-by-step programmatic guide. For a deeper understanding of AutoE evaluation pipelines, please refer to the [AutoE CLI Documentation](cli/autoe.md) for available commands and the [AutoE Notebook Example](notebooks/autoe.ipynb) for hands-on examples. - - diff --git a/docs/notebooks/cache_parallelism_benchmark.ipynb b/docs/notebooks/cache_parallelism_benchmark.ipynb new file mode 100644 index 0000000..bd8c268 --- /dev/null +++ b/docs/notebooks/cache_parallelism_benchmark.ipynb @@ -0,0 +1,653 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "overview", + "metadata": {}, + "source": [ + "# Cache and parallelism benchmark\n", + "\n", + "This notebook demonstrates the efficiency and reliability improvements of the SQLite/WAL cache compared with the original JSONL snapshot design. It uses deterministic simulated work instead of paid LLM calls.\n", + "\n", + "The demonstrations cover:\n", + "\n", + "1. incremental write scaling;\n", + "2. multiprocess integrity;\n", + "3. duplicate LLM-call avoidance;\n", + "4. interrupted-worker recovery;\n", + "5. configuration mismatch diagnostics; and\n", + "6. provenance inspection.\n", + "\n", + "> Performance varies by filesystem and machine. Run each benchmark several times and report medians, call counts, and integrity results rather than relying on one timing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "# Copyright (c) 2025 Microsoft Corporation.\n", + "from __future__ import annotations\n", + "\n", + "import contextlib\n", + "import json\n", + "import multiprocessing as mp\n", + "import shutil\n", + "import statistics\n", + "import subprocess # noqa: S404 - invokes the local benchmark-qed CLI only\n", + "import sys\n", + "import tempfile\n", + "import time\n", + "from concurrent.futures import ProcessPoolExecutor\n", + "from pathlib import Path\n", + "from typing import Any\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "from IPython.display import display as show\n", + "\n", + "from benchmark_qed.autoe.chunk_assertion.cache import (\n", + " ContentAddressedCache,\n", + " build_cache_metadata,\n", + " compute_cache_key,\n", + " compute_config_fingerprint,\n", + " compute_logical_key,\n", + ")\n", + "from benchmark_qed.cache import SQLiteCache, inspect_cache\n", + "\n", + "RUNS = 5\n", + "BATCHES = 30\n", + "BATCH_SIZE = 100\n", + "WORKERS = 4\n", + "SHARED_KEYS = 30\n", + "SIMULATED_CALL_SECONDS = 0.01\n", + "\n", + "workspace = Path(tempfile.mkdtemp(prefix=\"benchmark-qed-cache-demo-\"))\n", + "mp_context = mp.get_context(\"fork\") if \"fork\" in mp.get_all_start_methods() else None\n", + "print(f\"Workspace: {workspace}\")\n", + "print(f\"Multiprocess demonstrations available: {mp_context is not None}\")" + ] + }, + { + "cell_type": "markdown", + "id": "legacy-baseline", + "metadata": {}, + "source": [ + "## Original JSONL baseline\n", + "\n", + "The class below reproduces the original chunk-cache persistence behavior: load the complete JSONL file, keep an in-memory snapshot, write the complete snapshot to one shared `.tmp` file, and replace the cache on every flush." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "legacy-cache", + "metadata": {}, + "outputs": [], + "source": [ + "class LegacyJSONLCache:\n", + " \"\"\"Reproduce the original full-snapshot JSONL cache behavior.\"\"\"\n", + "\n", + " def __init__(self, path: Path) -> None:\n", + " self.path = path\n", + " self.path.parent.mkdir(parents=True, exist_ok=True)\n", + " self.data: dict[str, str] = {}\n", + " if self.path.exists():\n", + " for line in self.path.read_text(encoding=\"utf-8\").splitlines():\n", + " if not line:\n", + " continue\n", + " try:\n", + " record = json.loads(line)\n", + " except json.JSONDecodeError:\n", + " continue\n", + " if record.get(\"key\") and record.get(\"grade\"):\n", + " self.data[str(record[\"key\"])] = str(record[\"grade\"])\n", + "\n", + " def put(self, key: str, value: str) -> None:\n", + " \"\"\"Store a value in the in-memory snapshot.\"\"\"\n", + " self.data.setdefault(key, value)\n", + "\n", + " def flush(self) -> None:\n", + " \"\"\"Rewrite and atomically replace the complete JSONL snapshot.\"\"\"\n", + " temporary_path = self.path.with_suffix(self.path.suffix + \".tmp\")\n", + " with temporary_path.open(\"w\", encoding=\"utf-8\") as file:\n", + " for key, value in self.data.items():\n", + " file.write(json.dumps({\"key\": key, \"grade\": value}) + \"\\n\")\n", + " temporary_path.replace(self.path)\n", + "\n", + "\n", + "def reset_path(path: Path) -> None:\n", + " \"\"\"Remove one benchmark cache and its SQLite sidecars.\"\"\"\n", + " if path.is_dir():\n", + " shutil.rmtree(path)\n", + " elif path.exists():\n", + " path.unlink()\n", + " for suffix in (\"-wal\", \"-shm\"):\n", + " sidecar = Path(f\"{path}{suffix}\")\n", + " if sidecar.exists():\n", + " sidecar.unlink()" + ] + }, + { + "cell_type": "markdown", + "id": "incremental-heading", + "metadata": {}, + "source": [ + "## 1. Incremental write scaling\n", + "\n", + "This workload adds fixed-size batches and persists after every batch. The original cache rewrites every preceding entry; SQLite inserts only the new rows." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "incremental-benchmark", + "metadata": {}, + "outputs": [], + "source": [ + "def benchmark_incremental_writes(kind: str, run: int) -> tuple[float, list[float]]:\n", + " \"\"\"Measure total and per-batch persistence latency.\"\"\"\n", + " path = (\n", + " workspace\n", + " / f\"incremental-{kind}-{run}.{'jsonl' if kind == 'legacy' else 'sqlite3'}\"\n", + " )\n", + " reset_path(path)\n", + " batch_times: list[float] = []\n", + " started = time.perf_counter()\n", + " if kind == \"legacy\":\n", + " cache = LegacyJSONLCache(path)\n", + " for batch in range(BATCHES):\n", + " batch_started = time.perf_counter()\n", + " for offset in range(BATCH_SIZE):\n", + " index = batch * BATCH_SIZE + offset\n", + " cache.put(f\"key-{index}\", \"full_support\")\n", + " cache.flush()\n", + " batch_times.append(time.perf_counter() - batch_started)\n", + " else:\n", + " cache = SQLiteCache(path, \"benchmark\")\n", + " for batch in range(BATCHES):\n", + " batch_started = time.perf_counter()\n", + " entries = [\n", + " (f\"key-{batch * BATCH_SIZE + offset}\", \"full_support\", {})\n", + " for offset in range(BATCH_SIZE)\n", + " ]\n", + " cache.put_many(entries)\n", + " batch_times.append(time.perf_counter() - batch_started)\n", + " return time.perf_counter() - started, batch_times\n", + "\n", + "\n", + "incremental_runs: list[dict[str, Any]] = []\n", + "incremental_curves: dict[str, list[float]] = {}\n", + "for implementation in (\"legacy\", \"sqlite\"):\n", + " totals = []\n", + " curves = []\n", + " for run in range(RUNS):\n", + " total, curve = benchmark_incremental_writes(implementation, run)\n", + " totals.append(total)\n", + " curves.append(curve)\n", + " incremental_runs.append({\n", + " \"implementation\": implementation,\n", + " \"median_seconds\": statistics.median(totals),\n", + " \"p95_batch_ms\": statistics.quantiles(\n", + " [value for curve in curves for value in curve], n=20\n", + " )[18]\n", + " * 1000,\n", + " })\n", + " incremental_curves[implementation] = [\n", + " statistics.median(curve[index] for curve in curves) for index in range(BATCHES)\n", + " ]\n", + "\n", + "incremental_df = pd.DataFrame(incremental_runs)\n", + "show(incremental_df)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "incremental-plot", + "metadata": {}, + "outputs": [], + "source": [ + "for implementation, curve in incremental_curves.items():\n", + " plt.plot(\n", + " range(1, BATCHES + 1), [value * 1000 for value in curve], label=implementation\n", + " )\n", + "plt.xlabel(\"Persisted batch\")\n", + "plt.ylabel(\"Median batch latency (ms)\")\n", + "plt.title(\"Incremental cache persistence\")\n", + "plt.legend()\n", + "plt.grid(alpha=0.3)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "parallel-heading", + "metadata": {}, + "source": [ + "## 2. Multiprocess integrity\n", + "\n", + "Several processes load an empty cache, wait for the same start time, and then write disjoint keys. The check reports process errors and the number of entries that survive." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "parallel-benchmark", + "metadata": {}, + "outputs": [], + "source": [ + "def legacy_parallel_writer(\n", + " path_text: str, worker: int, start_at: float, count: int\n", + ") -> str | None:\n", + " \"\"\"Write disjoint entries through the original snapshot cache.\"\"\"\n", + " cache = LegacyJSONLCache(Path(path_text))\n", + " while time.time() < start_at:\n", + " time.sleep(0.001)\n", + " for index in range(count):\n", + " cache.put(f\"worker-{worker}-key-{index}\", \"full_support\")\n", + " try:\n", + " cache.flush()\n", + " except OSError as exc:\n", + " return str(exc)\n", + " return None\n", + "\n", + "\n", + "def sqlite_parallel_writer(\n", + " path_text: str, worker: int, start_at: float, count: int\n", + ") -> str | None:\n", + " \"\"\"Write disjoint entries through the SQLite cache.\"\"\"\n", + " cache = SQLiteCache(path_text, \"parallel-benchmark\")\n", + " while time.time() < start_at:\n", + " time.sleep(0.001)\n", + " cache.put_many([\n", + " (f\"worker-{worker}-key-{index}\", \"full_support\", {}) for index in range(count)\n", + " ])\n", + " return None\n", + "\n", + "\n", + "def count_legacy_entries(path: Path) -> int:\n", + " \"\"\"Count valid distinct keys surviving in a JSONL cache.\"\"\"\n", + " if not path.exists():\n", + " return 0\n", + " keys = set()\n", + " for line in path.read_text(encoding=\"utf-8\").splitlines():\n", + " with contextlib.suppress(json.JSONDecodeError, KeyError):\n", + " keys.add(json.loads(line)[\"key\"])\n", + " return len(keys)\n", + "\n", + "\n", + "parallel_rows = []\n", + "if mp_context is None:\n", + " print(\"Skipped: this demonstration requires a platform with the fork start method.\")\n", + "else:\n", + " expected = WORKERS * BATCH_SIZE\n", + " for implementation, writer in (\n", + " (\"legacy\", legacy_parallel_writer),\n", + " (\"sqlite\", sqlite_parallel_writer),\n", + " ):\n", + " path = (\n", + " workspace\n", + " / f\"parallel-{implementation}.{'jsonl' if implementation == 'legacy' else 'sqlite3'}\"\n", + " )\n", + " reset_path(path)\n", + " start_at = time.time() + 1\n", + " started = time.perf_counter()\n", + " with ProcessPoolExecutor(\n", + " max_workers=WORKERS, mp_context=mp_context\n", + " ) as executor:\n", + " futures = [\n", + " executor.submit(writer, str(path), worker, start_at, BATCH_SIZE)\n", + " for worker in range(WORKERS)\n", + " ]\n", + " errors = [future.result(timeout=30) for future in futures]\n", + " elapsed = time.perf_counter() - started\n", + " surviving = (\n", + " count_legacy_entries(path)\n", + " if implementation == \"legacy\"\n", + " else SQLiteCache(path, \"parallel-benchmark\").count()\n", + " )\n", + " parallel_rows.append({\n", + " \"implementation\": implementation,\n", + " \"workers\": WORKERS,\n", + " \"expected_entries\": expected,\n", + " \"surviving_entries\": surviving,\n", + " \"write_errors\": sum(error is not None for error in errors),\n", + " \"seconds\": elapsed,\n", + " \"integrity\": surviving == expected and not any(errors),\n", + " })\n", + "\n", + "show(pd.DataFrame(parallel_rows))" + ] + }, + { + "cell_type": "markdown", + "id": "coalescing-heading", + "metadata": {}, + "source": [ + "## 3. Duplicate expensive-call avoidance\n", + "\n", + "Every process requests the same keys. The legacy baseline performs every simulated call. The leased cache permits only the claim owner to perform each call; other processes reuse its published result." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "coalescing-benchmark", + "metadata": {}, + "outputs": [], + "source": [ + "def legacy_duplicate_worker(key_count: int, delay: float) -> int:\n", + " \"\"\"Perform every simulated expensive call without coalescing.\"\"\"\n", + " calls = 0\n", + " for _index in range(key_count):\n", + " time.sleep(delay)\n", + " calls += 1\n", + " return calls\n", + "\n", + "\n", + "def leased_duplicate_worker(\n", + " path_text: str, owner: str, key_count: int, delay: float\n", + ") -> int:\n", + " \"\"\"Perform only simulated calls transactionally claimed by this worker.\"\"\"\n", + " store = SQLiteCache(path_text, \"llm-call-benchmark\")\n", + " calls = 0\n", + " for index in range(key_count):\n", + " key = f\"shared-{index}\"\n", + " while store.get(key) is None:\n", + " if store.try_acquire(key, owner, ttl_seconds=5):\n", + " time.sleep(delay)\n", + " store.publish(\n", + " key,\n", + " \"result\",\n", + " {\"simulated\": True},\n", + " owner_id=owner,\n", + " logical_key=key,\n", + " config_fingerprint=\"demo\",\n", + " )\n", + " calls += 1\n", + " break\n", + " time.sleep(0.002)\n", + " return calls\n", + "\n", + "\n", + "coalescing_rows = []\n", + "if mp_context is None:\n", + " print(\"Skipped: this demonstration requires a platform with the fork start method.\")\n", + "else:\n", + " legacy_started = time.perf_counter()\n", + " with ProcessPoolExecutor(max_workers=WORKERS, mp_context=mp_context) as executor:\n", + " legacy_calls = sum(\n", + " executor.map(\n", + " legacy_duplicate_worker,\n", + " [SHARED_KEYS] * WORKERS,\n", + " [SIMULATED_CALL_SECONDS] * WORKERS,\n", + " )\n", + " )\n", + " legacy_seconds = time.perf_counter() - legacy_started\n", + "\n", + " leased_path = workspace / \"leased-duplicate.sqlite3\"\n", + " reset_path(leased_path)\n", + " leased_started = time.perf_counter()\n", + " with ProcessPoolExecutor(max_workers=WORKERS, mp_context=mp_context) as executor:\n", + " leased_calls = sum(\n", + " executor.map(\n", + " leased_duplicate_worker,\n", + " [str(leased_path)] * WORKERS,\n", + " [f\"worker-{index}\" for index in range(WORKERS)],\n", + " [SHARED_KEYS] * WORKERS,\n", + " [SIMULATED_CALL_SECONDS] * WORKERS,\n", + " )\n", + " )\n", + " leased_seconds = time.perf_counter() - leased_started\n", + "\n", + " logical_requests = WORKERS * SHARED_KEYS\n", + " coalescing_rows = [\n", + " {\n", + " \"implementation\": \"legacy\",\n", + " \"logical_requests\": logical_requests,\n", + " \"simulated_llm_calls\": legacy_calls,\n", + " \"duplicate_calls_avoided\": logical_requests - legacy_calls,\n", + " \"seconds\": legacy_seconds,\n", + " },\n", + " {\n", + " \"implementation\": \"sqlite_leases\",\n", + " \"logical_requests\": logical_requests,\n", + " \"simulated_llm_calls\": leased_calls,\n", + " \"duplicate_calls_avoided\": logical_requests - leased_calls,\n", + " \"seconds\": leased_seconds,\n", + " },\n", + " ]\n", + "\n", + "show(pd.DataFrame(coalescing_rows))" + ] + }, + { + "cell_type": "markdown", + "id": "recovery-heading", + "metadata": {}, + "source": [ + "## 4. Interrupted-worker recovery\n", + "\n", + "A short-lived process claims work and exits without publishing. Another worker cannot steal an active lease, but can claim it after expiration—without deleting lock files manually." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "recovery-demo", + "metadata": {}, + "outputs": [], + "source": [ + "def claim_and_exit(path_text: str) -> bool:\n", + " \"\"\"Claim work and exit without publishing to simulate interruption.\"\"\"\n", + " return SQLiteCache(path_text, \"recovery-demo\").try_acquire(\n", + " \"abandoned-key\", \"terminated-worker\", ttl_seconds=0.5\n", + " )\n", + "\n", + "\n", + "recovery_path = workspace / \"recovery.sqlite3\"\n", + "reset_path(recovery_path)\n", + "if mp_context is None:\n", + " print(\"Skipped: this demonstration requires a platform with the fork start method.\")\n", + "else:\n", + " with ProcessPoolExecutor(max_workers=1, mp_context=mp_context) as executor:\n", + " original_claim = executor.submit(claim_and_exit, str(recovery_path)).result(\n", + " timeout=30\n", + " )\n", + " recovery_store = SQLiteCache(recovery_path, \"recovery-demo\")\n", + " blocked_before_expiry = not recovery_store.try_acquire(\n", + " \"abandoned-key\", \"recovery-worker\", ttl_seconds=1\n", + " )\n", + " time.sleep(0.55)\n", + " recovered_after_expiry = recovery_store.try_acquire(\n", + " \"abandoned-key\", \"recovery-worker\", ttl_seconds=1\n", + " )\n", + " show(\n", + " pd.DataFrame([\n", + " {\n", + " \"original_worker_claimed\": original_claim,\n", + " \"protected_before_expiry\": blocked_before_expiry,\n", + " \"recovered_after_expiry\": recovered_after_expiry,\n", + " \"manual_cleanup_required\": False,\n", + " }\n", + " ])\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "metadata-heading", + "metadata": {}, + "source": [ + "## 5. Configuration mismatch and provenance\n", + "\n", + "The logical input fingerprint is independent from the model configuration. This allows the cache to detect that a result exists for the same input while correctly refusing to reuse it under different settings." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "metadata-demo", + "metadata": {}, + "outputs": [], + "source": [ + "provenance_path = workspace / \"provenance.sqlite3\"\n", + "cache = ContentAddressedCache(provenance_path)\n", + "assertion = \"The answer is supported by the passage.\"\n", + "chunk = \"A representative passage.\"\n", + "prompt_template = \"{assertion}\\n{chunk}\" # noqa: RUF027 - literal format template\n", + "first_metadata = build_cache_metadata(\n", + " model=\"model-a\",\n", + " call_args={\"temperature\": 0, \"api_key\": \"not-persisted\"},\n", + " system_prompt=\"Judge support.\",\n", + " user_prompt=prompt_template,\n", + ")\n", + "first_key = compute_cache_key(\n", + " assertion,\n", + " chunk,\n", + " model=\"model-a\",\n", + " call_args={\"temperature\": 0, \"api_key\": \"not-persisted\"},\n", + " system_prompt=\"Judge support.\",\n", + " user_prompt=prompt_template,\n", + ")\n", + "logical_key = compute_logical_key(assertion, chunk)\n", + "first_config = compute_config_fingerprint(\n", + " model=\"model-a\",\n", + " call_args={\"temperature\": 0},\n", + " system_prompt=\"Judge support.\",\n", + " user_prompt=prompt_template,\n", + ")\n", + "cache.claim(first_key, \"demo-owner\")\n", + "cache.publish(\n", + " first_key,\n", + " \"full_support\",\n", + " first_metadata,\n", + " owner_id=\"demo-owner\",\n", + " logical_key=logical_key,\n", + " config_fingerprint=first_config,\n", + ")\n", + "\n", + "second_metadata = build_cache_metadata(\n", + " model=\"model-b\",\n", + " call_args={\"temperature\": 0},\n", + " system_prompt=\"Judge support.\",\n", + " user_prompt=prompt_template,\n", + ")\n", + "mismatch_count, changed_fields = cache.find_configuration_mismatches(\n", + " [\n", + " (\n", + " logical_key,\n", + " compute_config_fingerprint(\n", + " model=\"model-b\",\n", + " call_args={\"temperature\": 0},\n", + " system_prompt=\"Judge support.\",\n", + " user_prompt=prompt_template,\n", + " ),\n", + " )\n", + " ],\n", + " second_metadata,\n", + ")\n", + "show(\n", + " pd.DataFrame([\n", + " {\n", + " \"matching_input\": True,\n", + " \"configuration_mismatches\": mismatch_count,\n", + " \"changed_fields\": \", \".join(changed_fields),\n", + " \"credential_persisted\": \"not-persisted\"\n", + " in json.dumps(inspect_cache(provenance_path)),\n", + " }\n", + " ])\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "inspection-demo", + "metadata": {}, + "outputs": [], + "source": [ + "inspection = inspect_cache(provenance_path)\n", + "print(json.dumps(inspection, indent=2, sort_keys=True))\n", + "\n", + "print(\"\\nEquivalent CLI command:\")\n", + "print(f\"uv run benchmark-qed cache inspect {provenance_path} --json\")\n", + "cli_result = subprocess.run( # noqa: S603 - fixed local module invocation\n", + " [\n", + " sys.executable,\n", + " \"-m\",\n", + " \"benchmark_qed\",\n", + " \"cache\",\n", + " \"inspect\",\n", + " str(provenance_path),\n", + " \"--json\",\n", + " ],\n", + " check=True,\n", + " capture_output=True,\n", + " text=True,\n", + ")\n", + "if json.loads(cli_result.stdout)[\"schema_version\"] != 2:\n", + " message = \"Cache inspection CLI returned an unexpected schema version\"\n", + " raise RuntimeError(message)" + ] + }, + { + "cell_type": "markdown", + "id": "reporting-heading", + "metadata": {}, + "source": [ + "## 6. Reporting guidance\n", + "\n", + "When sharing results, report the machine, filesystem, Python version, benchmark constants, median timings, p95 batch latency, surviving entry count, write errors, and actual simulated LLM-call count.\n", + "\n", + "The defensible conclusion is:\n", + "\n", + "> The SQLite/WAL cache is more efficient and reliable for repeated incremental writes and concurrent workloads. It prevents lost updates, coalesces duplicate model requests across processes, supports concurrent readers, warns about incompatible configurations, and recovers automatically from interrupted workers.\n", + "\n", + "Do not claim that SQLite wins every single-key microbenchmark; the primary gains are scaling, integrity, duplicate-call avoidance, and recoverability." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "summary", + "metadata": {}, + "outputs": [], + "source": [ + "summary_rows = []\n", + "summary_rows.extend(incremental_runs)\n", + "summary_rows.extend(parallel_rows)\n", + "summary_rows.extend(coalescing_rows)\n", + "show(pd.DataFrame(summary_rows))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cleanup", + "metadata": {}, + "outputs": [], + "source": [ + "# Run this after exporting any results you want to keep.\n", + "shutil.rmtree(workspace)\n", + "print(f\"Removed temporary workspace: {workspace}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mkdocs.yaml b/mkdocs.yaml index 0112ac5..338edc7 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -48,6 +48,7 @@ nav: - AutoD: notebooks/autod.ipynb - Assertion Generation: notebooks/assertion_gen.ipynb - Retrieval Metrics: notebooks/retrieval_metrics.ipynb + - Cache and Parallelism Benchmark: notebooks/cache_parallelism_benchmark.ipynb - Datasets: datasets.md repo_url: https://github.com/microsoft/benchmark-qed diff --git a/tests/autoe/chunk_assertion/cache_test.py b/tests/autoe/chunk_assertion/cache_test.py index 9ce9756..e41036a 100644 --- a/tests/autoe/chunk_assertion/cache_test.py +++ b/tests/autoe/chunk_assertion/cache_test.py @@ -2,12 +2,55 @@ # Licensed under the MIT License. """Tests for the content-addressed (assertion, chunk) cache.""" +import json +import sqlite3 +import time +from concurrent.futures import ProcessPoolExecutor from pathlib import Path from benchmark_qed.autoe.chunk_assertion.cache import ( ContentAddressedCache, + build_cache_metadata, compute_cache_key, + compute_config_fingerprint, + compute_logical_key, ) +from benchmark_qed.cache import SQLiteCache, inspect_cache + + +def _write_cache_entry(cache_path: str, cache_key: str, grade: str) -> None: + cache = ContentAddressedCache(cache_path) + cache.put(cache_key, grade, {"writer": cache_key}) + cache.flush() + + +def _try_claim(cache_path: str, owner_id: str) -> bool: + store = SQLiteCache(cache_path, "lease-test") + return store.try_acquire("shared", owner_id, ttl_seconds=30) + + +def _compute_once(cache_path: str, owner_id: str) -> bool: + store = SQLiteCache(cache_path, "coalescing-test") + while True: + if store.get("shared") is not None: + return False + if store.try_acquire("shared", owner_id, ttl_seconds=5): + time.sleep(0.1) + store.publish( + "shared", + "result", + {"owner": owner_id}, + owner_id=owner_id, + logical_key="logical", + config_fingerprint="config", + ) + return True + time.sleep(0.02) + + +def _acquire_then_exit(cache_path: str) -> bool: + store = SQLiteCache(cache_path, "interruption-test") + return store.try_acquire("abandoned", "crashed", ttl_seconds=0.2) class TestComputeCacheKey: @@ -43,20 +86,48 @@ def test_distinguishes_prompts(self) -> None: "a", "b", user_prompt="u2" ) + def test_credentials_do_not_affect_key(self) -> None: + """Credentials are neither persisted nor included in cache identity.""" + first = compute_cache_key("a", "b", call_args={"api_key": "first"}) + second = compute_cache_key("a", "b", call_args={"api_key": "second"}) + assert first == second + + +def test_build_cache_metadata_redacts_credentials() -> None: + metadata = build_cache_metadata( + model="test-model", + call_args={ + "temperature": 0, + "max_tokens": 100, + "api_key": "secret", + "headers": {"Authorization": "Bearer secret"}, + }, + system_prompt="system", + user_prompt="user", + ) + + assert metadata["model"] == "test-model" + assert metadata["call_args"] == { + "temperature": 0, + "max_tokens": 100, + "api_key": "", + "headers": {"Authorization": ""}, + } + class TestContentAddressedCache: """Tests for ContentAddressedCache persistence semantics.""" def test_put_get_roundtrip(self, tmp_path: Path) -> None: """Stored grades are retrievable and missing keys return None.""" - cache = ContentAddressedCache(tmp_path / "cache.jsonl") + cache = ContentAddressedCache(tmp_path / "cache.sqlite3") cache.put("k1", "full_support") assert cache.get("k1") == "full_support" assert cache.get("missing") is None def test_put_is_idempotent_for_new_count(self, tmp_path: Path) -> None: """Re-putting an existing key does not increment the new-entry count.""" - cache = ContentAddressedCache(tmp_path / "cache.jsonl") + cache = ContentAddressedCache(tmp_path / "cache.sqlite3") cache.put("k1", "full_support") cache.put("k1", "no_support") assert cache.new_count == 1 @@ -64,7 +135,7 @@ def test_put_is_idempotent_for_new_count(self, tmp_path: Path) -> None: def test_flush_persists_and_reloads(self, tmp_path: Path) -> None: """Flushed entries survive a reload from disk.""" - cache_path = tmp_path / "cache.jsonl" + cache_path = tmp_path / "cache.sqlite3" cache = ContentAddressedCache(cache_path) cache.put("k1", "full_support") cache.put("k2", "partial_support") @@ -75,8 +146,8 @@ def test_flush_persists_and_reloads(self, tmp_path: Path) -> None: assert reloaded.get("k2") == "partial_support" def test_repeated_flush_does_not_duplicate(self, tmp_path: Path) -> None: - """Incremental flushes rewrite the file without duplicating entries.""" - cache_path = tmp_path / "cache.jsonl" + """Incremental flushes preserve one row per cache key.""" + cache_path = tmp_path / "cache.sqlite3" cache = ContentAddressedCache(cache_path) cache.put("k1", "full_support") @@ -84,18 +155,216 @@ def test_repeated_flush_does_not_duplicate(self, tmp_path: Path) -> None: cache.put("k2", "no_support") cache.flush() - line_count = sum( - 1 for line in cache_path.read_text(encoding="utf-8").splitlines() if line - ) - assert line_count == 2 + with sqlite3.connect(cache_path) as connection: + row_count = connection.execute( + "SELECT COUNT(*) FROM cache_entries" + ).fetchone() + assert row_count == (2,) reloaded = ContentAddressedCache(cache_path) assert reloaded.get("k1") == "full_support" assert reloaded.get("k2") == "no_support" def test_flush_noop_when_no_new_entries(self, tmp_path: Path) -> None: - """Flushing with no new entries leaves the file untouched.""" - cache_path = tmp_path / "cache.jsonl" + """Flushing with no new entries leaves the database empty.""" + cache_path = tmp_path / "cache.sqlite3" cache = ContentAddressedCache(cache_path) + assert cache.flush() == 0 + with sqlite3.connect(cache_path) as connection: + row_count = connection.execute( + "SELECT COUNT(*) FROM cache_entries" + ).fetchone() + assert row_count == (0,) + + def test_metadata_persists(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + metadata = build_cache_metadata(model="gpt-test", call_args={"temperature": 0}) + cache = ContentAddressedCache(cache_path) + cache.put("k1", "full_support", metadata) cache.flush() - assert not cache_path.exists() + + assert ContentAddressedCache(cache_path).get_metadata("k1") == metadata + + def test_imports_legacy_jsonl_cache(self, tmp_path: Path) -> None: + legacy_path = tmp_path / "cache.jsonl" + legacy_path.write_text( + json.dumps({"key": "legacy", "grade": "partial_support"}) + "\n", + encoding="utf-8", + ) + + cache = ContentAddressedCache(legacy_path) + + assert cache.cache_path == tmp_path / "cache.sqlite3" + assert cache.get("legacy") == "partial_support" + assert cache.get_metadata("legacy") == { + "schema_version": 1, + "source": "legacy_jsonl", + } + + def test_concurrent_processes_preserve_all_entries(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + entries = [(f"k{index}", f"grade-{index}") for index in range(12)] + + with ProcessPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(_write_cache_entry, str(cache_path), key, grade) + for key, grade in entries + ] + for future in futures: + future.result(timeout=30) + + cache = ContentAddressedCache(cache_path) + assert {key: cache.get(key) for key, _grade in entries} == dict(entries) + + def test_concurrent_same_key_uses_first_writer(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + grades = [f"grade-{index}" for index in range(8)] + + with ProcessPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(_write_cache_entry, str(cache_path), "shared", grade) + for grade in grades + ] + for future in futures: + future.result(timeout=30) + + assert ContentAddressedCache(cache_path).get("shared") in grades + with sqlite3.connect(cache_path) as connection: + row_count = connection.execute( + "SELECT COUNT(*) FROM cache_entries WHERE key = 'shared'" + ).fetchone() + assert row_count == (1,) + + def test_only_one_process_acquires_same_work(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + + with ProcessPoolExecutor(max_workers=4) as executor: + acquired = list( + executor.map( + _try_claim, + [str(cache_path)] * 8, + [f"owner-{index}" for index in range(8)], + ) + ) + + assert acquired.count(True) == 1 + assert inspect_cache(cache_path)["active_leases"] == 1 + + def test_concurrent_processes_compute_same_work_once(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + + with ProcessPoolExecutor(max_workers=4) as executor: + computed = list( + executor.map( + _compute_once, + [str(cache_path)] * 4, + [f"owner-{index}" for index in range(4)], + ) + ) + + assert computed.count(True) == 1 + store = SQLiteCache(cache_path, "coalescing-test") + assert store.get("shared") is not None + assert inspect_cache(cache_path)["active_leases"] == 0 + + def test_expired_lease_is_recovered_after_writer_interruption( + self, tmp_path: Path + ) -> None: + store = SQLiteCache(tmp_path / "cache.sqlite3", "lease-test") + + assert store.try_acquire("k1", "crashed", ttl_seconds=1, now=100) + assert not store.try_acquire("k1", "waiting", ttl_seconds=1, now=100.5) + assert store.try_acquire("k1", "recovered", ttl_seconds=1, now=101) + + def test_recovers_lease_abandoned_by_exited_process(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + with ProcessPoolExecutor(max_workers=1) as executor: + assert executor.submit(_acquire_then_exit, str(cache_path)).result( + timeout=30 + ) + + time.sleep(0.25) + store = SQLiteCache(cache_path, "interruption-test") + assert store.try_acquire("abandoned", "recovered", ttl_seconds=1) + + def test_publish_releases_lease_atomically(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + store = SQLiteCache(cache_path, "lease-test") + assert store.try_acquire("k1", "owner", ttl_seconds=30) + + inserted = store.publish( + "k1", + "value", + {"model": "test"}, + owner_id="owner", + logical_key="logical", + config_fingerprint="config", + ) + + assert inserted + assert store.get("k1") == ("value", {"model": "test"}) + assert inspect_cache(cache_path)["active_leases"] == 0 + + def test_migrates_v1_schema(self, tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + with sqlite3.connect(cache_path) as connection: + connection.execute( + "CREATE TABLE cache_properties (name TEXT PRIMARY KEY, value TEXT NOT NULL)" + ) + connection.execute( + "INSERT INTO cache_properties VALUES ('schema_version', '1')" + ) + connection.execute( + """ + CREATE TABLE cache_entries ( + namespace TEXT NOT NULL, + key TEXT NOT NULL, + value_json TEXT NOT NULL, + metadata_json TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (namespace, key) + ) + """ + ) + connection.execute( + """ + INSERT INTO cache_entries(namespace, key, value_json, metadata_json) + VALUES ('chunk_assertion', 'old', '"full_support"', '{}') + """ + ) + + store = SQLiteCache(cache_path, "chunk_assertion") + + assert store.get("old") == ("full_support", {}) + assert inspect_cache(cache_path)["schema_version"] == 2 + with sqlite3.connect(cache_path) as connection: + columns = { + row[1] for row in connection.execute("PRAGMA table_info(cache_entries)") + } + assert {"logical_key", "config_fingerprint"} <= columns + + def test_reports_alternate_configuration(self, tmp_path: Path) -> None: + cache = ContentAddressedCache(tmp_path / "cache.sqlite3") + assertion = "assertion" + chunk = "chunk" + first_metadata = build_cache_metadata(model="first") + first_key = compute_cache_key(assertion, chunk, model="first") + logical_key = compute_logical_key(assertion, chunk) + first_fingerprint = compute_config_fingerprint(model="first") + assert cache.claim(first_key, "owner") + cache.publish( + first_key, + "full_support", + first_metadata, + owner_id="owner", + logical_key=logical_key, + config_fingerprint=first_fingerprint, + ) + + mismatch_count, fields = cache.find_configuration_mismatches( + [(logical_key, compute_config_fingerprint(model="second"))], + build_cache_metadata(model="second"), + ) + + assert mismatch_count == 1 + assert fields == ["model"] diff --git a/tests/autoe/chunk_assertion/scoring_test.py b/tests/autoe/chunk_assertion/scoring_test.py index f2ed9f4..eeccdc5 100644 --- a/tests/autoe/chunk_assertion/scoring_test.py +++ b/tests/autoe/chunk_assertion/scoring_test.py @@ -59,6 +59,7 @@ async def _run( cache_path: Path, *, k_list: list[int] | None = None, + model: str = "test-model", ) -> dict[str, Any]: """Invoke run_assertion_eval_chunk_mode with test-friendly defaults.""" output_dir.mkdir(parents=True, exist_ok=True) @@ -73,7 +74,7 @@ async def _run( llm_client=cast("LLMCompletion", object()), llm_config=cast( "LLMConfig", - SimpleNamespace(concurrent_requests=2, call_args={}, model="test-model"), + SimpleNamespace(concurrent_requests=2, call_args={}, model=model), ), output_storage=FileStorage(base_dir=str(output_dir)), pass_threshold=0.5, @@ -208,3 +209,83 @@ async def test_second_run_uses_cache( await _run(eval_results, question_set, tmp_path / "out2", cache_path) assert call_counter[0] == calls_after_first # fully served from cache + + +@pytest.mark.usefixtures("patched_chat") +async def test_duplicate_pairs_share_one_llm_call( + tmp_path: Path, grade_by_chunk: dict[str, str], call_counter: list[int] +) -> None: + """Identical uncached pairs across questions share one in-process request.""" + grade_by_chunk["same chunk"] = "full_support" + eval_results = [ + { + "question_id": question_id, + "text": "Q", + "context": [{"text": "same chunk", "chunk_id": "c1", "rank": 1}], + } + for question_id in ("q1", "q2") + ] + question_set = { + "assertions": [ + { + "question_id": question_id, + "question_text": "Q", + "assertions": [{"statement": "same assertion"}], + } + for question_id in ("q1", "q2") + ] + } + + summaries = await _run( + eval_results, + question_set, + tmp_path / "out", + tmp_path / "cache.sqlite3", + ) + + assert call_counter[0] == 1 + assert summaries["all"].successful_calls == 2 + + +@pytest.mark.usefixtures("patched_chat") +async def test_warns_when_cached_configuration_differs( + tmp_path: Path, + grade_by_chunk: dict[str, str], + capsys: pytest.CaptureFixture[str], +) -> None: + grade_by_chunk["chunk"] = "full_support" + eval_results = [ + { + "question_id": "q1", + "text": "Q", + "context": [{"text": "chunk", "chunk_id": "c1", "rank": 1}], + } + ] + question_set = { + "assertions": [ + { + "question_id": "q1", + "question_text": "Q", + "assertions": [{"statement": "assertion"}], + } + ] + } + cache_path = tmp_path / "cache.sqlite3" + await _run( + eval_results, + question_set, + tmp_path / "out1", + cache_path, + model="first-model", + ) + capsys.readouterr() + + await _run( + eval_results, + question_set, + tmp_path / "out2", + cache_path, + model="second-model", + ) + + assert "same inputs with different model" in capsys.readouterr().out diff --git a/tests/autoe/relevance_cache_test.py b/tests/autoe/relevance_cache_test.py new file mode 100644 index 0000000..267d61e --- /dev/null +++ b/tests/autoe/relevance_cache_test.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Tests for retrieval relevance caching.""" + +import json +from pathlib import Path +from typing import Any + +import pytest + +from benchmark_qed.autod.data_model.text_unit import TextUnit +from benchmark_qed.autoe.data_model.relevance import ( + RelevanceAssessmentItem, + RelevanceAssessmentResponse, +) +from benchmark_qed.autoe.retrieval_metrics.relevance_assessment.base import ( + RelevanceRater, +) +from benchmark_qed.cache import SQLiteCache + + +class StubRelevanceRater(RelevanceRater): + def __init__( + self, + cache_dir: Path | None, + *, + rater_params: dict[str, Any] | None = None, + ) -> None: + self.calls: list[list[str]] = [] + self.rater_params = rater_params or {"model": "test-model"} + super().__init__(cache_dir=cache_dir) + + def _get_cache_relevant_params(self) -> dict[str, Any]: + return self.rater_params + + async def _rate_relevance_impl( + self, query: str, text_units: list[TextUnit] + ) -> RelevanceAssessmentResponse: + self.calls.append([unit.id for unit in text_units]) + return RelevanceAssessmentResponse( + assessment=[ + RelevanceAssessmentItem( + text_unit=unit, + reasoning=f"{query}:{unit.text}", + score=len(unit.text) % 4, + ) + for unit in text_units + ] + ) + + +def _text_unit(unit_id: str, text: str) -> TextUnit: + return TextUnit(id=unit_id, short_id=None, text=text) + + +async def test_second_rater_uses_sqlite_cache(tmp_path: Path) -> None: + units = [_text_unit("one", "first"), _text_unit("two", "second")] + first = StubRelevanceRater(tmp_path) + expected = await first.rate_relevance("query", units) + + second = StubRelevanceRater(tmp_path) + actual = await second.rate_relevance("query", units) + + assert len(first.calls) == 1 + assert second.calls == [] + assert actual == expected + assert second.cache_hits == 2 + assert second.get_cache_stats()["cache_files"] == 2 + + +async def test_duplicate_pairs_share_one_assessment(tmp_path: Path) -> None: + rater = StubRelevanceRater(tmp_path) + + result = await rater.rate_relevance( + "query", + [_text_unit("one", "same"), _text_unit("two", " SAME ")], + ) + + assert rater.calls == [["one"]] + assert len(result.assessment) == 2 + assert result.assessment[0] == result.assessment[1] + assert rater.cache_misses == 2 + assert rater.get_cache_stats()["cache_files"] == 1 + + +async def test_configuration_metadata_is_redacted(tmp_path: Path) -> None: + params = { + "model": "test-model", + "call_args": {"temperature": 0, "api_key": "secret"}, + } + rater = StubRelevanceRater(tmp_path, rater_params=params) + unit = _text_unit("one", "text") + await rater.rate_relevance("query", [unit]) + + key = rater._generate_cache_key("query", unit, params) + entry = SQLiteCache(tmp_path / "relevance_cache.sqlite3", "StubRelevanceRater").get( + key + ) + + assert entry is not None + assert entry[1]["rater_params"]["call_args"] == { + "api_key": "", + "temperature": 0, + } + assert "secret" not in json.dumps(entry) + + +async def test_imports_legacy_per_key_json(tmp_path: Path) -> None: + unit = _text_unit("one", "legacy") + key_builder = StubRelevanceRater(None) + key = key_builder._generate_cache_key( + "query", unit, key_builder._get_cache_relevant_params() + ) + assessment = RelevanceAssessmentItem( + text_unit=unit, + reasoning="from legacy cache", + score=3, + ) + (tmp_path / f"{key}.json").write_text( + json.dumps({ + "timestamp": "2026-09-11T00:00:00+00:00", + "query": "query", + "assessment_item": assessment.model_dump(), + }), + encoding="utf-8", + ) + + rater = StubRelevanceRater(tmp_path) + result = await rater.rate_relevance("query", [unit]) + + assert rater.calls == [] + assert result.assessment == [assessment] + + +async def test_clear_cache_removes_sqlite_and_legacy_entries(tmp_path: Path) -> None: + rater = StubRelevanceRater(tmp_path) + await rater.rate_relevance("query", [_text_unit("one", "text")]) + legacy_file = tmp_path / "legacy.json" + legacy_file.write_text("{}", encoding="utf-8") + + rater.clear_cache() + + assert rater.get_cache_stats()["cache_files"] == 0 + assert not legacy_file.exists() + assert rater.cache_hits == 0 + assert rater.cache_misses == 0 + + +async def test_invalid_result_count_is_reported(tmp_path: Path) -> None: + class InvalidRater(StubRelevanceRater): + async def _rate_relevance_impl( + self, query: str, text_units: list[TextUnit] + ) -> RelevanceAssessmentResponse: + return RelevanceAssessmentResponse(assessment=[]) + + rater = InvalidRater(tmp_path) + + with pytest.raises(RuntimeError, match="returned 0 results for 1"): + await rater.rate_relevance("query", [_text_unit("one", "text")]) + + +async def test_warns_when_relevance_configuration_differs( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + unit = _text_unit("one", "text") + first = StubRelevanceRater(tmp_path, rater_params={"model": "first"}) + await first.rate_relevance("query", [unit]) + + second = StubRelevanceRater(tmp_path, rater_params={"model": "second"}) + with caplog.at_level("WARNING"): + await second.rate_relevance("query", [unit]) + + assert "same inputs with different rater_params" in caplog.text diff --git a/tests/test_cache_cli.py b/tests/test_cache_cli.py new file mode 100644 index 0000000..147a08e --- /dev/null +++ b/tests/test_cache_cli.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Tests for cache inspection commands.""" + +import json +from pathlib import Path + +from typer.testing import CliRunner + +from benchmark_qed.__main__ import app +from benchmark_qed.cache import SQLiteCache + + +def test_cache_inspect_json(tmp_path: Path) -> None: + cache_path = tmp_path / "cache.sqlite3" + store = SQLiteCache(cache_path, "test") + store.put_many( + [("key", {"score": 3}, {"model": "test-model"})], + identities={"key": ("logical", "config")}, + ) + + result = CliRunner().invoke(app, ["cache", "inspect", str(cache_path), "--json"]) + + assert result.exit_code == 0 + details = json.loads(result.stdout) + assert details["schema_version"] == 2 + assert details["active_leases"] == 0 + assert details["namespaces"] == [ + {"namespace": "test", "entries": 1, "configurations": 1} + ] + assert details["recent_provenance"][0]["metadata"]["model"] == "test-model"