tests/test_inversion_v2_runner.py

"""CPU-only evidence-schema tests for the v2 production runner."""
from __future__ import annotations

import copy
import hashlib
import os
from pathlib import Path
import sys
from types import SimpleNamespace

import pytest
import torch
from safetensors.torch import load as load_safetensors

from music3lab import inversion_v2_runner as runner
from music3lab.fd_io import AnchoredDirectory
from music3lab.inversion_v2 import FIXED_RULER, load_v2_config
from music3lab.inversion_v2_runner import (
    EvaluatorMetrics,
    MetricsPayload,
    RulerSnapshot,
    TraceEntry,
    TracePayload,
    evaluate_gates,
    load_parent_inputs,
)
from music3lab.manifests import semantic_digest


def _ruler(value: float) -> RulerSnapshot:
    terms = {
        "time_nmse": (value,) * 4,
        "complex_stft_nmse": (value,) * 4,
        "legacy_mrstft": (value,) * 4,
        "mid_side_nmse": (value,) * 4,
        "relative_envelope": (value,) * 4,
    }
    score = sum(value * FIXED_RULER[key] for key in FIXED_RULER)
    return RulerSnapshot(**terms, fixed_scores=(score,) * 4)


def _entry(global_step: int, stage: str, stage_step: int, value: float, best: float) -> TraceEntry:
    return TraceEntry(
        global_step=global_step,
        stage=stage,
        stage_step=stage_step,
        learning_rate=0.0 if global_step == 0 else 0.01,
        prior_coefficient=0.0 if global_step == 0 else 0.005,
        audio_weights=dict(FIXED_RULER),
        ruler=_ruler(value),
        training_objective=(value,) * 4,
        best_fixed_scores=(best,) * 4,
        best_latent_sha256=(hashlib.sha256(str(global_step).encode()).hexdigest(),) * 4,
        gradient_norm=(0.0 if global_step == 0 else 1.0,) * 4,
    )


def test_trace_schema_recomputes_digest_and_rejects_truncation_and_forged_best() -> None:
    first = _ruler(2.0).fixed_scores[0]
    second = _ruler(1.0).fixed_scores[0]
    trace = TracePayload.create(
        experiment_id="P2-E1",
        restart_seeds=(101, 103, 107, 109),
        fp32_steps=1,
        bf16_steps=1,
        trajectory_stride_steps=1,
        selection_rule="fixed_exact_bf16_audio_ruler",
        entries=(
            _entry(0, "initial", 0, 2.0, first),
            _entry(1, "fp32", 1, 1.0, second),
            _entry(2, "bf16", 1, 1.5, second),
        ),
    )
    forged = trace.model_dump(mode="json")
    forged["entries"][1]["best_fixed_scores"][0] = 0.0
    forged["semantic_digest"] = semantic_digest(
        {key: value for key, value in forged.items() if key != "semantic_digest"}
    )
    with pytest.raises(ValueError, match="best"):
        TracePayload.model_validate(forged)
    truncated = trace.model_dump(mode="json")
    truncated["entries"].pop()
    truncated["semantic_digest"] = semantic_digest(
        {key: value for key, value in truncated.items() if key != "semantic_digest"}
    )
    with pytest.raises(ValueError, match="truncated"):
        TracePayload.model_validate(truncated)


def test_metric_claims_are_derived_from_primitives_not_stored_booleans() -> None:
    config = load_v2_config(Path(__file__).parents[1] / "configs" / "inversion-v2.yaml").config
    thresholds = config.experiments[1].thresholds
    evaluator = EvaluatorMetrics(
        waveform_mae=1.0,
        correlation=0.0,
        si_sdr_db=-20.0,
        unscaled_snr_db=-20.0,
        loudness_error_db=3.0,
        stereo_correlation_error=1.0,
        latent_rmse=None,
    )
    initial = (2.0, 2.0, 2.0, 2.0)
    final = (1.9, 1.9, 1.9, 1.9)
    gates = evaluate_gates(initial, final, 0, evaluator, thresholds)
    metrics = MetricsPayload.create(
        experiment_id="P2-E2",
        selected_restart_index=0,
        selected_restart_seed=101,
        selection_tie_rule="lowest_restart_index",
        evaluator_computed_after_lock=True,
        initial_fixed_scores=initial,
        best_fixed_scores=final,
        evaluator=evaluator,
        thresholds=thresholds,
        gates=gates,
        status="FAIL",
        high_fidelity_status="NOT_HIGH_FIDELITY",
        quality_claim="NO_QUALITY_CLAIM",
    )
    forged = metrics.model_dump(mode="json")
    forged["status"] = "FEASIBILITY_PASS"
    forged["quality_claim"] = "feasibility_only"
    forged["semantic_digest"] = semantic_digest(
        {key: value for key, value in forged.items() if key != "semantic_digest"}
    )
    with pytest.raises(ValueError, match="claims"):
        MetricsPayload.model_validate(forged)


