neuron-tai/packages/node/meshnet_node/doctor.py

"""`meshnet-node doctor` — prove the selected shard actually runs.

The doctor answers one question: *would the model/shard/recipe this node is
configured to serve really execute here?* It answers it the only way that is
not a guess — by loading the selection through the production backend path and
pushing a bounded, real forward through the selected layers. Generic hardware
probing (is there a GPU, can Torch allocate a tensor) proves nothing about a
shard and is deliberately not what this reports on.

Two shapes of probe, chosen by where the shard sits, never by which model it is:

* head shard — tokenize a short prompt, embed it, run this shard's layers.
* mid/tail shard — synthesize a small hidden-state tensor in the same wire
  format peers send, and push it through `forward_bytes`. A tail shard decodes
  it, which also exercises the final norm and `lm_head`.

Everything here is model-agnostic: `model_id` is opaque, and no vendor or kernel
name is a branch. Failures are reported as a category plus an actionable hint
(never a raw traceback, unless the caller asks for one) and produce a *failed*
capability report — a failure is evidence too, and NCA-003 refuses to register
without a fresh passing one.
"""

from __future__ import annotations

import base64
import struct
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence

from .capability import (
    STATUS_FAILED,
    STATUS_PASSED,
    CapabilityReport,
    build_capability_report,
)
from .recipe_manifest import (
    DEFAULT_RECIPE_ID,
    Recipe,
    RecipeManifest,
    RecipeManifestError,
    load_recipe_manifest,
)

# The probe is deliberately tiny: enough tokens to drive every layer in the
# shard once, small enough that `doctor` costs seconds beyond the model load.
PROBE_TOKENS = 4
PROBE_PROMPT = "meshnet capability probe"

# Failure categories. These are what an operator acts on, so they name the thing
# to fix, not the exception that surfaced it.
CATEGORY_NO_MODEL = "no-model-selected"
CATEGORY_MISSING_DEPENDENCY = "missing-dependency"
CATEGORY_MODEL_UNAVAILABLE = "model-unavailable"
CATEGORY_INSUFFICIENT_MEMORY = "insufficient-memory"
CATEGORY_INVALID_SHARD = "invalid-shard"
CATEGORY_UNSUPPORTED_RECIPE = "unsupported-recipe"
CATEGORY_LOAD_FAILED = "load-failed"
CATEGORY_FORWARD_FAILED = "forward-failed"

CATEGORY_HINTS: Mapping[str, str] = {
    CATEGORY_NO_MODEL: (
        "No model is selected. Pass --model <repo-or-name>, or run `meshnet-node` "
        "once to save a config."
    ),
    CATEGORY_MISSING_DEPENDENCY: (
        "The model runtime is not installed. Install the node's model extras "
        "(torch, transformers, safetensors, accelerate, bitsandbytes)."
    ),
    CATEGORY_MODEL_UNAVAILABLE: (
        "The model files could not be read. Check the model id, --download-dir, "
        "and that the artifact is downloaded or reachable."
    ),
    CATEGORY_INSUFFICIENT_MEMORY: (
        "This shard does not fit in memory. Serve fewer layers (--shard-start / "
        "--shard-end) or use a smaller quantization (-q int8, -q nf4)."
    ),
    CATEGORY_INVALID_SHARD: (
        "The requested layer range does not exist in this model. Check "
        "--shard-start / --shard-end against the model's layer count."
    ),
    CATEGORY_UNSUPPORTED_RECIPE: (
        "The recipe asks for an execution setting this backend cannot apply. "
        "Select a different recipe with --recipe."
    ),
    CATEGORY_LOAD_FAILED: (
        "The shard could not be loaded. Re-run with --debug for the full traceback."
    ),
    CATEGORY_FORWARD_FAILED: (
        "The shard loaded but could not execute a forward pass. This node cannot "
        "serve this model/shard; re-run with --debug for the full traceback."
    ),
}


class DoctorError(RuntimeError):
    """A validation failure with an operator-facing category and hint."""

    def __init__(self, category: str, message: str) -> None:
        super().__init__(message)
        self.category = category

    @property
    def hint(self) -> str:
        return CATEGORY_HINTS.get(self.category, "")


