scripts/replay_vocoder_oracles.py

"""Replay pinned short/multichunk Phase-0 latents through the frozen vocoder."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import shutil
import tempfile

import torch

from music3lab.manifests import (
    atomic_write_bytes,
    canonical_json_bytes,
    semantic_digest,
)
from music3lab.vocoder import (
    build_replay_manifest,
    load_frozen_vocoder,
    load_phase0_vocoder_oracle,
    module_state_sha256,
    probe_latent_gradient,
    publish_replay_bundle,
    publish_replay_session_root,
    replay_vocoder_oracle,
    secure_replay_output_parent,
    verify_replay_bundle,
    verify_replay_session_tree,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--snapshot", type=Path, required=True)
    parser.add_argument("--base-manifest", type=Path, required=True)
    parser.add_argument("--diffusers-root", type=Path, required=True)
    parser.add_argument("--phase0-artifacts", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--device", default="cuda")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    device = torch.device(args.device)
    if device.type != "cuda" or not torch.cuda.is_available():
        raise RuntimeError(
            "oracle replay requires the explicitly authorized CUDA GPU"
        )
    device_name = torch.cuda.get_device_name(device)
    if "H100" not in device_name:
        raise RuntimeError(f"oracle replay requires H100, observed {device_name}")
    device_capability = tuple(torch.cuda.get_device_capability(device))
    cuda_runtime = torch.version.cuda
    if not cuda_runtime:
        raise RuntimeError("PyTorch does not expose a CUDA runtime identity")

    output_root = Path(
        os.path.abspath(os.fspath(args.output_root.expanduser()))
    )
    if os.path.lexists(output_root):
        raise FileExistsError(f"replay output already exists: {output_root}")
    output_root.parent.mkdir(parents=True, exist_ok=True)
    secure_replay_output_parent(output_root.parent)

    oracles = {
        kind: load_phase0_vocoder_oracle(args.phase0_artifacts, kind)
        for kind in ("short", "multi")
    }
    adapter = load_frozen_vocoder(
        snapshot=args.snapshot,
        base_manifest=args.base_manifest,
        diffusers_root=args.diffusers_root,
    )
    if adapter.report.project_git_dirty:
        raise RuntimeError("oracle replay refuses dirty project source")

    pending = []
    summaries = []
    for mode, dtype in (
        ("fp32_reference", torch.float32),
        ("bf16_pipeline", torch.bfloat16),
    ):
        adapter.to(device=device, dtype=dtype)
        torch.cuda.reset_peak_memory_stats(device)
        mode_weight_sha256 = module_state_sha256(adapter.model)
        gradient_probe = probe_latent_gradient(
            adapter,
            oracles["short"].latents[0],
            device=device,
            dtype=dtype,
        )
        torch.cuda.synchronize(device)
        if gradient_probe.weight_state_sha256_before != mode_weight_sha256:
            raise RuntimeError("gradient probe started from different weights")

        for kind, oracle in oracles.items():
            with torch.inference_mode():
                result = replay_vocoder_oracle(
                    adapter,
                    oracle,
                    device=device,
                    dtype=dtype,
                )
            torch.cuda.synchronize(device)
            peak_allocated = torch.cuda.max_memory_allocated(device)
            peak_reserved = torch.cuda.max_memory_reserved(device)
            manifest, audio_bytes, wav_bytes = build_replay_manifest(
                mode=mode,
                adapter=adapter,
                oracle=oracle,
                result=result,
                weight_state_sha256_before=mode_weight_sha256,
                gradient_probe=gradient_probe,
                device_name=device_name,
                device_capability=device_capability,
                cuda_runtime=cuda_runtime,
                peak_cuda_allocated_bytes=peak_allocated,
                peak_cuda_reserved_bytes=peak_reserved,
            )
            if mode == "bf16_pipeline" and not manifest.exact_oracle_audio:
                raise RuntimeError(
                    f"BF16 official replay differs for {kind}; stopping"
                )
            relative = Path(kind) / mode
            pending.append((relative, manifest, audio_bytes, wav_bytes))
            summary = {
                "oracle": kind,
                "mode": mode,
                "output": str(output_root / relative),
                "manifest_semantic_digest": manifest.semantic_digest,
                "output_content_sha256": manifest.output_content_sha256,
                "expected_content_sha256": manifest.expected_content_sha256,
                "output_wav_sha256": manifest.output_wav_artifact.sha256,
                "exact_oracle_audio": manifest.exact_oracle_audio,
                "max_abs_error": manifest.max_abs_error,
                "mean_abs_error": manifest.mean_abs_error,
                "weight_state_sha256": manifest.weight_state_sha256_after,
                "latent_gradient_sha256": (
                    manifest.gradient_probe.gradient_content_sha256
                ),
                "peak_cuda_allocated_bytes": peak_allocated,
                "peak_cuda_reserved_bytes": peak_reserved,
            }
            print(json.dumps(summary, sort_keys=True), flush=True)
            summaries.append(summary)
            del result
        if module_state_sha256(adapter.model) != mode_weight_sha256:
            raise RuntimeError("vocoder weights changed across replay mode")

    temporary_root = Path(
        tempfile.mkdtemp(
            prefix=f".{output_root.name}.",
            suffix=".tmp",
            dir=output_root.parent,
        )
    )
    published = False
    try:
        verified_summaries = []
        for relative, manifest, audio_bytes, wav_bytes in pending:
            leaf = temporary_root / relative
            publish_replay_bundle(
                leaf,
                manifest=manifest,
                audio_bytes=audio_bytes,
                wav_bytes=wav_bytes,
            )
            verified = verify_replay_bundle(
                leaf,
                expected_adapter_semantic_digest=adapter.report.semantic_digest,
                expected_oracle_semantic_digest=manifest.oracle_semantic_digest,
            )
            verified_summaries.append(
                {
                    "path": relative.as_posix(),
                    "manifest_file_sha256": verified.manifest_file_sha256,
                    "manifest_semantic_digest": manifest.semantic_digest,
                    "audio_file_sha256": manifest.output_artifact.sha256,
                    "wav_file_sha256": manifest.output_wav_artifact.sha256,
                    "audio_content_sha256": manifest.output_content_sha256,
                }
            )
        session_payload = {
            "schema_version": "music3lab.vocoder-replay-session.v2",
            "status": "PASS",
            "adapter_report": adapter.report.model_dump(mode="json"),
            "device": str(device),
            "device_name": device_name,
            "device_capability": list(device_capability),
            "cuda_runtime": cuda_runtime,
            "replays": verified_summaries,
        }
        session_payload["semantic_digest"] = semantic_digest(session_payload)
        atomic_write_bytes(
            temporary_root / "session.json",
            canonical_json_bytes(session_payload),
            mode=0o644,
        )
        publish_replay_session_root(temporary_root, output_root)
        published = True
    finally:
        if not published and temporary_root.exists():
            shutil.rmtree(temporary_root)

    for relative, manifest, _audio_bytes, _wav_bytes in pending:
        verify_replay_bundle(
            output_root / relative,
            expected_adapter_semantic_digest=adapter.report.semantic_digest,
            expected_oracle_semantic_digest=manifest.oracle_semantic_digest,
        )
    verify_replay_session_tree(output_root)
    print(
        json.dumps(
            {
                "status": "PASS",
                "adapter_semantic_digest": adapter.report.semantic_digest,
                "project_git_commit": adapter.report.project_git_commit,
                "project_source_sha256": adapter.report.project_source_sha256,
                "session_semantic_digest": session_payload["semantic_digest"],
                "output_root": str(output_root),
                "replays": summaries,
            },
            sort_keys=True,
        )
    )
    return 0


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