bundle/v0/scripts/build_hf_cpu_job.py

"""Build the self-contained Hugging Face CPU job from a pinned Git commit."""

from __future__ import annotations

import argparse
import base64
import gzip
import hashlib
import io
import subprocess
import tarfile
from pathlib import Path


PINNED_COMMIT = "870ead3e8203ad6524656ed383543185d34bdc3d"
LANE_PREFIX = "reproductions/papers/loss-aware-dro-ot/"
INCLUDED_PREFIXES = (
    LANE_PREFIX + "src/loss_aware_dro_repro/",
    LANE_PREFIX + "configs/",
    LANE_PREFIX + "environment/",
)
INCLUDED_FILES = (LANE_PREFIX + "pyproject.toml",)


def repository_root() -> Path:
    return Path(
        subprocess.check_output(
            ["git", "rev-parse", "--show-toplevel"], text=True
        ).strip()
    )


def pinned_paths(root: Path) -> list[str]:
    paths = subprocess.check_output(
        ["git", "ls-tree", "-r", "--name-only", PINNED_COMMIT],
        cwd=root,
        text=True,
    ).splitlines()
    selected = [
        path
        for path in paths
        if path in INCLUDED_FILES or any(path.startswith(prefix) for prefix in INCLUDED_PREFIXES)
    ]
    if not selected or LANE_PREFIX + "configs/paper_scale.json" not in selected:
        raise RuntimeError("pinned scientific source selection is incomplete")
    return sorted(selected)


def source_bundle(root: Path) -> bytes:
    raw = io.BytesIO()
    with gzip.GzipFile(filename="", mode="wb", fileobj=raw, mtime=0, compresslevel=9) as compressed:
        with tarfile.open(fileobj=compressed, mode="w") as archive:
            for path in pinned_paths(root):
                payload = subprocess.check_output(["git", "show", f"{PINNED_COMMIT}:{path}"], cwd=root)
                relative = path.removeprefix(LANE_PREFIX)
                info = tarfile.TarInfo(relative)
                info.size = len(payload)
                info.mode = 0o644
                info.uid = info.gid = 0
                info.uname = info.gname = ""
                info.mtime = 0
                archive.addfile(info, io.BytesIO(payload))
    return raw.getvalue()


def build(template: Path, output: Path) -> dict[str, str | int]:
    root = repository_root()
    bundle = source_bundle(root)
    bundle_hash = hashlib.sha256(bundle).hexdigest()
    text = template.read_text(encoding="utf-8")
    if text.count("__SOURCE_BUNDLE_SHA256__") != 1 or text.count("__SOURCE_BUNDLE_BASE85__") != 1:
        raise RuntimeError("template must contain each source placeholder exactly once")
    generated = text.replace("__SOURCE_BUNDLE_SHA256__", bundle_hash).replace(
        "__SOURCE_BUNDLE_BASE85__", base64.b85encode(bundle).decode("ascii")
    )
    output.write_text(generated, encoding="utf-8", newline="\n")
    return {
        "pinned_commit": PINNED_COMMIT,
        "source_bundle_sha256": bundle_hash,
        "source_bundle_bytes": len(bundle),
        "generated_script_sha256": hashlib.sha256(generated.encode()).hexdigest(),
        "generated_script_bytes": len(generated.encode()),
    }


def main() -> int:
    lane = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--template",
        type=Path,
        default=lane / "jobs" / "hf_cpu_empirical_a4_a6.template.py",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=lane / "jobs" / "hf_cpu_empirical_a4_a6.py",
    )
    args = parser.parse_args()
    receipt = build(args.template, args.output)
    print(receipt)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())