bundle/v0/tests/test_trace_streams.py

from __future__ import annotations

import hashlib
import gzip
import json
from pathlib import Path

import pytest

from loss_aware_dro_repro.core import ContractError, canonical_bytes
from loss_aware_dro_repro.trace_streams import (
    CHECKPOINT_STATES_FILENAME,
    ITERATION_TRACE_FILENAME,
    IncrementalJsonlWriter,
    LEGACY_TRACE_STREAM_CONTRACT,
    checkpoint_reasons,
    checkpoint_state_record,
    scalar_iteration_record,
    trace_stream_binding,
    trace_stream_contract,
    validate_trace_streams,
)


def _stopping(reason=None):
    diagnostic = {"signed_relative_improvement": None}
    return {
        "reason": reason,
        "paper_total_phi_literal": diagnostic,
        "paper_total_phi_abs_denominator": diagnostic,
        "released_lower_objective_abs_denominator": diagnostic,
    }


def _bundle(root: Path, iterations: int = 205, store_every: int = 100):
    run_identity = "sha256:" + "1" * 64
    task_id = "suite/d001/r00/n010"
    binding = trace_stream_binding(run_identity, task_id, store_every)
    iteration_path = root / ITERATION_TRACE_FILENAME
    checkpoint_path = root / CHECKPOINT_STATES_FILENAME
    iteration_writer = IncrementalJsonlWriter(iteration_path)
    checkpoint_writer = IncrementalJsonlWriter(checkpoint_path)
    with iteration_writer, checkpoint_writer:
        for iteration in range(iterations):
            reason = (
                "maximum_outer_iterations_reached"
                if iteration == iterations - 1
                else None
            )
            stopping = _stopping(reason)
            reasons = checkpoint_reasons(
                iteration, store_every=store_every, stop_reason=reason
            )
            scalar = scalar_iteration_record(
                stream_binding=binding,
                run_identity=run_identity,
                task_id=task_id,
                iteration=iteration,
                lower_objective=-float(iteration),
                total_objective=-float(iteration),
                penalty=0.0,
                raw_gradient_fro=1.0,
                applied_gradient_fro=1.0,
                solver_status="optimal",
                solver_residual_maximum=1e-10,
                stopping=stopping,
                iteration_seconds=0.1,
                elapsed_seconds=0.1 * (iteration + 1),
                checkpoint=bool(reasons),
            )
            iteration_writer.write(scalar)
            if reasons:
                checkpoint_writer.write(
                    checkpoint_state_record(
                        stream_binding=binding,
                        run_identity=run_identity,
                        task_id=task_id,
                        iteration=iteration,
                        reasons=reasons,
                        full_state={
                            "lower_objective": scalar["lower_objective"],
                            "total_objective": scalar["total_objective"],
                            "penalty": 0.0,
                            "raw_gradient_fro": 1.0,
                            "applied_gradient_fro": 1.0,
                            "solver_status": "optimal",
                            "solver_residual_maximum": 1e-10,
                            "metric_factor": [[1.0]],
                            "decision": [float(iteration)],
                            "stopping": stopping,
                        },
                    )
                )
        checkpoint_writer.write(
            checkpoint_state_record(
                stream_binding=binding,
                run_identity=run_identity,
                task_id=task_id,
                iteration=iterations,
                reasons=checkpoint_reasons(
                    iterations,
                    store_every=store_every,
                    stop_reason=None,
                    terminal=True,
                ),
                full_state={
                    "lower_objective": -float(iterations),
                    "total_objective": -float(iterations),
                    "penalty": 0.0,
                    "raw_gradient_fro": 1.0,
                    "applied_gradient_fro": 1.0,
                    "solver_status": "optimal",
                    "solver_residual_maximum": 1e-10,
                    "metric_factor": [[2.0]],
                    "decision": [float(iterations)],
                    "stopping": _stopping("maximum_outer_iterations_reached"),
                },
            )
        )
    hashes = {
        "iteration_trace": hashlib.sha256(iteration_path.read_bytes()).hexdigest(),
        "checkpoint_states": hashlib.sha256(checkpoint_path.read_bytes()).hexdigest(),
    }
    optimization = {
        "iterations": iterations,
        "stop_reason": "maximum_outer_iterations_reached",
        "stopping": _stopping("maximum_outer_iterations_reached"),
        "final_stationarity": 1.0,
        "trace_stream_binding": binding,
        "trace_stream_contract": trace_stream_contract(),
        "checkpoint_store_every": store_every,
        "iteration_trace_path": ITERATION_TRACE_FILENAME,
        "iteration_trace_sha256": hashes["iteration_trace"],
        "iteration_trace_records": iterations,
        "iteration_trace_bytes": iteration_writer.byte_count,
        "iteration_trace_uncompressed_bytes": iteration_writer.uncompressed_byte_count,
        "checkpoint_state_path": CHECKPOINT_STATES_FILENAME,
        "checkpoint_state_sha256": hashes["checkpoint_states"],
        "checkpoint_state_records": 5,
        "checkpoint_state_bytes": checkpoint_writer.byte_count,
        "checkpoint_state_uncompressed_bytes": checkpoint_writer.uncompressed_byte_count,
    }
    arguments = {
        "artifact_root": root,
        "optimization": optimization,
        "artifact_hashes": hashes,
        "run_identity": run_identity,
        "task_id": task_id,
        "metrics": {"worst_case_initial": 0.0, "worst_case_final": -float(iterations)},
        "scientific_contract": {
            "initial_metric_factor": [[1.0]],
            "final_metric_factor": [[2.0]],
            "initial_decision": [0.0],
            "final_decision": [float(iterations)],
        },
        "solver": {"status": "optimal"},
    }
    return arguments