@dataclass(frozen=True)
class DoctorSelection:
    """The one model/shard/config combination startup would load."""

    model_id: str
    shard_start: int
    shard_end: int
    quantization: str = "auto"
    cache_dir: Path | None = None
    force_cpu: bool = False

    @property
    def shard_label(self) -> str:
        return f"layers {self.shard_start}–{self.shard_end}"


@dataclass(frozen=True)
class RecipeResult:
    """One recipe's validation outcome, with the report it produced."""

    recipe: Recipe
    report: CapabilityReport
    category: str | None = None
    error: BaseException | None = None

    @property
    def passed(self) -> bool:
        return self.report.passed

    @property
    def hint(self) -> str:
        return CATEGORY_HINTS.get(self.category or "", "")


@dataclass(frozen=True)
class DoctorResult:
    """The outcome of a doctor run over one or more recipes."""

    selection: DoctorSelection
    results: tuple[RecipeResult, ...] = ()

    @property
    def passed(self) -> bool:
        return bool(self.results) and all(r.passed for r in self.results)

    @property
    def reports(self) -> tuple[CapabilityReport, ...]:
        return tuple(r.report for r in self.results)

    @property
    def exit_code(self) -> int:
        return 0 if self.passed else 1


# --- selection: the same resolution startup performs ------------------------


def resolve_selection(
    cfg: Mapping[str, Any],
    *,
    detect_layers: Callable[[str, Path | None], int | None] | None = None,
) -> DoctorSelection:
    """Resolve config + flags into the selection startup would load.

    This mirrors `startup.run_startup`: the same model id, the same
    `bf16`→`bfloat16` quantization normalization, and the same shard default of
    the whole model when no range is pinned. It deliberately does *not* ask the
    tracker for a gap assignment — the doctor is an offline check of what this
    node can run, and startup re-validates whatever range it is finally given.
    """
    model_id = _selected_model_id(cfg)
    if not model_id:
        raise DoctorError(
            CATEGORY_NO_MODEL, "no model is selected in config or flags"
        )

    cache_dir = Path(cfg["download_dir"]) if cfg.get("download_dir") else None
    quantization = str(cfg.get("quantization") or "auto").replace("bf16", "bfloat16")

    shard_start = cfg.get("shard_start")
    shard_end = cfg.get("shard_end")
    if shard_start is None or shard_end is None:
        detect = detect_layers or _detect_layers
        total = detect(model_id, cache_dir)
        if total is None:
            raise DoctorError(
                CATEGORY_MODEL_UNAVAILABLE,
                f"could not read the layer count from the {model_id} config; "
                "pass --shard-start and --shard-end explicitly",
            )
        shard_start = 0 if shard_start is None else shard_start
        shard_end = total - 1 if shard_end is None else shard_end

    if shard_start < 0 or shard_end < shard_start:
        raise DoctorError(
            CATEGORY_INVALID_SHARD,
            f"invalid shard range {shard_start}–{shard_end}: start must be "
            "non-negative and not greater than end",
        )

    return DoctorSelection(
        model_id=model_id,
        shard_start=int(shard_start),
        shard_end=int(shard_end),
        quantization=quantization,
        cache_dir=cache_dir,
        force_cpu=bool(cfg.get("force_cpu", False)),
    )


def _selected_model_id(cfg: Mapping[str, Any]) -> str | None:
    """The HF repo startup would load, resolving a catalog alias if needed."""
    hf_repo = str(cfg.get("model_hf_repo") or "").strip()
    if hf_repo:
        return hf_repo
    name = str(cfg.get("model_name") or "").strip()
    if not name:
        return None
    from .model_catalog import resolve_model_alias

    preset = resolve_model_alias(name)
    if preset is not None and preset.hf_repo:
        return preset.hf_repo
    return name if "/" in name else None


def _detect_layers(model_id: str, cache_dir: Path | None) -> int | None:
    from .startup import _detect_num_layers

    return _detect_num_layers(model_id, cache_dir=cache_dir)


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


@dataclass(frozen=True)
class ProbeInput:
    """A synthetic hidden-state payload in the same wire format peers send."""

    body: bytes
    shape: list[int]
    attention_mask_header: str | None
    position_ids_header: str | None


