File size: 3,683 Bytes
0693474 | 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 | from __future__ import annotations
import os
import json
from pathlib import Path
from types import SimpleNamespace
import torch
from offload.dspark.d2.train.reference import taps_window
from offload.dspark.d2.train.replay import capture_inputs, load_capture_set
from offload.dspark.d2.train.replay import g1_drafts
from offload.dspark.d2.train.tests.test_data import write_shard
from offload.dspark.d2.train.data import D2Blocks, sha256
def test_d2_data_reader_same_input_contract_and_taps_only_window(tmp_path: Path) -> None:
manifest = write_shard(tmp_path)
reader = os.environ.get("D2_DATA_READER", "/opt/dsv41/repos/.worktrees/dsv41-d2-data/offload/dspark/d2/shards.py")
cs = load_capture_set(str(manifest), reader, "train")
assert cs.plumbing == []
sample = cs.samples[0]
inp, meta = capture_inputs(cs, sample, "cpu")
assert inp["anchor"].item() == 4
assert inp["positions"].tolist() == [[0, 1, 2, 3]]
assert inp["taps"].shape == (1, 4, 15360)
assert inp["has_taps"].all() and not inp["engine_mx"].any()
assert meta["engine_mx_rows"] == 0
saved = []
ref = SimpleNamespace(main_x=lambda x: x[:5120].float(), stage_kv=lambda p, x: torch.zeros(3, 512),
reset_window=lambda: None, fill_window=lambda p, kv: saved.append((p, kv)))
cache = SimpleNamespace(mx_native={}, kv={})
positions, sources, mx, source = taps_window(ref, cs, sample, cache)
assert positions == [0, 1, 2, 3] and sources == ["taps"]*4 and source == "taps"
assert saved[0][1].shape == (4, 3, 512)
assert mx.shape == (5120,)
taps_window(ref, cs, cs.samples[1], cache)
assert len(cache.mx_native) == len(cache.kv) == 5
def test_empty_sequence_keeps_v1_sequence_key_numbering(tmp_path: Path) -> None:
entries = []
for name, gen in (("empty", 9), ("full", 4)):
root = tmp_path/name
root.mkdir()
manifest = write_shard(root, gen)
spec = json.loads(manifest.read_text())["splits"]["train"]["files"][0]
sidecar = root/spec["sidecar"]
side = json.loads(sidecar.read_text())
side["sequences"][0]["seq_id"] = name
sidecar.write_text(json.dumps(side))
spec["sidecar_sha256"] = sha256(sidecar)
spec.update(file=name+"/"+spec["file"], sidecar=name+"/"+spec["sidecar"])
entries.append(spec)
manifest = tmp_path/"MANIFEST.json"
manifest.write_text(json.dumps({"schema": "d2.manifest.v1", "checks": {"plumbing_errors": 0}, "splits": {"train": {"files": entries}}}))
data = D2Blocks(manifest, "train")
assert data.sequence_b == [1]
reader = os.environ.get("D2_DATA_READER", "/opt/dsv41/repos/.worktrees/dsv41-d2-data/offload/dspark/d2/shards.py")
captures = load_capture_set(str(manifest), reader, "train")
assert captures.plumbing == []
assert captures.samples[0].b == data.sequence_b[0]
def test_actual_g1_jsonl_mapping_and_manifest_without_optional_error_counter(tmp_path: Path) -> None:
manifest = write_shard(tmp_path)
spec = json.loads(manifest.read_text())
spec["checks"] = {"sequences": 1, "skipped_sequences": 0} # actual dataset.py shape
manifest.write_text(json.dumps(spec))
d = D2Blocks(manifest, "train")
path = tmp_path/"train.k2_drafts.jsonl"
path.write_text(json.dumps({"seq_id": "p1", "drafts": {str(P): [0]*5 for P in range(3, 8)}})+"\n")
assert len(g1_drafts(path, d)["p1"]) == 5
reader = os.environ.get("D2_DATA_READER", "/opt/dsv41/repos/.worktrees/dsv41-d2-data/offload/dspark/d2/shards.py")
cs = load_capture_set(str(manifest), reader, "train", str(path))
assert all(not s.engine_drafts_placeholder for s in cs.samples)
|