Files
neuron-tai/tests/test_node_admission.py
2026-07-12 11:17:03 +03:00

397 lines
14 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""NCA-003: startup fails closed — no registration without a fresh matching proof.
Two layers are covered here:
* the gate itself (`meshnet_node.admission.admit`) — which reports admit a
selection, and which are refused as failed, stale, or about something else;
* `run_startup` — that a refused report means the tracker is never called, and
that an admitted one travels with the registration payload.
Torch is a stub in the dev venv, so the backend is faked by duck-typing
`TorchModelShard` (see `_FakeBackend`); the *production* validator still runs a
real `doctor` forward against it, so the fail-closed path is exercised without a
bypass. Tests that cannot supply an executable backend pass the explicit
test-safe validator from `meshnet_node.testing`.
"""
import base64
import struct
import time
import pytest
from meshnet_node.admission import (
REASON_BACKEND_MISMATCH,
REASON_MODEL_MISMATCH,
REASON_NO_REPORT,
REASON_NOT_PASSED,
REASON_RECIPE_MISMATCH,
REASON_SHARD_MISMATCH,
REASON_STALE,
AdmissionRequirement,
CapabilityAdmissionError,
CapabilityContext,
admit,
probe_capability,
)
from meshnet_node.capability import STATUS_FAILED, STATUS_SKIPPED
from meshnet_node.doctor import DoctorSelection
from meshnet_node.recipe_manifest import DEFAULT_RECIPE_ID, load_recipe_manifest
from meshnet_node.startup import run_startup
from meshnet_node.testing import capability_report_for, capability_stub
MODEL = "acme/opaque-model-7b"
class _FakeDevice:
def __init__(self, type_: str = "cpu"):
self.type = type_
class _FakeOutput:
def __init__(self, hidden_size: int, tokens: int = 4):
self.body = b"\x00" * (tokens * hidden_size * 2)
self.shape = [1, tokens, hidden_size]
self.attention_mask_header = _int64_header([[1] * tokens])
self.position_ids_header = _int64_header([list(range(tokens))])
def _int64_header(rows):
flat = [int(v) for row in rows for v in row]
raw = struct.pack(f"<{len(flat)}q", *flat)
return f"{len(rows)},{len(rows[0])}:{base64.b64encode(raw).decode('ascii')}"
class _FakeBackend:
"""Duck-types the parts of `TorchModelShard` the doctor probe touches."""
total_layers = 24
hidden_size = 8
def __init__(self, *, shard_start=0, shard_end=23, device="cpu", forward_error=None):
self.shard_start = shard_start
self.shard_end = shard_end
self.is_head = shard_start == 0
self.is_tail = shard_end == self.total_layers - 1
self.device = _FakeDevice(device)
self.model_id = MODEL
self._forward_error = forward_error
def encode_prompt(self, _prompt):
if self._forward_error:
raise self._forward_error
return _FakeOutput(self.hidden_size)
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
if self._forward_error:
raise self._forward_error
return _FakeOutput(self.hidden_size)
def _context(backend=None, *, model_id=MODEL, shard_start=0, shard_end=23, device="cpu"):
manifest = load_recipe_manifest()
return CapabilityContext(
backend=backend,
selection=DoctorSelection(
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization="bfloat16",
),
recipe=manifest.require(DEFAULT_RECIPE_ID),
manifest=manifest,
device=device,
)
# ---------------------------------------------------------------------------
# The gate: which reports admit a selection
# ---------------------------------------------------------------------------
def test_a_fresh_matching_passing_report_admits_the_selection():
"The proof covers exactly what is about to be advertised, so the node may register.\n\nTags: node, admission"
ctx = _context()
report = capability_report_for(ctx)
assert admit(AdmissionRequirement.for_context(ctx), report) is report
def test_a_missing_report_is_refused():
"No proof at all is the default state, and it must not register.\n\nTags: node, admission"
ctx = _context()
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), None)
assert exc.value.reason == REASON_NO_REPORT
@pytest.mark.parametrize("status", [STATUS_FAILED, STATUS_SKIPPED])
def test_a_report_that_did_not_pass_is_refused(status):
"A failed or skipped validation is evidence against admission, not for it.\n\nTags: node, admission"
ctx = _context()
report = capability_report_for(
ctx, status=status, diagnostics=["the shard forward returned no output"]
)
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_NOT_PASSED
assert "the shard forward returned no output" in str(exc.value)
def test_a_passing_report_for_another_model_is_refused():
"A proof about one model says nothing about another — no reuse across models.\n\nTags: node, admission"
ctx = _context()
report = capability_report_for(ctx, model_id="other/model-1b")
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_MODEL_MISMATCH
assert "other/model-1b" in str(exc.value)
def test_a_passing_report_for_another_shard_range_is_refused():
"Layers 011 running is no proof that layers 1223 fit.\n\nTags: node, admission"
ctx = _context(shard_start=12, shard_end=23)
report = capability_report_for(ctx, shard_start=0, shard_end=11)
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_SHARD_MISMATCH
def test_a_passing_report_for_another_recipe_version_is_refused():
"A recipe's execution params changed, so its old proof no longer applies.\n\nTags: node, admission"
ctx = _context()
report = capability_report_for(ctx, recipe_version="0")
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_RECIPE_MISMATCH
def test_a_passing_report_from_another_backend_or_device_is_refused():
"A CUDA proof does not admit a node that would serve the shard on CPU.\n\nTags: node, admission"
ctx = _context(device="cpu")
report = capability_report_for(ctx, device="cuda")
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_BACKEND_MISMATCH
other_backend = capability_report_for(ctx, backend_id="some-other-runtime")
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), other_backend)
assert exc.value.reason == REASON_BACKEND_MISMATCH
def test_a_report_older_than_the_freshness_window_is_refused():
"Hardware, drivers and weights move; an old proof is not a current one.\n\nTags: node, admission"
ctx = _context()
requirement = AdmissionRequirement.for_context(ctx, max_age_seconds=900)
report = capability_report_for(ctx, age_seconds=901)
with pytest.raises(CapabilityAdmissionError) as exc:
admit(requirement, report)
assert exc.value.reason == REASON_STALE
still_fresh = capability_report_for(ctx, age_seconds=899)
assert admit(requirement, still_fresh) is still_fresh
def test_a_future_dated_report_is_refused():
"A report from the future is a broken clock, not a fresh proof.\n\nTags: node, admission"
ctx = _context()
report = capability_report_for(ctx, validated_at=time.time() + 3600)
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_STALE
# ---------------------------------------------------------------------------
# The production validator: a real forward through the loaded backend
# ---------------------------------------------------------------------------
def test_the_production_validator_proves_a_working_backend_with_a_real_forward():
"probe_capability runs the doctor's bounded forward on the backend that would serve.\n\nTags: node, admission"
ctx = _context(_FakeBackend())
report = probe_capability(ctx)
assert report.passed
assert admit(AdmissionRequirement.for_context(ctx), report) is report
def test_the_production_validator_fails_a_backend_that_cannot_execute():
"A shard that loads but cannot run a forward yields a failed report, not a pass.\n\nTags: node, admission"
ctx = _context(_FakeBackend(forward_error=RuntimeError("CUDA out of memory")))
report = probe_capability(ctx)
assert not report.passed
with pytest.raises(CapabilityAdmissionError) as exc:
admit(AdmissionRequirement.for_context(ctx), report)
assert exc.value.reason == REASON_NOT_PASSED
def test_the_production_validator_fails_a_node_with_no_backend_at_all():
"A server with no model backend cannot prove anything, so it is not admitted.\n\nTags: node, admission"
ctx = _context(None)
assert not probe_capability(ctx).passed
# ---------------------------------------------------------------------------
# run_startup: a refused report means the tracker is never called
# ---------------------------------------------------------------------------
class _FakeTorchNodeServer:
started = False
def __init__(self, **kwargs):
self.kwargs = kwargs
self.backend = kwargs.pop("_backend", None) or _FakeBackend()
def start(self):
type(self).started = True
return 7099
def stop(self):
pass
@pytest.fixture
def startup_env(monkeypatch):
"""Fake hardware, wallet and tracker; records registrations *for this model*.
Heartbeat threads that other tests leave running are daemon threads that
outlive their test and re-register through this same patched `_post_json`, so
a plain call log would be polluted by whatever ran before. Only posts naming
this test's model can have come from this test's node.
"""
import meshnet_node.startup as startup_mod
posted: list[tuple[str, dict]] = []
_FakeTorchNodeServer.started = False
def _record(url, payload, timeout=10.0):
if url.endswith("/v1/nodes/register") and payload.get("hf_repo") == MODEL:
posted.append((url, payload))
return {"node_id": "node-nca"}
monkeypatch.setattr(
startup_mod,
"detect_hardware",
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
)
monkeypatch.setattr(
startup_mod, "benchmark_throughput_checked", lambda _device: (10.0, True, None)
)
monkeypatch.setattr(startup_mod, "TorchNodeServer", _FakeTorchNodeServer)
monkeypatch.setattr(
startup_mod, "load_or_create_wallet", lambda **_kw: (b"", b"", "wallet-nca")
)
monkeypatch.setattr(
startup_mod, "_get_json", lambda _url, timeout=10.0: {"relay_url": None, "nodes": []}
)
monkeypatch.setattr(startup_mod, "_post_json", _record)
monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)
return posted
def _start(**kwargs):
return run_startup(
tracker_url="http://127.0.0.1:8080",
model_id=MODEL,
shard_start=0,
shard_end=23,
**kwargs,
)
def test_backend_validation_failure_registers_nothing(startup_env, monkeypatch):
"A shard that cannot run a forward must never reach /v1/nodes/register.\n\nTags: node, admission, startup"
import meshnet_node.startup as startup_mod
broken = _FakeBackend(forward_error=RuntimeError("CUDA out of memory"))
monkeypatch.setattr(
startup_mod,
"TorchNodeServer",
lambda **kw: _FakeTorchNodeServer(_backend=broken, **kw),
)
with pytest.raises(CapabilityAdmissionError) as exc:
_start() # no validator: the production real-forward path
assert exc.value.reason == REASON_NOT_PASSED
assert startup_env == [], "the tracker was called despite a failed validation"
assert not _FakeTorchNodeServer.started, "a failed node still opened an endpoint"
def test_a_report_for_a_different_model_cannot_be_reused_to_register(startup_env):
"A success for one model must not admit another — the shape of a replayed proof.\n\nTags: node, admission, startup"
with pytest.raises(CapabilityAdmissionError) as exc:
_start(capability_validator=capability_stub(model_id="other/model-1b"))
assert exc.value.reason == REASON_MODEL_MISMATCH
assert startup_env == []
def test_a_stale_report_cannot_be_reused_to_register(startup_env):
"An aged-out proof is refused before the node advertises itself.\n\nTags: node, admission, startup"
with pytest.raises(CapabilityAdmissionError) as exc:
_start(capability_validator=capability_stub(age_seconds=86_400))
assert exc.value.reason == REASON_STALE
assert startup_env == []
def test_a_matching_passing_report_registers_and_travels_with_the_payload(startup_env):
"Registration carries the proof for exactly the model/shard/recipe it advertises.\n\nTags: node, admission, startup"
node = _start() # production validator against a working fake backend
node.stop()
assert len(startup_env) == 1
url, payload = startup_env[0]
assert url.endswith("/v1/nodes/register")
report = payload["capability_report"]
assert report["status"] == "passed"
assert report["model"]["model_id"] == MODEL
assert (report["shard"]["start"], report["shard"]["end"]) == (0, 23)
assert report["recipe"]["recipe_id"] == DEFAULT_RECIPE_ID
assert report["backend"]["device"] == "cpu"
def test_the_served_backend_is_loaded_with_the_recipe_that_was_validated(startup_env):
"The recipe named in the report is the one the serving backend actually ran.\n\nTags: node, admission, startup"
node = _start(recipe_id="eager-attention")
node.stop()
assert node.kwargs["recipe_params"] == {"attn_implementation": "eager"}
report = startup_env[0][1]["capability_report"]
assert report["recipe"]["recipe_id"] == "eager-attention"
def test_an_unknown_recipe_fails_before_any_weights_are_loaded(startup_env):
"A typo'd recipe id is caught at resolution, not after a multi-minute load.\n\nTags: node, admission, startup"
from meshnet_node.recipe_manifest import RecipeManifestError
with pytest.raises(RecipeManifestError):
_start(recipe_id="does-not-exist")
assert startup_env == []
assert not _FakeTorchNodeServer.started