def _int64_header(rows: Sequence[Sequence[int]]) -> str:
    """Encode an int64 tensor as `shape:base64`, matching the backend's format."""
    flat = [int(v) for row in rows for v in row]
    raw = struct.pack(f"<{len(flat)}q", *flat)
    shape = f"{len(rows)},{len(rows[0])}" if rows else "0"
    return f"{shape}:{base64.b64encode(raw).decode('ascii')}"


def build_probe_input(hidden_size: int, tokens: int = PROBE_TOKENS) -> ProbeInput:
    """Build a bounded mid-shard probe: `tokens` positions of bfloat16 zeros.

    Zeros are a legitimate hidden state; what is being proven is that the
    layers execute on this device, not that the output means anything. The
    payload is built with plain bytes so callers need no Torch import.
    """
    if hidden_size <= 0:
        raise DoctorError(
            CATEGORY_FORWARD_FAILED,
            "the backend reports no hidden size, so no probe tensor can be built",
        )
    ones = [[1] * tokens]
    positions = [list(range(tokens))]
    return ProbeInput(
        body=b"\x00" * (tokens * hidden_size * 2),  # bfloat16 == 2 bytes
        shape=[1, tokens, hidden_size],
        attention_mask_header=_int64_header(ones),
        position_ids_header=_int64_header(positions),
    )


def probe_forward(backend: Any, *, tokens: int = PROBE_TOKENS) -> dict:
    """Run one bounded real forward through the shard `backend` holds.

    Returns a small detail dict for the human summary. Raises `DoctorError`
    (category `forward-failed`) if the shard cannot execute or returns nothing.
    """
    is_head = bool(getattr(backend, "is_head", False))
    is_tail = bool(getattr(backend, "is_tail", False))

    try:
        if is_head:
            output = backend.encode_prompt(PROBE_PROMPT)
            kind = "prompt"
            if is_tail:
                # A head+tail shard owns the lm_head too. Re-entering above the
                # last layer runs no layer again — it only decodes — so the whole
                # selected shard is covered without a second forward through it.
                output = backend.forward_bytes(
                    output.body,
                    output.shape,
                    output.attention_mask_header,
                    output.position_ids_header,
                    start_layer=int(getattr(backend, "shard_end", 0)) + 1,
                )
                kind = "prompt+decode"
        else:
            probe = build_probe_input(int(getattr(backend, "hidden_size", 0) or 0))
            output = backend.forward_bytes(
                probe.body,
                probe.shape,
                probe.attention_mask_header,
                probe.position_ids_header,
                start_layer=getattr(backend, "shard_start", None),
            )
            kind = "hidden-states"
    except DoctorError:
        raise
    except Exception as exc:
        raise DoctorError(CATEGORY_FORWARD_FAILED, _describe(exc)) from exc

    return {"probe": kind, "tokens": tokens, **_describe_output(output)}


def _describe_output(output: Any) -> dict:
    """Validate the forward produced real output, and summarize it."""
    if output is None:
        raise DoctorError(
            CATEGORY_FORWARD_FAILED, "the shard forward returned no output"
        )

    token_id = getattr(output, "token_id", None)
    if token_id is not None:  # tail shard: decoded a token
        return {"output": "token", "token_id": int(token_id)}

    body = getattr(output, "body", None)
    shape = list(getattr(output, "shape", []) or [])
    if not body or not shape:
        raise DoctorError(
            CATEGORY_FORWARD_FAILED,
            "the shard forward returned an empty hidden-state payload",
        )
    return {"output": "hidden-states", "shape": shape}


# --- running the doctor -----------------------------------------------------


def default_load_backend(
    selection: DoctorSelection,
    recipe: Recipe,
) -> Any:
    """Load the shard through the exact path startup uses."""
    from .torch_server import _load_backend

    return _load_backend(
        selection.model_id,
        selection.shard_start,
        selection.shard_end,
        selection.quantization,
        selection.cache_dir,
        force_cpu=selection.force_cpu,
        recipe_params=recipe.params,
    )


