giant/tests/test_cli_predict.py

import uuid

import torch
import yaml
from typer.testing import CliRunner

from giant.cli import (
    _CEPH_PREDICTIONS,
    _build_predict_timing,
    _resolve_prediction_output,
    _write_prediction_ref,
    app,
)

runner = CliRunner()


# ---------------------------------------------------------------------------
# _resolve_prediction_output
# ---------------------------------------------------------------------------


def test_non_ceph_path_goes_to_data_parent(tmp_path):
    data = tmp_path / "pools" / "pbwo4" / "full.manifest"
    out, dataset_path, pred_uuid = _resolve_prediction_output(data, None)

    assert out.parent == data.parent
    assert out.name == f"{pred_uuid}.parquet"
    assert dataset_path == data.resolve()


def test_ceph_path_goes_to_central_store(tmp_path, monkeypatch):
    # Patch resolve() so /ceph/... exists on any machine running the tests.
    ceph_data = _CEPH_PREDICTIONS.parent / "pools" / "pbwo4" / "full.manifest"
    monkeypatch.setattr(
        "giant.cli.Path.resolve",
        lambda self: ceph_data if self == ceph_data else self.absolute(),
    )
    out, _, pred_uuid = _resolve_prediction_output(ceph_data, None)

    assert out.parent == _CEPH_PREDICTIONS
    assert out.name == f"{pred_uuid}.parquet"


def test_explicit_out_is_used_as_is(tmp_path):
    data = tmp_path / "data.parquet"
    explicit = tmp_path / "my_output.parquet"
    out, _, _ = _resolve_prediction_output(data, explicit)

    assert out == explicit


def test_uuid_is_valid(tmp_path):
    data = tmp_path / "data.parquet"
    _, _, pred_uuid = _resolve_prediction_output(data, None)
    parsed = uuid.UUID(pred_uuid)
    assert parsed.version == 4


def test_each_call_produces_a_distinct_uuid(tmp_path):
    data = tmp_path / "data.parquet"
    _, _, uuid1 = _resolve_prediction_output(data, None)
    _, _, uuid2 = _resolve_prediction_output(data, None)
    assert uuid1 != uuid2


def test_dataset_path_is_resolved(tmp_path):
    data = tmp_path / "data.parquet"
    _, dataset_path, _ = _resolve_prediction_output(data, None)
    assert dataset_path.is_absolute()


# ---------------------------------------------------------------------------
# _write_prediction_ref
# ---------------------------------------------------------------------------


def test_ref_file_created_in_checkpoint_dir(tmp_path):
    ckpt_dir = tmp_path / "checkpoints" / "run1"
    ckpt_dir.mkdir(parents=True)
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    out = tmp_path / "predictions" / "abc.parquet"
    dataset = tmp_path / "pools" / "pbwo4" / "full.manifest"
    pred_uuid = str(uuid.uuid4())

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset)

    assert ref_path == ckpt_dir / f"{pred_uuid}.yaml"
    assert ref_path.exists()


def test_ref_yaml_contains_expected_fields(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    out = tmp_path / "pred.parquet"
    dataset = tmp_path / "full.manifest"
    pred_uuid = str(uuid.uuid4())

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset)
    data = yaml.safe_load(ref_path.read_text())

    assert data["kind"] == "prediction"
    assert data["prediction_id"] == pred_uuid
    assert data["output"] == str(out)
    assert data["dataset"] == str(dataset)
    assert data["checkpoint"] == str(checkpoint.resolve())
    assert "timestamp" in data
    assert "comment" not in data


