mirror of
https://github.com/Dynamis-Labs/spectralquant.git
synced 2026-10-21 09:49:45 -05:00
Ashwin GopinathandAnirudh B. Vangara
19ebecbc4d
Upgrades SpectralQuant from uniform bit allocation to greedy water-filling inside the semantic subspace, with per-dimension Lloyd-Max codebooks fit to the empirical scalar marginal. Adds full three-way evaluation harness (perplexity, generation, latency, LongBench) on Modal H200 GPUs, comprehensive documentation, and the consolidated arxiv-ready manuscript. v1 is preserved on the 'v1' branch. Co-Authored-By: Anirudh B. Vangara <anirudh@sentra.app>
420 lines
14 KiB
Python
420 lines
14 KiB
Python
"""Tests for ``run_longbench``'s per-example generation progress and the
|
|
safe partial-checkpoint discipline added in response to the 2026-04-30
|
|
post-calibration generation-stage stall.
|
|
|
|
Coverage:
|
|
|
|
* ``GenerationProgress.start/step/end`` emits ``eval_task_start`` /
|
|
``eval_progress`` / ``eval_task_end`` events with the expected
|
|
``method``, ``task``, ``completed``, ``total``, and ``elapsed_s`` keys
|
|
in ``details``.
|
|
* Per-example ``eval_progress`` actually fires every ``progress_every_n``
|
|
examples, plus on the first and last steps.
|
|
* The per-(method, task) partial JSON shard is rewritten on every step
|
|
and is tagged ``paper_valid: false, partial: true``.
|
|
* The run-level ``partial_status.json`` records planned/completed
|
|
methods + tasks and is tagged ``paper_valid: false, partial: true``.
|
|
* Partial artifacts are NEVER written into the canonical result path.
|
|
* ``GenerationProgress`` is a no-op when ``status`` is None.
|
|
* The new stages appear in ``run_status.KNOWN_STAGES``.
|
|
|
|
These tests do not load any HF model — they exercise the progress
|
|
instrumentation directly and via small in-memory fakes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _add_paths(monkeypatch):
|
|
monkeypatch.syspath_prepend(str(REPO_ROOT))
|
|
monkeypatch.syspath_prepend(str(REPO_ROOT / "src"))
|
|
|
|
|
|
def _make_status(tmp_path: Path):
|
|
from experiments import run_status
|
|
|
|
return run_status.StatusWriter(
|
|
status_dir=tmp_path / "status_root" / "rid",
|
|
run_id="rid",
|
|
commit="abc",
|
|
model="m",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stage registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_eval_task_start_and_end_are_known_stages():
|
|
from experiments import run_status
|
|
|
|
assert "eval_task_start" in run_status.KNOWN_STAGES
|
|
assert "eval_task_end" in run_status.KNOWN_STAGES
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GenerationProgress event emission
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _read_events(status) -> List[Dict[str, Any]]:
|
|
p = Path(status.events_path)
|
|
if not p.exists():
|
|
return []
|
|
return [json.loads(line) for line in p.read_text().splitlines()]
|
|
|
|
|
|
def test_generation_progress_emits_start_per_step_and_end(tmp_path):
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="spectralquant_v2", task="narrativeqa", total=3,
|
|
)
|
|
prog.start()
|
|
prog.step(completion="answer-1")
|
|
prog.step(completion="answer-2")
|
|
prog.step(completion="answer-3")
|
|
prog.end(score=0.42)
|
|
|
|
events = _read_events(status)
|
|
stages = [e["stage"] for e in events]
|
|
assert "eval_task_start" in stages
|
|
assert stages.count("eval_progress") == 3 # one per example
|
|
assert "eval_task_end" in stages
|
|
|
|
# Every progress event carries method/task/completed/total/elapsed_s.
|
|
for e in events:
|
|
if e["stage"] != "eval_progress":
|
|
continue
|
|
det = e["details"]
|
|
assert det["method"] == "spectralquant_v2"
|
|
assert det["task"] == "narrativeqa"
|
|
assert isinstance(det["completed"], int)
|
|
assert det["total"] == 3
|
|
assert isinstance(det["elapsed_s"], (int, float))
|
|
|
|
# eval_task_end carries the score.
|
|
end = next(e for e in events if e["stage"] == "eval_task_end")
|
|
assert end["details"]["score"] == pytest.approx(0.42)
|
|
assert end["details"]["completed"] == 3
|
|
|
|
|
|
def test_generation_progress_progress_every_n_thinning(tmp_path):
|
|
"""progress_every_n=5 should still fire at step 1 (first), step 5
|
|
(boundary), and step 10 (last) — but not on every step."""
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="fp16", task="qasper", total=10,
|
|
progress_every_n=5,
|
|
)
|
|
prog.start()
|
|
for i in range(10):
|
|
prog.step(completion=f"c{i}")
|
|
prog.end()
|
|
|
|
events = _read_events(status)
|
|
progress_completed = [
|
|
e["details"]["completed"] for e in events
|
|
if e["stage"] == "eval_progress"
|
|
]
|
|
# First (1), 5 (boundary), 10 (last) — and 10 is also a multiple-of-5
|
|
# boundary, so we get exactly 3 events.
|
|
assert progress_completed == [1, 5, 10]
|
|
|
|
|
|
def test_generation_progress_no_status_is_noop(tmp_path):
|
|
from experiments import run_longbench
|
|
|
|
prog = run_longbench.GenerationProgress(
|
|
status=None, method="fp16", task="t", total=2,
|
|
)
|
|
prog.start()
|
|
prog.step(completion="x")
|
|
prog.step(completion="y")
|
|
prog.end()
|
|
# No partial dir created, no events.
|
|
assert not (tmp_path / "status_root").exists()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-(method, task) partial JSON shard
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_partial_task_shard_written_after_every_step(tmp_path):
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="spectralquant_v2", task="narrativeqa", total=4,
|
|
)
|
|
prog.start()
|
|
prog.step(completion="alpha")
|
|
shard_path = (
|
|
Path(status.status_dir) / run_longbench.PARTIAL_SUBDIR
|
|
/ "spectralquant_v2__narrativeqa.json"
|
|
)
|
|
assert shard_path.exists()
|
|
after_one = json.loads(shard_path.read_text())
|
|
assert after_one["completed"] == 1
|
|
assert after_one["total"] == 4
|
|
assert after_one["paper_valid"] is False
|
|
assert after_one["partial"] is True
|
|
assert after_one["status"] == "running"
|
|
assert after_one["kind"] == "longbench_partial_task"
|
|
assert after_one["last_completion_chars"] == len("alpha")
|
|
|
|
prog.step(completion="beta")
|
|
after_two = json.loads(shard_path.read_text())
|
|
assert after_two["completed"] == 2
|
|
|
|
prog.end(score=0.7)
|
|
final = json.loads(shard_path.read_text())
|
|
assert final["status"] == "ended"
|
|
assert final["score"] == pytest.approx(0.7)
|
|
|
|
|
|
def test_partial_task_shard_filename_is_sanitized(tmp_path):
|
|
"""A method or task containing path separators must not write
|
|
outside the partial directory."""
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="../escape", task="../t", total=1,
|
|
)
|
|
prog.start()
|
|
prog.step(completion="x")
|
|
pd = Path(status.status_dir) / run_longbench.PARTIAL_SUBDIR
|
|
files = list(pd.iterdir())
|
|
# Exactly one shard, sitting in the partial dir — not anywhere else.
|
|
assert len(files) == 1
|
|
assert files[0].parent.resolve() == pd.resolve()
|
|
# The shard name has no path separator, so it cannot escape pd even
|
|
# if a future filename change reintroduces "..".
|
|
assert "/" not in files[0].name
|
|
assert "\\" not in files[0].name
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Run-level partial_status.json
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_level_partial_status_paper_valid_false(tmp_path):
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
args = run_longbench.parse_args([
|
|
"--model", "m", "--output-dir", str(tmp_path),
|
|
"--methods", "fp16", "spectralquant_v2",
|
|
"--subset", "minimal", "--n-per-task", "5",
|
|
])
|
|
run_longbench._write_run_partial_status(
|
|
status, args=args, mode="full",
|
|
methods_planned=["fp16", "spectralquant_v2"],
|
|
tasks_planned=["narrativeqa", "qasper"],
|
|
current_method="spectralquant_v2",
|
|
current_task="narrativeqa",
|
|
completed_methods=["fp16"],
|
|
completed_tasks_per_method={"spectralquant_v2": []},
|
|
note="starting method=spectralquant_v2",
|
|
)
|
|
snap_path = (
|
|
Path(status.status_dir)
|
|
/ run_longbench.PARTIAL_SUBDIR
|
|
/ run_longbench.PARTIAL_STATUS_FILENAME
|
|
)
|
|
assert snap_path.exists()
|
|
snap = json.loads(snap_path.read_text())
|
|
assert snap["paper_valid"] is False
|
|
assert snap["partial"] is True
|
|
assert snap["kind"] == "longbench_partial_run"
|
|
assert snap["methods_planned"] == ["fp16", "spectralquant_v2"]
|
|
assert snap["completed_methods"] == ["fp16"]
|
|
assert snap["current_method"] == "spectralquant_v2"
|
|
assert snap["current_task"] == "narrativeqa"
|
|
assert snap["note"] == "starting method=spectralquant_v2"
|
|
|
|
|
|
def test_partial_directory_separate_from_canonical_result_path(tmp_path):
|
|
"""Partial artifacts live under <status_dir>/partial/, never at the
|
|
canonical <output_dir>/<run_id>.json result path. A partial JSON
|
|
must never be mistaken for a paper-valid result."""
|
|
from experiments import eval_common, run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="fp16", task="t", total=1,
|
|
)
|
|
prog.start()
|
|
prog.step(completion="x")
|
|
prog.end()
|
|
|
|
args = run_longbench.parse_args([
|
|
"--model", "m", "--output-dir", str(tmp_path / "out"),
|
|
"--methods", "fp16", "--subset", "minimal", "--n-per-task", "5",
|
|
])
|
|
canonical = eval_common.output_path(
|
|
args.output_dir, run_longbench.derive_run_id(args, "full"),
|
|
)
|
|
# Canonical result path was never created (partials only).
|
|
assert not canonical.exists()
|
|
# Partials are under status_dir / partial / and NOT inside output_dir.
|
|
pd = Path(status.status_dir) / run_longbench.PARTIAL_SUBDIR
|
|
assert pd.is_dir()
|
|
assert canonical not in pd.iterdir()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _generate_for_task threads progress through to per-example events.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeIds:
|
|
def __init__(self, n: int = 1) -> None:
|
|
self._n = n
|
|
|
|
def to(self, *_a, **_k):
|
|
return self
|
|
|
|
def size(self, _dim: int) -> int:
|
|
return self._n
|
|
|
|
|
|
class _FakeBatch:
|
|
def __init__(self, ids: _FakeIds) -> None:
|
|
self.input_ids = ids
|
|
|
|
|
|
class _FakeTokenizer:
|
|
def __call__(self, *_a, **_k):
|
|
return _FakeBatch(_FakeIds(1))
|
|
|
|
def decode(self, ids, **_k):
|
|
# Return a deterministic tag using the integer state of the ids.
|
|
return f"completion-{getattr(ids, '_tag', 0)}"
|
|
|
|
|
|
class _FakeOutIds:
|
|
def __init__(self, tag: int) -> None:
|
|
self._tag = tag
|
|
|
|
def __getitem__(self, _slice):
|
|
return self # pretend slicing returns same fake
|
|
|
|
|
|
class _FakeModel:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def generate(self, *_a, **_k):
|
|
self.calls += 1
|
|
return _FakeOutIds(self.calls)
|
|
|
|
|
|
class _FakeRow:
|
|
def __init__(self, prompt: str) -> None:
|
|
self.prompt = prompt
|
|
self.expected_max_new_tokens = 8
|
|
|
|
|
|
def test_generate_for_task_emits_per_example_events(tmp_path, monkeypatch):
|
|
"""``_generate_for_task`` should drive the progress object once per
|
|
row. We avoid loading HF/torch by stubbing the import."""
|
|
fake_torch = type(sys)("torch")
|
|
|
|
class _NoGrad:
|
|
def __enter__(self): return self
|
|
def __exit__(self, *a): return False
|
|
fake_torch.no_grad = _NoGrad
|
|
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
|
|
|
from experiments import run_longbench
|
|
|
|
status = _make_status(tmp_path)
|
|
args = run_longbench.parse_args([
|
|
"--model", "m", "--output-dir", str(tmp_path),
|
|
"--methods", "fp16", "--subset", "minimal", "--n-per-task", "3",
|
|
"--max-input-tokens", "128", "--max-new-tokens", "8",
|
|
])
|
|
rows = [_FakeRow(f"p{i}") for i in range(3)]
|
|
prog = run_longbench.GenerationProgress(
|
|
status, method="fp16", task="narrativeqa", total=len(rows),
|
|
)
|
|
completions = run_longbench._generate_for_task(
|
|
_FakeModel(), _FakeTokenizer(), device="cpu", args=args, rows=rows,
|
|
progress=prog,
|
|
)
|
|
prog.end()
|
|
assert len(completions) == 3
|
|
|
|
events = _read_events(status)
|
|
progress_events = [e for e in events if e["stage"] == "eval_progress"]
|
|
completeds = [e["details"]["completed"] for e in progress_events]
|
|
assert completeds == [1, 2, 3]
|
|
|
|
# Per-task partial JSON exists and shows ended.
|
|
shard = (
|
|
Path(status.status_dir) / run_longbench.PARTIAL_SUBDIR
|
|
/ "fp16__narrativeqa.json"
|
|
)
|
|
assert shard.exists()
|
|
final = json.loads(shard.read_text())
|
|
assert final["completed"] == 3
|
|
assert final["total"] == 3
|
|
assert final["status"] == "ended"
|
|
assert final["paper_valid"] is False
|
|
assert final["partial"] is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Paper-valid discipline: partial outputs do NOT bypass paper-valid gates.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_partial_artifacts_do_not_make_payload_paper_valid(tmp_path):
|
|
"""A run that left only partial artifacts must not be able to claim
|
|
paper_valid=true. We simulate by feeding a result JSON's methods
|
|
through ``build_payload`` and asserting that even without partials
|
|
the gate logic continues to work — i.e. partial artifacts on disk
|
|
are independent of the schema-validated payload."""
|
|
from experiments import run_longbench, eval_common
|
|
|
|
args = run_longbench.parse_args([
|
|
"--model", "m", "--output-dir", str(tmp_path),
|
|
"--methods", "fp16",
|
|
"--n-per-task", "50", "--subset", "minimal",
|
|
])
|
|
methods = {
|
|
"fp16": {
|
|
"label": "fp16",
|
|
"per_task": [{"task": "narrativeqa", "metric": "qa_f1_en",
|
|
"score": 0.4, "n_examples": 50}],
|
|
"aggregate": {"macro_score": 0.4},
|
|
"evidence_ids": ["x"],
|
|
"dataset_source": "huggingface_thudm",
|
|
},
|
|
}
|
|
payload = run_longbench.build_payload(
|
|
args, ["x"], eval_common.MODE_FULL, methods,
|
|
)
|
|
# Without spectralquant_v2 partials, FP16-only minimal is paper-valid
|
|
# subject to its own caveats (transparent subset). The test confirms
|
|
# that partial-status code paths did not silently flip the gate.
|
|
assert payload["paper_valid"] is True
|