def _gzip_rows(path: Path) -> list[dict]:
    with gzip.open(path, "rb") as handle:
        return [json.loads(line) for line in handle]


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)


def _rebind(arguments: dict, label: str, path: Path) -> None:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    if label == "iteration_trace":
        prefix = "iteration_trace"
    else:
        prefix = "checkpoint_state"
    arguments["optimization"][f"{prefix}_sha256"] = digest
    arguments["optimization"][f"{prefix}_bytes"] = path.stat().st_size
    try:
        with gzip.open(path, "rb") as handle:
            arguments["optimization"][f"{prefix}_uncompressed_bytes"] = len(
                handle.read()
            )
    except (OSError, EOFError):
        pass
    arguments["artifact_hashes"][label] = digest


def test_validates_periodic_schedule_and_mutual_stream_binding(tmp_path):
    arguments = _bundle(tmp_path)
    validate_trace_streams(**arguments)
    checkpoints = _gzip_rows(tmp_path / CHECKPOINT_STATES_FILENAME)
    assert [row["iteration"] for row in checkpoints] == [0, 100, 200, 204, 205]
    assert checkpoints[-2]["checkpoint_reasons"] == ["iteration_cap"]
    assert checkpoints[-1]["checkpoint_reasons"] == ["terminal"]


def test_rejects_noncanonical_stream_even_when_hashes_are_rebound(tmp_path):
    arguments = _bundle(tmp_path)
    path = tmp_path / ITERATION_TRACE_FILENAME
    rows = _gzip_rows(path)
    _write_fixed_gzip(
        path,
        ("\n".join(json.dumps(row, sort_keys=True, indent=None) for row in rows) + "\n").encode(),
    )
    _rebind(arguments, "iteration_trace", path)
    with pytest.raises(ContractError, match="not canonical"):
        validate_trace_streams(**arguments)


def test_rejects_missing_terminal_checkpoint_after_rehash(tmp_path):
    arguments = _bundle(tmp_path)
    path = tmp_path / CHECKPOINT_STATES_FILENAME
    rows = _gzip_rows(path)[:-1]
    _write_fixed_gzip(path, b"".join(canonical_bytes(row) + b"\n" for row in rows))
    _rebind(arguments, "checkpoint_states", path)
    arguments["optimization"]["checkpoint_state_records"] = 4
    with pytest.raises(ContractError, match="frozen schedule|terminal checkpoint"):
        validate_trace_streams(**arguments)