def test_ref_yaml_includes_comment_when_provided(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    out = tmp_path / "pred.parquet"
    dataset = tmp_path / "full.manifest"
    pred_uuid = str(uuid.uuid4())

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
    data = yaml.safe_load(ref_path.read_text())

    assert data["comment"] == "baseline sweep run 3"


def test_ref_timestamp_is_iso_format(tmp_path):
    from datetime import datetime

    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    pred_uuid = str(uuid.uuid4())
    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
    data = yaml.safe_load(ref_path.read_text())

    # Must parse without error and be timezone-aware (UTC).
    ts = datetime.fromisoformat(data["timestamp"])
    assert ts.tzinfo is not None


def test_ref_yaml_merges_extra_after_base_fields(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    out = tmp_path / "pred.parquet"
    dataset = tmp_path / "full.manifest"
    pred_uuid = str(uuid.uuid4())

    ref_path = _write_prediction_ref(
        checkpoint,
        pred_uuid,
        out,
        dataset,
        extra={"coord": "global", "n_rows": 42, "timing": {"setup_s": 1.0}},
    )
    data = yaml.safe_load(ref_path.read_text())

    # Base fields untouched, extras layered on top.
    assert data["kind"] == "prediction"
    assert data["prediction_id"] == pred_uuid
    assert data["coord"] == "global"
    assert data["n_rows"] == 42
    assert data["timing"] == {"setup_s": 1.0}


def test_ref_yaml_without_extra_matches_today(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    pred_uuid = str(uuid.uuid4())
    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
    data = yaml.safe_load(ref_path.read_text())

    assert set(data) == {"kind", "prediction_id", "output", "dataset", "checkpoint", "timestamp"}


def test_ref_checkpoint_path_is_absolute(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
    checkpoint = ckpt_dir / "best.pt"
    checkpoint.touch()

    pred_uuid = str(uuid.uuid4())
    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
    data = yaml.safe_load(ref_path.read_text())

    assert data["checkpoint"].startswith("/")


# ---------------------------------------------------------------------------
# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
# not just at the giant.checkpoint_io unit level).
# ---------------------------------------------------------------------------


def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
    checkpoint = tmp_path / "bad.pt"
    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)

    result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])

    assert result.exit_code == 1
    assert "checkpoint has no model_config" in result.output


# ---------------------------------------------------------------------------
# --set (gitea #87)
# ---------------------------------------------------------------------------


def test_predict_set_flag_without_equals_exits_1(tmp_path):
    checkpoint = tmp_path / "missing.pt"

    result = runner.invoke(
        app,
        ["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
    )

    assert result.exit_code == 1
    assert "must be 'dotted.path=value'" in result.output


def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
    checkpoint = tmp_path / "ckpt.pt"
    torch.save(
        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
        checkpoint,
    )

    result = runner.invoke(
        app,
        [
            "predict",
            "dummy.parquet",
            "--checkpoint",
            str(checkpoint),
            "--set",
            "stage1_model.hidden_dim=999",
        ],
    )

    assert result.exit_code == 1
    assert "not an inference-safe override" in result.output


# ---------------------------------------------------------------------------
# schema v3 constants (truth-column tagging)
# ---------------------------------------------------------------------------


def test_predict_schema_version_is_v3():
    from giant.constants import PREDICT_SCHEMA_VERSION

    assert PREDICT_SCHEMA_VERSION == "3"


def test_predict_truth_metadata_key_exists():
    from giant.constants import PREDICT_TRUTH_METADATA_KEY

    assert PREDICT_TRUTH_METADATA_KEY == "giant.predict.has_truth"


# ---------------------------------------------------------------------------
# _build_predict_timing
# ---------------------------------------------------------------------------


def test_build_predict_timing_computes_per_step_cost():
    timing = _build_predict_timing(
        setup_s=1.0,
        predict_s=10.0,
        write_s=2.0,
        n_rows=100,
        device="cpu",
        torch_threads=4,
    )
    assert timing["n_rows"] == 100
    assert timing["sample_s"] == 8.0  # predict_s - write_s
    assert timing["us_per_step"] == 8.0 / 100 * 1e6
    assert timing["write_us_per_step"] == 2.0 / 100 * 1e6
    assert timing["rows_per_s"] == 10.0
    assert timing["device"] == "cpu" and timing["torch_threads"] == 4


def test_build_predict_timing_handles_zero_rows():
    timing = _build_predict_timing(
        setup_s=1.0,
        predict_s=0.0,
        write_s=0.0,
        n_rows=0,
        device="cpu",
        torch_threads=1,
    )
    assert timing["us_per_step"] is None
    assert timing["write_us_per_step"] is None
    assert timing["rows_per_s"] is None


def test_predict_has_truth_flag_default_on():
    # Inspecting rendered --help text is brittle across terminal
    # widths/color settings (wraps or re-colors mid-flag); go straight to
    # the underlying click command's registered option instead.
    from typing import cast

    import typer
    from click import Group

    predict_cmd = cast(Group, typer.main.get_command(app)).commands["predict"]
    truth_param = next(p for p in predict_cmd.params if p.name == "truth")
    assert truth_param.opts == ["--truth"]
    assert truth_param.secondary_opts == ["--no-truth"]
    assert truth_param.default is True