bundle/v0/tests/test_oos_timing_harness.py

from __future__ import annotations

import hashlib
import json

import numpy as np
import pytest

from loss_aware_dro_repro.core import ContractError, load_plan
from loss_aware_dro_repro.datasets import generate_dataset
from loss_aware_dro_repro.empirical_task import _regression_oos_losses, select_empirical_task
from loss_aware_dro_repro.oos_timing_harness import (
    BLAS_ENVIRONMENT_VARIABLES,
    OOS_CLI_BOOTSTRAP_VARIABLE,
    PAPER_SAMPLE_SIZES,
    PAPER_TASK_GROUPS,
    evaluate_grouped_regression_oos,
    normalize_oos_timing_config,
    run_oos_timing_sample,
)


def _diagnostic_config(sample_count: int = 257, chunk_size: int = 31):
    return {
        "schema_version": 1,
        "execution_scale": "diagnostic",
        "groups": [
            {
                **group,
                "sample_sizes": list(PAPER_SAMPLE_SIZES),
                "sample_count": sample_count,
                "chunk_size": chunk_size,
            }
            for group in PAPER_TASK_GROUPS
        ],
        "decision_pair_policy": "representative_five_pair_slopes_v1",
        "claim_eligible": False,
    }


def test_paper_contract_is_exact_and_cannot_be_scaled_down():
    plan = load_plan()
    exact = normalize_oos_timing_config(plan)
    assert [group["loss"] for group in exact["groups"]] == ["absolute", "squared"]
    assert all(group["sample_count"] == 10_000_000 for group in exact["groups"])
    assert all(group["chunk_size"] == 100_000 for group in exact["groups"])
    assert all(group["sample_sizes"] == [10, 20, 30, 40, 50] for group in exact["groups"])
    with pytest.raises(ContractError, match="may not be overridden"):
        normalize_oos_timing_config(
            plan,
            {
                **exact,
                "groups": [{**exact["groups"][0], "sample_count": 100}, exact["groups"][1]],
            },
        )


@pytest.mark.parametrize(
    "task_id",
    [
        "regression_absolute_main/d010/r09/n050",
        "regression_squared/d010/r09/n050",
    ],
)
def test_grouped_evaluator_matches_scalar_contract_and_is_repeatable(task_id):
    task = select_empirical_task(task_id)
    _, metadata = generate_dataset(task)
    pairs = [
        {"task_id": "a", "initial": 0.25, "final": 0.75},
        {"task_id": "b", "initial": -0.5, "final": 1.5},
    ]
    grouped, _ = evaluate_grouped_regression_oos(
        loss=task["loss"],
        weight=float(metadata["parameters"]["weight"]),
        noise_variance=float(metadata["parameters"]["noise_variance"]),
        oos_seed=task["seeds"]["oos"],
        decision_pairs=pairs,
        sample_count=257,
        chunk_size=31,
    )
    repeated, _ = evaluate_grouped_regression_oos(
        loss=task["loss"],
        weight=float(metadata["parameters"]["weight"]),
        noise_variance=float(metadata["parameters"]["noise_variance"]),
        oos_seed=task["seeds"]["oos"],
        decision_pairs=pairs,
        sample_count=257,
        chunk_size=31,
    )
    assert grouped == repeated
    for pair, result in zip(pairs, grouped["results"], strict=True):
        scalar = _regression_oos_losses(
            task,
            metadata,
            np.asarray([pair["initial"]]),
            np.asarray([pair["final"]]),
            sample_count=257,
            chunk_size=31,
        )
        np.testing.assert_allclose(
            [result["initial_loss"], result["final_loss"]],
            scalar[:2],
            rtol=0.0,
            atol=1e-12,
        )


