From fb8d9b2e6796cb3d155d65efc908e76d8518ea16 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 21 Apr 2026 14:20:22 +0200 Subject: [PATCH] fix(estimator): check index version and harden meta/per_prefix parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ensure_index_fresh now rebuilds on ESTIMATOR_INDEX_VERSION mismatch, so a schema bump actually invalidates the cached pickle. - Corrupt meta.json rebuilds instead of crashing. - Index builder no longer falls back to the whole cache dict when per_prefix is missing — that previously treated top-level keys (algo_version, created_at) as prefix stats. Co-Authored-By: Claude Opus 4.7 --- lora_estimator.py | 16 +++++++++++- scripts/build_estimator_index.py | 4 +-- tests/test_estimator_freshness.py | 41 +++++++++++++++++++++++++------ 3 files changed, 51 insertions(+), 10 deletions(-) diff --git a/lora_estimator.py b/lora_estimator.py index 8fe5415..43d2f11 100644 --- a/lora_estimator.py +++ b/lora_estimator.py @@ -169,12 +169,26 @@ def ensure_index_fresh( if not meta_path.exists(): rebuild_fn() return + try: + cached = json.loads(meta_path.read_text()) + except (json.JSONDecodeError, OSError) as e: + _log.warning("[LoRA Estimator] meta.json unreadable (%s) — rebuilding", e) + rebuild_fn() + return + # Schema-version mismatch trumps SHA: a new feature layout can't read an old pickle. + cached_version = cached.get("estimator_index_version") + if cached_version != ESTIMATOR_INDEX_VERSION: + _log.info( + "[LoRA Estimator] Index schema changed (%s → %s), rebuilding…", + cached_version or "?", ESTIMATOR_INDEX_VERSION, + ) + rebuild_fn() + return try: current_sha = _fetch_hf_head_sha(repo_id) except Exception as e: _log.warning("[LoRA Estimator] HF SHA check failed — using cached index: %s", e) return - cached = json.loads(meta_path.read_text()) cached_sha = cached.get("hf_commit_sha") if cached_sha != current_sha: _log.info( diff --git a/scripts/build_estimator_index.py b/scripts/build_estimator_index.py index 0399211..1f71a05 100644 --- a/scripts/build_estimator_index.py +++ b/scripts/build_estimator_index.py @@ -64,7 +64,7 @@ def load_phase1_for_config(dataset_dir: Path, sorted_hashes) -> tuple[dict, dict if lora_path.exists(): raw = _load_json(lora_path) # The cache format stores stats under `per_prefix`. - lora_stats[i] = {"per_prefix": raw.get("per_prefix", raw)} + lora_stats[i] = {"per_prefix": raw.get("per_prefix", {})} for i, ha in enumerate(sorted_hashes): for j, hb in enumerate(sorted_hashes): if i >= j: @@ -73,7 +73,7 @@ def load_phase1_for_config(dataset_dir: Path, sorted_hashes) -> tuple[dict, dict pair_path = dataset_dir / "pair" / f"{key_a}_{key_b}.pair.json" if pair_path.exists(): raw = _load_json(pair_path) - pair_stats[(i, j)] = {"per_prefix": raw.get("per_prefix", raw)} + pair_stats[(i, j)] = {"per_prefix": raw.get("per_prefix", {})} return pair_stats, lora_stats diff --git a/tests/test_estimator_freshness.py b/tests/test_estimator_freshness.py index 62d9de9..0d8c135 100644 --- a/tests/test_estimator_freshness.py +++ b/tests/test_estimator_freshness.py @@ -6,12 +6,19 @@ from pathlib import Path from unittest import mock +def _write_meta(dir_path: Path, sha: str = "abc", version: str | None = None): + """Write a meta.json with a valid version by default (so version-check passes).""" + from lora_estimator import ESTIMATOR_INDEX_VERSION + meta = {"hf_commit_sha": sha, + "estimator_index_version": version or ESTIMATOR_INDEX_VERSION} + (dir_path / "meta.json").write_text(json.dumps(meta)) + + class TestFreshness(unittest.TestCase): def test_stale_sha_triggers_rebuild(self): from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - meta_path = Path(tmp) / "meta.json" - meta_path.write_text(json.dumps({"hf_commit_sha": "abc"})) + _write_meta(Path(tmp), sha="abc") with mock.patch("lora_estimator._fetch_hf_head_sha", return_value="xyz"): rebuild_fn = mock.Mock() ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="auto") @@ -20,7 +27,7 @@ class TestFreshness(unittest.TestCase): def test_fresh_sha_no_rebuild(self): from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - (Path(tmp) / "meta.json").write_text(json.dumps({"hf_commit_sha": "abc"})) + _write_meta(Path(tmp), sha="abc") with mock.patch("lora_estimator._fetch_hf_head_sha", return_value="abc"): rebuild_fn = mock.Mock() ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="auto") @@ -29,7 +36,7 @@ class TestFreshness(unittest.TestCase): def test_force_mode_always_rebuilds(self): from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - (Path(tmp) / "meta.json").write_text(json.dumps({"hf_commit_sha": "abc"})) + _write_meta(Path(tmp)) rebuild_fn = mock.Mock() ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="force") rebuild_fn.assert_called_once() @@ -37,7 +44,7 @@ class TestFreshness(unittest.TestCase): def test_skip_mode_never_rebuilds(self): from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - (Path(tmp) / "meta.json").write_text(json.dumps({"hf_commit_sha": "abc"})) + _write_meta(Path(tmp)) with mock.patch("lora_estimator._fetch_hf_head_sha", return_value="xyz"): rebuild_fn = mock.Mock() ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="skip") @@ -46,7 +53,7 @@ class TestFreshness(unittest.TestCase): def test_no_internet_falls_back_to_cache(self): from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - (Path(tmp) / "meta.json").write_text(json.dumps({"hf_commit_sha": "abc"})) + _write_meta(Path(tmp)) with mock.patch("lora_estimator._fetch_hf_head_sha", side_effect=ConnectionError("no network")): rebuild_fn = mock.Mock() @@ -57,7 +64,27 @@ class TestFreshness(unittest.TestCase): """No meta.json on disk → must rebuild regardless of network state.""" from lora_estimator import ensure_index_fresh with tempfile.TemporaryDirectory() as tmp: - # No meta.json written. + rebuild_fn = mock.Mock() + ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="auto") + rebuild_fn.assert_called_once() + + def test_stale_version_triggers_rebuild_without_network(self): + """Index-version mismatch rebuilds unconditionally (schema change must invalidate pickle).""" + from lora_estimator import ensure_index_fresh + with tempfile.TemporaryDirectory() as tmp: + _write_meta(Path(tmp), sha="abc", version="0.0.1-old") + # Network should not even be consulted when the schema is stale. + with mock.patch("lora_estimator._fetch_hf_head_sha") as m_sha: + rebuild_fn = mock.Mock() + ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="auto") + rebuild_fn.assert_called_once() + m_sha.assert_not_called() + + def test_corrupt_meta_triggers_rebuild(self): + """Malformed meta.json should rebuild rather than crash.""" + from lora_estimator import ensure_index_fresh + with tempfile.TemporaryDirectory() as tmp: + (Path(tmp) / "meta.json").write_text("{not valid json") rebuild_fn = mock.Mock() ensure_index_fresh(Path(tmp), rebuild_fn=rebuild_fn, mode="auto") rebuild_fn.assert_called_once()