File size: 6,735 Bytes
201dc77 | 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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """Regression tests for audit findings U09-1..U09-4 (align/data/store.py).
Run: uv run pytest tests/test_audit_g5_data_store_store.py -m "not gpu" -v
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pytest
from align.data.store import ShardDataset, ShardReader, ShardWriter, read_shard_receipt
def _add(w, ids, pid, **extra):
w.add_row(ids, problem_id=pid, subject="Algebra", boundary_positions=[], **extra)
def _make_shard(d, name="s", rows=((1, 1, 1), (2, 2, 2, 2), (3, 3))):
with ShardWriter(d, name) as w:
for i, ids in enumerate(rows):
_add(w, list(ids), f"p{i}")
return w.receipt
# -- U09-1 --------------------------------------------------------------------
def test_u09_1_failed_add_row_leaves_no_partial_row(tmp_workdir):
"""A row rejected by schema/JSON validation must not leave its tokens in
the blobs: the caller catches and continues (build.py's pattern), and the
NEXT row's metadata would be paired with the PREVIOUS row's tokens."""
w = ShardWriter(tmp_workdir, "s")
_add(w, [1, 1, 1], "a", k=1)
with pytest.raises(ValueError):
_add(w, [2, 2, 2, 2], "b", k=1, extra_key=5) # schema mismatch vs row 0
with pytest.raises(TypeError):
_add(w, [9, 9, 9, 9, 9], "c", k=np.int64(3)) # non-JSON extra
_add(w, [3, 3], "d", k=2)
w.finalize()
rows = list(ShardReader(tmp_workdir, "s")[i] for i in range(2))
assert [r.problem_id for r in rows] == ["a", "d"]
assert rows[0].input_ids.tolist() == [1, 1, 1]
assert rows[1].input_ids.tolist() == [3, 3]
assert w.receipt.n_tokens == 5
def test_u09_1_non_json_extra_on_first_row_does_not_poison_schema(tmp_workdir):
w = ShardWriter(tmp_workdir, "s")
with pytest.raises(TypeError):
_add(w, [1, 1], "a", bad=np.int64(3))
_add(w, [2, 2], "b", good=1) # schema is decided by the first ACCEPTED row
w.finalize()
r = ShardReader(tmp_workdir, "s")
assert len(r) == 1 and r[0].problem_id == "b" and r[0].extras == {"good": 1}
def test_u09_1_ragged_field_failure_is_atomic_across_fields(tmp_workdir):
w = ShardWriter(tmp_workdir, "s", {"input_ids": "int32", "rej": "int32"})
w.add_row(
[1, 1], problem_id="a", subject="x", boundary_positions=[], ragged={"rej": [5, 5]}
)
with pytest.raises(ValueError):
w.add_row(
[2, 2],
problem_id="b",
subject="x",
boundary_positions=[],
ragged={"rej": [[1, 2], [3, 4]]}, # 2-D: fails after input_ids passes
)
w.add_row(
[3, 3, 3], problem_id="c", subject="x", boundary_positions=[], ragged={"rej": [6]}
)
w.finalize()
r = ShardReader(tmp_workdir, "s")
assert [r[i].problem_id for i in range(len(r))] == ["a", "c"]
assert r[1].input_ids.tolist() == [3, 3, 3] and r[1].ragged["rej"].tolist() == [6]
# -- U09-2 --------------------------------------------------------------------
def test_u09_2_exception_in_with_block_leaves_no_receipt(tmp_workdir):
with pytest.raises(KeyboardInterrupt):
with ShardWriter(tmp_workdir, "s") as w:
_add(w, [1, 1, 1], "a")
raise KeyboardInterrupt
assert w.receipt is None
assert list(Path(tmp_workdir).glob("s.*")) == []
with pytest.raises(FileNotFoundError):
ShardDataset(tmp_workdir)
def test_u09_2_clean_exit_still_finalizes(tmp_workdir):
r = _make_shard(tmp_workdir)
assert r.n_rows == 3 and (Path(tmp_workdir) / "s.receipt.json").exists()
def test_u09_2_finalized_shard_survives_a_later_exception(tmp_workdir):
with pytest.raises(RuntimeError):
with ShardWriter(tmp_workdir, "s") as w:
_add(w, [1, 1], "a")
w.finalize()
raise RuntimeError("after finalize")
assert (Path(tmp_workdir) / "s.receipt.json").exists()
def test_u09_2_poisoned_writer_refuses_to_finalize(tmp_workdir):
"""A write that failed AFTER the first byte hit disk cannot be undone."""
w = ShardWriter(tmp_workdir, "s")
_add(w, [1, 1], "a")
w._handles["input_ids"].close() # next write raises ValueError mid-row
with pytest.raises(ValueError):
_add(w, [2, 2], "b")
with pytest.raises(RuntimeError, match="poison|abort|partial"):
_add(w, [3, 3], "c")
with pytest.raises(RuntimeError, match="poison|abort|partial"):
w.finalize()
# -- U09-3 --------------------------------------------------------------------
def test_u09_3_tampered_idx_is_detected(tmp_workdir):
_make_shard(tmp_workdir)
idx = Path(tmp_workdir) / "s.input_ids.idx.npy"
np.save(idx, np.array([0, 3, 7], dtype=np.int64)) # missing last offset
with pytest.raises((RuntimeError, ValueError)):
r = ShardReader(tmp_workdir, "s")
r[0]
def test_u09_3_trailing_bytes_in_bin_are_detected(tmp_workdir):
_make_shard(tmp_workdir)
with open(Path(tmp_workdir) / "s.input_ids.bin", "ab") as f:
f.write(b"\0" * 8)
with pytest.raises((RuntimeError, ValueError)):
r = ShardReader(tmp_workdir, "s")
r[0]
def test_u09_3_shifted_row_boundaries_change_the_hash(tmp_workdir):
"""Same length, same last offset, different row boundaries: only a hash
that covers the idx files notices."""
_make_shard(tmp_workdir)
r = ShardReader(tmp_workdir, "s")
r.verify() # untampered shard verifies
idx = Path(tmp_workdir) / "s.input_ids.idx.npy"
np.save(idx, np.array([0, 2, 9, 9], dtype=np.int64))
with pytest.raises(RuntimeError, match="hash"):
ShardReader(tmp_workdir, "s").verify()
def test_u09_3_legacy_v1_receipt_still_verifies(tmp_workdir):
"""Shards built before the idx files were hashed carry format_version 1 and
must keep verifying against the file set they were hashed over."""
from align.data.store import _hash_files
_make_shard(tmp_workdir)
d = Path(tmp_workdir)
legacy = _hash_files(
[d / f"s.{f}.bin" for f in ("boundary_positions", "input_ids")] + [d / "s.extras.jsonl"]
)
rj = json.loads((d / "s.receipt.json").read_text())
rj["content_hash"], rj["format_version"] = legacy, 1
(d / "s.receipt.json").write_text(json.dumps(rj))
ShardReader(tmp_workdir, "s").verify()
# -- U09-4 --------------------------------------------------------------------
def test_u09_4_rewriting_a_shard_drops_its_stale_receipt_immediately(tmp_workdir):
"""Between `open(wb)` truncating the blobs and finalize(), the OLD receipt
must not be left describing bytes that no longer exist."""
_make_shard(tmp_workdir)
w = ShardWriter(tmp_workdir, "s")
assert not (Path(tmp_workdir) / "s.receipt.json").exists()
w.abort()
|