@pytest.mark.skipif("MINIMAX_V13_ROOT" not in os.environ, reason="real v1.3 authority not configured")
def test_real_v13_parent_loader_binds_all_four_restarts() -> None:
    root = Path(os.environ["MINIMAX_V13_ROOT"])
    config_path = Path(os.environ.get(
        "MINIMAX_V2_CONFIG",
        Path(__file__).parents[1] / "configs" / "inversion-v2.yaml",
    ))
    authority = AnchoredDirectory.open_absolute(
        root,
        label="test v1.3 parent",
        require_readonly=False,
    )
    try:
        parent = load_parent_inputs(authority, load_v2_config(config_path).config)
        assert tuple(parent.e1_initial.shape) == (4, 128, 86)
        assert tuple(parent.e2_initial.shape) == (4, 128, 86)
        assert tuple(parent.target_audio.shape) == (1, 2, 44032)
        assert parent.session_semantic_digest == "1554159fc7a5e364de063e3c8345c3cf05426288aa23bd0039957026aabc3a7c"
    finally:
        authority.close(validate=True)



def test_v2_config_bytes_are_code_pinned_even_when_yaml_semantics_match(
    tmp_path: Path,
) -> None:
    canonical = Path(__file__).parents[1] / "configs" / "inversion-v2.yaml"
    forged = tmp_path / "inversion-v2.yaml"
    forged.write_bytes(canonical.read_bytes() + b"\n# semantically inert forgery\n")
    forged.chmod(0o644)
    with pytest.raises(ValueError, match="frozen preregistration"):
        load_v2_config(forged)


