runtime/release_contract.py

"""Pure-Python release path, identity, and lineage helpers."""

from __future__ import annotations

import hashlib
import importlib.metadata
import json
import os
import platform
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
CONFIG_PATH = ROOT / "config" / "deployment.json"


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


def load_config() -> dict[str, Any]:
    return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))


def release_path(relative: str) -> Path:
    candidate = (ROOT / relative).resolve()
    if candidate != ROOT and ROOT not in candidate.parents:
        raise ValueError(f"release path escapes package root: {relative}")
    return candidate


def adapter_path(config: dict[str, Any] | None = None) -> Path:
    config = config or load_config()
    return release_path(config["adapter"]["path"])


def tokenizer_path(config: dict[str, Any] | None = None) -> Path:
    config = config or load_config()
    return release_path(config["foundation"]["tokenizer_path"])


def resolve_model_root(
    explicit: str | Path | None = None,
    config: dict[str, Any] | None = None,
) -> Path:
    config = config or load_config()
    env_name = config["foundation"]["model_root_env"]
    value = explicit if explicit is not None else os.environ.get(env_name)
    if not value:
        raise ValueError(
            f"Foundation weights are external. Set {env_name} to the model-root "
            "directory described in config/foundation_models.lock.json."
        )
    return Path(value).expanduser().resolve()


def foundation_files(
    model_root: Path,
    config: dict[str, Any] | None = None,
) -> list[tuple[dict[str, Any], Path]]:
    config = config or load_config()
    return [
        (row, (model_root / row["relative_path"]).resolve())
        for row in config["foundation"]["required_files"]
    ]


def foundation_model_dir(
    model_root: Path,
    config: dict[str, Any] | None = None,
) -> Path:
    config = config or load_config()
    return (model_root / config["foundation"]["model_id"]).resolve()


def export_lineage(
    model_root: Path,
    *,
    config: dict[str, Any] | None = None,
) -> dict[str, Any]:
    config = config or load_config()
    adapter = adapter_path(config)
    base_dit = next(
        path
        for row, path in foundation_files(model_root, config)
        if row["relative_path"].endswith("diffusion_pytorch_model.safetensors")
    )
    exporter = (
        ROOT
        / "vendor"
        / "raven-image3f-joint"
        / "control"
        / "export_image3f_block_stack_onnx.py"
    )
    source_manifest = ROOT / "config" / "SOURCE_MANIFEST.sha256"
    profile = {
        "batch_size": config["generation_contract"]["batch_size"],
        "views": config["generation_contract"]["views"],
        "per_view_width": config["generation_contract"]["per_view_width"],
        "per_view_height": config["generation_contract"]["per_view_height"],
        "frames": config["generation_contract"]["frames"],
        "reference_mode": "lag256_rgb_all_views",
        "onnx_opset": 20,
    }
    toolchain = {
        "python": platform.python_version(),
        "torch": importlib.metadata.version("torch"),
        "onnx": importlib.metadata.version("onnx"),
        "onnxscript": importlib.metadata.version("onnxscript"),
        "safetensors": importlib.metadata.version("safetensors"),
    }
    payload = {
        "adapter_sha256": sha256(adapter),
        "base_dit_sha256": sha256(base_dit),
        "exporter_sha256": sha256(exporter),
        "release_source_manifest_sha256": sha256(source_manifest),
        "export_toolchain": toolchain,
        "profile": profile,
    }
    canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
    payload["export_key"] = hashlib.sha256(canonical).hexdigest()
    return payload


def generated_onnx_dir(
    model_root: Path,
    config: dict[str, Any] | None = None,
) -> Path:
    config = config or load_config()
    lineage = export_lineage(model_root, config=config)
    return release_path(config["generated_layout"]["onnx_root"]) / lineage[
        "export_key"
    ][:24]


def current_engine_dir(config: dict[str, Any] | None = None) -> Path:
    config = config or load_config()
    return release_path(config["generated_layout"]["engine_root"]) / "current"