neuron-tai/tests/test_node_doctor.py

"""NCA-002 tests for `meshnet-node doctor`.

The unit tests inject a fake backend, so none of them download a model, import
Torch, or need a GPU. The one test that runs a real model is `integration`-marked
and takes its model identity from the environment — it has no model default, on
purpose: the doctor is model-agnostic and so is its test.
"""

import base64
import json
import os
import struct
from pathlib import Path

import pytest

from meshnet_node import doctor
from meshnet_node.capability import STATUS_FAILED, STATUS_PASSED, CapabilityReport
from meshnet_node.doctor import (
    CATEGORY_FORWARD_FAILED,
    CATEGORY_INSUFFICIENT_MEMORY,
    CATEGORY_INVALID_SHARD,
    CATEGORY_MISSING_DEPENDENCY,
    CATEGORY_NO_MODEL,
    CATEGORY_UNSUPPORTED_RECIPE,
    PROBE_TOKENS,
    DoctorError,
    DoctorSelection,
    build_probe_input,
    classify_failure,
    probe_forward,
    render_result,
    resolve_selection,
    run_doctor,
    select_recipes,
    write_reports,
)
from meshnet_node.model_backend import (
    InsufficientVRAMError,
    MissingModelDependencyError,
    UnsupportedRecipeParam,
    validate_recipe_params,
)
from meshnet_node.recipe_manifest import parse_recipe_manifest

# Deliberately not a model this project ships against: nothing here may special-case it.
FIXTURE_MODEL = "acme-labs/Widget-9000-Instruct"

MANIFEST = parse_recipe_manifest(
    {
        "schema_version": 1,
        "catalogue_version": "test-1",
        "recipes": [
            {"id": "baseline", "version": "1", "backend_id": "torch-transformers"},
            {
                "id": "stateless",
                "version": "2",
                "backend_id": "torch-transformers",
                "params": {"use_cache": False},
            },
        ],
    },
    source="<test manifest>",
)


class _Payload:
    """Stands in for model_backend.TensorPayload."""

    def __init__(self, body: bytes, shape: list[int]) -> None:
        self.body = body
        self.shape = shape
        self.attention_mask_header = None
        self.position_ids_header = None


class _TailToken:
    """Stands in for model_backend.TailTokenResult."""

    def __init__(self, token_id: int = 7) -> None:
        self.token_id = token_id
        self.text = "ok"


class _Device:
    def __init__(self, type_: str = "cpu") -> None:
        self.type = type_


class _FakeBackend:
    """A backend that loads but records exactly how it was driven."""

    hidden_size = 8

    def __init__(
        self,
        *,
        is_head: bool = True,
        is_tail: bool = False,
        shard_start: int = 0,
        shard_end: int = 3,
        forward_error: Exception | None = None,
    ) -> None:
        self.model_id = FIXTURE_MODEL
        self.is_head = is_head
        self.is_tail = is_tail
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.device = _Device("cpu")
        self.forward_error = forward_error
        self.encoded_prompts: list[str] = []
        self.forwards: list[dict] = []

    def encode_prompt(self, prompt: str):
        if self.forward_error is not None:
            raise self.forward_error
        self.encoded_prompts.append(prompt)
        return _Payload(b"\x00" * (PROBE_TOKENS * self.hidden_size * 2),
                        [1, PROBE_TOKENS, self.hidden_size])

    def forward_bytes(
        self,
        body,
        shape,
        attention_mask_header,
        position_ids_header,
        start_layer=None,
        **kwargs,
    ):
        if self.forward_error is not None:
            raise self.forward_error
        self.forwards.append(
            {
                "body_len": len(body),
                "shape": shape,
                "start_layer": start_layer,
                "attention_mask_header": attention_mask_header,
                "position_ids_header": position_ids_header,
            }
        )
        if self.is_tail:
            return _TailToken()
        return _Payload(body, shape)


def _loader(backend=None, *, error: Exception | None = None):
    """A load_backend stub that records the (selection, recipe) pairs it saw."""
    calls: list[tuple[DoctorSelection, object]] = []

    def load(selection, recipe):
        calls.append((selection, recipe))
        if error is not None:
            raise error
        return backend if backend is not None else _FakeBackend()

    load.calls = calls  # type: ignore[attr-defined]
    return load


