| """Multi-source training data (``--records path[:weight]``, repeatable): parsing, per-source |
| split handling and reloads, and the per-micro-batch source shares (CPU, no model).""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import random |
| from collections import Counter |
| from pathlib import Path |
|
|
| import pytest |
|
|
| pytest.importorskip("torch") |
|
|
| from deg.model import EncodedContext |
| from deg.train import ( |
| DEG_RECORDS, |
| Example, |
| Source, |
| active_sources, |
| batches, |
| load_sources, |
| make_sources, |
| parse_records_arg, |
| ) |
|
|
|
|
| def _ex(uid: str, n: int = 50) -> Example: |
| return Example("deg", uid, EncodedContext(tuple(range(n)), 0, ()), []) |
|
|
|
|
| def _write(path: Path, rows: list[dict]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("a") as fh: |
| for r in rows: |
| fh.write(json.dumps(r) + "\n") |
|
|
|
|
| def test_parse_records_weights_and_names(): |
| s = parse_records_arg("/data/deg/datagen/v3-long/records.jsonl:0.2") |
| assert s.path == Path("/data/deg/datagen/v3-long/records.jsonl") and s.weight == 0.2 |
| assert s.name == "v3-long" |
| s = parse_records_arg("/x/y/records.jsonl") |
| assert s.weight == 1.0 and s.name == "y" |
| assert parse_records_arg("/x/a:b/records.jsonl").path == Path("/x/a:b/records.jsonl") |
| assert parse_records_arg("other.jsonl:3").name == "other" |
| with pytest.raises(ValueError): |
| parse_records_arg("/x/records.jsonl:0") |
| assert [s.path for s in make_sources(None)] == [DEG_RECORDS] |
| a, b = make_sources(["/a/records.jsonl", "/b/a/records.jsonl:2"]) |
| assert (a.name, b.name, b.weight) == ("a", "a#2", 2.0) |
|
|
|
|
| def test_load_sources_keeps_each_split_and_reloads_all(tmp_path): |
| v2, v3 = tmp_path / "v2" / "records.jsonl", tmp_path / "v3" / "records.jsonl" |
| _write(v2, [{"uid": f"a{i}", "split": s} for i, s in enumerate( |
| ["train", "val", "test", "train"])]) |
| sources = make_sources([str(v2), f"{v3}:0.2"]) |
| built = [] |
|
|
| def build(rec): |
| built.append(rec["uid"]) |
| return None if rec["uid"] == "a3" else _ex(rec["uid"]) |
|
|
| assert load_sources(sources, build, 0, 1) == {"v2": 1, "v3": 0} |
| assert built == ["a0", "a3"] |
| _write(v3, [{"uid": "v3-d1", "split": "train"}, {"uid": "v3-d2", "split": "val"}]) |
| _write(v2, [{"uid": "a9", "split": "train"}]) |
| assert load_sources(sources, build, 0, 1) == {"v2": 1, "v3": 1} |
| assert built == ["a0", "a3", "a9", "v3-d1"] |
| assert [e.uid for e in sources[1].pool] == ["v3-d1"] |
|
|
| def broken(rec): |
| raise KeyError("schema") |
|
|
| logged = [] |
| _write(v3, [{"uid": "v3-d3", "split": "train"}]) |
| assert load_sources(sources, broken, 0, 1, logged.append) == {"v2": 0, "v3": 0} |
| assert logged[0]["skip"] == "v3-d3" and logged[0]["source"] == "v3" |
|
|
|
|
| def test_load_sources_shards_by_uid(tmp_path): |
| path = tmp_path / "r" / "records.jsonl" |
| _write(path, [{"uid": f"u{i}", "split": "train"} for i in range(40)]) |
| pools = [] |
| for rank in range(2): |
| srcs = make_sources([str(path)]) |
| load_sources(srcs, lambda r: _ex(r["uid"]), rank, 2) |
| pools.append({e.uid for e in srcs[0].pool}) |
| assert pools[0].isdisjoint(pools[1]) and len(pools[0] | pools[1]) == 40 |
|
|
|
|
| def _shares(sources, n=4000, min_train=50): |
| rng = random.Random(0) |
| gen = batches(sources, None, rng, budget=10_000, jev_frac=0.0, max_records=8, |
| min_train=min_train) |
| got = Counter() |
| for _ in range(n): |
| batch = next(gen) |
| srcs = {e.uid.split("-")[0] for e in batch} |
| assert len(srcs) == 1 |
| got[srcs.pop()] += 1 |
| return {k: v / n for k, v in got.items()} |
|
|
|
|
| def test_small_source_gets_fixed_share_once_it_has_min_train(): |
| big = Source(Path("/v2/records.jsonl"), 0.8) |
| big.pool = [_ex(f"v2-{i}", 20 + i % 300) for i in range(5000)] |
| small = Source(Path("/v3/records.jsonl"), 0.2) |
| small.pool = [_ex(f"v3-{i}", 3000 + i) for i in range(49)] |
| assert active_sources([big, small], 50) == [big] |
| assert _shares([big, small]) == {"v2": 1.0} |
| small.pool.append(_ex("v3-49", 4000)) |
| assert active_sources([big, small], 50) == [big, small] |
| share = _shares([big, small]).get("v3", 0.0) |
| assert abs(share - 0.2) < 0.03, share |
|
|
|
|
| def test_no_source_over_threshold_still_trains(): |
| a = Source(Path("/a/records.jsonl")) |
| a.pool = [_ex("a-1")] |
| assert active_sources([a, Source(Path("/b/records.jsonl"))], 50) == [a] |
| assert _shares([a], n=10) == {"a": 1.0} |
| with pytest.raises(RuntimeError): |
| next(batches([Source(Path("/c/records.jsonl"))], None, random.Random(0), 100, 0.0, 4)) |
|
|
|
|
| def test_lr_scale_restart_continues_the_schedule(): |
| from deg.train import lr_scale |
|
|
| full = [lr_scale(s, 20000, 30) for s in range(1, 20001)] |
| assert full[0] == pytest.approx(1 / 30 * (0.5 * (1 + 1) * 0.9 + 0.1), rel=1e-6) |
| assert full[-1] == pytest.approx(0.1) |
| |
| assert lr_scale(1001, 20000, 30, 1000) == pytest.approx(full[1000] / 30) |
| for s in (1030, 5000, 20000): |
| assert lr_scale(s, 20000, 30, 1000) == pytest.approx(full[s - 1]) |
|
|
|
|
| def test_heads_only_steps_then_backbone_warms_up(): |
| from deg.train import backbone_live, lr_scale |
|
|
| assert not backbone_live(1001, 1000, 500) |
| assert not backbone_live(1500, 1000, 500) |
| assert backbone_live(1501, 1000, 500) |
| assert backbone_live(1, 0, 0) |
| |
| assert lr_scale(1501, 20000, 30, 1500) == pytest.approx(lr_scale(1501, 20000, 30) / 30, |
| rel=1e-6) |
|
|
