bundle/v0/tests/test_batch_recovery.py

from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path

import pytest

from loss_aware_dro_repro.batch_control import DENIED_AUTHORITY
from loss_aware_dro_repro.batch_recovery import (
    ALLOWED_FAILURE,
    _failure_index,
    _file_hash,
    _recovery_command,
    _scientific_projection,
    _validate_original_control,
    validate_recovery_manifest,
)
from loss_aware_dro_repro.core import ContractError, sha256_value


def _contracts(tmp_path):
    recovery_lane = Path(__file__).resolve().parents[1]
    source_lane = (tmp_path / "source-lane").resolve()
    recovery_output_root = (tmp_path / "recovery-output").resolve()
    python_executable = Path(sys.executable).resolve()
    source_task = {
        "task_id": "suite/d001/r00/n010",
        "scientific_run_identity": "sha256:" + "1" * 64,
        "normalized_execution_config": {"schema_version": 2, "max_outer_iterations": 5000},
        "execution_config_hash": "sha256:" + "3" * 64,
        "output_relpath": "tasks/example",
    }
    source = {
        "manifest_hash": "sha256:" + "2" * 64,
        "execution": {
            "requested_execution_config": {
                "schema_version": 2,
                "execution_scale": "stopping_storage_sample",
                "max_outer_iterations": 5000,
            }
        },
        "tasks": [source_task],
    }
    failure = {
        "event": "failed",
        "schema_version": 1,
        "manifest_hash": source["manifest_hash"],
        "task_id": source_task["task_id"],
        "scientific_run_identity": source_task["scientific_run_identity"],
        "reason": ALLOWED_FAILURE,
    }
    failure["failure_id"] = sha256_value(failure)
    task = {
        "task_id": source_task["task_id"],
        "scientific_run_identity": source_task["scientific_run_identity"],
        "output_relpath": source_task["output_relpath"],
        "requested_execution_config": source["execution"]["requested_execution_config"],
        "normalized_execution_config": source_task["normalized_execution_config"],
        "execution_config_hash": source_task["execution_config_hash"],
        "original_failure": failure,
    }
    command = _recovery_command(
        python_executable, source_lane, recovery_output_root, source_task
    )
    task["command"] = command
    launch_binding = {
        "phase": "controller_exception_recovery_v1",
        "recovery_code_commit": "a" * 40,
        "source_manifest_hash": source["manifest_hash"],
        "task_id": source_task["task_id"],
        "scientific_run_identity": source_task["scientific_run_identity"],
        "execution_config_hash": source_task["execution_config_hash"],
        "source_commit": "b" * 40,
        "python_executable": python_executable.as_posix(),
        "python_executable_sha256": _file_hash(python_executable),
        "recovery_output_root": recovery_output_root.as_posix(),
        "command": command,
    }
    task["recovery_run_key"] = sha256_value(launch_binding)
    body = {
        "schema_version": 1,
        "kind": "controller_exception_recovery_v1",
        "source_manifest_hash": source["manifest_hash"],
        "source_manifest_path": (tmp_path / "manifest.json").resolve().as_posix(),
        "source_manifest_file_sha256": "0" * 64,
        "source_output_root": (tmp_path / "source-output").resolve().as_posix(),
        "source_commit": "b" * 40,
        "source_lane": source_lane.as_posix(),
        "recovery_code_commit": "a" * 40,
        "recovery_lane": recovery_lane.as_posix(),
        "python_executable": python_executable.as_posix(),
        "python_executable_sha256": _file_hash(python_executable),
        "recovery_output_root": recovery_output_root.as_posix(),
        "validator_source_hashes": {
            name: _file_hash(recovery_lane / "src" / "loss_aware_dro_repro" / name)
            for name in ("artifacts.py", "batch_control.py", "batch_recovery.py")
        },
        "allowed_failure_reason": ALLOWED_FAILURE,
        "task_count": 1,
        "tasks": [task],
        "authority": DENIED_AUTHORITY,
    }
    recovery = {**body, "recovery_manifest_hash": sha256_value(body)}
    return source, recovery


def _rehash(recovery):
    body = {key: value for key, value in recovery.items() if key != "recovery_manifest_hash"}
    recovery["recovery_manifest_hash"] = sha256_value(body)


def test_recovery_manifest_accepts_exact_new_launch_identity(tmp_path):
    source, recovery = _contracts(tmp_path)
    validate_recovery_manifest(recovery, source)