def _selection(**overrides) -> DoctorSelection:
    kwargs = dict(model_id=FIXTURE_MODEL, shard_start=0, shard_end=3)
    kwargs.update(overrides)
    return DoctorSelection(**kwargs)


# --- selection resolves the same as startup ---------------------------------


def test_resolve_selection_uses_the_configured_repo_shard_and_quantization():
    selection = resolve_selection(
        {
            "model_hf_repo": FIXTURE_MODEL,
            "model_name": "Widget-9000-Instruct",
            "shard_start": 4,
            "shard_end": 11,
            "quantization": "bf16",  # startup normalizes this to bfloat16
            "download_dir": "/models",
        }
    )

    assert selection.model_id == FIXTURE_MODEL
    assert (selection.shard_start, selection.shard_end) == (4, 11)
    assert selection.quantization == "bfloat16"
    assert selection.cache_dir == Path("/models")


def test_resolve_selection_defaults_to_the_whole_model_like_startup():
    """With no pinned shard, startup serves layers 0..n-1 — so doctor validates that."""
    seen: list[tuple[str, Path | None]] = []

    def detect(model_id, cache_dir):
        seen.append((model_id, cache_dir))
        return 24

    selection = resolve_selection(
        {"model_hf_repo": FIXTURE_MODEL}, detect_layers=detect
    )

    assert (selection.shard_start, selection.shard_end) == (0, 23)
    assert seen == [(FIXTURE_MODEL, None)]


def test_resolve_selection_without_a_model_is_actionable():
    with pytest.raises(DoctorError) as exc:
        resolve_selection({"model_hf_repo": "", "model_name": ""})

    assert exc.value.category == CATEGORY_NO_MODEL
    assert "--model" in exc.value.hint


def test_resolve_selection_rejects_an_inverted_shard_range():
    with pytest.raises(DoctorError) as exc:
        resolve_selection(
            {"model_hf_repo": FIXTURE_MODEL, "shard_start": 9, "shard_end": 2}
        )

    assert exc.value.category == CATEGORY_INVALID_SHARD


def test_resolve_selection_reports_an_unreadable_model_config():
    with pytest.raises(DoctorError) as exc:
        resolve_selection(
            {"model_hf_repo": FIXTURE_MODEL}, detect_layers=lambda *_: None
        )

    assert exc.value.category == doctor.CATEGORY_MODEL_UNAVAILABLE
    assert "--shard-start" in str(exc.value)


# --- the bounded real forward ------------------------------------------------


def test_a_pass_requires_a_real_forward_through_the_selected_shard():
    """Hardware being fine is not the bar: the shard itself has to execute."""
    backend = _FakeBackend(is_head=True)

    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=_loader(backend), now=lambda: 1000.0
    )

    assert result.passed
    assert backend.encoded_prompts == [doctor.PROBE_PROMPT]
    report = result.reports[0]
    assert report.status == STATUS_PASSED
    assert report.model.model_id == FIXTURE_MODEL
    assert (report.shard.start, report.shard.end) == (0, 3)


def test_a_backend_that_loads_but_cannot_forward_never_passes():
    """The regression this story exists for: a load is not a validation."""
    backend = _FakeBackend(forward_error=RuntimeError("kernel exploded"))

    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=_loader(backend), now=lambda: 1.0
    )

    assert not result.passed
    assert result.exit_code == 1
    report = result.reports[0]
    assert report.status == STATUS_FAILED
    assert result.results[0].category == CATEGORY_FORWARD_FAILED
    assert any("kernel exploded" in d for d in report.diagnostics)


def test_a_mid_shard_is_probed_with_peer_shaped_hidden_states():
    backend = _FakeBackend(is_head=False, shard_start=4, shard_end=7)

    detail = probe_forward(backend)

    assert detail["probe"] == "hidden-states"
    assert backend.encoded_prompts == []
    forward = backend.forwards[0]
    assert forward["shape"] == [1, PROBE_TOKENS, backend.hidden_size]
    # bfloat16 == 2 bytes per element, and the probe stays bounded to PROBE_TOKENS.
    assert forward["body_len"] == PROBE_TOKENS * backend.hidden_size * 2
    assert forward["start_layer"] == 4


