aituner/tests/test_oracle_gap.py

from __future__ import annotations

import json
import sys
from pathlib import Path


ORACLE_GAP = Path(__file__).resolve().parents[1] / "scripts/oracle_gap"
sys.path.insert(0, str(ORACLE_GAP))

from analyze import score_trial, summarize_trials  # noqa: E402
from run_frontier import (  # noqa: E402
    AMENDMENT,
    UP_EXTENSIONS,
    boundary_closure_needs,
    resume_compatible,
)


def _request(
    request_id: str,
    admitted: float,
    first: float,
    completed: float,
    *,
    input_tokens: int = 512,
    output_tokens: int = 64,
) -> dict:
    return {
        "request_id": request_id,
        "scheduled_s": admitted - 0.001,
        "admitted_s": admitted,
        "first_token_s": first,
        "completed_s": completed,
        "input_tokens": input_tokens,
        "requested_output_tokens": output_tokens,
        "actual_output_tokens": output_tokens,
        "success": True,
        "error_kind": None,
    }


def test_score_trial_uses_clean_admission_cohort_and_both_slos(tmp_path: Path) -> None:
    result_path = tmp_path / "result.json"
    result_path.write_text(
        json.dumps(
            {
                "warmup_seconds": 10,
                "clean_segment_seconds": 20,
                "num_clean_segments": 1,
            }
        )
    )
    rows = [
        _request("warmup", 9.0, 9.1, 10.0),
        _request("pass", 11.0, 11.1, 13.0),
        # (16.0 - 12.1) / 63 = 61.9 ms: TPOT miss.
        _request("tpot-fail", 12.0, 12.1, 16.0),
        _request("after-clean", 31.0, 31.1, 32.0),
    ]
    request_path = tmp_path / "requests.jsonl"
    request_path.write_text("".join(json.dumps(row) + "\n" for row in rows))

    scored = score_trial(
        request_path,
        result_path,
        phase="P01",
        config="C00",
        target_rate=0.1,
        repetition=0,
        role="test",
    )

    assert scored["cohort_n"] == 2
    assert scored["pass_n"] == 1
    assert scored["pass_rate"] == 0.5
    assert scored["slo_goodput_rps"] == 0.05
    assert scored["failure_reasons"] == {"tpot_slo": 1}
    assert scored["offered_rate_valid"]
    assert not scored["feasible"]


def _frontier_row(
    phase: str,
    config: str,
    rate: float,
    feasible: bool,
) -> dict:
    cohort = 100
    passes = 99 if feasible else 80
    return {
        "phase": phase,
        "config": config,
        "target_rate_rps": rate,
        "repetition": 0,
        "role": "test",
        "cohort_n": cohort,
        "pass_n": passes,
        "pass_rate": passes / cohort,
        "slo_goodput_rps": rate * passes / cohort,
        "feasible": feasible,
        "invariants": {"test": True},
    }


def test_conservative_oracle_bound_can_refute_ten_percent_gate() -> None:
    brackets = {
        "P01": {
            "C00": (30.0, 32.0),
            "C10": (28.0, 30.0),
            "C01": (29.0, 31.0),
            "C11": (27.0, 29.0),
        },
        "P06": {
            "C00": (1.6, 1.7),
            "C10": (1.7, 1.8),
            "C01": (1.55, 1.65),
            "C11": (1.6, 1.7),
        },
    }
    rows = []
    for phase, configs in brackets.items():
        for config, (lower, upper) in configs.items():
            for _ in range(3):
                rows.append(_frontier_row(phase, config, lower, True))
                rows.append(_frontier_row(phase, config, upper, False))

    summary = summarize_trials(rows)

    assert summary["verdict"] == "REFUTED"
    assert summary["worst_mixture_conservative"]["gap"] < 0.10
    assert summary["sanity"]["invariants"]["all_cells_bracketed"]
    assert summary["sanity"]["invariants"]["all_frontiers_monotone"]


def test_nonmonotone_frontier_blocks_inference() -> None:
    rows = []
    for phase in ("P01", "P06"):
        for config in ("C00", "C10", "C01", "C11"):
            for _ in range(3):
                rows.append(_frontier_row(phase, config, 1.0, True))
                rows.append(_frontier_row(phase, config, 2.0, False))
    for _ in range(3):
        rows.append(_frontier_row("P01", "C00", 3.0, True))

    summary = summarize_trials(rows)

    assert summary["verdict"] == "INCONCLUSIVE"
    assert not summary["sanity"]["invariants"]["all_frontiers_monotone"]


def test_a_og_1_extends_only_mutable_resume_fields() -> None:
    immutable = {
        "p5_client_sha256": "p5",
        "p3_client_sha256": "p3",
        "vllm_commit": "vllm",
        "model": "model",
        "manifests": {"P01": "a", "P06": "b"},
        "runtime": "runtime",
        "driver": "driver",
        "analyzer_sha256": "analyzer",
        "config_details": {"C00": {}},
        "base_rates": {"P01": [26], "P06": [1.4]},
        "up_extensions": {"P01": [38.0], "P06": [2.1, 2.2, 2.3]},
    }
    amended = {
        **immutable,
        "controller_sha256": "new",
        "repo_commit": "new",
        "up_extensions": {
            "P01": [38.0],
            "P06": [2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.8, 3.0],
        },
    }

    assert UP_EXTENSIONS["P06"][-5:] == (2.4, 2.5, 2.6, 2.8, 3.0)
    assert resume_compatible(immutable, amended)
    assert not resume_compatible(
        immutable, {**amended, "manifests": {"P01": "changed"}}
    )


def test_a_og_4_requires_the_amended_grid_to_stay_fixed() -> None:
    old = {
        "analyzer_sha256": "analyzer",
        "p5_client_sha256": "p5",
        "p3_client_sha256": "p3",
        "vllm_commit": "vllm",
        "model": "model",
        "manifests": {"P01": "a", "P06": "b"},
        "runtime": "runtime",
        "driver": "driver",
        "config_details": {"C00": {}},
        "base_rates": {"P01": [26], "P06": [1.4]},
        "up_extensions": {"P01": [38.0], "P06": [2.1, 2.2, 2.3, 2.4]},
    }

    assert AMENDMENT == "A-OG-4"
    assert resume_compatible(old, {**old, "controller_sha256": "new"})
    assert not resume_compatible(
        old,
        {**old, "up_extensions": {"P01": [38.0], "P06": [2.1, 2.2, 2.3, 2.5]}},
    )


def test_boundary_closure_follows_a_majority_shift() -> None:
    rows = []
    rows.append(_frontier_row("P06", "C00", 2.3, True))
    rows.extend(_frontier_row("P06", "C00", 2.4, value) for value in (True, False, False))
    rows.extend(_frontier_row("P06", "C00", 2.5, False) for _ in range(3))

    assert boundary_closure_needs(rows) == [2.3]

    rows.extend(_frontier_row("P06", "C00", 2.3, True) for _ in range(2))
    assert boundary_closure_needs(rows) == []