def test_diagnostic_run_writes_atomic_bound_timing_receipt(tmp_path, monkeypatch):
    commit = "1" * 40
    monkeypatch.setattr(
        "loss_aware_dro_repro.oos_timing_harness._repository_state",
        lambda: (commit, True),
    )
    for variable in BLAS_ENVIRONMENT_VARIABLES:
        monkeypatch.setenv(variable, "1")
    monkeypatch.setenv(OOS_CLI_BOOTSTRAP_VARIABLE, "before_scientific_imports_v1")
    output = tmp_path / "oos-timing"
    receipt = run_oos_timing_sample(
        output, requested_config=_diagnostic_config()
    )
    assert set(path.name for path in output.iterdir()) == {
        "receipt.json",
        "result_payload.json",
    }
    persisted = json.loads((output / "receipt.json").read_text(encoding="utf-8"))
    assert persisted == receipt
    assert receipt["identity"]["implementation_commit"] == commit
    identity_groups = receipt["identity"]["groups"]
    assert [group["loss"] for group in identity_groups] == ["absolute", "squared"]
    assert all(group["sample_count"] == 257 for group in identity_groups)
    assert all(group["chunk_size"] == 31 for group in identity_groups)
    assert len({group["oos_seed"] for group in identity_groups}) == 2
    assert all(group["group_identity"].startswith("sha256:") for group in identity_groups)
    assert all(group["parameter_hash"].startswith("sha256:") for group in identity_groups)
    assert all(set(group["generating_parameters"]) >= {"weight", "noise_variance"} for group in identity_groups)
    assert receipt["identity"]["config_hash"].startswith("sha256:")
    assert receipt["identity"]["source_tree_hash"].startswith("sha256:")
    assert receipt["claim_eligible"] is False
    assert receipt["cost_freeze_eligible"] is False
    assert set(receipt["scientific_verdicts"].values()) == {"HOLD"}
    assert receipt["authority"]["paid_compute"] is False
    assert receipt["authority"]["portfolio_cost_projection"] is False
    assert receipt["runtime"]["cost_usd"] == 0.0
    assert receipt["timing_seconds"]["sampling_evaluation"] >= 0.0
    assert receipt["timing_seconds"]["aggregation"] >= 0.0
    assert receipt["timing_seconds"]["result_serialization"] >= 0.0
    timing_groups = receipt["timing_seconds"]["groups"]
    assert set(timing_groups) == {
        "absolute_regression_d010_r09",
        "squared_regression_d010_r09",
    }
    assert all(group["sampling_evaluation"] >= 0.0 for group in timing_groups.values())
    payload = json.loads((output / "result_payload.json").read_text(encoding="utf-8"))
    assert [group["result"]["loss"] for group in payload["groups"]] == [
        "absolute",
        "squared",
    ]
    payload_path = output / receipt["artifacts"]["result_payload"]
    assert hashlib.sha256(payload_path.read_bytes()).hexdigest() == receipt["artifacts"][
        "result_payload_sha256"
    ]
    with pytest.raises(ContractError, match="already exists"):
        run_oos_timing_sample(output, requested_config=_diagnostic_config())


def test_run_refuses_dirty_snapshot(tmp_path, monkeypatch):
    monkeypatch.setattr(
        "loss_aware_dro_repro.oos_timing_harness._repository_state",
        lambda: ("1" * 40, False),
    )
    with pytest.raises(ContractError, match="paper lane is not clean"):
        run_oos_timing_sample(
            tmp_path / "dirty", requested_config=_diagnostic_config(sample_count=7, chunk_size=3)
        )


def test_run_refuses_unbound_blas_threads(tmp_path, monkeypatch):
    monkeypatch.setattr(
        "loss_aware_dro_repro.oos_timing_harness._repository_state",
        lambda: ("1" * 40, True),
    )
    for variable in BLAS_ENVIRONMENT_VARIABLES:
        monkeypatch.setenv(variable, "1")
    monkeypatch.setenv(OOS_CLI_BOOTSTRAP_VARIABLE, "before_scientific_imports_v1")
    monkeypatch.setenv("OMP_NUM_THREADS", "2")
    with pytest.raises(ContractError, match="requires every supported BLAS"):
        run_oos_timing_sample(
            tmp_path / "threads", requested_config=_diagnostic_config(sample_count=7, chunk_size=3)
        )


def test_run_refuses_source_change_before_atomic_publish(tmp_path, monkeypatch):
    states = iter([("1" * 40, True), ("2" * 40, True)])
    monkeypatch.setattr(
        "loss_aware_dro_repro.oos_timing_harness._repository_state",
        lambda: next(states),
    )
    for variable in BLAS_ENVIRONMENT_VARIABLES:
        monkeypatch.setenv(variable, "1")
    monkeypatch.setenv(OOS_CLI_BOOTSTRAP_VARIABLE, "before_scientific_imports_v1")
    output = tmp_path / "changed"
    with pytest.raises(ContractError, match="source snapshot changed"):
        run_oos_timing_sample(
            output, requested_config=_diagnostic_config(sample_count=7, chunk_size=3)
        )
    assert not output.exists()