def test_a_head_and_tail_shard_also_decodes_so_the_lm_head_is_covered():
    backend = _FakeBackend(is_head=True, is_tail=True, shard_end=5)

    detail = probe_forward(backend)

    assert detail["probe"] == "prompt+decode"
    assert detail["output"] == "token"
    # Re-entering above the last layer decodes without re-running any layer.
    assert backend.forwards[0]["start_layer"] == 6


def test_a_tail_shard_that_decodes_a_token_passes():
    backend = _FakeBackend(is_head=False, is_tail=True, shard_start=8, shard_end=11)

    detail = probe_forward(backend)

    assert detail == {
        "probe": "hidden-states",
        "tokens": PROBE_TOKENS,
        "output": "token",
        "token_id": 7,
    }


def test_an_empty_forward_result_is_a_failure_not_a_pass():
    backend = _FakeBackend(is_head=False)
    backend.forward_bytes = lambda *a, **k: _Payload(b"", [])  # type: ignore[assignment]

    with pytest.raises(DoctorError) as exc:
        probe_forward(backend)

    assert exc.value.category == CATEGORY_FORWARD_FAILED


def test_a_backend_with_no_hidden_size_cannot_be_probed():
    with pytest.raises(DoctorError) as exc:
        build_probe_input(0)

    assert exc.value.category == CATEGORY_FORWARD_FAILED


def test_probe_headers_decode_as_int64_tensors():
    probe = build_probe_input(hidden_size=8, tokens=3)

    shape, encoded = probe.position_ids_header.split(":", 1)
    raw = base64.b64decode(encoded)
    assert shape == "1,3"
    assert list(struct.unpack("<3q", raw)) == [0, 1, 2]


# --- recipes -----------------------------------------------------------------


def test_the_default_run_validates_only_the_selected_recipe():
    """Onboarding must not pay to validate recipes the node was not asked to serve."""
    load = _loader()

    result = run_doctor(_selection(), manifest=MANIFEST, load_backend=load)

    assert [r.recipe.id for r in result.results] == ["baseline"]
    assert len(load.calls) == 1


def test_all_recipes_is_explicit_and_validates_every_recipe():
    load = _loader()

    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=load, all_recipes=True
    )

    assert [r.recipe.id for r in result.results] == ["baseline", "stateless"]
    assert len(load.calls) == 2
    assert result.passed


def test_each_recipe_reaches_the_backend_that_runs_it():
    """A recipe that never reaches the loader was not really validated."""
    load = _loader()

    run_doctor(_selection(), manifest=MANIFEST, load_backend=load, all_recipes=True)

    params = [recipe.params for _, recipe in load.calls]
    assert params == [{}, {"use_cache": False}]


def test_an_unknown_recipe_names_the_ones_that_exist():
    with pytest.raises(DoctorError) as exc:
        select_recipes(MANIFEST, recipe_id="does-not-exist")

    assert exc.value.category == CATEGORY_UNSUPPORTED_RECIPE
    assert "baseline" in str(exc.value)


def test_recipe_and_all_recipes_are_mutually_exclusive():
    with pytest.raises(DoctorError):
        select_recipes(MANIFEST, recipe_id="baseline", all_recipes=True)


def test_a_recipe_the_backend_cannot_apply_is_a_failure_not_a_silent_pass():
    validate_recipe_params({"use_cache": False, "attn_implementation": "eager"})

    with pytest.raises(UnsupportedRecipeParam) as exc:
        validate_recipe_params({"sparkle_mode": True})

    assert "sparkle_mode" in str(exc.value)
    assert classify_failure(exc.value) == CATEGORY_UNSUPPORTED_RECIPE


def test_the_shipped_recipes_are_all_applicable_by_the_backend():
    """recipes.json and the backend's supported params must not drift apart."""
    from meshnet_node.recipe_manifest import load_recipe_manifest

    for recipe in load_recipe_manifest().recipes:
        validate_recipe_params(recipe.params)


# --- failure reporting -------------------------------------------------------


@pytest.mark.parametrize(
    "exc, category",
    [
        (MissingModelDependencyError("no torch"), CATEGORY_MISSING_DEPENDENCY),
        (InsufficientVRAMError("too big"), CATEGORY_INSUFFICIENT_MEMORY),
        (UnsupportedRecipeParam("nope"), CATEGORY_UNSUPPORTED_RECIPE),
        (ValueError("shard_end 99 exceeds last layer index 23"), CATEGORY_INVALID_SHARD),
        (FileNotFoundError("config.json"), doctor.CATEGORY_MODEL_UNAVAILABLE),
        (RuntimeError("something else"), doctor.CATEGORY_LOAD_FAILED),
    ],
)
def test_load_failures_are_classified_into_actionable_categories(exc, category):
    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=_loader(error=exc)
    )

    assert not result.passed
    item = result.results[0]
    assert item.category == category
    assert item.hint  # every category tells the operator what to do next
    assert item.report.status == STATUS_FAILED


