| """ACC-SPLIT tests on synthetic files (no campaign data, no text): every layout, de-duplication, the role filter, the |
| estimators and the shift decomposition.""" |
| from __future__ import annotations |
|
|
| import json |
| import math |
| from collections.abc import Mapping |
| from pathlib import Path |
| from typing import cast |
|
|
| import pytest |
|
|
| from offload.drafter.acc_split import __main__ as cli |
| from offload.drafter.acc_split import analysis as a |
| from offload.drafter.acc_split import sources as s |
|
|
| SHA_A, SHA_B = 'a' * 64, 'b' * 64 |
| NAME = 'DeepSeek' |
|
|
|
|
| def _spec(rounds: float, drafts: float, accepted: float, pos: list[float]) -> dict[str, object]: |
| return {'rounds': rounds, 'proposed_tokens': drafts, 'accepted_tokens': accepted, |
| 'accepted_per_position': {str(i): v for i, v in enumerate(pos)}} |
|
|
|
|
| def _bench2(path: Path, rid: str, sha: str, workload: str, ctx: int, max_tokens: int, rounds: float, accepted: float, |
| *, ts: str = '2026-10-23T10:00:00Z', label: str = 'X', error: object = None, answer: int = 0, |
| tps: float = 50.0, extra: dict[str, object] | None = None) -> None: |
| run = {'run_id': rid, 'config_label': label, 'type': 'run', 'ts': ts, 'target_ctx': ctx, 'workload': workload, |
| 'concurrency': 1, 'cache_state': 'cold', 'spec': _spec(rounds, 3 * rounds, accepted, [accepted * .6, accepted * .4])} |
| req = {'run_id': rid, 'config_label': label, 'type': 'request', 'request_index': 0, 'prompt_sha256': sha, |
| 'prompt_tokens': 100, 'completion_tokens': max_tokens, 'reasoning_tokens': max_tokens - answer, |
| 'answer_tokens': answer, 'decode_tps': tps, 'error': error, 'step_events': int(rounds), |
| 'step_gap_ms': {'mean': 40.0, 'p50': 39.0}, 'settings': {'max_tokens': max_tokens}} |
| req.update(extra or {}) |
| with path.open('a') as f: |
| f.write(json.dumps(run) + '\n' + json.dumps(req) + '\n') |
|
|
|
|
| def test_bench2_dedupes_copies_and_drops_errors(tmp_path: Path) -> None: |
| one = tmp_path / 'r1' / 'A-short.jsonl' |
| one.parent.mkdir() |
| _bench2(one, 'r1', SHA_A, 'legal', 0, 1500, 700, 700) |
| _bench2(one, 'r2', SHA_A, 'legal', 0, 1500, 700, 350, error='boom') |
| copy = tmp_path / 'checkout' / 'A-short.jsonl' |
| copy.parent.mkdir() |
| copy.write_bytes(one.read_bytes()) |
| files = s.bench2_files([tmp_path]) |
| assert len(files) == 1 |
| obs = s.load_bench2(files + [copy]) |
| assert [o.where for o in obs] == ['r1'] |
| assert obs[0].accepted_per_round == 1.0 and obs[0].depth == 3.0 and obs[0].cell == 'legal-0' |
|
|
|
|
| def _metrics(drafts: float, draft_tokens: float, accepted: float, pos: list[float]) -> dict[str, float]: |
| m = {f'{s.DRAFTS}{{engine="0",model_name="{NAME}"}}': drafts, f'{s.DRAFT_TOKENS}{{engine="0"}}': draft_tokens, |
| f'{s.ACCEPTED}{{engine="0"}}': accepted} |
| m.update({f'{s.PER_POS}{{engine="0",position="{i}"}}': v for i, v in enumerate(pos)}) |
| return m |
|
|
|
|
| def _windows(root: Path) -> None: |
| sp1 = root / 'SP1-W-20261026T144153Z' / 'cells' |
| sp1.mkdir(parents=True) |
| stream = {'prompt': 'short', 'prompt_tokens': 1014, 'prompt_sha256': SHA_B, 'reasoning_tokens': 1500, |
| 'answer_and_control_tokens': 0.0, 'sampling': {'max_tokens': 1500}} |
| (sp1 / 'ctl-short-b1.json').write_text(json.dumps({'streams': [stream], 'spec': _spec(600, 1800, 690, [400, 200, 90])})) |
| (sp1 / 'warmup.json').write_text(json.dumps({'streams': [stream], 'spec': _spec(10, 30, 10, [10])})) |
| pc1 = root / 'PC1-x-W-20261027T100908Z' |
| pc1.mkdir() |
| (pc1 / 'summary.json').write_text(json.dumps({'starts': {'s1-control': {'role': 'control'}, |
| 's2-cand': {'role': 'candidate'}}})) |
| rep_stream = {'prompt_sha256': SHA_B, 'prompt_tokens': 1014, 'reasoning_tokens': 512, 'answer_and_control_tokens': 0.0, |
| 'sampling': {'max_tokens': 512}, 'decode_tps': 55.0} |
| for start, acc in (('s1-control', 260.0), ('s2-cand', 100.0)): |
| reps = [{'streams': [rep_stream], 'metrics_before': _metrics(100, 300, 50, [30, 15, 5]), |
| 'metrics_after': _metrics(300, 1000, 50 + acc, [30 + acc * .6, 15 + acc * .3, 5 + acc * .1]), |
| 'server_decode_ms': 7800.0, 'spec_request_rounds': 200.0}] |
| (pc1 / f'{start}-short-b1.json').write_text(json.dumps({'cell': 'short-b1', 'repetitions': reps})) |
| (pc1 / f'{start}-greedy.json').write_text(json.dumps({'cell': 'greedy'})) |
|
|
|
|
| def test_window_layouts_and_role_filter(tmp_path: Path) -> None: |
| _windows(tmp_path) |
| obs = s.load_windows(s.window_dirs(tmp_path)) |
| by = {(o.label[:3], o.role): o for o in obs} |
| assert len(obs) == 3 |
| sp1 = by[('SP1', 'control')] |
| assert sp1.cell == 'ctl-short-b1' and sp1.accepted_per_round == pytest.approx(1.15) and sp1.date == '2026-10-26' |
| pc1 = by[('PC1', 'control')] |
| assert pc1.rounds == 200 and pc1.accepted_per_round == pytest.approx(1.3) and pc1.depth == pytest.approx(3.5) |
| assert pc1.ms_per_round == pytest.approx(39.0) and pc1.max_tokens == 512 and pc1.per_position[0] == pytest.approx(156) |
| rows = a.by_prompt(obs) |
| assert {(r['max_tokens'], record_n(r)) for r in rows} == {(512, 1), (1500, 1)} |
|
|
|
|
| def record_n(row: Mapping[str, object]) -> int: |
| acc = row['accepted_per_round'] |
| assert isinstance(acc, dict) |
| n = acc['n'] |
| assert isinstance(n, int) |
| return n |
|
|
|
|
| def test_summarize_interval() -> None: |
| sm = a.summarize([1.0, 1.2, 1.4]) |
| assert sm.mean == pytest.approx(1.2) and sm.sd == pytest.approx(.2) |
| assert sm.hi - sm.mean == pytest.approx(4.303 * .2 / math.sqrt(3)) |
| assert a.summarize([2.0]).lo == 2.0 |
| with pytest.raises(ValueError): |
| a.summarize([]) |
|
|
|
|
| def test_decompose_prompt_then_length(tmp_path: Path) -> None: |
| f = tmp_path / 'b.jsonl' |
| for i, acc in enumerate((1000.0, 1020.0, 980.0)): |
| _bench2(f, f'o{i}', SHA_A, 'legal', 0, 1500, 1000, acc) |
| obs = s.load_bench2([f]) |
| _windows(tmp_path / 'w') |
| obs += s.load_windows(s.window_dirs(tmp_path / 'w')) |
| d = a.decompose(obs, 'bench2:legal-0', 'corpus:short-b1', 1500, 512) |
| assert d['prompt_step'] == pytest.approx(0.15) and d['length_step'] == pytest.approx(0.15) |
| assert d['total'] == pytest.approx(0.30) |
| with pytest.raises(ValueError): |
| a.decompose(obs, 'bench2:legal-0', 'corpus:nothing', 1500, 512) |
|
|
|
|
| def test_segment_estimators(tmp_path: Path) -> None: |
| f = tmp_path / 'b.jsonl' |
| |
| times = [i * 40.0 for i in range(700)] |
| _bench2(f, 'r', SHA_A, 'legal', 0, 1400, 700, 700, extra={ |
| 'step_times_ms': times, 'first_by_kind_s': {'reasoning': 0.0}, 'last_by_kind_s': {'reasoning': times[-1] / 1000}}) |
| |
| |
| _bench2(f, 'm', SHA_B, 'summary', 32000, 16384, 450, 600, answer=450, extra={ |
| 'completion_tokens': 1050, 'reasoning_tokens': 600, 'reasoning_decode_tps': 50.0, 'answer_decode_tps': 75.0}) |
| seg = a.segment_split(s.load_bench2([f])) |
| assert seg['validation_runs'] == 1 |
| (m,) = cast(list[dict[str, object]], seg['mixed_runs']) |
| assert m['method'] == 'tps_equal_round_time' |
| assert m['round_s_estimated'] == pytest.approx(0.04) |
| assert m['reasoning_estimated'] == pytest.approx(1.0) and m['answer_estimated'] == pytest.approx(2.0) |
|
|
|
|
| def test_metric_reconciliation(tmp_path: Path) -> None: |
| f = tmp_path / 'b.jsonl' |
| for i in range(3): |
| _bench2(f, f'r{i}', SHA_A, 'legal', 0, 1500, 750, 750, tps=50.0) |
| (row,) = a.metric_reconciliation(s.load_bench2([f])) |
| assert row['server'] == 1.0 and row['tps_x_p50_gap'] == pytest.approx(0.95) and row['tps_x_mean_gap'] == 1.0 |
| assert row['tokens_per_sse_event'] == 1.0 and row['rounds_per_sse_event'] == 1.0 |
|
|
|
|
| def test_cli_writes_and_refuses_overwrite(tmp_path: Path) -> None: |
| f = tmp_path / 'bench' / 'b.jsonl' |
| f.parent.mkdir() |
| for i, acc in enumerate((1000.0, 1020.0)): |
| _bench2(f, f'o{i}', SHA_A, 'legal', 0, 1500, 1000, acc) |
| _bench2(f, f's{i}', 'c' * 64, 'summary', 32000, 1500, 1000, 1600 + acc / 100) |
| _bench2(f, f'y{i}', 'd' * 64, 'synth', 150000, 1500, 1000, 1200 + acc / 100) |
| camp = tmp_path / 'camp' |
| _windows(camp) |
| for name, sha in (('32k-b1', 'e' * 64), ('150k-b1', 'f' * 64)): |
| for w in ('SP1-W-20261026T144153Z', 'PC1-x-W-20261027T100908Z'): |
| src = camp / w / ('cells/ctl-short-b1.json' if w.startswith('SP1') else 's1-control-short-b1.json') |
| dst = src.with_name(src.name.replace('short-b1', name)) |
| dst.write_text(src.read_text().replace(SHA_B, sha)) |
| out = tmp_path / 'out' |
| assert cli.main(['--campaign', str(camp), '--bench2-root', str(f.parent), '--out', str(out)]) == 0 |
| rep = json.loads((out / 'acc-split.json').read_text()) |
| assert set(rep['shifts']) == {'32K', 'short', '150K'} and (out / 'acc-split.md').read_text().startswith('# ACC-SPLIT') |
| with pytest.raises(SystemExit): |
| cli.main(['--campaign', str(camp), '--bench2-root', str(f.parent), '--out', str(out)]) |
|
|
