scripts/build_phase21a_hf_dataset_bundle.py

#!/usr/bin/env python3
"""Materialize the minimal portable Phase 2.1A dataset bundle.

The source view may contain symlinks into cold storage. This script copies only
frame archives referenced by the frozen manifest and preserves their relative
paths so ``--processed-root <bundle>/processed_datasets`` remains valid.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--processed-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--execute",
        action="store_true",
        help="Copy files and write inventory; without this flag only print the plan.",
    )
    args = parser.parse_args()

    manifest = args.manifest.resolve()
    # Preserve the lexical parent because ``processed_root`` is commonly a
    # symlink named ``processed_datasets`` and manifest paths include that name.
    processed_root = args.processed_root.absolute()
    source_root = processed_root.parent
    output = args.output_dir.resolve()
    resolved_processed_root = processed_root.resolve()
    if output == resolved_processed_root or resolved_processed_root in output.parents:
        raise ValueError("output directory must not be inside the source dataset view")

    rows = [json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines() if line]
    references: set[str] = set()
    split_counts: dict[str, int] = {}
    for row in rows:
        split = str(row["split"])
        split_counts[split] = split_counts.get(split, 0) + 1
        for role in ("chosen", "rejected"):
            value = str(row[role]["source_frame_file"])
            relative = Path(value)
            if relative.is_absolute() or ".." in relative.parts:
                raise ValueError(f"non-portable frame path: {value}")
            if relative.parts and relative.parts[0] == ".":
                relative = Path(*relative.parts[1:])
            if not relative.parts or relative.parts[0] != "processed_datasets":
                raise ValueError(f"frame path is outside processed_datasets: {value}")
            references.add(relative.as_posix())

    sources: list[tuple[str, Path, int]] = []
    for relative in sorted(references):
        relative_path = Path(relative)
        source = resolved_processed_root / Path(*relative_path.parts[1:])
        if not source.is_file():
            raise FileNotFoundError(source)
        sources.append((relative, source, source.stat().st_size))

    plan = {
        "status": "ready" if args.execute else "dry_run",
        "manifest": str(manifest),
        "manifest_sha256": sha256(manifest),
        "pair_rows": len(rows),
        "split_counts": split_counts,
        "unique_frame_files": len(sources),
        "frame_bytes": sum(size for _, _, size in sources),
        "frame_gib": sum(size for _, _, size in sources) / 2**30,
        "output_dir": str(output),
    }
    if not args.execute:
        print(json.dumps(plan, indent=2, sort_keys=True))
        return

    manifest_target = output / "manifests" / manifest.name
    manifest_target.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(manifest, manifest_target)
    summary_source = manifest.with_name(f"{manifest.stem}_summary.json")
    if summary_source.is_file():
        shutil.copy2(summary_source, manifest_target.parent / summary_source.name)

    inventory = []
    checksum_lines = []
    for index, (relative, source, size) in enumerate(sources, start=1):
        target = output / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        if not target.is_file() or target.stat().st_size != size:
            shutil.copy2(source, target)
        checksum = sha256(target)
        inventory.append({"path": relative, "bytes": size, "sha256": checksum})
        checksum_lines.append(f"{checksum}  {relative}\n")
        if index % 100 == 0:
            print(f"materialized {index}/{len(sources)}", flush=True)

    inventory_payload = dict(plan)
    inventory_payload["status"] = "complete"
    inventory_payload["files"] = inventory
    (output / "phase21a_data_inventory.json").write_text(
        json.dumps(inventory_payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    (output / "SHA256SUMS").write_text("".join(checksum_lines), encoding="utf-8")
    print(json.dumps({key: value for key, value in inventory_payload.items() if key != "files"}, indent=2))


if __name__ == "__main__":
    main()