"""Unit tests for precision_baseline_store — no server, no model loading, no HF network.""" from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") import json import os import tempfile import unittest from datetime import datetime, timezone from pathlib import Path from unittest.mock import MagicMock, patch from sglang.test import precision_baseline_store as hfs from sglang.test.test_utils import CustomTestCase # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_config() -> hfs.HfStoreConfig: return hfs.HfStoreConfig(repo="test/repo", revision="main") def _make_rows(n: int, *, model: str = "org/model", base_index: int = 0) -> list[dict]: return [ { "model": model, "run_path": f"org__model/2025/01/{i:02d}/run-abc123{i}", "date": f"2025-01-{i + base_index:02d}", "push_index": (i + base_index) * 1000, } for i in range(n) ] # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- class TestHfStoreConfig(CustomTestCase): def test_from_env_reads_required_var(self): with patch.dict( os.environ, {"SGLANG_PRECISION_HF_REPO": "my/repo"}, clear=False ): cfg = hfs.HfStoreConfig.from_env() self.assertEqual(cfg.repo, "my/repo") self.assertEqual(cfg.revision, "main") def test_from_env_reads_optional_revision(self): with patch.dict( os.environ, { "SGLANG_PRECISION_HF_REPO": "my/repo", "SGLANG_PRECISION_HF_REVISION": "dev", }, clear=False, ): cfg = hfs.HfStoreConfig.from_env() self.assertEqual(cfg.revision, "dev") def test_from_env_raises_when_missing(self): with patch.dict(os.environ, {}, clear=True): with self.assertRaises(RuntimeError): hfs.HfStoreConfig.from_env() class TestSanitizeModelName(CustomTestCase): def test_slashes_and_spaces(self): self.assertEqual(hfs._sanitize_model_name("org/model name"), "org__model_name") def test_no_changes_needed(self): self.assertEqual(hfs._sanitize_model_name("simple"), "simple") class TestRowRecencyKey(CustomTestCase): def test_uses_explicit_push_index(self): row = {"push_index": 100} self.assertEqual(hfs._row_recency_key(row, 5), (100, 5)) def test_falls_back_to_index(self): row = {} self.assertEqual(hfs._row_recency_key(row, 5), (-1, 5)) def test_invalid_push_index(self): row = {"push_index": "bad"} self.assertEqual(hfs._row_recency_key(row, 5), (-1, 5)) def test_none_push_index(self): row = {"push_index": None} self.assertEqual(hfs._row_recency_key(row, 5), (-1, 5)) class TestSelectLatestRun(CustomTestCase): def test_picks_highest_recency(self): rows = _make_rows(3) result = hfs._select_latest_run(rows, model="org/model") self.assertEqual(result, rows[-1]["run_path"]) def test_filters_by_model(self): rows = [ {"model": "a/model", "run_path": "a", "push_index": 1}, {"model": "b/model", "run_path": "b", "push_index": 2}, ] self.assertEqual(hfs._select_latest_run(rows, model="a/model"), "a") def test_filters_by_capture_signature(self): rows = [ { "model": "org/m", "run_path": "old", "capture_signature": "abc123", "push_index": 1, }, { "model": "org/m", "run_path": "new", "capture_signature": "def456", "push_index": 2, }, ] self.assertEqual( hfs._select_latest_run(rows, model="org/m", capture_signature="def456"), "new", ) def test_returns_none_on_empty(self): self.assertIsNone(hfs._select_latest_run([], model="org/m")) def test_skips_rows_without_run_path(self): rows = [ {"model": "org/m", "push_index": 1}, {"model": "org/m", "run_path": "good", "push_index": 2}, ] self.assertEqual(hfs._select_latest_run(rows, model="org/m"), "good") def test_returns_none_when_signature_mismatch(self): rows = [ { "model": "org/m", "run_path": "old", "capture_signature": "abc123", "push_index": 1, }, ] self.assertIsNone( hfs._select_latest_run(rows, model="org/m", capture_signature="zzz") ) def test_prefers_older_passed_over_newer_failed(self): # A failed run must not shadow an older good baseline, or a persistent # regression is masked after one night. rows = [ { "model": "org/m", "run_path": "good", "pass_label": "passed", "push_index": 1, }, { "model": "org/m", "run_path": "bad", "pass_label": "failed", "push_index": 2, }, ] self.assertEqual(hfs._select_latest_run(rows, model="org/m"), "good") def test_prefers_baseline_established_over_newer_failed(self): rows = [ { "model": "org/m", "run_path": "seed", "pass_label": "baseline_established", "push_index": 1, }, { "model": "org/m", "run_path": "bad", "pass_label": "failed", "push_index": 2, }, ] self.assertEqual(hfs._select_latest_run(rows, model="org/m"), "seed") def test_falls_back_to_failed_when_only_failed(self): rows = [ { "model": "org/m", "run_path": "bad1", "pass_label": "failed", "push_index": 1, }, { "model": "org/m", "run_path": "bad2", "pass_label": "failed", "push_index": 2, }, ] self.assertEqual(hfs._select_latest_run(rows, model="org/m"), "bad2") def test_missing_pass_label_treated_as_usable(self): # Legacy rows without pass_label stay usable as baselines. rows = [ {"model": "org/m", "run_path": "legacy", "push_index": 1}, { "model": "org/m", "run_path": "bad", "pass_label": "failed", "push_index": 2, }, ] self.assertEqual(hfs._select_latest_run(rows, model="org/m"), "legacy") class TestReadManifest(CustomTestCase): @patch("sglang.test.precision_baseline_store.hf_hub_download") def test_parses_valid_manifest(self, mock_download): content = ( '{"model":"a","run_path":"p1","push_index":1}\n' '{"model":"b","run_path":"p2","push_index":2}\n' ) tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) try: tmp.write(content) tmp.close() mock_download.return_value = tmp.name rows, text = hfs._read_manifest(_make_config()) finally: os.unlink(tmp.name) self.assertEqual(len(rows), 2) self.assertEqual(rows[0]["model"], "a") self.assertEqual(text, content) @patch("sglang.test.precision_baseline_store.hf_hub_download") def test_skips_blank_and_corrupt_lines(self, mock_download): content = ( '{"model":"a","run_path":"p1"}\n\nnot-json\n{"model":"b","run_path":"p2"}\n' ) tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) try: tmp.write(content) tmp.close() mock_download.return_value = tmp.name rows, _ = hfs._read_manifest(_make_config()) finally: os.unlink(tmp.name) self.assertEqual(len(rows), 2) @patch("sglang.test.precision_baseline_store.hf_hub_download") def test_returns_empty_on_not_found(self, mock_download): from huggingface_hub.errors import EntryNotFoundError mock_download.side_effect = EntryNotFoundError("not found") rows, text = hfs._read_manifest(_make_config()) self.assertEqual(rows, []) self.assertEqual(text, "") class TestFetchLatestBaseline(CustomTestCase): @patch("sglang.test.precision_baseline_store.snapshot_download") @patch.object(hfs, "_read_manifest") def test_downloads_and_copies_tensors(self, mock_manifest, mock_snapshot): rows = [ { "model": "org/m", "run_path": "org__m/2025/01/01/run-abc", "push_index": 1, } ] mock_manifest.return_value = (rows, "") with tempfile.TemporaryDirectory() as snap_dir: tensors = Path(snap_dir) / "org__m/2025/01/01/run-abc/tensors" tensors.mkdir(parents=True) (tensors / "layer0.pt").write_bytes(b"\x00") mock_snapshot.return_value = snap_dir with tempfile.TemporaryDirectory() as target: result = hfs.fetch_latest_baseline( config=_make_config(), model="org/m", target_tensors_dir=Path(target), ) self.assertEqual(result, "org__m/2025/01/01/run-abc") @patch.object(hfs, "_read_manifest") def test_returns_none_when_no_runs(self, mock_manifest): mock_manifest.return_value = ([], "") with tempfile.TemporaryDirectory() as target: result = hfs.fetch_latest_baseline( config=_make_config(), model="org/m", target_tensors_dir=Path(target), ) self.assertIsNone(result) @patch("sglang.test.precision_baseline_store.snapshot_download") @patch.object(hfs, "_read_manifest") def test_passes_capture_signature(self, mock_manifest, mock_snapshot): rows = [ { "model": "org/m", "run_path": "run_new", "capture_signature": "sig2", "push_index": 2, }, { "model": "org/m", "run_path": "run_old", "capture_signature": "sig1", "push_index": 1, }, ] mock_manifest.return_value = (rows, "") with tempfile.TemporaryDirectory() as snap_dir: tensors = Path(snap_dir) / "run_new/tensors" tensors.mkdir(parents=True) (tensors / "layer0.pt").write_bytes(b"\x00") mock_snapshot.return_value = snap_dir with tempfile.TemporaryDirectory() as target: result = hfs.fetch_latest_baseline( config=_make_config(), model="org/m", target_tensors_dir=Path(target), capture_signature="sig2", ) self.assertEqual(result, "run_new") class TestPushRun(CustomTestCase): """push_run deletes its temp manifest file in a finally block, so tests that inspect the manifest content must capture it via a side_effect on the mock upload_file *before* push_run cleans up.""" @staticmethod def _make_push_mocks(mock_manifest, mock_api_cls): mock_manifest.return_value = ([], "") mock_api = MagicMock() mock_api_cls.return_value = mock_api # Capture manifest text before push_run's finally block deletes it. captured = [] mock_api.upload_file.side_effect = lambda *a, **kw: captured.append( Path(kw["path_or_fileobj"]).read_text() ) return mock_api, captured @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_uploads_tensors_and_manifest(self, mock_manifest, mock_api_cls): mock_api, captured = self._make_push_mocks(mock_manifest, mock_api_cls) with tempfile.TemporaryDirectory() as tensor_dir: (Path(tensor_dir) / "layer0.pt").write_bytes(b"\x01") meta = {"tp_size": 8, "capture_signature": "abc", "hardware": "H200"} run_path = hfs.push_run( config=_make_config(), model="org/m", sglang_commit="abc1234567", today_tensors_dir=Path(tensor_dir), meta=meta, ) mock_api.upload_folder.assert_called_once() mock_api.upload_file.assert_called_once() row = json.loads(captured[0].strip().splitlines()[-1]) self.assertEqual(row["model"], "org/m") self.assertEqual(row["capture_signature"], "abc") self.assertEqual(row["tp_size"], 8) self.assertTrue(run_path.startswith("org__m/")) @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_skips_existing_tensors_unless_force(self, mock_manifest, mock_api_cls): # The run_path must match what push_run generates: model/date/sha7. # _today_path() returns today's date, so build the path accordingly. today_date, today_date_path = hfs._today_path() existing_run_path = f"org__m/{today_date_path}/run-abc1234" existing_row = { "model": "org/m", "run_path": existing_run_path, "date": today_date, "push_index": 1, } mock_manifest.return_value = ([existing_row], json.dumps(existing_row) + "\n") mock_api = MagicMock() mock_api_cls.return_value = mock_api # Capture pt file count before push_run cleans up the temp staging dir. captured_pt_count = [] mock_api.upload_folder.side_effect = lambda *a, **kw: captured_pt_count.append( len(list(Path(kw["folder_path"]).rglob("*.pt"))) ) with tempfile.TemporaryDirectory() as tensor_dir: (Path(tensor_dir) / "layer0.pt").write_bytes(b"\x01") hfs.push_run( config=_make_config(), model="org/m", sglang_commit="abc1234567", today_tensors_dir=Path(tensor_dir), meta={"tp_size": 8}, ) self.assertEqual(captured_pt_count[0], 0) @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_force_re_uploads(self, mock_manifest, mock_api_cls): # Use today's date so the run_path matches what push_run generates. today_date, today_date_path = hfs._today_path() existing_run_path = f"org__m/{today_date_path}/run-abc1234" existing_row = { "model": "org/m", "run_path": existing_run_path, "date": today_date, "push_index": 1, } mock_manifest.return_value = ([existing_row], json.dumps(existing_row) + "\n") mock_api = MagicMock() mock_api_cls.return_value = mock_api # Capture pt file count before push_run cleans up the temp staging dir. captured_pt_count = [] mock_api.upload_folder.side_effect = lambda *a, **kw: captured_pt_count.append( len(list(Path(kw["folder_path"]).rglob("*.pt"))) ) with tempfile.TemporaryDirectory() as tensor_dir: (Path(tensor_dir) / "layer0.pt").write_bytes(b"\x01") hfs.push_run( config=_make_config(), model="org/m", sglang_commit="abc1234567", today_tensors_dir=Path(tensor_dir), meta={"tp_size": 8}, force=True, ) self.assertGreater(captured_pt_count[0], 0) @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_manifest_row_promotes_keys(self, mock_manifest, mock_api_cls): _mock_api, captured = self._make_push_mocks(mock_manifest, mock_api_cls) with tempfile.TemporaryDirectory() as tensor_dir: (Path(tensor_dir) / "layer0.pt").write_bytes(b"\x01") meta = { "tp_size": 4, "hardware": "H100", "capture_signature": "sig1", "num_layers_compared": 10, "num_layers_passed": 10, "num_layers_failed": 0, "max_rel_diff": 0.001, "ci_run_id": "12345", "extra_key_not_promoted": True, } hfs.push_run( config=_make_config(), model="org/m", sglang_commit="abc1234567", today_tensors_dir=Path(tensor_dir), meta=meta, ) row = json.loads(captured[0].strip().splitlines()[-1]) for key in hfs._MANIFEST_PROMOTE_KEYS: if key in meta: self.assertEqual( row.get(key), meta[key], f"manifest missing promoted key: {key}", ) self.assertNotIn("extra_key_not_promoted", row) @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_includes_comparator_report(self, mock_manifest, mock_api_cls): mock_manifest.return_value = ([], "") mock_api = MagicMock() mock_api_cls.return_value = mock_api # Capture file existence before push_run cleans up the temp staging dir. captured_files = [] mock_api.upload_folder.side_effect = lambda *a, **kw: captured_files.append( list(Path(kw["folder_path"]).iterdir()) ) with tempfile.TemporaryDirectory() as tensor_dir: (Path(tensor_dir) / "layer0.pt").write_bytes(b"\x01") report_path = Path(tensor_dir) / "report.jsonl" report_path.write_text('{"type":"comparison_tensor"}\n') hfs.push_run( config=_make_config(), model="org/m", sglang_commit="abc1234567", today_tensors_dir=Path(tensor_dir), meta={"tp_size": 8}, comparator_report=report_path, ) staged_names = [f.name for f in captured_files[0]] self.assertIn("comparator_report.jsonl", staged_names) class TestPruneOldRuns(CustomTestCase): @patch.object(hfs, "_read_manifest") def test_keeps_recent_runs(self, mock_manifest): today = datetime.now(timezone.utc).strftime("%Y-%m-%d") rows = [ {"model": "org/m", "run_path": "recent", "date": today}, ] mock_manifest.return_value = (rows, json.dumps(rows[0]) + "\n") result = hfs.prune_old_runs(config=_make_config(), keep_days=30) self.assertIn("recent", result["kept"]) self.assertEqual(result["pruned"], []) @patch.object(hfs, "_read_manifest") def test_archives_one_per_week(self, mock_manifest): rows = [ {"model": "org/m", "run_path": "old1", "date": "2020-01-06"}, {"model": "org/m", "run_path": "old2", "date": "2020-01-07"}, {"model": "org/m", "run_path": "old3", "date": "2020-01-08"}, ] mock_manifest.return_value = (rows, "") result = hfs.prune_old_runs( config=_make_config(), keep_days=0, weekly_archive=True, dry_run=True ) self.assertEqual(len(result["kept"]), 1) self.assertEqual(result["kept"][0], "old3") self.assertEqual(len(result["pruned"]), 2) @patch.object(hfs, "_read_manifest") def test_prune_without_archive(self, mock_manifest): rows = [ {"model": "org/m", "run_path": "old1", "date": "2020-01-06"}, {"model": "org/m", "run_path": "old2", "date": "2020-01-07"}, ] mock_manifest.return_value = (rows, "") result = hfs.prune_old_runs( config=_make_config(), keep_days=0, weekly_archive=False, dry_run=True ) self.assertEqual(result["kept"], []) self.assertEqual(len(result["pruned"]), 2) @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_dry_run_does_not_delete(self, mock_manifest, mock_api_cls): rows = [ {"model": "org/m", "run_path": "old1", "date": "2020-01-06"}, ] mock_manifest.return_value = (rows, "") mock_api_cls.return_value = MagicMock() hfs.prune_old_runs(config=_make_config(), keep_days=0, dry_run=True) mock_api_cls.return_value.upload_file.assert_not_called() mock_api_cls.return_value.delete_folder.assert_not_called() @patch("sglang.test.precision_baseline_store.HfApi") @patch.object(hfs, "_read_manifest") def test_live_mode_deletes(self, mock_manifest, mock_api_cls): rows = [ {"model": "org/m", "run_path": "old1", "date": "2020-01-06"}, {"model": "org/m", "run_path": "old2", "date": "2020-01-07"}, ] mock_manifest.return_value = (rows, "") mock_api = MagicMock() mock_api_cls.return_value = mock_api result = hfs.prune_old_runs( config=_make_config(), keep_days=0, weekly_archive=True, dry_run=False ) self.assertEqual(len(result["kept"]), 1) self.assertEqual(len(result["pruned"]), 1) mock_api.upload_file.assert_called_once() mock_api.delete_folder.assert_called_once() @patch.object(hfs, "_read_manifest") def test_filters_by_model(self, mock_manifest): rows = [ {"model": "org/m1", "run_path": "m1_old", "date": "2020-01-06"}, {"model": "org/m2", "run_path": "m2_old", "date": "2020-01-06"}, ] mock_manifest.return_value = (rows, "") result = hfs.prune_old_runs( config=_make_config(), model="org/m1", keep_days=0, weekly_archive=False, dry_run=True, ) self.assertIn("m2_old", result["kept"]) self.assertIn("m1_old", result["pruned"]) class TestWithRetries(CustomTestCase): @patch("sglang.test.precision_baseline_store.time") def test_succeeds_on_first_attempt(self, mock_time): result = hfs._with_retries(lambda: 42, what="test") self.assertEqual(result, 42) mock_time.sleep.assert_not_called() @patch("sglang.test.precision_baseline_store.time") def test_retries_on_429(self, mock_time): from huggingface_hub.errors import HfHubHTTPError resp_429 = MagicMock() resp_429.status_code = 429 exc_429 = HfHubHTTPError("rate limited", response=resp_429) mock_op = MagicMock(side_effect=[exc_429, "ok"]) result = hfs._with_retries(mock_op, what="test", base_delay=0.01) self.assertEqual(result, "ok") mock_time.sleep.assert_called_once() @patch("sglang.test.precision_baseline_store.time") def test_raises_on_auth_error(self, mock_time): from huggingface_hub.errors import HfHubHTTPError resp_401 = MagicMock() resp_401.status_code = 401 exc_401 = HfHubHTTPError("unauthorized", response=resp_401) mock_op = MagicMock(side_effect=exc_401) with self.assertRaises(HfHubHTTPError): hfs._with_retries(mock_op, what="test") mock_time.sleep.assert_not_called() @patch("sglang.test.precision_baseline_store.time") def test_raises_after_max_attempts(self, mock_time): from huggingface_hub.errors import HfHubHTTPError resp_500 = MagicMock() resp_500.status_code = 500 exc = HfHubHTTPError("server error", response=resp_500) mock_op = MagicMock(side_effect=exc) with self.assertRaises(HfHubHTTPError): hfs._with_retries(mock_op, what="test", attempts=2, base_delay=0.001) self.assertEqual(mock_time.sleep.call_count, 1) if __name__ == "__main__": unittest.main()