def test_a_failure_report_carries_the_hint_and_no_traceback():
    result = run_doctor(
        _selection(),
        manifest=MANIFEST,
        load_backend=_loader(error=InsufficientVRAMError("insufficient VRAM to load")),
    )

    diagnostics = " ".join(result.reports[0].diagnostics)
    assert "insufficient VRAM to load" in diagnostics
    assert "--shard-start" in diagnostics  # the actionable next step
    assert "Traceback" not in diagnostics
    assert ".py" not in diagnostics  # no file/line noise from a stack


def test_a_failure_report_still_identifies_what_was_being_validated():
    """NCA-003 refuses to register without a matching report — including a failed one."""
    result = run_doctor(
        _selection(shard_start=4, shard_end=9, quantization="int8"),
        manifest=MANIFEST,
        load_backend=_loader(error=RuntimeError("boom")),
        now=lambda: 4242.0,
    )

    report = result.reports[0]
    assert report.identity_key() == (
        FIXTURE_MODEL, 4, 9, "baseline", "1", "torch-transformers", "unknown",
    )
    assert report.validated_at == 4242.0
    assert report.recipe.catalogue_version == "test-1"


def test_the_report_records_the_device_the_forward_actually_ran_on():
    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=_loader(_FakeBackend())
    )

    assert result.reports[0].backend.device == "cpu"
    assert result.reports[0].backend.backend_id == "torch-transformers"


def test_reports_round_trip_through_the_written_json(tmp_path):
    result = run_doctor(
        _selection(), manifest=MANIFEST, load_backend=_loader(), all_recipes=True
    )

    path = write_reports(result.reports, tmp_path / "nested" / "capability.json")
    payload = json.loads(path.read_text())

    assert [CapabilityReport.from_dict(d).recipe.recipe_id for d in payload] == [
        "baseline",
        "stateless",
    ]


def test_a_single_report_is_written_as_one_object():
    """One selected recipe writes one report — the shape NCA-003 will read."""
    import tempfile

    result = run_doctor(_selection(), manifest=MANIFEST, load_backend=_loader())
    with tempfile.TemporaryDirectory() as tmp:
        path = write_reports(result.reports, Path(tmp) / "capability.json")
        report = CapabilityReport.from_json(path.read_text())

    assert report.passed


def test_the_summary_tells_a_failing_operator_what_to_fix():
    result = run_doctor(
        _selection(),
        manifest=MANIFEST,
        load_backend=_loader(error=MissingModelDependencyError("torch is not installed")),
    )

    text = render_result(result, report_path=Path("/tmp/capability.json"))

    assert "FAIL" in text
    assert CATEGORY_MISSING_DEPENDENCY in text
    assert "torch is not installed" in text
    assert "/tmp/capability.json" in text
    assert "Traceback" not in text


def test_the_summary_names_the_shard_that_passed():
    result = run_doctor(_selection(), manifest=MANIFEST, load_backend=_loader())

    text = render_result(result)

    assert "PASS" in text
    assert FIXTURE_MODEL in text
    assert "layers 0–3" in text


# --- the CLI wiring ----------------------------------------------------------


def _run_cli(monkeypatch, argv, backend=None, error=None):
    """Drive `meshnet-node doctor` end to end with an injected backend."""
    import sys

    from meshnet_node import cli, config

    monkeypatch.setattr(
        config, "load_config", lambda *a, **k: {
            "model_hf_repo": FIXTURE_MODEL,
            "shard_start": 0,
            "shard_end": 3,
            "quantization": "auto",
        }
    )
    monkeypatch.setattr(
        doctor, "default_load_backend", _loader(backend, error=error)
    )
    monkeypatch.setattr(doctor, "load_recipe_manifest", lambda *a, **k: MANIFEST)
    monkeypatch.setattr(sys, "argv", ["meshnet-node", *argv])

    with pytest.raises(SystemExit) as exit_info:
        cli.main()
    return exit_info.value.code