def test_live_bf16_replay_rejects_coherently_rehashed_forged_ruler(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class FakeAdapter:
        def __init__(self) -> None:
            self.model = torch.nn.Linear(1, 1, bias=False)
            self.model.requires_grad_(False)

    checkpoints = torch.stack((
        torch.full((4, 128, 86), 2.0),
        torch.full((4, 128, 86), 1.0),
    ))
    initial_score = _ruler(7.0).fixed_scores[0]
    final_score = _ruler(6.0).fixed_scores[0]
    forged = TracePayload.create(
        experiment_id="P2-E1",
        restart_seeds=(101, 103, 107, 109),
        fp32_steps=1,
        bf16_steps=0,
        trajectory_stride_steps=1,
        selection_rule="fixed_exact_bf16_audio_ruler",
        entries=(
            _entry(0, "initial", 0, 7.0, initial_score),
            _entry(1, "fp32", 1, 6.0, final_score),
        ),
    )

    def fake_decode(_adapter, latents: torch.Tensor, _dtype) -> torch.Tensor:
        return latents[:, :2, :1].expand(-1, -1, 44032).float()

    def fake_components(
        audio: torch.Tensor,
        _target: torch.Tensor,
        _config,
    ) -> dict[str, torch.Tensor]:
        value = audio[:, 0, 0]
        return {key: value for key in FIXED_RULER}

    monkeypatch.setattr(runner, "_decode", fake_decode)
    monkeypatch.setattr(runner, "ruler_components", fake_components)
    config = load_v2_config(
        Path(__file__).parents[1] / "configs" / "inversion-v2.yaml"
    ).config
    target = torch.zeros((1, 2, 44032))
    actual_initial_ruler = runner._ruler_snapshot(
        fake_components(fake_decode(FakeAdapter(), checkpoints[0], torch.bfloat16), target, config.loss)
    )
    actual_final_ruler = runner._ruler_snapshot(
        fake_components(fake_decode(FakeAdapter(), checkpoints[1], torch.bfloat16), target, config.loss)
    )
    legitimate = TracePayload.create(
        experiment_id="P2-E1",
        restart_seeds=(101, 103, 107, 109),
        fp32_steps=1,
        bf16_steps=0,
        trajectory_stride_steps=1,
        selection_rule="fixed_exact_bf16_audio_ruler",
        entries=(
            _entry(
                0,
                "initial",
                0,
                2.0,
                _ruler(2.0).fixed_scores[0],
            ).model_copy(update={
                "ruler": actual_initial_ruler,
                "best_fixed_scores": actual_initial_ruler.fixed_scores,
                "best_latent_sha256": runner._hash_each(checkpoints[0]),
            }),
            _entry(
                1,
                "fp32",
                1,
                1.0,
                _ruler(1.0).fixed_scores[0],
            ).model_copy(update={
                "ruler": actual_final_ruler,
                "best_fixed_scores": actual_final_ruler.fixed_scores,
                "best_latent_sha256": runner._hash_each(checkpoints[1]),
            }),
        ),
    )
    initial_audio, best_latents, best_audio, handoff = runner._replay_exact_trajectory(
        FakeAdapter(),
        checkpoints,
        torch.zeros((1, 2, 44032)),
        legitimate,
        config,
    )
    assert torch.equal(best_latents, checkpoints[1])
    assert torch.equal(handoff, checkpoints[1])
    assert tuple(initial_audio.shape) == (4, 2, 44032)
    assert tuple(best_audio.shape) == (4, 2, 44032)
    with pytest.raises(RuntimeError, match="live BF16 ruler differs"):
        runner._replay_exact_trajectory(
            FakeAdapter(),
            checkpoints,
            torch.zeros((1, 2, 44032)),
            forged,
            load_v2_config(Path(__file__).parents[1] / "configs" / "inversion-v2.yaml").config,
        )


def test_isolated_v13_dependency_bootstrap_uses_lexical_venv_and_rejects_wrong_origin(
    tmp_path: Path,
) -> None:
    executable = Path(sys.executable).absolute()
    assert executable.is_symlink()
    site = runner._venv_site_packages(executable)
    expected = (
        executable.parent.parent
        / "lib"
        / f"python{sys.version_info.major}.{sys.version_info.minor}"
        / "site-packages"
    ).resolve(strict=True)
    assert site == expected

    payload = runner.probe_isolated_v13_dependencies()
    assert payload["isolated"] is True
    assert payload["no_site"] is True
    assert payload["dont_write_bytecode"] is True
    assert payload["pythonpath_present"] is False
    assert set(payload["origins"]) == {"pydantic", "safetensors", "yaml"}
    for origin in payload["origins"].values():
        Path(origin).resolve(strict=True).relative_to(site)

    missing_executable = tmp_path / "missing-venv" / "bin" / "python"
    with pytest.raises(RuntimeError, match="site-packages is missing"):
        runner._venv_site_packages(missing_executable)

    wrong = tmp_path / "wrong-yaml.py"
    wrong.write_text("# wrong origin\n", encoding="utf-8")
    wrong.chmod(0o644)
    forged = copy.deepcopy(payload)
    forged["origins"]["yaml"] = str(wrong)
    with pytest.raises(RuntimeError, match="wrong origin: yaml"):
        runner._validate_dependency_probe(forged, site)


def test_nested_vocoder_loader_uses_canonical_snapshot_while_outer_authority_stays_live(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    base_id = "a" * 64
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir(mode=0o755)
    captured: list[Path] = []

    class FakeEvidence:
        def __init__(self, observed_base_id: str) -> None:
            self.frozen_base = SimpleNamespace(
                loader_path=Path("/proc/self/fd/37"),
                manifest=SimpleNamespace(semantic_digest=observed_base_id),
            )
            self.stable_calls = 0

        def assert_stable(self) -> None:
            self.stable_calls += 1

    def fake_load_frozen_vocoder(*, snapshot: Path, **_kwargs):
        captured.append(snapshot)
        return SimpleNamespace(report=SimpleNamespace(base_id=base_id))

    monkeypatch.setattr(runner, "load_frozen_vocoder", fake_load_frozen_vocoder)
    evidence = FakeEvidence(base_id)
    adapters = runner._load_nested_v2_vocoders(
        snapshot=snapshot,
        base_manifest=tmp_path / "base.json",
        diffusers_root=tmp_path / "diffusers",
        evidence=evidence,
        expected_base_id=base_id,
    )
    assert len(adapters) == 2
    assert captured == [snapshot.resolve(strict=True), snapshot.resolve(strict=True)]
    assert all(not str(path).startswith("/proc/") for path in captured)
    assert evidence.stable_calls == 2
    assert evidence.frozen_base.loader_path == Path("/proc/self/fd/37")

    wrong_outer = FakeEvidence("b" * 64)
    with pytest.raises(RuntimeError, match="outer frozen-base authority digest"):
        runner._load_nested_v2_vocoders(
            snapshot=snapshot,
            base_manifest=tmp_path / "base.json",
            diffusers_root=tmp_path / "diffusers",
            evidence=wrong_outer,
            expected_base_id=base_id,
        )

    linked = tmp_path / "linked-snapshot"
    linked.symlink_to(snapshot, target_is_directory=True)
    with pytest.raises(RuntimeError, match="ordinary/canonical"):
        runner._load_nested_v2_vocoders(
            snapshot=linked,
            base_manifest=tmp_path / "base.json",
            diffusers_root=tmp_path / "diffusers",
            evidence=FakeEvidence(base_id),
            expected_base_id=base_id,
        )


def test_selected_artifact_tensors_are_detached_and_round_trip_exactly() -> None:
    best_latents = torch.arange(4 * 3 * 2, dtype=torch.float32).reshape(4, 3, 2)
    best_audio = torch.arange(4 * 2 * 5, dtype=torch.float32).reshape(4, 2, 5)
    selected_latent = runner._detached_selected_tensor(best_latents, 2)
    selected_audio = runner._detached_selected_tensor(best_audio, 2)

    assert torch.equal(selected_latent, best_latents[2:3])
    assert torch.equal(selected_audio, best_audio[2:3])
    assert (
        selected_latent.untyped_storage().data_ptr()
        != best_latents.untyped_storage().data_ptr()
    )
    assert (
        selected_audio.untyped_storage().data_ptr()
        != best_audio.untyped_storage().data_ptr()
    )

    encoded = runner.save_safetensors({
        "best_audio": best_audio,
        "best_latents": best_latents,
        "selected_audio": selected_audio,
        "selected_latent": selected_latent,
    })
    decoded = load_safetensors(encoded)


def test_public_verifier_rejects_rehashed_selected_seed_manifest_and_session_attack(
    tmp_path: Path,
) -> None:
    loaded = load_v2_config(
        Path(__file__).parents[1] / "configs" / "inversion-v2.yaml"
    )
    config = loaded.config
    report = SimpleNamespace(
        project_git_dirty=False,
        project_git_commit="1" * 40,
        project_source_sha256="2" * 64,
        semantic_digest="3" * 64,
        base_id=config.expected_base_id,
        diffusers_revision=config.diffusers_revision,
    )
    model = torch.nn.Linear(1, 1, bias=False).requires_grad_(False)
    adapter = SimpleNamespace(report=report, model=model)
    parent = SimpleNamespace(
        session_file_sha256="4" * 64,
        session_semantic_digest="5" * 64,
        manifest_semantic_digests=("6" * 64, "7" * 64),
    )

    def make_manifest(
        experiment_id: str,
        parent_digest: str,
    ) -> runner.V2ExperimentManifest:
        return runner.V2ExperimentManifest.create(
            experiment_id=experiment_id,
            config_file_sha256=loaded.file_sha256,
            config_semantic_digest=loaded.semantic_digest,
            project_git_commit=report.project_git_commit,
            project_source_sha256=report.project_source_sha256,
            diffusers_revision=config.diffusers_revision,
            base_id=config.expected_base_id,
            adapter_semantic_digest=report.semantic_digest,
            oracle_semantic_digest=config.v1_authority.oracle_semantic_digest,
            v1_session_semantic_digest=parent.session_semantic_digest,
            v1_parent_manifest_semantic_digest=parent_digest,
            restart_seeds=config.execution.restart_seeds,
            latent_shape=(4, 128, 86),
            audio_shape=(4, 2, 44032),
            trace_semantic_digest="8" * 64,
            metrics_semantic_digest="9" * 64,
            selected_restart_index=2,
            selected_restart_seed=107,
            initial_latents_sha256="a" * 64,
            best_latents_sha256="b" * 64,
            selected_latent_sha256="c" * 64,
            selected_audio_sha256="d" * 64,
            target_audio_sha256="e" * 64,
            trajectory_latents_sha256="f" * 64,
            weight_state_sha256_before=runner.module_state_sha256(model),
            weight_state_sha256_after=runner.module_state_sha256(model),
            trainable_vocoder_parameter_count=0,
            vocoder_parameter_gradient_count=0,
            device="cuda",
            device_name="NVIDIA H100 80GB HBM3",
            device_capability=(9, 0),
            cuda_runtime="12.8",
            elapsed_seconds=1.0,
            peak_cuda_allocated_bytes=1,
            peak_cuda_reserved_bytes=1,
            fp32_handoff_latents_sha256="0" * 64,
            status="FAIL",
            high_fidelity_status="NOT_HIGH_FIDELITY",
            quality_claim="NO_QUALITY_CLAIM",
            artifacts={},
        )

    manifests = (
        make_manifest("P2-E1", parent.manifest_semantic_digests[0]),
        make_manifest("P2-E2", parent.manifest_semantic_digests[1]),
    )
    legitimate_bytes = tuple(runner.canonical_json_bytes(item) for item in manifests)
    references = tuple(
        runner.SessionExperimentReference(
            experiment_id=manifest.experiment_id,
            path=manifest.experiment_id,
            manifest_file_sha256=hashlib.sha256(data).hexdigest(),
            manifest_semantic_digest=manifest.semantic_digest,
            selected_restart_index=manifest.selected_restart_index,
            selected_restart_seed=manifest.selected_restart_seed,
            status=manifest.status,
            high_fidelity_status=manifest.high_fidelity_status,
        )
        for manifest, data in zip(manifests, legitimate_bytes, strict=True)
    )
    session = runner.V2SessionManifest.create(
        config_file_sha256=loaded.file_sha256,
        config_semantic_digest=loaded.semantic_digest,
        project_git_commit=report.project_git_commit,
        project_source_sha256=report.project_source_sha256,
        adapter_semantic_digest=report.semantic_digest,
        oracle_semantic_digest=config.v1_authority.oracle_semantic_digest,
        v1_session_file_sha256=parent.session_file_sha256,
        v1_session_semantic_digest=parent.session_semantic_digest,
        restart_seeds=config.execution.restart_seeds,
        experiments=references,
        all_feasibility_pass=False,
        all_high_fidelity_pass=False,
        quality_claim="NO_QUALITY_CLAIM",
    )

    forged_manifest = manifests[0].model_dump(mode="json")
    forged_manifest["selected_restart_seed"] = 101
    forged_manifest["semantic_digest"] = semantic_digest({
        key: value for key, value in forged_manifest.items() if key != "semantic_digest"
    })
    forged_manifest_bytes = runner.canonical_json_bytes(forged_manifest)
    forged_session = session.model_dump(mode="json")
    forged_session["experiments"][0]["manifest_file_sha256"] = hashlib.sha256(
        forged_manifest_bytes
    ).hexdigest()
    forged_session["experiments"][0]["manifest_semantic_digest"] = forged_manifest[
        "semantic_digest"
    ]
    forged_session["semantic_digest"] = semantic_digest({
        key: value for key, value in forged_session.items() if key != "semantic_digest"
    })

    coherently_forged_session = copy.deepcopy(forged_session)
    coherently_forged_session["experiments"][0]["selected_restart_seed"] = 101
    coherently_forged_session["semantic_digest"] = semantic_digest({
        key: value
        for key, value in coherently_forged_session.items()
        if key != "semantic_digest"
    })
    with pytest.raises(ValueError, match="selected seed/index"):
        runner.V2SessionManifest.model_validate(coherently_forged_session)

    evidence = tmp_path / "v2-evidence"
    evidence.mkdir()
    evidence.chmod(config.publication.root_mode)
    for experiment_id in runner.EXPERIMENT_IDS:
        leaf = evidence / experiment_id
        leaf.mkdir()
        leaf.chmod(config.publication.experiment_directory_mode)
        for name in runner.EXPERIMENT_FILES:
            (leaf / name).write_bytes(b"")
            (leaf / name).chmod(config.publication.file_mode)
    (evidence / "P2-E1" / "manifest.json").write_bytes(forged_manifest_bytes)
    (evidence / "P2-E2" / "manifest.json").write_bytes(legitimate_bytes[1])
    (evidence / "session.json").write_bytes(runner.canonical_json_bytes(forged_session))
    (evidence / "session.json").chmod(config.publication.file_mode)

    authority = AnchoredDirectory.open_absolute(
        evidence,
        label="forged v2 evidence",
        require_readonly=False,
    )
    try:
        with pytest.raises(ValueError, match="selected seed/index"):
            runner.verify_v2_session_authority(
                authority,
                loaded=loaded,
                parent=parent,
                adapter_fp32=adapter,
                adapter_bf16=adapter,
            )
    finally:
        authority.close(validate=True)