def select_recipes(
    manifest: RecipeManifest,
    *,
    recipe_id: str | None = None,
    all_recipes: bool = False,
) -> tuple[Recipe, ...]:
    """The recipes to validate: the selected one, or every one on request.

    `--all-recipes` is the only way to pay for validating recipes the node was
    not asked to serve; ordinary onboarding validates exactly one.
    """
    if all_recipes:
        if recipe_id is not None:
            raise DoctorError(
                CATEGORY_UNSUPPORTED_RECIPE,
                "--recipe and --all-recipes are mutually exclusive",
            )
        return manifest.recipes
    try:
        return (manifest.require(recipe_id or DEFAULT_RECIPE_ID),)
    except RecipeManifestError as exc:
        raise DoctorError(CATEGORY_UNSUPPORTED_RECIPE, str(exc)) from exc


def run_doctor(
    selection: DoctorSelection,
    *,
    manifest: RecipeManifest | None = None,
    recipe_id: str | None = None,
    all_recipes: bool = False,
    load_backend: Callable[[DoctorSelection, Recipe], Any] | None = None,
    now: Callable[[], float] | None = None,
) -> DoctorResult:
    """Validate the selection, one bounded real forward per recipe.

    Never raises for a validation failure: every recipe yields a report, passed
    or failed, so the caller can write the evidence out either way. `DoctorError`
    only escapes for input the caller got wrong (an unknown recipe id).
    """
    manifest = manifest or load_recipe_manifest()
    recipes = select_recipes(manifest, recipe_id=recipe_id, all_recipes=all_recipes)
    clock = now or time.time
    load = load_backend or default_load_backend

    results = [
        _validate_recipe(selection, recipe, manifest, load, clock)
        for recipe in recipes
    ]
    return DoctorResult(selection=selection, results=tuple(results))


def validate_loaded_backend(
    backend: Any,
    selection: DoctorSelection,
    recipe: Recipe,
    manifest: RecipeManifest,
    *,
    now: Callable[[], float] | None = None,
) -> RecipeResult:
    """Validate a shard that is already loaded, without loading it a second time.

    Startup calls this on the very backend that would serve traffic, so the proof
    it produces is about that object, not about a re-load that might have landed
    on a different device.
    """
    return _validate_recipe(
        selection, recipe, manifest, lambda *_: backend, now or time.time
    )


def _validate_recipe(
    selection: DoctorSelection,
    recipe: Recipe,
    manifest: RecipeManifest,
    load_backend: Callable[[DoctorSelection, Recipe], Any],
    clock: Callable[[], float],
) -> RecipeResult:
    started = time.monotonic()
    backend: Any = None
    category: str | None = None
    error: BaseException | None = None
    diagnostics: list[str] = []
    detail: dict = {}

    try:
        backend = load_backend(selection, recipe)
        detail = probe_forward(backend)
    except DoctorError as exc:
        category, error = exc.category, exc
        diagnostics = [str(exc), exc.hint]
    except Exception as exc:  # noqa: BLE001 — every failure becomes a report
        category = classify_failure(exc)
        error = exc
        diagnostics = [_describe(exc), CATEGORY_HINTS.get(category, "")]
    duration_ms = int((time.monotonic() - started) * 1000)

    device = _backend_device(backend, selection)
    report = build_capability_report(
        model_id=selection.model_id,
        shard_start=selection.shard_start,
        shard_end=selection.shard_end,
        recipe_id=recipe.id,
        recipe_version=recipe.version,
        catalogue_version=manifest.catalogue_version,
        backend_id=recipe.backend_id,
        device=device,
        device_name=_backend_device_name(device),
        quantization=selection.quantization,
        runtime=_runtime_versions(),
        model_config=_model_config(backend),
        status=STATUS_FAILED if category else STATUS_PASSED,
        duration_ms=duration_ms,
        diagnostics=[d for d in diagnostics if d] or None,
        validated_at=clock(),
    )
    if category:
        return RecipeResult(
            recipe=recipe, report=report, category=category, error=error
        )
    return RecipeResult(recipe=recipe, report=report)


