bundle/v0/tests/test_paper_task_runner.py

from __future__ import annotations

import hashlib
import gzip
import json
from pathlib import Path

import pytest

import loss_aware_dro_repro.paper_task as paper_task_module
from loss_aware_dro_repro.artifacts import validate_result
from loss_aware_dro_repro.core import ContractError, canonical_bytes, load_plan, plan_hash
from loss_aware_dro_repro.matrix import expand_tasks
from loss_aware_dro_repro.paper_task import (
    normalize_execution_config,
    require_execution_snapshot,
    run_gaussian_task,
    select_gaussian_task,
    verify_execution_snapshot_unchanged,
)


TEST_EXECUTION = {
    "schema_version": 2,
    "execution_scale": "test",
    "max_outer_iterations": 1,
}

STOPPING_STORAGE_EXECUTION = {
    "schema_version": 2,
    "execution_scale": "stopping_storage_sample",
    "max_outer_iterations": 5000,
}


def _load(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def _sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _scientific_trace(path: Path):
    with gzip.open(path, "rb") as handle:
        rows = [json.loads(line) for line in handle]
    for row in rows:
        row.pop("timing")
    return rows


def _write_fixed_gzip(path: Path, payload: bytes) -> None:
    with path.open("wb") as raw:
        with gzip.GzipFile(
            filename="", mode="wb", compresslevel=9, fileobj=raw, mtime=0
        ) as compressed:
            compressed.write(payload)


@pytest.mark.parametrize(
    "task_id,suite",
    [
        ("portfolio_gaussian_main/d001/r00/n010", "portfolio_gaussian_main"),
        (
            "portfolio_gaussian_coverage_ablation/d001/r00/n010",
            "portfolio_gaussian_coverage_ablation",
        ),
        ("portfolio_gaussian_highdim/d001/r00/n010", "portfolio_gaussian_highdim"),
    ],
)
def test_selects_each_supported_gaussian_suite(task_id, suite):
    assert select_gaussian_task(task_id)["suite"] == suite


def test_rejects_non_gaussian_suite():
    with pytest.raises(ContractError, match="not supported"):
        select_gaussian_task("regression_absolute_main/d001/r00/n010")


def test_paper_scale_config_rejects_overrides():
    plan = load_plan()
    exact = normalize_execution_config(plan)
    assert exact["max_outer_iterations"] == 1_000_000
    assert exact["optimizer"] == "paper_algorithm_plain_gradient_descent"
    with pytest.raises(ContractError, match="may not override"):
        normalize_execution_config(plan, {**exact, "max_outer_iterations": 1})
    with pytest.raises(ContractError, match=r"\[1, 5\]"):
        normalize_execution_config(
            plan,
            {"schema_version": 2, "execution_scale": "test", "max_outer_iterations": 6},
        )
    changed_optimizer = {
        **plan,
        "paper_hyperparameters": {
            **plan["paper_hyperparameters"],
            "optimizer": "adam",
        },
    }
    with pytest.raises(ContractError, match="frozen paper optimizer"):
        normalize_execution_config(changed_optimizer)


def test_stopping_storage_sample_normalizes_to_frozen_paper_mechanics():
    plan = load_plan()
    normalized = normalize_execution_config(plan, STOPPING_STORAGE_EXECUTION)
    assert normalized == {
        **normalize_execution_config(plan),
        "execution_scale": "stopping_storage_sample",
        "max_outer_iterations": 5000,
        "claim_eligible": False,
    }
    assert normalized["store_every"] == 100
    assert normalized["relative_objective_improvement_tolerance"] == 1e-6
    assert normalized["risk_coefficient_mode"] == plan["numerics"][
        "gaussian_risk_estimands"
    ]["primary"]


@pytest.mark.parametrize(
    "override",
    [
        {"max_outer_iterations": 4999},
        {"max_outer_iterations": 5000.0},
        {"schema_version": 2.0},
        {"store_every": 1},
        {"optimizer": "adam"},
        {"risk_coefficient_mode": "empirical_bootstrap"},
    ],
)
def test_stopping_storage_sample_rejects_caps_and_scientific_overrides(override):
    with pytest.raises(ContractError, match="exact frozen 5000-step request"):
        normalize_execution_config(
            load_plan(),
            {**STOPPING_STORAGE_EXECUTION, **override},
        )


def test_paper_scale_requires_clean_snapshot_but_test_scale_does_not():
    paper = normalize_execution_config(load_plan())
    with pytest.raises(ContractError, match="paper lane is not clean"):
        require_execution_snapshot(paper, False)
    sample = normalize_execution_config(load_plan(), STOPPING_STORAGE_EXECUTION)
    with pytest.raises(ContractError, match="paper lane is not clean"):
        require_execution_snapshot(sample, False)
    require_execution_snapshot(TEST_EXECUTION, False)


def test_stopping_storage_sample_rejects_a_changed_execution_snapshot(monkeypatch):
    sample = normalize_execution_config(load_plan(), STOPPING_STORAGE_EXECUTION)
    monkeypatch.setattr(
        paper_task_module,
        "_repository_state",
        lambda: ("b" * 40, True),
    )
    with pytest.raises(ContractError, match="source snapshot changed"):
        verify_execution_snapshot_unchanged(sample, "a" * 40, True)


@pytest.mark.parametrize(
    "task_id",
    [
        "portfolio_gaussian_main/d001/r00/n010",
        "portfolio_gaussian_coverage_ablation/d001/r00/n010",
        "portfolio_gaussian_highdim/d001/r00/n010",
    ],
)
def test_executes_supported_task_and_binds_all_receipts(tmp_path, task_id):
    output = tmp_path / task_id.split("/")[0]
    result = run_gaussian_task(
        task_id, output, requested_execution_config=TEST_EXECUTION
    )
    expected = next(task for task in expand_tasks() if task["task_id"] == task_id)
    validate_result(
        result,
        expected,
        plan_hash(load_plan()),
        artifact_root=output,
        solver_residual_max=load_plan()["numerics"]["solver_residual_max"],
    )
    assert set(path.name for path in output.iterdir()) == {
        "lineage_receipt.json",
        "result.json",
        "solver_receipt.json",
        "stopping_receipt.json",
        "oos_receipt.json",
        "iteration_trace.jsonl.gz",
        "checkpoint_states.jsonl.gz",
    }
    lineage = _load(output / "lineage_receipt.json")
    assert lineage["task_hash"] == expected["task_hash"]
    assert lineage["execution_config_hash"].startswith("sha256:")
    assert lineage["source_tree_hash"].startswith("sha256:")
    assert lineage["claim_eligible"] is False
    assert set(lineage["scientific_verdicts"].values()) == {"HOLD"}
    assert result["artifact_hashes"]["iteration_trace"] == _sha256(
        output / "iteration_trace.jsonl.gz"
    )
    assert result["artifact_hashes"]["checkpoint_states"] == _sha256(
        output / "checkpoint_states.jsonl.gz"
    )
    assert result["optimization"]["iteration_trace_records"] == 1
    assert result["optimization"]["checkpoint_state_records"] == 2
    solver = _load(output / "solver_receipt.json")
    stopping = _load(output / "stopping_receipt.json")
    for receipt in (solver, stopping):
        assert receipt["task_hash"] == expected["task_hash"]
        assert receipt["execution_config_hash"] == lineage["execution_config_hash"]
        assert receipt["source_tree_hash"] == lineage["source_tree_hash"]
    assert solver["all_pass"] is True
    assert sum(solver["primary_status_counts"].values()) == 2
    assert solver["primary_residual_maximum"] >= 0.0
    assert solver["solver_call_count"] == 2 + solver["accuracy_refinement_count"]
    oos = _load(output / "oos_receipt.json")
    assert oos["run_identity"] == result["run_identity"]
    assert oos["evaluation"] == "analytic_true_gaussian_cvar"
    assert oos["initial_cvar"] == result["metrics"]["oos_initial"]
    assert oos["final_cvar"] == result["metrics"]["oos_final"]
    assert result["artifact_hashes"]["oos_receipt"] == _sha256(
        output / "oos_receipt.json"
    )
    assert stopping["stop_reason"] == "maximum_outer_iterations_reached"
    assert stopping["iterations"] == 1
    assert stopping["primary_stopping_rule"] == "paper_total_phi_literal"
    assert stopping["safety_diagnostic"] == "paper_total_phi_abs_denominator"
    assert stopping["released_source_diagnostic"] == "released_lower_objective_abs_denominator"
    assert stopping["terminal_diagnostics"] == result["optimization"]["stopping"]


def test_same_task_has_deterministic_scientific_payload_and_is_immutable(tmp_path):
    task_id = "portfolio_gaussian_main/d001/r00/n010"
    first, second = tmp_path / "first", tmp_path / "second"
    result_a = run_gaussian_task(
        task_id, first, requested_execution_config=TEST_EXECUTION
    )
    result_b = run_gaussian_task(
        task_id, second, requested_execution_config=TEST_EXECUTION
    )
    assert _scientific_trace(first / "iteration_trace.jsonl.gz") == _scientific_trace(
        second / "iteration_trace.jsonl.gz"
    )
    assert (first / "checkpoint_states.jsonl.gz").read_bytes() == (
        second / "checkpoint_states.jsonl.gz"
    ).read_bytes()
    assert result_a["run_identity"] == result_b["run_identity"]
    assert result_a["metrics"] == result_b["metrics"]
    assert result_a["solver"] == result_b["solver"]
    assert _load(first / "lineage_receipt.json")["deterministic_payload_hash"] == _load(
        second / "lineage_receipt.json"
    )["deterministic_payload_hash"]
    with pytest.raises(ContractError, match="already exists"):
        run_gaussian_task(
            task_id, first, requested_execution_config=TEST_EXECUTION
        )


def test_validator_rejects_tampered_receipt(tmp_path):
    task_id = "portfolio_gaussian_main/d001/r00/n010"
    output = tmp_path / "tampered"
    result = run_gaussian_task(task_id, output, requested_execution_config=TEST_EXECUTION)
    (output / "solver_receipt.json").write_text("{}\n", encoding="utf-8")
    expected = next(task for task in expand_tasks() if task["task_id"] == task_id)
    with pytest.raises(ContractError, match="artifact hash mismatch"):
        validate_result(
            result,
            expected,
            plan_hash(load_plan()),
            artifact_root=output,
            solver_residual_max=load_plan()["numerics"]["solver_residual_max"],
        )


def test_validator_rejects_noncanonical_trace_with_rebound_hash(tmp_path):
    task_id = "portfolio_gaussian_main/d001/r00/n010"
    output = tmp_path / "noncanonical-trace"
    result = run_gaussian_task(task_id, output, requested_execution_config=TEST_EXECUTION)
    trace_path = output / "iteration_trace.jsonl.gz"
    with gzip.open(trace_path, "rb") as handle:
        row = json.loads(handle.read())
    payload = (json.dumps(row, sort_keys=True) + "\n").encode()
    _write_fixed_gzip(trace_path, payload)
    digest = hashlib.sha256(trace_path.read_bytes()).hexdigest()
    result["optimization"]["iteration_trace_sha256"] = digest
    result["optimization"]["iteration_trace_bytes"] = trace_path.stat().st_size
    result["optimization"]["iteration_trace_uncompressed_bytes"] = len(payload)
    result["artifact_hashes"]["iteration_trace"] = digest
    expected = next(task for task in expand_tasks() if task["task_id"] == task_id)
    with pytest.raises(ContractError, match="not canonical"):
        validate_result(
            result,
            expected,
            plan_hash(load_plan()),
            artifact_root=output,
            solver_residual_max=load_plan()["numerics"]["solver_residual_max"],
        )


def test_validator_rejects_rebound_checkpoint_schedule_tamper(tmp_path):
    task_id = "portfolio_gaussian_main/d001/r00/n010"
    output = tmp_path / "schedule-tamper"
    result = run_gaussian_task(task_id, output, requested_execution_config=TEST_EXECUTION)
    checkpoint_path = output / "checkpoint_states.jsonl.gz"
    with gzip.open(checkpoint_path, "rb") as handle:
        rows = [json.loads(line) for line in handle]
    rows[0]["checkpoint_reasons"].append("periodic")
    payload = b"".join(canonical_bytes(row) + b"\n" for row in rows)
    _write_fixed_gzip(checkpoint_path, payload)
    digest = hashlib.sha256(checkpoint_path.read_bytes()).hexdigest()
    result["optimization"]["checkpoint_state_sha256"] = digest
    result["optimization"]["checkpoint_state_bytes"] = checkpoint_path.stat().st_size
    result["optimization"]["checkpoint_state_uncompressed_bytes"] = len(payload)
    result["artifact_hashes"]["checkpoint_states"] = digest
    expected = next(task for task in expand_tasks() if task["task_id"] == task_id)
    with pytest.raises(ContractError, match="checkpoint reasons"):
        validate_result(
            result,
            expected,
            plan_hash(load_plan()),
            artifact_root=output,
            solver_residual_max=load_plan()["numerics"]["solver_residual_max"],
        )


def test_failed_gaussian_run_removes_unpublished_streams(tmp_path, monkeypatch):
    import loss_aware_dro_repro.paper_task as module

    output = tmp_path / "failed"
    monkeypatch.setattr(
        module,
        "evaluate_outer_state",
        lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("injected")),
    )
    with pytest.raises(RuntimeError, match="injected"):
        run_gaussian_task(
            "portfolio_gaussian_main/d001/r00/n010",
            output,
            requested_execution_config=TEST_EXECUTION,
        )
    assert not output.exists()
    assert not list(tmp_path.glob(".paper-task-*"))