@pytest.mark.parametrize(
    ("mutation", "message"),
    [
        ("forged_failure", "failure binding"),
        ("identity_drift", "identity drift"),
        ("omitted_task", "task count"),
        ("duplicate_task", "duplicate"),
        ("execution_drift", "configuration drift"),
        ("output_drift", "output path drift"),
        ("execution_hash_drift", "configuration hash drift"),
        ("launch_key_drift", "launch identity"),
        ("command_drift", "command binding"),
        ("failure_id_drift", "failure identity"),
    ],
)
def test_recovery_manifest_rejects_contract_drift(mutation, message, tmp_path):
    source, recovery = _contracts(tmp_path)
    if mutation == "forged_failure":
        recovery["tasks"][0]["original_failure"]["reason"] = "invented"
    elif mutation == "identity_drift":
        recovery["tasks"][0]["scientific_run_identity"] = "sha256:" + "9" * 64
    elif mutation == "omitted_task":
        recovery["tasks"] = []
    elif mutation == "duplicate_task":
        recovery["tasks"].append(dict(recovery["tasks"][0]))
        recovery["task_count"] = 2
    elif mutation == "execution_drift":
        recovery["tasks"][0]["normalized_execution_config"] = {"schema_version": 2, "max_outer_iterations": 1}
    elif mutation == "output_drift":
        recovery["tasks"][0]["output_relpath"] = "tasks/copied"
    elif mutation == "execution_hash_drift":
        recovery["tasks"][0]["execution_config_hash"] = "sha256:" + "8" * 64
    elif mutation == "launch_key_drift":
        recovery["tasks"][0]["recovery_run_key"] = "sha256:" + "8" * 64
    elif mutation == "command_drift":
        recovery["tasks"][0]["command"][-1] = "forged.json"
    else:
        recovery["tasks"][0]["original_failure"]["failure_id"] = "sha256:" + "8" * 64
    _rehash(recovery)
    with pytest.raises(ContractError, match=message):
        validate_recovery_manifest(recovery, source)


def test_failure_index_rejects_duplicate_identity(tmp_path):
    source, recovery = _contracts(tmp_path)
    control = tmp_path / "control"
    control.mkdir()
    row = recovery["tasks"][0]["original_failure"]
    (control / "failure-ledger.jsonl").write_text(
        json.dumps(row) + "\n" + json.dumps(row) + "\n",
        encoding="utf-8",
    )
    with pytest.raises(ContractError, match="duplicate"):
        _failure_index(tmp_path, source)


def test_scientific_projection_ignores_only_runtime_and_artifact_locations():
    left = {
        "task_id": "task",
        "metrics": {"loss": 1.25},
        "solver": {"status": "Solved"},
        "optimization": {
            "state": {"iteration": 2},
            "iteration_trace_bytes": 100,
            "iteration_trace_sha256": "a",
            "iteration_trace_uncompressed_bytes": 500,
        },
        "runtime": {"duration_seconds": 1.0, "blas_threads": 1},
        "artifact_hashes": {"iteration_trace": "a", "solver_receipt": "same"},
    }
    right = {
        **left,
        "runtime": {"duration_seconds": 9.0, "blas_threads": 1},
        "artifact_hashes": {"iteration_trace": "b", "solver_receipt": "same"},
        "optimization": {
            "state": {"iteration": 2},
            "iteration_trace_bytes": 101,
            "iteration_trace_sha256": "b",
            "iteration_trace_uncompressed_bytes": 501,
        },
    }
    assert _scientific_projection(left) == _scientific_projection(right)
    right["metrics"] = {"loss": 1.24}
    assert _scientific_projection(left) != _scientific_projection(right)
    right["metrics"] = left["metrics"]
    right["runtime"] = {"duration_seconds": 9.0, "blas_threads": 2}
    assert _scientific_projection(left) != _scientific_projection(right)
    right["runtime"] = {"duration_seconds": 9.0, "blas_threads": 1}
    right["artifact_hashes"]["solver_receipt"] = "changed"
    assert _scientific_projection(left) != _scientific_projection(right)


def test_original_control_requires_exact_membership_and_reservation(tmp_path):
    identity = "sha256:" + "1" * 64
    manifest = {
        "manifest_hash": "sha256:" + "2" * 64,
        "tasks": [{"task_id": "suite/task", "scientific_run_identity": identity}],
    }
    control = tmp_path / "control"
    reservations = control / "reservations"
    reservations.mkdir(parents=True)
    reservation = {
        "manifest_hash": manifest["manifest_hash"],
        "task_id": "suite/task",
        "scientific_run_identity": identity,
    }
    (reservations / f"{identity.removeprefix('sha256:')}.json").write_text(
        json.dumps(reservation), encoding="utf-8"
    )
    events = [
        {**reservation, "event": "reserved"},
        {**reservation, "event": "success", "result_sha256": "a" * 64},
    ]
    (control / "launch-ledger.jsonl").write_text(
        "".join(json.dumps(row) + "\n" for row in events), encoding="utf-8"
    )
    (control / "failure-ledger.jsonl").write_text("", encoding="utf-8")
    _validate_original_control(tmp_path, manifest)
    events.append(
        {
            "manifest_hash": manifest["manifest_hash"],
            "task_id": "foreign/task",
            "scientific_run_identity": "sha256:" + "9" * 64,
            "event": "reserved",
        }
    )
    (control / "launch-ledger.jsonl").write_text(
        "".join(json.dumps(row) + "\n" for row in events), encoding="utf-8"
    )
    with pytest.raises(ContractError, match="membership"):
        _validate_original_control(tmp_path, manifest)


def test_recovery_cli_bootstraps_sibling_source_without_pythonpath():
    lane = Path(__file__).resolve().parents[1]
    environment = os.environ.copy()
    environment.pop("PYTHONPATH", None)
    completed = subprocess.run(
        [sys.executable, str(lane / "scripts" / "recover_failed_batch.py"), "--help"],
        cwd=lane,
        env=environment,
        text=True,
        capture_output=True,
        check=False,
    )
    assert completed.returncode == 0, completed.stderr
    assert "controller exceptions" in completed.stdout