def classify_failure(exc: BaseException) -> str:
    """Map a backend exception to an operator-facing category.

    Matches on the backend's own error types, never on model or vendor names.
    """
    from .model_backend import (
        InsufficientVRAMError,
        MissingModelDependencyError,
        PartialModelLoadUnsupported,
        UnsupportedRecipeParam,
    )

    if isinstance(exc, MissingModelDependencyError):
        return CATEGORY_MISSING_DEPENDENCY
    if isinstance(exc, InsufficientVRAMError):
        return CATEGORY_INSUFFICIENT_MEMORY
    if isinstance(exc, UnsupportedRecipeParam):
        return CATEGORY_UNSUPPORTED_RECIPE
    if isinstance(exc, PartialModelLoadUnsupported):
        return CATEGORY_LOAD_FAILED
    if isinstance(exc, ValueError):  # shard range vs. the model's real layers
        return CATEGORY_INVALID_SHARD
    if isinstance(exc, (FileNotFoundError, OSError)):
        return CATEGORY_MODEL_UNAVAILABLE
    return CATEGORY_LOAD_FAILED


def _describe(exc: BaseException) -> str:
    """A one-line, traceback-free description. Sanitized by the report."""
    text = str(exc).strip()
    return f"{type(exc).__name__}: {text}" if text else type(exc).__name__


def _backend_device(backend: Any, selection: DoctorSelection) -> str:
    device = getattr(backend, "device", None)
    if device is None:
        # The load failed, so no device was chosen — record the one that was asked for.
        return "cpu" if selection.force_cpu else "unknown"
    return str(getattr(device, "type", device))


def _backend_device_name(device: str) -> str | None:
    """The accelerator's name, when the shard actually landed on one."""
    if device != "cuda":
        return None
    from .hardware import detect_hardware

    try:
        return detect_hardware().get("gpu_name") or None
    except Exception:
        return None


def _model_config(backend: Any) -> Any:
    """The loaded model's config, for the report's fingerprint."""
    config = getattr(getattr(backend, "model", None), "config", None)
    to_dict = getattr(config, "to_dict", None)
    if not callable(to_dict):
        return None
    try:
        return to_dict()
    except Exception:
        return None


def _runtime_versions() -> dict[str, str]:
    """Versions of the stack that ran the forward — opaque labels, never branches."""
    versions: dict[str, str] = {}
    for name in ("torch", "transformers"):
        try:
            module = __import__(name)
        except Exception:
            continue
        version = getattr(module, "__version__", None)
        if version:
            versions[name] = str(version)
    return versions


# --- output -----------------------------------------------------------------

DEFAULT_REPORT_FILENAME = "capability.json"


def default_report_path() -> Path:
    from .config import config_path

    return config_path().parent / DEFAULT_REPORT_FILENAME


def write_reports(reports: Sequence[CapabilityReport], path: Path) -> Path:
    """Write the capability report(s) as JSON. A failed run writes too."""
    import json

    path.parent.mkdir(parents=True, exist_ok=True)
    if len(reports) == 1:
        path.write_text(reports[0].to_json(indent=2) + "\n", encoding="utf-8")
    else:
        payload = [r.to_dict() for r in reports]
        path.write_text(
            json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
    return path


def render_result(result: DoctorResult, *, report_path: Path | None = None) -> str:
    """The human summary: what was validated, what to do if it failed."""
    selection = result.selection
    lines = [
        "meshnet-node doctor",
        f"  Model:        {selection.model_id}",
        f"  Shard:        {selection.shard_label}",
        f"  Quantization: {selection.quantization}",
        "",
    ]

    for item in result.results:
        mark = "PASS" if item.passed else "FAIL"
        device = item.report.backend.device
        lines.append(
            f"  [{mark}] recipe {item.recipe.id} (v{item.recipe.version}) "
            f"on {device} — {item.report.duration_ms} ms"
        )
        if not item.passed:
            for diagnostic in item.report.diagnostics:
                lines.append(f"         {diagnostic}")

    lines.append("")
    if result.passed:
        count = len(result.results)
        what = "recipe" if count == 1 else "recipes"
        lines.append(
            f"  OK — the selected shard ran a real forward for {count} {what}."
        )
    else:
        failed = [r for r in result.results if not r.passed]
        categories = ", ".join(dict.fromkeys(r.category or "unknown" for r in failed))
        lines.append(f"  FAILED — {categories}. This node cannot serve this shard.")

    if report_path is not None:
        lines.append(f"  Capability report: {report_path}")
    return "\n".join(lines)