"""Arm B shard dedup: the permanent invariant, and the dedup script's logic. The ref-logprob cache build refused the original Arm B shards with `ValueError: duplicate pair_id`: 87 train and 2 held_out pair_ids were present exactly twice (curation's depth-round top-up re-added already selected pairs; the per-problem merge in `prm800k_pairs.curate` never deduped by pair_id). `scripts/dedup_arm_b_shards.py` removed the exact re-adds. From now on, unique pair_ids per split -- and zero train/held overlap -- is a data invariant, not a hope. """ import importlib.util import os import sys from pathlib import Path import numpy as np import pytest REPO = Path("/apps/development/rlhf/nano-mistral-o1") ARM_B = REPO / "data" / "shards" / "arm_b" def _load_dedup_module(): """Import the script by path; scripts/ is not a package.""" path = REPO / "scripts" / "dedup_arm_b_shards.py" spec = importlib.util.spec_from_file_location("dedup_arm_b_shards", path) mod = importlib.util.module_from_spec(spec) # The script pins CUDA_VISIBLE_DEVICES="" for its own runs; do not let # importing it here clobber the test session's environment. saved = os.environ.get("CUDA_VISIBLE_DEVICES") try: spec.loader.exec_module(mod) finally: if saved is None: os.environ.pop("CUDA_VISIBLE_DEVICES", None) else: os.environ["CUDA_VISIBLE_DEVICES"] = saved sys.modules.pop("dedup_arm_b_shards", None) return mod # --------------------------------------------------------------------------- # (a) Post-condition on the LIVE shards: the permanent invariant # --------------------------------------------------------------------------- def _split_pair_ids(split: str) -> list[str]: from align.data.store import ShardDataset ds = ShardDataset(ARM_B / split) return [ds[i].extras["pair_id"] for i in range(len(ds))] @pytest.mark.skipif(not ARM_B.is_dir(), reason=f"arm_b shards not present at {ARM_B}") def test_live_arm_b_shards_have_unique_pair_ids_and_disjoint_splits(): """STRICT: every pair_id unique within each split, zero overlap across. A duplicate here means either the dedup regressed or a rebuild reintroduced the curation top-up defect; both must fail at CPU-test time, before the ref cache refuses them after GPUs have been claimed. """ by_split = {} for split in ("train", "held_out"): ids = _split_pair_ids(split) dupes = sorted({p for p in ids if ids.count(p) > 1}) if len(set(ids)) < len(ids) else [] assert len(set(ids)) == len(ids), ( f"arm_b/{split}: {len(ids) - len(set(ids))} duplicate pair_id " f"occurrences (e.g. {dupes[:3]}). The dedup invariant is broken; " "see scripts/dedup_arm_b_shards.py." ) by_split[split] = set(ids) overlap = by_split["train"] & by_split["held_out"] assert not overlap, ( f"{len(overlap)} pair_ids present in BOTH arm_b splits (e.g. " f"{sorted(overlap)[:3]}); held-out preference accuracy would measure " "memorization while looking like generalization." ) @pytest.mark.skipif( not (REPO / "data" / "manifest.json").is_file(), reason="build manifest not present" ) def test_manifest_arm_b_count_matches_live_shards(): """counts.arm_b_curated must equal what is actually on disk: the manifest gate derives the optimizer-step horizon from it.""" import json from align.data.store import ShardDataset if not ARM_B.is_dir(): pytest.skip(f"arm_b shards not present at {ARM_B}") manifest = json.loads((REPO / "data" / "manifest.json").read_text()) n = sum(len(ShardDataset(ARM_B / s)) for s in ("train", "held_out")) assert manifest["counts"]["arm_b_curated"] == n # --------------------------------------------------------------------------- # (b) The dedup script's kept-first / content-identity logic, synthetically # --------------------------------------------------------------------------- def _write_tiny_shard(tmp_path: Path, rows: list[dict]) -> Path: from align.data.store import ShardWriter fields = { "input_ids": "int32", "rejected_input_ids": "int32", "rejected_boundary_positions": "int32", } with ShardWriter(tmp_path, "tiny-0000", fields) as w: for r in rows: w.add_row( r["input_ids"], problem_id=r["problem_id"], subject="algebra", boundary_positions=r["boundary_positions"], ragged={ "rejected_input_ids": r["rejected_input_ids"], "rejected_boundary_positions": [0], }, chosen_prompt_len=2, rejected_prompt_len=2, shared_prefix_len=3, step_index=r.get("step_index", 1), pair_id=r["pair_id"], source="tiny", ) return tmp_path def _row(pair_id, input_ids=(5, 6, 7, 8), rejected=(5, 6, 9), problem_id="p1", **kw): n = len(input_ids) return { "pair_id": pair_id, "input_ids": list(input_ids), "rejected_input_ids": list(rejected), "boundary_positions": [1, n - 1] if n > 2 else [n - 1], "problem_id": problem_id, **kw, } def test_dedup_keeps_first_and_drops_identical_re_adds(tmp_path): from align.data.store import ShardDataset dedup = _load_dedup_module() # a, b, a-again (identical), c, b-again (identical) rows = [ _row("pair_a"), _row("pair_b", input_ids=(1, 2, 3, 4)), _row("pair_a"), _row("pair_c", input_ids=(9, 9, 9)), _row("pair_b", input_ids=(1, 2, 3, 4)), ] ds = ShardDataset(_write_tiny_shard(tmp_path, rows)) kept, dropped = dedup.dedup_row_indices(ds) assert kept == [0, 1, 3], "first occurrence wins, in original order" assert dropped == [2, 4] assert [ds[i].extras["pair_id"] for i in kept] == ["pair_a", "pair_b", "pair_c"] def test_dedup_refuses_non_identical_duplicate(tmp_path): """A shared pair_id with DIFFERENT content is a collision, not a re-add; the script must refuse rather than resolve it by row order.""" from align.data.store import ShardDataset dedup = _load_dedup_module() rows = [ _row("pair_a"), _row("pair_a", input_ids=(5, 6, 7, 8, 8)), # same id, different tokens ] ds = ShardDataset(_write_tiny_shard(tmp_path, rows)) with pytest.raises(ValueError, match="DIFFERENT content"): dedup.dedup_row_indices(ds) def test_rows_identical_covers_every_stored_field(tmp_path): """The identity check must see extras and the rejected side, not just input_ids -- otherwise a real divergence hides behind a shared pair_id.""" from align.data.store import ShardDataset dedup = _load_dedup_module() ds = ShardDataset( _write_tiny_shard( tmp_path, [ _row("pair_a"), _row("pair_b"), # same tokens as pair_a, different pair_id extra _row("pair_c", rejected=(5, 6, 1)), # differs only on the rejected side ], ) ) assert dedup.rows_identical(ds[0], ds[0]) assert not dedup.rows_identical(ds[0], ds[1]), "extras differ (pair_id)" assert not dedup.rows_identical(ds[0], ds[2]), "rejected_input_ids differ" assert np.array_equal(ds[0].input_ids, ds[1].input_ids), "sanity: tokens do match"