code/tests/test_data_store.py · viai957/align at main

align / code /tests /test_data_store.py

Add files using upload-large-folder tool

201dc77 verified 15 days ago

7.82 kB

"""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)

	"""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)