def test_gzip_bytes_are_deterministic_and_lossless(tmp_path):
    first = _bundle(tmp_path / "first", iterations=205)
    second = _bundle(tmp_path / "second", iterations=205)
    for filename in (ITERATION_TRACE_FILENAME, CHECKPOINT_STATES_FILENAME):
        assert (tmp_path / "first" / filename).read_bytes() == (
            tmp_path / "second" / filename
        ).read_bytes()
    assert first["optimization"]["trace_stream_contract"] == trace_stream_contract()
    compressed = sum(
        first["optimization"][key]
        for key in ("iteration_trace_bytes", "checkpoint_state_bytes")
    )
    uncompressed = sum(
        first["optimization"][key]
        for key in (
            "iteration_trace_uncompressed_bytes",
            "checkpoint_state_uncompressed_bytes",
        )
    )
    assert compressed < uncompressed


@pytest.mark.parametrize("corruption", ["truncate", "flip"])
def test_rejects_corrupt_gzip_after_hash_rebinding(tmp_path, corruption):
    arguments = _bundle(tmp_path)
    path = tmp_path / ITERATION_TRACE_FILENAME
    payload = bytearray(path.read_bytes())
    if corruption == "truncate":
        del payload[-7:]
    else:
        payload[len(payload) // 2] ^= 0xFF
    path.write_bytes(payload)
    _rebind(arguments, "iteration_trace", path)
    with pytest.raises(ContractError, match="cannot read trace stream|invalid JSONL"):
        validate_trace_streams(**arguments)


def test_writer_removes_partial_stream_on_failure(tmp_path):
    path = tmp_path / ITERATION_TRACE_FILENAME
    with pytest.raises(RuntimeError, match="injected"):
        with IncrementalJsonlWriter(path) as writer:
            writer.write({"record": 1})
            raise RuntimeError("injected")
    assert not path.exists()


def test_legacy_uncompressed_streams_remain_explicitly_interpretable(tmp_path):
    arguments = _bundle(tmp_path)
    binding = trace_stream_binding(
        arguments["run_identity"],
        arguments["task_id"],
        arguments["optimization"]["checkpoint_store_every"],
        LEGACY_TRACE_STREAM_CONTRACT,
    )
    for compressed_name, legacy_name, label, path_key, hash_key in (
        (
            ITERATION_TRACE_FILENAME,
            "iteration_trace.jsonl",
            "iteration_trace",
            "iteration_trace_path",
            "iteration_trace_sha256",
        ),
        (
            CHECKPOINT_STATES_FILENAME,
            "checkpoint_states.jsonl",
            "checkpoint_states",
            "checkpoint_state_path",
            "checkpoint_state_sha256",
        ),
    ):
        rows = _gzip_rows(tmp_path / compressed_name)
        for row in rows:
            row["trace_contract_version"] = 1
            row["stream_binding"] = binding
        legacy_path = tmp_path / legacy_name
        legacy_path.write_bytes(
            b"".join(canonical_bytes(row) + b"\n" for row in rows)
        )
        digest = hashlib.sha256(legacy_path.read_bytes()).hexdigest()
        arguments["optimization"][path_key] = legacy_name
        arguments["optimization"][hash_key] = digest
        arguments["artifact_hashes"][label] = digest
    arguments["optimization"].pop("trace_stream_contract")
    for key in (
        "iteration_trace_bytes",
        "iteration_trace_uncompressed_bytes",
        "checkpoint_state_bytes",
        "checkpoint_state_uncompressed_bytes",
    ):
        arguments["optimization"].pop(key)
    arguments["optimization"]["trace_stream_binding"] = binding
    validate_trace_streams(**arguments)