File size: 7,823 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 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | """Ragged shard store + receipts.
Run: uv run pytest tests/test_data_store.py -m "not gpu" -v
The store is the only place training data enters the run, so it is where the
"missing value silently replaced by a plausible one" failure class has to be
made impossible: `problem_id`, `subject` and the per-row boundary receipt are
MANDATORY keyword arguments with no defaults, validated at write time.
"""
from __future__ import annotations
import json
import os
import numpy as np
import pytest
from align.data import tokenization as tk
from align.data.store import (
ShardDataset,
ShardReader,
ShardWriter,
read_shard_receipt,
)
NL = 781 # test-side constant; the PACKAGE derives it, never hardcodes it
def _spec() -> tk.BoundarySpec:
return tk.BoundarySpec(newline_id=NL, tokenizer_hash="test", vocab_absorbers=0)
def _row_ids(n_steps: int) -> list[int]:
"""A synthetic token row with `n_steps - 1` newline-bigram boundaries."""
out: list[int] = []
for k in range(n_steps):
if k:
out += [NL, NL]
out += [100 + k, 200 + k, 300 + k]
return out
def test_shard_metadata_carries_problem_id_and_subject(tmp_workdir):
"""Every row written to a shard round-trips problem_id and subject via
metadata extras.
Incident: the predecessor added these post-hoc -- subject was `null` on all
10,606 RL prompts and Stage 4 would have aborted AFTER loading a 7B model.
Here they are build-time mandatory and travel with the row."""
spec = _spec()
with ShardWriter(tmp_workdir, "sft-0000") as w:
for i in range(5):
ids = _row_ids(i + 2)
w.add_row(
ids,
problem_id=f"pid{i:04d}",
subject=["Algebra", "Geometry"][i % 2],
boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
)
receipt = w.receipt
assert receipt.n_rows == 5
r = ShardReader(tmp_workdir, "sft-0000")
assert len(r) == 5
for i in range(5):
row = r[i]
assert row.problem_id == f"pid{i:04d}"
assert row.subject == ["Algebra", "Geometry"][i % 2]
assert row.input_ids.tolist() == _row_ids(i + 2)
# The receipt is on disk and reloads identically -- the manifest quotes it.
assert read_shard_receipt(tmp_workdir, "sft-0000") == receipt
def test_receipts_roundtrip_against_fresh_scan(tmp_workdir):
"""Boundary positions stored in shard extras match a fresh integer scan of
the stored input_ids for EVERY row of a small built shard (uses the
section-02 scan; no decode anywhere).
Incident: a receipt that drifts from the ids it describes is undetectable
downstream -- interp would patch the wrong positions and report a real
number about the wrong tokens."""
spec = _spec()
with ShardWriter(tmp_workdir, "sft-0000") as w:
for i in range(8):
ids = _row_ids(i + 1)
w.add_row(
ids,
problem_id=f"p{i}",
subject="Algebra",
boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
)
ds = ShardDataset(tmp_workdir)
assert len(ds) == 8
checked = 0
for i in range(len(ds)):
row = ds[i]
fresh = tk.find_step_boundaries(row.input_ids.tolist(), spec.newline_id)
assert row.boundary_positions.tolist() == fresh
checked += 1
assert checked == 8
def test_row_missing_mandatory_extras_raises(tmp_workdir):
"""Writing a row without problem_id/subject/boundary receipts raises; the
store never backfills a plausible default (the silent-substitution failure
class). Both directions: a complete row writes silently."""
spec = _spec()
ids = _row_ids(3)
good = tk.find_step_boundaries(ids, spec.newline_id)
with ShardWriter(tmp_workdir, "s0") as w:
with pytest.raises(TypeError):
w.add_row(ids, subject="Algebra", boundary_positions=good)
with pytest.raises(TypeError):
w.add_row(ids, problem_id="p", boundary_positions=good)
with pytest.raises(TypeError):
w.add_row(ids, problem_id="p", subject="Algebra")
for bad in ("", None):
with pytest.raises(ValueError):
w.add_row(ids, problem_id=bad, subject="Algebra", boundary_positions=good)
with pytest.raises(ValueError):
w.add_row(ids, problem_id="p", subject=bad, boundary_positions=good)
# A receipt pointing outside the row is a corrupt receipt, not a row
# to be quietly clipped.
with pytest.raises(ValueError):
w.add_row(ids, problem_id="p", subject="Algebra", boundary_positions=[len(ids) + 5])
with pytest.raises(ValueError):
w.add_row(
ids, problem_id="p", subject="Algebra", boundary_positions=[good[-1], good[0]]
) # not increasing
# Silent direction: the complete row goes through.
w.add_row(ids, problem_id="p", subject="Algebra", boundary_positions=good)
assert w.receipt.n_rows == 1
def test_content_hash_detects_a_flipped_token(tmp_workdir):
"""The shard content hash is what the build manifest pins, so it must
change when a single stored token changes. Both directions: an identical
rebuild reproduces the hash exactly."""
spec = _spec()
def build(dirname: str, tweak: int) -> str:
d = os.path.join(tmp_workdir, dirname)
with ShardWriter(d, "s") as w:
ids = _row_ids(3)
ids[0] += tweak
w.add_row(
ids,
problem_id="p",
subject="Algebra",
boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
)
return w.receipt.content_hash
a = build("a", 0)
b = build("b", 0)
c = build("c", 1)
assert a == b # deterministic rebuild
assert a != c # one flipped token changes the hash
def test_extras_sidecar_is_one_json_object_per_row(tmp_workdir):
"""The extras sidecar is line-delimited so a shard streams; a row count
mismatch between the binary blob and the sidecar raises at open time
rather than mis-pairing metadata with tokens."""
spec = _spec()
with ShardWriter(tmp_workdir, "s") as w:
for i in range(3):
ids = _row_ids(2)
w.add_row(
ids,
problem_id=f"p{i}",
subject="Algebra",
boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
)
path = os.path.join(tmp_workdir, "s.extras.jsonl")
lines = [json.loads(ln) for ln in open(path) if ln.strip()]
assert len(lines) == 3
assert {"problem_id", "subject", "boundary_positions"} <= set(lines[0])
with open(path, "a") as f:
f.write(
json.dumps({"problem_id": "ghost", "subject": "Algebra", "boundary_positions": []})
+ "\n"
)
with pytest.raises(RuntimeError, match="row count"):
ShardReader(tmp_workdir, "s")
def test_lengths_read_from_index_only(tmp_workdir):
"""Length-grouped batching scans every row length each epoch; that scan
must not touch the data pages (it would page in the whole corpus on a box
with no swap)."""
spec = _spec()
with ShardWriter(tmp_workdir, "s") as w:
for i in range(6):
ids = _row_ids(i + 1)
w.add_row(
ids,
problem_id=f"p{i}",
subject="Algebra",
boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
)
ds = ShardDataset(tmp_workdir)
lengths = ds.lengths()
assert lengths.tolist() == [len(_row_ids(i + 1)) for i in range(6)]
assert isinstance(lengths, np.ndarray)
|