File size: 5,402 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """Regression tests for audit findings U09-4..U09-7 (align/data/build.py,
align/data/prm800k_pairs.py).
Run: uv run pytest tests/test_audit_g5_data_store_build.py -m "not gpu" -v
"""
from __future__ import annotations
import json
from pathlib import Path
import pytest
from align.data import build as bd
from align.data import prm800k_pairs as pb
from align.data.build import _PREF_FIELDS, _ShardSet, _held_out_ids
from align.data.store import ShardDataset
# -- U09-4 --------------------------------------------------------------------
def _fill(set_, pids):
for pid in pids:
set_.add(input_ids=[1, 2, 3], problem_id=pid, subject="s", boundary_positions=[])
return set_.finalize()
def test_u09_4_rebuild_into_reused_dir_does_not_leave_stale_shards(tmp_workdir):
old = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3)
_fill(old, ["o0", "o1", "o2", "o3", "o4"]) # shards 0000 + 0001
new = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3)
_fill(new, ["n0"]) # only 0000 rewritten
ids = [r.problem_id for r in ShardDataset(tmp_workdir).iter_rows()]
assert ids == ["n0"]
def test_u09_4_purge_is_scoped_to_its_own_prefix(tmp_workdir):
other = _ShardSet(tmp_workdir, "sft-held")
_fill(other, ["h0"])
_fill(_ShardSet(tmp_workdir, "sft-train"), ["t0"])
assert [r.problem_id for r in ShardDataset(tmp_workdir, names=["sft-held-0000"]).iter_rows()] == ["h0"]
# -- U09-5 --------------------------------------------------------------------
def _ids(n):
return [f"problem-{i:06d}" for i in range(n)]
def test_u09_5_held_out_assignment_is_independent_of_the_id_set():
ids = _ids(4000)
held_full = _held_out_ids(ids, 0.05, 0)
# Arm-B-like subset: every id's membership must be the same as in the
# full set, otherwise Arm B's held-out set leaks SFT-train problems.
subset = ids[::3]
held_sub = _held_out_ids(subset, 0.05, 0)
assert held_sub == held_full & set(subset)
def test_u09_5_adding_one_id_changes_at_most_that_id():
ids = _ids(4000)
base = _held_out_ids(ids, 0.05, 0)
more = _held_out_ids(ids + ["problem-new"], 0.05, 0)
assert base - more == set()
assert more - base <= {"problem-new"}
def test_u09_5_fraction_is_close_to_requested_and_nonempty_for_tiny_inputs():
held = _held_out_ids(_ids(20000), 0.05, 0)
assert 0.04 * 20000 < len(held) < 0.06 * 20000
assert len(_held_out_ids(["only-one"], 0.05, 0)) == 1
assert len(_held_out_ids(_ids(5), 0.05, 0)) >= 1
assert _held_out_ids([], 0.05, 0) == set()
def test_u09_5_seed_changes_the_split():
assert _held_out_ids(_ids(4000), 0.05, 0) != _held_out_ids(_ids(4000), 0.05, 1)
def test_u09_5_rule_is_recorded_in_pins():
assert bd.PinnedYields().held_out_rule == "hash-v1"
# -- U09-6 --------------------------------------------------------------------
def _pair(pid="p0", step=0, chosen="good", rejected="bad", conf=1.0):
return pb.PrefPair(
problem="Q", problem_id=pid, step_index=step, prefix_steps=(),
chosen_step=chosen, rejected_step=rejected, confidence=conf,
)
def test_u09_6_curate_never_emits_duplicate_pair_ids():
dup = [_pair(), _pair(), _pair(rejected="worse")]
selected, remainder = pb.curate(dup, set(), {"p0": "Algebra"}, target=1)
ids = [p.pair_id for p in selected] + [p.pair_id for p in remainder]
assert len(ids) == len(set(ids)) == 2
assert not ({p.pair_id for p in selected} & {p.pair_id for p in remainder})
def test_u09_6_curate_dedupe_is_order_independent():
a, b = _pair(), _pair(rejected="worse")
s1, r1 = pb.curate([a, a, b], set(), {"p0": "Algebra"}, target=1)
s2, r2 = pb.curate([b, a, a], set(), {"p0": "Algebra"}, target=1)
assert [p.pair_id for p in s1] == [p.pair_id for p in s2]
assert [p.pair_id for p in r1] == [p.pair_id for p in r2]
# -- U09-7 --------------------------------------------------------------------
class _FakeSpec:
newline_id = 1
vocab_absorbers = 0
tokenizer_hash = "t"
def to_json(self, path):
Path(path).write_text("{}")
def test_u09_7_aborted_rebuild_invalidates_the_previous_manifest(tmp_workdir, monkeypatch):
out = Path(tmp_workdir) / "build"
out.mkdir()
(out / "manifest.json").write_text(json.dumps({"pins_verified": True, "build_version": 2}))
monkeypatch.setattr(bd, "derive_boundary_spec", lambda tok: _FakeSpec())
monkeypatch.setattr(bd, "math_subject_index", lambda src: {})
def boom(*a, **k):
raise RuntimeError("strict pin failure")
monkeypatch.setattr(bd, "build_eval_stage", boom)
with pytest.raises(RuntimeError, match="strict pin"):
bd.build_all(object(), out, object(), bd.PinnedYields())
assert not (out / "manifest.json").exists()
def test_u09_7_atomic_write_never_exposes_a_partial_file(tmp_workdir, monkeypatch):
from align.data import store
target = Path(tmp_workdir) / "m.json"
target.write_text("OLD")
def fail_replace(src, dst):
raise OSError("killed before rename")
monkeypatch.setattr(store.os, "replace", fail_replace)
with pytest.raises(OSError):
store.atomic_write_text(target, "NEW")
assert target.read_text() == "OLD"
assert [p.name for p in Path(tmp_workdir).iterdir()] == ["m.json"]
monkeypatch.undo()
store.atomic_write_text(target, "NEW")
assert target.read_text() == "NEW"
|