403 lines
14 KiB
Python
403 lines
14 KiB
Python
"""Fail-closed admission: no routable registration without a fresh matching proof.
|
||
|
||
This module does not *produce* proof — `doctor` does that, by pushing a bounded
|
||
real forward through the selected shard (NCA-002). This module *decides whether a
|
||
proof covers what is about to be advertised*, and startup calls it immediately
|
||
before it registers with the tracker.
|
||
|
||
A capability report proves one combination: model artifact, shard range, recipe,
|
||
backend and device. Reusing it for anything else is the exact hole this closes —
|
||
a report that failed, aged out, or describes a different model, shard, recipe or
|
||
device is rejected here, and the node exits without ever registering an endpoint.
|
||
|
||
Nothing in here branches on a model, vendor or kernel name: identity fields are
|
||
opaque labels that are compared, never interpreted.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import Any, Callable
|
||
|
||
from . import __version__ as _PACKAGE_VERSION
|
||
from .capability import CapabilityReport, config_fingerprint
|
||
from .doctor import DoctorSelection
|
||
from .recipe_manifest import Recipe, RecipeManifest
|
||
from .runtime_recipe import (
|
||
build_artifact_identity,
|
||
build_runtime_recipe_identity,
|
||
compatibility_fingerprint,
|
||
fingerprint_payload,
|
||
)
|
||
from .gguf_ownership import authoritative_dense_llama_ownership
|
||
|
||
# How long a passing report stays usable. Startup normally validates in-process
|
||
# (age ≈ 0); this bounds how far a report written by an earlier `doctor` run can
|
||
# be carried forward, after which the hardware, drivers or weights may have moved.
|
||
DEFAULT_MAX_REPORT_AGE_SECONDS = 900.0
|
||
|
||
# A report timestamped this far in the future is not fresh, it is wrong.
|
||
_MAX_CLOCK_SKEW_SECONDS = 60.0
|
||
|
||
REASON_NO_REPORT = "no-report"
|
||
REASON_NOT_PASSED = "not-passed"
|
||
REASON_STALE = "stale"
|
||
REASON_MODEL_MISMATCH = "model-mismatch"
|
||
REASON_SHARD_MISMATCH = "shard-mismatch"
|
||
REASON_RECIPE_MISMATCH = "recipe-mismatch"
|
||
REASON_BACKEND_MISMATCH = "backend-mismatch"
|
||
REASON_COMPATIBILITY_MISMATCH = "compatibility-mismatch"
|
||
|
||
|
||
class CapabilityAdmissionError(RuntimeError):
|
||
"""This node may not advertise the selection: the proof does not cover it."""
|
||
|
||
def __init__(self, reason: str, message: str) -> None:
|
||
super().__init__(message)
|
||
self.reason = reason
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CapabilityContext:
|
||
"""What is about to be advertised, and the loaded backend that would serve it."""
|
||
|
||
backend: Any
|
||
selection: DoctorSelection
|
||
recipe: Recipe
|
||
manifest: RecipeManifest
|
||
device: str
|
||
|
||
|
||
# A validator turns the context into the report the gate then judges. Production
|
||
# uses `probe_capability`; tests pass an explicit test-safe one (see
|
||
# `meshnet_node.testing`) rather than switching this module into a lenient mode.
|
||
CapabilityValidator = Callable[[CapabilityContext], CapabilityReport]
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class AdmissionRequirement:
|
||
"""The one capability a report must prove for this node to register."""
|
||
|
||
model_id: str
|
||
shard_start: int
|
||
shard_end: int
|
||
recipe_id: str
|
||
recipe_version: str
|
||
backend_id: str
|
||
device: str
|
||
compatibility_fingerprint: str
|
||
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS
|
||
|
||
@classmethod
|
||
def for_context(
|
||
cls,
|
||
context: CapabilityContext,
|
||
*,
|
||
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS,
|
||
) -> AdmissionRequirement:
|
||
return cls(
|
||
model_id=context.selection.model_id,
|
||
shard_start=context.selection.shard_start,
|
||
shard_end=context.selection.shard_end,
|
||
recipe_id=context.recipe.id,
|
||
recipe_version=context.recipe.version,
|
||
backend_id=context.recipe.backend_id,
|
||
device=context.device,
|
||
compatibility_fingerprint=_compatibility_fingerprint_for_context(
|
||
context
|
||
),
|
||
max_age_seconds=max_age_seconds,
|
||
)
|
||
|
||
@property
|
||
def shard_label(self) -> str:
|
||
return f"layers {self.shard_start}–{self.shard_end}"
|
||
|
||
|
||
def admit(
|
||
requirement: AdmissionRequirement,
|
||
report: CapabilityReport | None,
|
||
*,
|
||
now: float | None = None,
|
||
) -> CapabilityReport:
|
||
"""Return `report` if it admits `requirement`; otherwise refuse to register.
|
||
|
||
Checks run selection-first, so the operator is told the report is about the
|
||
wrong thing before being told it is old.
|
||
"""
|
||
if report is None:
|
||
raise CapabilityAdmissionError(
|
||
REASON_NO_REPORT,
|
||
f"no capability report for {requirement.model_id} "
|
||
f"{requirement.shard_label}: this node has not proven it can serve it",
|
||
)
|
||
|
||
if report.model.model_id != requirement.model_id:
|
||
raise _mismatch(
|
||
REASON_MODEL_MISMATCH,
|
||
requirement,
|
||
"model",
|
||
report.model.model_id,
|
||
requirement.model_id,
|
||
)
|
||
|
||
if (report.shard.start, report.shard.end) != (
|
||
requirement.shard_start,
|
||
requirement.shard_end,
|
||
):
|
||
raise _mismatch(
|
||
REASON_SHARD_MISMATCH,
|
||
requirement,
|
||
"shard",
|
||
f"layers {report.shard.start}–{report.shard.end}",
|
||
requirement.shard_label,
|
||
)
|
||
|
||
if (report.recipe.recipe_id, report.recipe.recipe_version) != (
|
||
requirement.recipe_id,
|
||
requirement.recipe_version,
|
||
):
|
||
raise _mismatch(
|
||
REASON_RECIPE_MISMATCH,
|
||
requirement,
|
||
"recipe",
|
||
f"{report.recipe.recipe_id} (v{report.recipe.recipe_version})",
|
||
f"{requirement.recipe_id} (v{requirement.recipe_version})",
|
||
)
|
||
|
||
if (report.backend.backend_id, report.backend.device) != (
|
||
requirement.backend_id,
|
||
requirement.device,
|
||
):
|
||
raise _mismatch(
|
||
REASON_BACKEND_MISMATCH,
|
||
requirement,
|
||
"backend",
|
||
f"{report.backend.backend_id} on {report.backend.device}",
|
||
f"{requirement.backend_id} on {requirement.device}",
|
||
)
|
||
|
||
if report.compatibility_fingerprint != requirement.compatibility_fingerprint:
|
||
raise CapabilityAdmissionError(
|
||
REASON_COMPATIBILITY_MISMATCH,
|
||
f"capability proof fingerprint {report.compatibility_fingerprint!r} "
|
||
f"does not match the expected compatibility fingerprint for "
|
||
f"{requirement.model_id} {requirement.shard_label}; the artifact, "
|
||
f"tokenizer, architecture, boundary schema, activation recipe or "
|
||
f"cache layout differs",
|
||
)
|
||
|
||
if not report.passed:
|
||
raise CapabilityAdmissionError(
|
||
REASON_NOT_PASSED,
|
||
f"capability validation {report.status} for {requirement.model_id} "
|
||
f"{requirement.shard_label} with recipe {requirement.recipe_id}"
|
||
+ _diagnostics_suffix(report),
|
||
)
|
||
|
||
now = time.time() if now is None else now
|
||
age = now - report.validated_at
|
||
if age > requirement.max_age_seconds:
|
||
raise CapabilityAdmissionError(
|
||
REASON_STALE,
|
||
f"capability report for {requirement.model_id} {requirement.shard_label} "
|
||
f"is {age / 60:.0f} min old (limit "
|
||
f"{requirement.max_age_seconds / 60:.0f} min); re-run `meshnet-node doctor`",
|
||
)
|
||
if age < -_MAX_CLOCK_SKEW_SECONDS:
|
||
raise CapabilityAdmissionError(
|
||
REASON_STALE,
|
||
f"capability report for {requirement.model_id} {requirement.shard_label} "
|
||
f"is timestamped {-age:.0f}s in the future; check this host's clock",
|
||
)
|
||
|
||
return report
|
||
|
||
|
||
def _mismatch(
|
||
reason: str,
|
||
requirement: AdmissionRequirement,
|
||
field_name: str,
|
||
reported: str,
|
||
required: str,
|
||
) -> CapabilityAdmissionError:
|
||
return CapabilityAdmissionError(
|
||
reason,
|
||
f"capability report proves a different {field_name}: it validated "
|
||
f"{reported}, but this node would serve {required}. A report is only "
|
||
"proof for the exact combination it ran.",
|
||
)
|
||
|
||
|
||
def _diagnostics_suffix(report: CapabilityReport) -> str:
|
||
if not report.diagnostics:
|
||
return ""
|
||
return " — " + " ".join(report.diagnostics)
|
||
|
||
|
||
def probe_capability(context: CapabilityContext) -> CapabilityReport:
|
||
"""Production validator: one bounded real forward through the loaded shard."""
|
||
from .doctor import validate_loaded_backend
|
||
|
||
return validate_loaded_backend(
|
||
context.backend,
|
||
context.selection,
|
||
context.recipe,
|
||
context.manifest,
|
||
).report
|
||
|
||
|
||
def _compatibility_fingerprint_for_context(context: CapabilityContext) -> str:
|
||
backend = context.backend
|
||
selection = context.selection
|
||
recipe = context.recipe
|
||
model_config = getattr(getattr(backend, "model", None), "config", None)
|
||
model_config_payload = (
|
||
model_config.to_dict() if hasattr(model_config, "to_dict") else model_config
|
||
)
|
||
runtime_versions = _runtime_versions()
|
||
runtime_version = _PACKAGE_VERSION
|
||
ownership = authoritative_dense_llama_ownership(backend, selection)
|
||
artifact = build_artifact_identity(
|
||
model_id=selection.model_id,
|
||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||
model_config=model_config_payload,
|
||
shard_start=ownership.start_layer,
|
||
shard_end=ownership.end_layer,
|
||
)
|
||
runtime_recipe = build_runtime_recipe_identity(
|
||
model_id=selection.model_id,
|
||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||
model_config=model_config_payload,
|
||
recipe_params=recipe.params,
|
||
weight_quantization=selection.quantization,
|
||
backend_id=recipe.backend_id,
|
||
runtime_version=runtime_version,
|
||
activation_dtype="bfloat16",
|
||
compute_dtype=_backend_compute_dtype(backend),
|
||
kv_dtype=_backend_kv_dtype(backend),
|
||
kv_layout=_backend_kv_layout(backend),
|
||
tokenizer_revision=_backend_tokenizer_revision(backend, selection),
|
||
architecture_adapter=_backend_architecture_adapter(backend, recipe.backend_id),
|
||
boundary_schema_version=1,
|
||
cache_layout=_backend_cache_layout(backend, recipe.params),
|
||
)
|
||
return compatibility_fingerprint(
|
||
fingerprint_payload(
|
||
model={
|
||
"model_id": selection.model_id,
|
||
"revision": getattr(getattr(backend, "model", None), "revision", None),
|
||
"config_fingerprint": config_fingerprint(model_config_payload),
|
||
},
|
||
shard={
|
||
"start": ownership.start_layer,
|
||
"end": ownership.end_layer,
|
||
"owns_embedding": ownership.owns_embedding,
|
||
"owns_final_head": ownership.owns_final_head,
|
||
},
|
||
recipe={
|
||
"recipe_id": recipe.id,
|
||
"recipe_version": recipe.version,
|
||
"catalogue_version": context.manifest.catalogue_version,
|
||
},
|
||
backend={
|
||
"backend_id": recipe.backend_id,
|
||
"device": context.device,
|
||
"device_name": _backend_device_name(context.device),
|
||
"quantization": selection.quantization,
|
||
"runtime": runtime_versions,
|
||
},
|
||
artifact=artifact.to_dict(),
|
||
runtime_recipe=runtime_recipe.to_dict(),
|
||
)
|
||
)
|
||
|
||
|
||
def _runtime_versions() -> dict[str, str]:
|
||
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
|
||
|
||
|
||
def _backend_compute_dtype(backend: Any) -> str:
|
||
config = getattr(getattr(backend, "model", None), "config", None)
|
||
for candidate in (config, getattr(config, "text_config", None)):
|
||
if candidate is None:
|
||
continue
|
||
for attr in ("dtype", "torch_dtype"):
|
||
value = getattr(candidate, attr, None)
|
||
if value is None:
|
||
continue
|
||
return str(value).removeprefix("torch.")
|
||
return "bfloat16"
|
||
|
||
|
||
def _backend_kv_dtype(backend: Any) -> str:
|
||
return _backend_compute_dtype(backend)
|
||
|
||
|
||
def _backend_kv_layout(backend: Any) -> str:
|
||
return "session-cache" if getattr(backend, "supports_kv_cache", False) else "stateless"
|
||
|
||
|
||
def _backend_tokenizer_revision(backend: Any, selection: DoctorSelection) -> str:
|
||
model = getattr(backend, "model", None)
|
||
revision = getattr(model, "revision", None)
|
||
if isinstance(revision, str) and revision.strip():
|
||
return revision
|
||
tokenizer = getattr(backend, "tokenizer", None)
|
||
for attr in ("revision", "model_id"):
|
||
value = getattr(tokenizer, attr, None)
|
||
if isinstance(value, str) and value.strip():
|
||
return value
|
||
return selection.model_id
|
||
|
||
|
||
def _backend_architecture_adapter(backend: Any, default: str) -> str:
|
||
config = getattr(getattr(backend, "model", None), "config", None)
|
||
for candidate in (config, getattr(config, "text_config", None)):
|
||
if candidate is None:
|
||
continue
|
||
for attr in ("architecture_adapter", "model_type"):
|
||
value = getattr(candidate, attr, None)
|
||
if isinstance(value, str) and value.strip():
|
||
return value
|
||
architectures = getattr(candidate, "architectures", None)
|
||
if isinstance(architectures, (list, tuple)) and architectures:
|
||
first = architectures[0]
|
||
if isinstance(first, str) and first.strip():
|
||
return first
|
||
return default
|
||
|
||
|
||
def _backend_device_name(device: str) -> str | None:
|
||
if device != "cuda":
|
||
return None
|
||
from .hardware import detect_hardware
|
||
|
||
try:
|
||
return detect_hardware().get("gpu_name") or None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _backend_cache_layout(backend: Any, recipe_params: dict[str, Any] | None) -> str:
|
||
if getattr(backend, "supports_kv_cache", False) is False:
|
||
return "stateless"
|
||
if recipe_params is None:
|
||
return "local-hot-kv"
|
||
if recipe_params.get("use_cache") is False:
|
||
return "stateless"
|
||
value = recipe_params.get("cache_layout")
|
||
if isinstance(value, str) and value.strip():
|
||
return value
|
||
return "local-hot-kv"
|