def test_cli_doctor_exits_zero_and_writes_a_passing_report(monkeypatch, capsys, tmp_path):
    report = tmp_path / "capability.json"

    code = _run_cli(monkeypatch, ["doctor", "--report", str(report)], backend=_FakeBackend())

    assert code == 0
    assert capsys.readouterr().out.count("PASS") == 1
    assert CapabilityReport.from_json(report.read_text()).passed


def test_cli_doctor_exits_non_zero_and_writes_the_failed_report(monkeypatch, capsys, tmp_path):
    report = tmp_path / "capability.json"

    code = _run_cli(
        monkeypatch,
        ["doctor", "--report", str(report)],
        error=InsufficientVRAMError("insufficient VRAM to load 24 layers"),
    )

    out = capsys.readouterr().out
    assert code == 1
    assert "FAIL" in out
    assert CATEGORY_INSUFFICIENT_MEMORY in out
    assert "Traceback" not in out  # no raw traceback by default
    assert CapabilityReport.from_json(report.read_text()).status == STATUS_FAILED


def test_cli_doctor_all_recipes_is_opt_in(monkeypatch, capsys, tmp_path):
    report = tmp_path / "capability.json"

    code = _run_cli(
        monkeypatch,
        ["doctor", "--all-recipes", "--report", str(report)],
        backend=_FakeBackend(),
    )

    assert code == 0
    assert capsys.readouterr().out.count("PASS") == 2
    assert len(json.loads(report.read_text())) == 2


def test_cli_doctor_json_prints_the_capability_report(monkeypatch, capsys, tmp_path):
    code = _run_cli(
        monkeypatch,
        ["doctor", "--json", "--report", str(tmp_path / "c.json")],
        backend=_FakeBackend(),
    )

    payload = json.loads(capsys.readouterr().out)
    assert code == 0
    assert payload[0]["model"]["model_id"] == FIXTURE_MODEL


def test_cli_doctor_flags_select_what_is_validated(monkeypatch, capsys, tmp_path):
    """`doctor --shard-start/--shard-end` validates the shard startup would load."""
    report = tmp_path / "capability.json"

    code = _run_cli(
        monkeypatch,
        ["doctor", "--shard-start", "2", "--shard-end", "5", "--report", str(report)],
        backend=_FakeBackend(),
    )

    written = CapabilityReport.from_json(report.read_text())
    assert code == 0
    assert (written.shard.start, written.shard.end) == (2, 5)


# --- the real-model smoke test ----------------------------------------------

# Model identity comes from the environment; there is no default, so this test
# never smuggles a vendor-specific assumption into the suite.
DOCTOR_MODEL = os.environ.get("MESHNET_DOCTOR_MODEL")
DOCTOR_SHARD_START = int(os.environ.get("MESHNET_DOCTOR_SHARD_START", "0"))
DOCTOR_SHARD_END = os.environ.get("MESHNET_DOCTOR_SHARD_END")


@pytest.mark.integration
@pytest.mark.skipif(
    not DOCTOR_MODEL,
    reason="set MESHNET_DOCTOR_MODEL (and optionally MESHNET_DOCTOR_SHARD_START/END) to run",
)
def test_doctor_smoke_runs_a_real_forward_on_a_real_model(tmp_path):
    cfg = {
        "model_hf_repo": DOCTOR_MODEL,
        "quantization": os.environ.get("MESHNET_DOCTOR_QUANTIZATION", "auto"),
        "download_dir": os.environ.get("MESHNET_DOWNLOAD_DIR") or None,
        "shard_start": DOCTOR_SHARD_START,
        "shard_end": int(DOCTOR_SHARD_END) if DOCTOR_SHARD_END else None,
        "force_cpu": os.environ.get("MESHNET_DOCTOR_CPU") == "1",
    }
    selection = resolve_selection(cfg)

    result = run_doctor(selection)

    report = result.reports[0]
    assert result.passed, f"doctor failed: {report.diagnostics}"
    assert report.status == STATUS_PASSED
    assert report.model.model_id == DOCTOR_MODEL
    assert report.duration_ms > 0
    assert report.model.config_fingerprint.startswith("sha256:")

    path = write_reports(result.reports, tmp_path / "capability.json")
    assert CapabilityReport.from_json(path.read_text()).passed