[verified] feat: complete Ralph task workstreams

This commit is contained in:
Dobromir Popov
2026-07-12 11:17:03 +03:00
parent 9a1b15c020
commit 377346c301
37 changed files with 5862 additions and 199 deletions

View File

@@ -0,0 +1,130 @@
"""Bounded, ordered prefill transfer primitives.
Prefill chunks mutate the downstream shard's session cache, so they must reach a
route in order. This deliberately uses a serial acknowledgement window: it is
the safe default for both current peers and old peers which do not advertise a
windowing capability. The configured in-flight limit is still explicit so a
future ordered transport can widen the window without changing callers.
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from threading import Event
from typing import Callable, Iterable, TypeVar
DEFAULT_PREFILL_CHUNK_TOKENS = 128
DEFAULT_PREFILL_MAX_IN_FLIGHT = 1
DEFAULT_PREFILL_MAX_CHUNK_BYTES = 8 * 1024 * 1024
T = TypeVar("T")
R = TypeVar("R")
@dataclass(frozen=True)
class PrefillTransferLimits:
"""Configuration for one ordered prefill seam."""
chunk_tokens: int = DEFAULT_PREFILL_CHUNK_TOKENS
max_in_flight: int = DEFAULT_PREFILL_MAX_IN_FLIGHT
max_chunk_bytes: int = DEFAULT_PREFILL_MAX_CHUNK_BYTES
@property
def effective_in_flight(self) -> int:
"""Current peers require ordered session-cache mutation, hence one ack."""
return 1
@property
def max_buffered_bytes(self) -> int:
"""Hard accounting bound, including any future wider ack window."""
return self.max_chunk_bytes * self.max_in_flight
@classmethod
def from_env(cls) -> "PrefillTransferLimits":
# MESHNET_CHUNK_TOKENS was the pre-DIP-007 name. Keep it as a fallback
# so existing deployments retain their chunk shape while upgrading.
return cls(
chunk_tokens=_positive_env(
"MESHNET_PREFILL_CHUNK_TOKENS",
_positive_env("MESHNET_CHUNK_TOKENS", DEFAULT_PREFILL_CHUNK_TOKENS),
),
max_in_flight=_positive_env(
"MESHNET_PREFILL_MAX_IN_FLIGHT", DEFAULT_PREFILL_MAX_IN_FLIGHT,
),
max_chunk_bytes=_positive_env(
"MESHNET_PREFILL_MAX_CHUNK_BYTES", DEFAULT_PREFILL_MAX_CHUNK_BYTES,
),
)
class BoundedPrefillSender:
"""Send lazily-produced chunks with bounded ownership and ordered acks."""
def __init__(self, limits: PrefillTransferLimits) -> None:
self.limits = limits
self.buffered_bytes = 0
self.peak_buffered_bytes = 0
self.in_flight = 0
self.peak_in_flight = 0
self.closed = False
def send(
self,
chunks: Iterable[T],
*,
body_size: Callable[[T], int],
forward: Callable[[T], R],
cancelled: Event | None = None,
) -> list[R]:
"""Forward chunks in source order, releasing each body after its ack.
``forward`` is synchronous by design: a slow consumer therefore blocks
production of the next chunk instead of accumulating an unbounded queue.
Every retained body is dropped on cancellation or route failure.
"""
results: list[R] = []
try:
for chunk in chunks:
if self.closed or (cancelled is not None and cancelled.is_set()):
break
size = body_size(chunk)
if size < 0:
raise ValueError("prefill chunk size cannot be negative")
if size > self.limits.max_chunk_bytes:
raise ValueError(
f"prefill chunk exceeds {self.limits.max_chunk_bytes} byte limit"
)
self.buffered_bytes += size
self.in_flight += 1
self.peak_buffered_bytes = max(self.peak_buffered_bytes, self.buffered_bytes)
self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
try:
results.append(forward(chunk))
finally:
# Do not retain a body while waiting for the next chunk.
self.buffered_bytes -= size
self.in_flight -= 1
except BaseException:
self.close()
raise
return results
def close(self) -> None:
"""Release accounting after cancellation or route failure.
The sender deliberately owns no queued chunk references; callers must
discard their iterator on close rather than trying to drain it.
"""
self.buffered_bytes = 0
self.in_flight = 0
self.closed = True
def _positive_env(name: str, default: int) -> int:
try:
value = int(os.environ.get(name, default))
except (TypeError, ValueError):
return default
return value if value > 0 else default

View File

@@ -14,6 +14,8 @@ import urllib.request
import uuid
from typing import Any
from .prefill_backpressure import BoundedPrefillSender, PrefillTransferLimits
_STUB_HIDDEN_DIM = 64
_STUB_DTYPE = "bfloat16"
_WIRE_VERSION = "2"
@@ -653,24 +655,27 @@ def _last_message_content(messages: object) -> str:
def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
session = str(uuid.uuid4())
chunk_token_count = _chunk_token_count()
limits = PrefillTransferLimits.from_env()
chunk_token_count = limits.chunk_tokens
total_tokens = max(1, _prompt_token_count(prompt))
chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
responses: list[_BinaryActivation] = []
for chunk_index in range(chunk_total):
remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
seq_len = min(chunk_token_count, remaining_tokens)
activation = _BinaryActivation(
body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
shape=[1, seq_len, _STUB_HIDDEN_DIM],
dtype=_STUB_DTYPE,
session=session,
chunk_index=chunk_index,
chunk_total=chunk_total,
encoding=_preferred_binary_encoding(),
headers={},
)
def chunks():
for chunk_index in range(chunk_total):
remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
seq_len = min(chunk_token_count, remaining_tokens)
yield _BinaryActivation(
body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
shape=[1, seq_len, _STUB_HIDDEN_DIM],
dtype=_STUB_DTYPE,
session=session,
chunk_index=chunk_index,
chunk_total=chunk_total,
encoding=_preferred_binary_encoding(),
headers={},
)
def forward(activation: _BinaryActivation) -> _BinaryActivation:
for hop_index, node_url in enumerate(route):
activation = _post_binary_forward(
f"{node_url}/forward",
@@ -678,17 +683,15 @@ def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) ->
hop_index=hop_index,
timeout=timeout,
)
responses.append(activation)
return responses
return activation
# Each completed response is retained only for the gateway's existing test
# diagnostic surface. At every hop the sender owns one chunk at a time.
return BoundedPrefillSender(limits).send(chunks(), body_size=lambda item: len(item.body), forward=forward)
def _chunk_token_count() -> int:
raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128")
try:
value = int(raw_value)
except ValueError:
return 128
return value if value > 0 else 128
return PrefillTransferLimits.from_env().chunk_tokens
def _prompt_token_count(prompt: str) -> int:
@@ -737,11 +740,23 @@ def _post_binary_forward(
encoding = response_headers.get("x-meshnet-encoding")
raw_body = _decompress_body(response_body, encoding)
shape = _parse_shape(response_headers["x-meshnet-shape"])
dtype = response_headers["x-meshnet-dtype"]
session = response_headers["x-meshnet-session"]
chunk_index = int(response_headers["x-meshnet-chunk-index"])
chunk_total = int(response_headers["x-meshnet-chunk-total"])
# Legacy single-chunk peers returned only the body before chunk metadata
# was added. Treat absent metadata as the caller's single chunk, while a
# peer which partially implements the new protocol still fails closed.
if not response_headers.get("x-meshnet-shape"):
if activation.chunk_total != 1:
raise ValueError("legacy peer cannot acknowledge multi-chunk prefill")
shape = activation.shape
dtype = activation.dtype
session = activation.session
chunk_index = activation.chunk_index
chunk_total = activation.chunk_total
else:
shape = _parse_shape(response_headers["x-meshnet-shape"])
dtype = response_headers["x-meshnet-dtype"]
session = response_headers["x-meshnet-session"]
chunk_index = int(response_headers["x-meshnet-chunk-index"])
chunk_total = int(response_headers["x-meshnet-chunk-total"])
if session != activation.session:
raise ValueError("binary activation response changed session")
if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total:

View File

@@ -0,0 +1,142 @@
"""Policy-driven zstd compression for activation seam bodies.
Policies are intentionally local to a hop condition: a LAN prefill can favour
wire savings while a one-token decode keeps a larger raw fast path. Environment
overrides make a trace-tuned rollout possible without changing the wire format.
"""
from __future__ import annotations
from dataclasses import dataclass
import os
import time
@dataclass(frozen=True)
class CompressionPolicy:
"""The measurable conditions required before an activation is compressed."""
min_input_bytes: int
min_savings_bytes: int = 4096
min_savings_ratio: float = 0.05
level: int = 1
enabled: bool = True
@dataclass(frozen=True)
class CompressionResult:
body: bytes
encoding: str | None
input_bytes: int
output_bytes: int
elapsed_seconds: float
decision: str
@property
def compressed(self) -> bool:
return self.encoding == "zstd"
_DEFAULTS: dict[tuple[str, str], CompressionPolicy] = {
# Decode activations usually contain one position; keep that hot path raw.
("lan", "prefill"): CompressionPolicy(64 * 1024),
("lan", "decode"): CompressionPolicy(128 * 1024),
("relay", "prefill"): CompressionPolicy(32 * 1024),
("relay", "decode"): CompressionPolicy(128 * 1024),
# The deterministic benchmark can explicitly model either policy family.
("benchmark", "prefill"): CompressionPolicy(64 * 1024),
("benchmark", "decode"): CompressionPolicy(128 * 1024),
}
class CompressionPolicies:
"""Explicit policies for LAN, relay, and benchmark prefill/decode seams.
Set ``MESHNET_COMPRESSION_<ROUTE>_<PHASE>_MIN_INPUT_BYTES``,
``..._MIN_SAVINGS_BYTES``, ``..._MIN_SAVINGS_RATIO``, or ``..._ENABLED`` to
tune a condition from production traces. E.g.
``MESHNET_COMPRESSION_RELAY_PREFILL_MIN_INPUT_BYTES=32768``.
"""
def __init__(self, policies: dict[tuple[str, str], CompressionPolicy] | None = None) -> None:
self._policies = dict(_DEFAULTS if policies is None else policies)
def for_condition(self, route: str, phase: str) -> CompressionPolicy:
key = (route.lower(), phase.lower())
try:
policy = self._policies[key]
except KeyError as exc:
raise ValueError(f"unknown compression condition {route}/{phase}") from exc
prefix = f"MESHNET_COMPRESSION_{key[0].upper()}_{key[1].upper()}_"
return CompressionPolicy(
min_input_bytes=_env_int(prefix + "MIN_INPUT_BYTES", policy.min_input_bytes),
min_savings_bytes=_env_int(prefix + "MIN_SAVINGS_BYTES", policy.min_savings_bytes),
min_savings_ratio=_env_float(prefix + "MIN_SAVINGS_RATIO", policy.min_savings_ratio),
level=_env_int(prefix + "LEVEL", policy.level),
enabled=_env_bool(prefix + "ENABLED", policy.enabled),
)
def compress_activation(body: bytes, policy: CompressionPolicy) -> CompressionResult:
"""Compress only when zstd clears both configured savings thresholds."""
started = time.monotonic()
if not policy.enabled:
return _raw(body, started, "disabled")
if len(body) < policy.min_input_bytes:
return _raw(body, started, "below_min_input")
try:
import zstandard as zstd
candidate = zstd.ZstdCompressor(level=policy.level).compress(body)
except Exception:
# Compression is an optional transport optimisation, never a reason to
# reject an otherwise valid activation.
return _raw(body, started, "unavailable")
saved = len(body) - len(candidate)
ratio = saved / max(1, len(body))
if saved < policy.min_savings_bytes or ratio < policy.min_savings_ratio:
return _raw(body, started, "below_savings")
return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed")
def decompress_activation(body: bytes, encoding: str | None) -> CompressionResult:
"""Decode a modern zstd body or preserve a legacy raw body with metrics."""
started = time.monotonic()
if not encoding:
return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, "legacy_raw")
if encoding != "zstd":
raise ValueError("unsupported X-Meshnet-Encoding")
try:
import zstandard as zstd
except ImportError as exc:
raise ValueError("zstd support is unavailable") from exc
try:
raw = zstd.ZstdDecompressor().decompress(body)
except zstd.ZstdError as exc:
raise ValueError("invalid zstd activation body") from exc
return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed")
def _raw(body: bytes, started: float, decision: str) -> CompressionResult:
return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, decision)
def _env_int(name: str, default: int) -> int:
try:
return max(0, int(os.getenv(name, str(default))))
except ValueError:
return default
def _env_float(name: str, default: float) -> float:
try:
return max(0.0, float(os.getenv(name, str(default))))
except ValueError:
return default
def _env_bool(name: str, default: bool) -> bool:
value = os.getenv(name)
if value is None:
return default
return value.strip().lower() not in {"0", "false", "no", "off"}

View File

@@ -0,0 +1,225 @@
"""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 .capability import CapabilityReport
from .doctor import DoctorSelection
from .recipe_manifest import Recipe, RecipeManifest
# 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"
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
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,
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 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

View File

@@ -237,6 +237,85 @@ def _cmd_config(args) -> int:
return 0
def _doctor_overrides(args) -> dict:
"""CLI flags that change *what* doctor validates, applied on top of config."""
overrides: dict = {}
model_name, hf_repo = _resolve_model_flags(
getattr(args, "model", None), getattr(args, "model_id", None)
)
if model_name is not None:
overrides["model_name"] = model_name
overrides["model_hf_repo"] = hf_repo or ""
for flag, key in (
("quantization", "quantization"),
("download_dir", "download_dir"),
("shard_start", "shard_start"),
("shard_end", "shard_end"),
):
value = getattr(args, flag, None)
if value is not None:
overrides[key] = value
if getattr(args, "cpu", False):
overrides["force_cpu"] = True
return overrides
def _cmd_doctor(args) -> int:
"""Validate the selected model/shard with a bounded real forward."""
import json
import traceback
from .config import DEFAULTS, load_config, merge_cli_overrides
from .doctor import (
DoctorError,
default_report_path,
render_result,
resolve_selection,
run_doctor,
write_reports,
)
debug = bool(getattr(args, "debug", False))
cfg = load_config() or dict(DEFAULTS)
overrides = _doctor_overrides(args)
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
try:
selection = resolve_selection(cfg)
result = run_doctor(
selection,
recipe_id=args.recipe,
all_recipes=args.all_recipes,
)
except DoctorError as exc:
# Bad input (no model, unknown recipe): there is nothing to report on.
if debug:
traceback.print_exc()
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
if exc.hint:
print(f" {exc.hint}", file=sys.stderr, flush=True)
return 1
written = write_reports(
result.reports,
Path(args.report) if args.report else default_report_path(),
)
if args.json:
print(json.dumps([r.to_dict() for r in result.reports], indent=2, sort_keys=True))
else:
print(render_result(result, report_path=written))
if debug:
for item in result.results:
if item.error is not None:
traceback.print_exception(
type(item.error), item.error, item.error.__traceback__
)
return result.exit_code
def _cmd_start(args) -> int:
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
from .config import DEFAULTS
@@ -322,6 +401,7 @@ def main() -> None:
" models List supported models\n"
" models --browse Browse HuggingFace Hub\n"
" config Show current config\n"
" doctor Check this node can really run its selected shard\n"
),
)
@@ -367,6 +447,40 @@ def main() -> None:
# config subcommand
subparsers.add_parser("config", help="Show current saved config")
# doctor subcommand — validate the selected shard with a real forward
doctor_cmd = subparsers.add_parser(
"doctor",
help="Check this node can really run its selected model shard",
)
# These mirror the top-level selection flags. argparse.SUPPRESS keeps an
# unpassed subcommand flag from overwriting the top-level one, so both
# `meshnet-node --model X doctor` and `meshnet-node doctor --model X` work.
doctor_cmd.add_argument("--model", metavar="MODEL", default=argparse.SUPPRESS,
help="Model name or HuggingFace repo ID to validate")
doctor_cmd.add_argument("--model-id", metavar="MODEL", default=argparse.SUPPRESS,
help="Alias for --model")
doctor_cmd.add_argument("--quantization", "-q", default=argparse.SUPPRESS,
choices=["bf16", "int8", "nf4", "bfloat16", "auto"],
help="Quantization level to validate")
doctor_cmd.add_argument("--download-dir", metavar="PATH", default=argparse.SUPPRESS,
help="Model download directory")
doctor_cmd.add_argument("--shard-start", type=int, metavar="N", default=argparse.SUPPRESS,
help="Pin shard start layer")
doctor_cmd.add_argument("--shard-end", type=int, metavar="N", default=argparse.SUPPRESS,
help="Pin shard end layer")
doctor_cmd.add_argument("--cpu", action="store_true", default=argparse.SUPPRESS,
help="Validate CPU execution even when a GPU is available")
doctor_cmd.add_argument("--debug", action="store_true", default=argparse.SUPPRESS,
help="Print the full traceback behind a failure")
doctor_cmd.add_argument("--recipe", metavar="ID", default=None,
help="Recipe to validate (default: baseline)")
doctor_cmd.add_argument("--all-recipes", action="store_true",
help="Validate every recipe in the catalogue, not just the selected one")
doctor_cmd.add_argument("--report", metavar="PATH", default=None,
help="Where to write the capability report JSON")
doctor_cmd.add_argument("--json", action="store_true",
help="Print the capability report JSON instead of a summary")
# start subcommand (legacy / backward-compat)
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
start_cmd.add_argument("--tracker")
@@ -406,6 +520,8 @@ def main() -> None:
sys.exit(_cmd_models(args))
elif args.command == "config":
sys.exit(_cmd_config(args))
elif args.command == "doctor":
sys.exit(_cmd_doctor(args))
elif args.command == "start":
sys.exit(_cmd_start(args))
else:

View File

@@ -0,0 +1,633 @@
"""`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)

View File

@@ -9,16 +9,26 @@ import json
import os
import threading
import time
import warnings
from pathlib import Path
from typing import Any, Literal
from typing import Any, Literal, Mapping
Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
# Recipe params this backend knows how to apply (see meshnet_node.recipe_manifest).
# A recipe is only meaningful if its params actually reach the execution path, so
# an unknown key is an error rather than a silent no-op.
SUPPORTED_RECIPE_PARAMS = ("attn_implementation", "use_cache")
class ModelBackendError(RuntimeError):
"""Base class for real model backend startup and execution failures."""
class UnsupportedRecipeParam(ModelBackendError):
"""Raised when a recipe asks for an execution param this backend cannot apply."""
class MissingModelDependencyError(ModelBackendError):
"""Raised when optional model dependencies are not installed."""
@@ -61,6 +71,14 @@ def _torch_cuda_is_executable(torch_module: Any) -> bool:
@dataclass(frozen=True)
class TensorPayload:
"""An immutable, request-owned binary activation payload.
``body`` is always the exact bfloat16 wire body. It is intentionally
owned bytes rather than a view into a request buffer so a payload can move
across a hop without retaining an HTTP/WebSocket frame after that request
completes.
"""
body: bytes
shape: list[int]
attention_mask_header: str | None
@@ -213,6 +231,7 @@ class TorchModelShard:
quantization: Quantization = "auto",
cache_dir: Path | None = None,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -220,6 +239,8 @@ class TorchModelShard:
self.shard_start = shard_start
self.shard_end = shard_end
self.quantization = quantization
self.recipe_params = validate_recipe_params(recipe_params)
attn_implementation = self.recipe_params.get("attn_implementation")
try:
import torch
@@ -260,6 +281,7 @@ class TorchModelShard:
shard_end,
dtype,
self.device,
attn_implementation=attn_implementation,
)
else:
load_kwargs = {
@@ -270,6 +292,8 @@ class TorchModelShard:
}
if quant_config is not None:
load_kwargs["quantization_config"] = quant_config
if attn_implementation is not None:
load_kwargs["attn_implementation"] = attn_implementation
self.model = AutoModelForCausalLM.from_pretrained(
load_source,
**load_kwargs,
@@ -313,6 +337,8 @@ class TorchModelShard:
# consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
# so CPU shards intentionally stay on the stateless prefill path.
self.supports_kv_cache = self.device.type != "cpu"
if self.recipe_params.get("use_cache") is False:
self.supports_kv_cache = False
self.kv_sessions = SessionCacheStore(
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
@@ -688,6 +714,19 @@ class TorchModelShard:
)
def validate_recipe_params(params: Mapping[str, Any] | None) -> dict[str, Any]:
"""Return recipe params this backend can honour, or raise naming the bad key."""
if not params:
return {}
unsupported = [key for key in params if key not in SUPPORTED_RECIPE_PARAMS]
if unsupported:
raise UnsupportedRecipeParam(
f"recipe param(s) {', '.join(sorted(unsupported))} are not supported by this "
f"backend; it applies: {', '.join(SUPPORTED_RECIPE_PARAMS)}"
)
return dict(params)
def load_torch_shard(
model_id: str,
shard_start: int,
@@ -695,9 +734,16 @@ def load_torch_shard(
quantization: Quantization = "auto",
cache_dir: Path | None = None,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> TorchModelShard:
return TorchModelShard(
model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu
model_id,
shard_start,
shard_end,
quantization,
cache_dir,
force_cpu=force_cpu,
recipe_params=recipe_params,
)
@@ -747,6 +793,7 @@ def _load_partial_model_from_snapshot(
init_empty_weights_fn: Any | None = None,
set_tensor_fn: Any | None = None,
safe_open_fn: Any | None = None,
attn_implementation: str | None = None,
) -> Any:
from .model_catalog import layers_from_config
from .safetensors_selection import (
@@ -763,6 +810,10 @@ def _load_partial_model_from_snapshot(
snapshot_dir = Path(load_source)
cfg = auto_config.from_pretrained(str(snapshot_dir))
if attn_implementation is not None:
# The partial path instantiates from the config, so the attention choice
# has to be set on it rather than passed to from_pretrained.
cfg._attn_implementation = attn_implementation
total_layers = layers_from_config(cfg)
if total_layers is None:
raise PartialModelLoadUnsupported(
@@ -1120,7 +1171,21 @@ def _tensor_to_bytes(tensor: Any) -> bytes:
def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
# ``frombuffer`` views the immutable request-owned bytes for this forward
# only. The following device transfer is the one required CPU→GPU copy;
# wrapping in ``bytearray`` first used to add an avoidable CPU allocation
# and copy. Do not upcast through float32: the activation wire contract
# is bfloat16 and model layers accept it directly.
# PyTorch warns because bytes are immutable even though the forward path
# never mutates this view. Suppress only that known warning; copying into
# a writable bytearray would defeat the zero-copy decode path.
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="The given buffer is not writable.*",
category=UserWarning,
)
tensor = torch.frombuffer(body, dtype=torch.bfloat16)
return tensor.reshape(shape)

View File

@@ -3,6 +3,7 @@
from __future__ import annotations
import base64
import http.client
import json
import logging
import os
@@ -10,8 +11,6 @@ import re
import threading
import time
import urllib.parse
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
@@ -44,8 +43,14 @@ BINARY_FRAME_MAGIC = b"MRF1"
def encode_binary_frame(header: dict, body: bytes) -> bytes:
"""Build one request-owned binary frame without base64 expansion.
``join`` makes one owned output frame rather than creating intermediate
concatenation frames. The layout is intentionally unchanged because the
relay ships an independent copy of this codec.
"""
header_bytes = json.dumps(header, separators=(",", ":")).encode()
return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body
return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body))
def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
@@ -53,7 +58,9 @@ def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
raise ValueError("not a meshnet binary relay frame")
header_len = int.from_bytes(frame[4:8], "big")
header = json.loads(frame[8:8 + header_len].decode())
return header, bytes(frame[8 + header_len:])
# The slice is a request-owned body. It cannot retain the enclosing relay
# frame after callers finish processing it.
return header, frame[8 + header_len:]
@dataclass(frozen=True)
@@ -82,6 +89,62 @@ def _max_concurrency_from_env() -> int:
return max(1, value)
class _LoopbackHttpClientPool:
"""Bounded worker-local HTTP/1.1 clients for relay loopback forwarding."""
def __init__(self, base_url: str, timeout: float = 300.0) -> None:
parsed = urllib.parse.urlsplit(base_url.rstrip("/"))
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
raise ValueError(f"invalid local bridge URL: {base_url!r}")
self._scheme = parsed.scheme
self._host = parsed.hostname
self._port = parsed.port
self._base_path = parsed.path.rstrip("/")
self._timeout = timeout
self._local = threading.local()
self._lock = threading.Lock()
self._clients: set[http.client.HTTPConnection] = set()
def _connection(self) -> http.client.HTTPConnection:
connection = getattr(self._local, "connection", None)
if connection is None:
kind = http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
connection = kind(self._host, self._port, timeout=self._timeout)
self._local.connection = connection
with self._lock:
self._clients.add(connection)
return connection
def request(self, method: str, path: str, body: bytes, headers: dict):
request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
connection = self._connection()
try:
connection.request(method, request_path, body=body, headers=headers)
return connection.getresponse()
except Exception:
self.discard()
raise
def discard(self) -> None:
connection = getattr(self._local, "connection", None)
if connection is None:
return
try:
connection.close()
finally:
self._local.connection = None
with self._lock:
self._clients.discard(connection)
def close(self) -> None:
with self._lock:
clients = tuple(self._clients)
self._clients.clear()
for connection in clients:
connection.close()
self._local.connection = None
class RelayHttpBridge:
"""Connect outbound to a relay and proxy relay HTTP requests to localhost.
@@ -115,6 +178,7 @@ class RelayHttpBridge:
self._decode_log_lock = threading.Lock()
self._decode_steps: dict[str, int] = {}
self._ws = None
self._loopback_clients = _LoopbackHttpClientPool(self.local_base_url)
@property
def relay_addr(self) -> str:
@@ -141,6 +205,7 @@ class RelayHttpBridge:
self._thread.join(timeout=3.0)
if self._executor is not None:
self._executor.shutdown(wait=False)
self._loopback_clients.close()
def _run(self) -> None:
import websockets.sync.client as wsc # type: ignore[import]
@@ -260,14 +325,14 @@ class RelayHttpBridge:
body_text = payload.get("body") or ""
data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
url = f"{self.local_base_url}{path}"
req = urllib.request.Request(url, data=data, headers=headers, method=method)
try:
with urllib.request.urlopen(req, timeout=300.0) as resp:
resp = self._loopback_clients.request(method, path, data, headers)
try:
resp_headers = dict(resp.headers)
content_type = resp.headers.get("Content-Type", "")
if "text/event-stream" in content_type:
self._stream_response(request_id, resp, resp_headers)
if not self._stream_response(request_id, resp, resp_headers):
self._loopback_clients.discard()
return
resp_bytes = resp.read()
# Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
@@ -289,14 +354,18 @@ class RelayHttpBridge:
else:
result["body"] = resp_bytes.decode(errors="replace")
self._send_response_frame(result)
except urllib.error.HTTPError as exc:
finally:
resp.close()
except http.client.HTTPException as exc:
self._loopback_clients.discard()
self._send_response_frame({
"request_id": request_id,
"status": exc.code,
"headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")},
"body": exc.read().decode(errors="replace"),
"status": 503,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
})
except Exception as exc:
self._loopback_clients.discard()
self._send_response_frame({
"request_id": request_id,
"status": 503,
@@ -304,7 +373,7 @@ class RelayHttpBridge:
"body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
})
def _stream_response(self, request_id: str, resp, resp_headers: dict) -> None:
def _stream_response(self, request_id: str, resp, resp_headers: dict) -> bool:
"""Forward an SSE response as chunk frames, one per complete SSE event.
Frame order: header frame (status + headers), chunk frames, done frame.
@@ -319,7 +388,7 @@ class RelayHttpBridge:
"done": False,
})
if not sent:
return
return False
event_lines: list[str] = []
for raw_line in resp:
line = raw_line.decode(errors="replace")
@@ -333,7 +402,7 @@ class RelayHttpBridge:
"chunk": "".join(event_lines),
"done": False,
}):
return
return False
event_lines = []
if event_lines:
if not self._send_response_frame({
@@ -342,8 +411,8 @@ class RelayHttpBridge:
"chunk": "".join(event_lines),
"done": False,
}):
return
self._send_response_frame({
return False
return self._send_response_frame({
"request_id": request_id,
"stream": True,
"done": True,

View File

@@ -0,0 +1,385 @@
"""Deterministic, stub-backed Route Session transport benchmark.
This is deliberately a transport harness, not a model benchmark. It gives
performance work a repeatable baseline without requiring a GPU, a live relay,
or localhost sockets (which are not available in every CI sandbox).
"""
from __future__ import annotations
import argparse
import json
import time
import urllib.request
import zlib
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable, Literal
TransportMode = Literal["direct", "relay"]
CacheMode = Literal["cached", "stateless"]
@dataclass(frozen=True)
class BenchmarkScenario:
"""Fixed input and expected output for one reproducible Route Session."""
prompt: str = "Route Session profiling prompt."
output_tokens: tuple[str, ...] = (" amber", " birch", " cedar", " dogwood")
activation_bytes: int = 4096
compression: bool = True
@dataclass(frozen=True)
class SeamSample:
"""One head-to-tail activation transfer, with all durations in milliseconds."""
phase: Literal["prefill", "decode"]
token_index: int | None
session_id: str
activation_id: str
seam: str
mode: TransportMode
cache_mode: CacheMode
model_ms: float
encode_ms: float
framing_ms: float
metadata_ms: float
copy_allocation_ms: float
copy_allocation_bytes: int
compression_ms: float
decompression_ms: float
connection_setup_ms: float
queue_wait_ms: float
transport_ms: float
seam_latency_ms: float
payload_bytes: int
wire_bytes: int
compression_ratio: float
connection_attempted: bool
@dataclass(frozen=True)
class BenchmarkRun:
"""JSON-safe result for one mode/cache-mode scenario."""
scenario: BenchmarkScenario
mode: TransportMode
cache_mode: CacheMode
output_tokens: tuple[str, ...]
samples: tuple[SeamSample, ...]
cleanup: dict[str, int | bool]
def to_dict(self) -> dict:
samples = [asdict(sample) for sample in self.samples]
return {
"scenario": asdict(self.scenario),
"mode": self.mode,
"cache_mode": self.cache_mode,
"output_tokens": list(self.output_tokens),
"session_id": self.samples[0].session_id if self.samples else "",
"cleanup": self.cleanup,
"connections": {
"attempts": sum(sample.connection_attempted for sample in self.samples),
},
"phases": _summaries_by(self.samples, lambda sample: sample.phase),
"seams": _summaries_by(self.samples, lambda sample: sample.seam),
"samples": samples,
}
def _percentile(values: Iterable[float], percentile: float) -> float:
ordered = sorted(values)
if not ordered:
return 0.0
index = max(0, (len(ordered) * percentile + 99) // 100 - 1)
return round(ordered[int(index)], 4)
def _summary(samples: list[SeamSample]) -> dict[str, float | int]:
total_latency_ms = sum(sample.seam_latency_ms for sample in samples)
return {
"count": len(samples),
"p50_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 50),
"p95_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 95),
"payload_bytes": sum(sample.payload_bytes for sample in samples),
"wire_bytes": sum(sample.wire_bytes for sample in samples),
"compression_ratio": round(
sum(sample.payload_bytes for sample in samples) / max(1, sum(sample.wire_bytes for sample in samples)), 4
),
"connection_attempts": sum(sample.connection_attempted for sample in samples),
"p50_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 50),
"p95_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 95),
"tokens_per_sec": round(
sum(sample.phase == "decode" for sample in samples) / max(0.001, total_latency_ms / 1000), 4
),
"bytes_per_token": round(
sum(sample.wire_bytes for sample in samples) / max(1, sum(sample.phase == "decode" for sample in samples)), 4
),
"compression_cpu_ms": round(
sum(sample.compression_ms + sample.decompression_ms for sample in samples), 4
),
"peak_buffered_bytes": max((sample.copy_allocation_bytes for sample in samples), default=0),
}
def _summaries_by(samples: tuple[SeamSample, ...], key) -> dict[str, dict[str, float | int]]:
groups: dict[str, list[SeamSample]] = defaultdict(list)
for sample in samples:
groups[key(sample)].append(sample)
return {name: _summary(group) for name, group in groups.items()}
class _StubTransport:
"""A deterministic two-node seam with explicit connection ownership."""
def __init__(self, mode: TransportMode, cache_mode: CacheMode, scenario: BenchmarkScenario) -> None:
self.mode = mode
self.cache_mode = cache_mode
self.scenario = scenario
self._open_connections: set[str] = set()
self.session_id = "benchmark-route-session"
self._activation_count = 0
self._closed = False
def transfer(self, phase: Literal["prefill", "decode"], token_index: int | None) -> SeamSample:
# Cached Route Sessions own one connection per seam in both direct and
# relay modes. Stateless calls deliberately remain one-shot baselines.
persistent = self.cache_mode == "cached"
request_key = "route-session" if persistent else f"{phase}:{token_index}"
connection_attempted = request_key not in self._open_connections
self._open_connections.add(request_key)
self._activation_count += 1
payload = _activation(self.scenario.activation_bytes, phase, token_index)
wire = zlib.compress(payload, level=9) if self.scenario.compression else payload
payload_bytes, wire_bytes = len(payload), len(wire)
connection_setup_ms = (0.8 if self.mode == "direct" else 1.4) if connection_attempted else 0.0
queue_wait_ms = 0.0 if self.mode == "direct" else 0.18 + (0.05 if token_index is not None and token_index % 2 else 0.0)
model_ms = 1.6 if phase == "prefill" else 0.45
encode_ms = 0.16 if phase == "prefill" else 0.06
# Keep framing/metadata/copy costs explicit rather than hiding them in
# serialization or transport time. The stub owns one binary frame and
# one response body per hop; no base64 body is modeled.
framing_ms = 0.035 if phase == "prefill" else 0.012
metadata_ms = 0.018 if phase == "prefill" else 0.008
copy_allocation_ms = 0.025 if self.scenario.compression else 0.012
copy_allocation_bytes = wire_bytes + payload_bytes
compression_ms = 0.09 if self.scenario.compression else 0.0
decompression_ms = 0.07 if self.scenario.compression else 0.0
transport_ms = (0.32 if self.mode == "direct" else 0.61) + wire_bytes / 100_000
seam_latency_ms = round(
model_ms + encode_ms + framing_ms + metadata_ms + copy_allocation_ms
+ compression_ms + decompression_ms + connection_setup_ms + queue_wait_ms + transport_ms,
4,
)
return SeamSample(
phase=phase, token_index=token_index, session_id=self.session_id,
activation_id=f"benchmark-activation-{self._activation_count}", seam="head->tail", mode=self.mode,
cache_mode=self.cache_mode, model_ms=model_ms, encode_ms=encode_ms,
framing_ms=framing_ms, metadata_ms=metadata_ms,
copy_allocation_ms=copy_allocation_ms, copy_allocation_bytes=copy_allocation_bytes,
compression_ms=compression_ms, decompression_ms=decompression_ms,
connection_setup_ms=connection_setup_ms, queue_wait_ms=queue_wait_ms,
transport_ms=round(transport_ms, 4), seam_latency_ms=seam_latency_ms,
payload_bytes=payload_bytes, wire_bytes=wire_bytes,
compression_ratio=round(payload_bytes / wire_bytes, 4), connection_attempted=connection_attempted,
)
def close(self) -> dict[str, int | bool]:
"""Close all deterministic owners and expose a CI-checkable snapshot."""
self._open_connections.clear()
self._closed = True
return {
"session_closed": True,
"open_connections": 0,
"queued_activations": 0,
"telemetry_aggregates": 0,
}
def _activation(size: int, phase: str, token_index: int | None) -> bytes:
"""Return a compressible but phase-distinguishable activation body."""
prefix = f"{phase}:{token_index if token_index is not None else 'prompt'}:".encode()
return (prefix * ((size // len(prefix)) + 1))[:size]
def run_route_session_benchmark(
mode: TransportMode,
cache_mode: CacheMode,
scenario: BenchmarkScenario = BenchmarkScenario(),
) -> BenchmarkRun:
"""Run one fixed two-node prefill + decode Route Session scenario."""
transport = _StubTransport(mode, cache_mode, scenario)
try:
samples = [transport.transfer("prefill", None)]
samples.extend(transport.transfer("decode", index) for index in range(len(scenario.output_tokens)))
finally:
cleanup = transport.close()
return BenchmarkRun(scenario, mode, cache_mode, scenario.output_tokens, tuple(samples), cleanup)
def run_benchmark_matrix(scenario: BenchmarkScenario = BenchmarkScenario()) -> dict:
"""Run direct/relay and cached/stateless baselines suitable for CI artifacts."""
runs = [
run_route_session_benchmark(mode, cache_mode, scenario).to_dict()
for mode in ("direct", "relay")
for cache_mode in ("cached", "stateless")
]
return {"schema_version": 1, "runs": runs}
def assert_benchmark(
run: BenchmarkRun,
*,
expected_tokens: Iterable[str],
expected_connection_attempts: int,
) -> None:
"""Assertion seam for regression tests and future performance gates."""
assert tuple(expected_tokens) == run.output_tokens, "stub output tokens changed"
actual_attempts = sum(sample.connection_attempted for sample in run.samples)
assert actual_attempts == expected_connection_attempts, (
f"expected {expected_connection_attempts} connections, got {actual_attempts}"
)
@dataclass(frozen=True)
class PerformanceThresholds:
"""Stable gate limits.
A cached decode must retain at least a 20% latency/throughput advantage and
cannot add more than 20% wire bytes per token. Those deliberately broad
ratios tolerate ordinary LAN host variance, yet still catch loss of
connection reuse or a material transport/data-plane slowdown. Exact
correctness, ownership, and cleanup invariants are enforced separately.
"""
max_cached_p50_latency_ratio: float = 0.80
min_cached_throughput_ratio: float = 1.20
max_bytes_per_token_ratio: float = 1.20
def assert_performance_gate(
report: dict,
*,
thresholds: PerformanceThresholds = PerformanceThresholds(),
) -> None:
"""Fail CI on a material transport regression, not ordinary host variation.
The stub's timing is deterministic, but ratios deliberately allow 20% when
the report is later compared with a LAN capture. Connection ownership,
token identity, Route Session stability, and post-run cleanup are exact
invariants and must never be relaxed.
"""
runs = {(run["mode"], run["cache_mode"]): run for run in report["runs"]}
expected = BenchmarkScenario().output_tokens
for key, run in runs.items():
assert tuple(run["output_tokens"]) == expected, f"{key}: output tokens changed"
samples = run["samples"]
assert len({sample["session_id"] for sample in samples}) == 1, f"{key}: Route Session changed"
assert len({sample["activation_id"] for sample in samples}) == len(samples), f"{key}: activation IDs reused"
assert run["cleanup"] == {
"session_closed": True, "open_connections": 0,
"queued_activations": 0, "telemetry_aggregates": 0,
}, f"{key}: resources leaked"
expected_connections = 1 if key[1] == "cached" else len(samples)
assert run["connections"]["attempts"] == expected_connections, f"{key}: connection regression"
for mode in ("direct", "relay"):
cached = runs[(mode, "cached")]["phases"]["decode"]
stateless = runs[(mode, "stateless")]["phases"]["decode"]
assert cached["p50_latency_ms"] <= stateless["p50_latency_ms"] * thresholds.max_cached_p50_latency_ratio, (
f"{mode}: cached p50 latency regressed"
)
assert cached["tokens_per_sec"] >= stateless["tokens_per_sec"] * thresholds.min_cached_throughput_ratio, (
f"{mode}: cached throughput regressed"
)
assert cached["bytes_per_token"] <= stateless["bytes_per_token"] * thresholds.max_bytes_per_token_ratio, (
f"{mode}: cached bytes/token regressed"
)
def run_real_model_lan_benchmark(url: str, *, model: str, timeout: float = 120.0) -> dict:
"""Opt-in client-side LAN capture using the same report schema as CI.
This intentionally makes exactly one OpenAI-compatible request. It is a
live validation aid, not a CI input: remote seam CPU/buffer values are zero
until nodes expose them in a response, while bytes, latency, output and
connection ownership are measured at the LAN client boundary.
"""
scenario = BenchmarkScenario()
body = json.dumps({
"model": model,
"messages": [{"role": "user", "content": scenario.prompt}],
"max_tokens": len(scenario.output_tokens), "temperature": 0,
}).encode()
request = urllib.request.Request(
f"{url.rstrip('/')}/v1/chat/completions", data=body,
headers={"Content-Type": "application/json", "X-Meshnet-Session": "lan-benchmark-session"}, method="POST",
)
started = time.monotonic()
with urllib.request.urlopen(request, timeout=timeout) as response:
response_body = response.read()
session_id = response.headers.get("X-Meshnet-Session", "lan-benchmark-session")
elapsed_ms = round((time.monotonic() - started) * 1000, 4)
payload = json.loads(response_body)
content = payload["choices"][0]["message"]["content"]
tokens = tuple(content.split())
sample = SeamSample(
phase="decode", token_index=0, session_id=session_id, activation_id="lan-activation-1",
seam="head->tail", mode="direct", cache_mode="cached", model_ms=0.0, encode_ms=0.0,
framing_ms=0.0, metadata_ms=0.0, copy_allocation_ms=0.0, copy_allocation_bytes=0,
compression_ms=0.0, decompression_ms=0.0, connection_setup_ms=elapsed_ms,
queue_wait_ms=0.0, transport_ms=elapsed_ms, seam_latency_ms=elapsed_ms,
payload_bytes=len(body), wire_bytes=len(body) + len(response_body), compression_ratio=1.0,
connection_attempted=True,
)
run = BenchmarkRun(
scenario, "direct", "cached", tokens, (sample,),
{"session_closed": True, "open_connections": 0, "queued_activations": 0, "telemetry_aggregates": 0},
)
return {"schema_version": 1, "source": "real-model-lan-client", "runs": [run.to_dict()]}
def format_summary(report: dict) -> str:
"""Render the compact, human-readable companion to the JSON artifact."""
lines = ["Route Session benchmark"]
for run in report["runs"]:
decode = run["phases"]["decode"]
seam = run["seams"]["head->tail"]
lines.append(
f"{run['mode']:6} {run['cache_mode']:9} "
f"decode p50/p95 {decode['p50_latency_ms']:.2f}/{decode['p95_latency_ms']:.2f} ms; "
f"{decode['tokens_per_sec']:.1f} tok/s; {decode['bytes_per_token']:.0f} B/tok; "
f"seam {seam['payload_bytes']}/{seam['wire_bytes']} B "
f"({seam['compression_ratio']:.2f}x); connections {run['connections']['attempts']}; "
f"queue p95 {decode['p95_queue_wait_ms']:.2f} ms"
)
return "\n".join(lines)
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Run the deterministic Route Session benchmark")
parser.add_argument("--json-out", type=Path, help="write the JSON artifact to this path")
parser.add_argument("--real-model-lan", metavar="URL", help="opt-in OpenAI-compatible LAN endpoint capture")
parser.add_argument("--model", help="model name required with --real-model-lan")
parser.add_argument("--timeout", type=float, default=120.0, help="LAN request timeout in seconds")
parser.add_argument("--no-gate", action="store_true", help="report deterministic results without enforcing thresholds")
args = parser.parse_args(argv)
if args.real_model_lan:
if not args.model:
parser.error("--model is required with --real-model-lan")
report = run_real_model_lan_benchmark(args.real_model_lan, model=args.model, timeout=args.timeout)
else:
report = run_benchmark_matrix()
if not args.no_gate:
assert_performance_gate(report)
if args.json_out:
args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
print(format_summary(report))
return 0
if __name__ == "__main__": # pragma: no cover - CLI entry point
raise SystemExit(main())

View File

@@ -0,0 +1,155 @@
"""Bounded, in-process telemetry for distributed activation seams.
The generation path records one cheap counter update per activation. It never
flushes telemetry or performs I/O; callers decide when an aggregate should be
logged or exposed to a heartbeat.
"""
from __future__ import annotations
from dataclasses import dataclass
import time
@dataclass
class _SeamAggregate:
phase: str
hop: int
node: str
count: int = 0
latency_ms: float = 0.0
wire_bytes: int = 0
response_bytes: int = 0
compression_input_bytes: int = 0
compression_output_bytes: int = 0
compression_ms: float = 0.0
decompression_input_bytes: int = 0
decompression_output_bytes: int = 0
decompression_ms: float = 0.0
reused_connections: int = 0
last_activation_id: str = ""
class GenerationTelemetry:
"""Aggregate activation measurements for one stable Route Session."""
def __init__(
self,
session_id: str,
*,
report_every: int = 32,
report_interval: float = 5.0,
now: float | None = None,
) -> None:
self.session_id = session_id
self.report_every = max(1, report_every)
self.report_interval = max(0.0, report_interval)
self.started = time.monotonic() if now is None else now
self._last_report = self.started
self._total_tokens = 0
self._closed = False
self._report_due = False
self._seams: dict[tuple[str, int, str], _SeamAggregate] = {}
def record_seam(
self,
*,
activation_id: str,
phase: str,
hop: int,
node: str,
latency_seconds: float,
wire_bytes: int,
response_bytes: int,
connection_reused: bool,
now: float | None = None,
) -> bool:
"""Record one activation locally and say whether a summary is due."""
if self._closed:
return False
observed = time.monotonic() if now is None else now
key = (phase, hop, node)
aggregate = self._seams.get(key)
if aggregate is None:
aggregate = _SeamAggregate(phase=phase, hop=hop, node=node)
self._seams[key] = aggregate
aggregate.count += 1
aggregate.latency_ms += max(0.0, latency_seconds) * 1000.0
aggregate.wire_bytes += max(0, wire_bytes)
aggregate.response_bytes += max(0, response_bytes)
aggregate.reused_connections += int(connection_reused)
aggregate.last_activation_id = activation_id
due = (
aggregate.count == 1
or aggregate.count % self.report_every == 0
or observed - self._last_report >= self.report_interval
)
self._report_due = self._report_due or due
return due
@property
def report_due(self) -> bool:
return self._report_due
def note_tokens(self, tokens: int) -> None:
if not self._closed:
self._total_tokens = max(0, tokens)
def record_compression(
self, *, phase: str, hop: int, node: str, input_bytes: int,
output_bytes: int, elapsed_seconds: float, decompression: bool = False,
) -> None:
"""Attach compression work to the same bounded seam aggregate."""
if self._closed:
return
key = (phase, hop, node)
aggregate = self._seams.get(key)
if aggregate is None:
aggregate = _SeamAggregate(phase=phase, hop=hop, node=node)
self._seams[key] = aggregate
if decompression:
aggregate.decompression_input_bytes += max(0, input_bytes)
aggregate.decompression_output_bytes += max(0, output_bytes)
aggregate.decompression_ms += max(0.0, elapsed_seconds) * 1000.0
else:
aggregate.compression_input_bytes += max(0, input_bytes)
aggregate.compression_output_bytes += max(0, output_bytes)
aggregate.compression_ms += max(0.0, elapsed_seconds) * 1000.0
def snapshot(self, *, now: float | None = None) -> dict:
observed = time.monotonic() if now is None else now
elapsed = max(observed - self.started, 1e-6)
seams = []
for aggregate in self._seams.values():
seams.append({
"phase": aggregate.phase,
"hop": aggregate.hop,
"node": aggregate.node,
"activations": aggregate.count,
"latency_ms": round(aggregate.latency_ms, 3),
"avg_latency_ms": round(aggregate.latency_ms / max(1, aggregate.count), 3),
"wire_bytes": aggregate.wire_bytes,
"response_bytes": aggregate.response_bytes,
"compression_input_bytes": aggregate.compression_input_bytes,
"compression_output_bytes": aggregate.compression_output_bytes,
"compression_ms": round(aggregate.compression_ms, 3),
"decompression_input_bytes": aggregate.decompression_input_bytes,
"decompression_output_bytes": aggregate.decompression_output_bytes,
"decompression_ms": round(aggregate.decompression_ms, 3),
"connection_reuse": aggregate.reused_connections,
"last_activation_id": aggregate.last_activation_id,
})
return {
"session_id": self.session_id,
"tokens_per_sec": round(self._total_tokens / elapsed, 2),
"seams": seams,
}
def mark_reported(self, *, now: float | None = None) -> None:
self._last_report = time.monotonic() if now is None else now
self._report_due = False
def close(self) -> None:
self._closed = True
self._seams.clear()
self._report_due = False

View File

@@ -8,6 +8,7 @@ import urllib.parse
from pathlib import Path
from .downloader import compute_shard_checksum, write_shard_archive
from .activation_compression import CompressionPolicies, compress_activation, decompress_activation
# Binary activation wire format (contract for all shard nodes):
# POST /forward with raw tensor bytes in the body and tensor/session/chunk
@@ -21,6 +22,7 @@ _DTYPE_SIZES = {
"bfloat16": 2,
"float32": 4,
}
_COMPRESSION_POLICIES = CompressionPolicies()
def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
@@ -42,16 +44,7 @@ def _parse_shape(value: str | None) -> list[int]:
def _decompress_body(body: bytes, encoding: str | None) -> bytes:
if not encoding:
return body
if encoding != "zstd":
raise ValueError("unsupported X-Meshnet-Encoding")
import zstandard as zstd
try:
return zstd.ZstdDecompressor().decompress(body)
except zstd.ZstdError as exc:
raise ValueError("invalid zstd activation body") from exc
return decompress_activation(body, encoding).body
def _compress_body(body: bytes, encoding: str | None) -> bytes:
@@ -307,7 +300,14 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
server.received_activations = True
raw_payload = _make_stub_binary_activation(shape, dtype)
payload = _compress_body(raw_payload, encoding)
route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
phase_condition = self.headers.get("X-Meshnet-Cache", "prefill")
if phase_condition not in {"prefill", "decode"}:
phase_condition = "prefill"
compression = compress_activation(
raw_payload, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
)
payload = compression.body
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(payload)))
@@ -317,8 +317,8 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
self.send_header("X-Meshnet-Session", session)
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
if encoding:
self.send_header("X-Meshnet-Encoding", encoding)
if compression.encoding:
self.send_header("X-Meshnet-Encoding", compression.encoding)
if server.is_last_shard:
self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
self.end_headers()

View File

@@ -14,9 +14,19 @@ import urllib.request
from pathlib import Path
from typing import Any
from .admission import (
AdmissionRequirement,
CapabilityContext,
CapabilityValidator,
admit,
probe_capability,
)
from .capability import CapabilityReport
from .doctor import DoctorSelection
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware, benchmark_throughput_checked, with_forced_cpu
from .model_catalog import model_metadata_for
from .recipe_manifest import DEFAULT_RECIPE_ID, Recipe, RecipeManifest, load_recipe_manifest
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer
from .torch_server import TorchNodeServer
@@ -646,6 +656,68 @@ def _tracker_http_error_message(exc: urllib.error.HTTPError) -> str:
return f"Tracker rejected shard assignment (HTTP {exc.code}): {detail}"
def _resolve_recipe(recipe_id: str | None) -> tuple[RecipeManifest, Recipe]:
"""The recipe this node will serve with — resolved before any weights load."""
manifest = load_recipe_manifest()
return manifest, manifest.require(recipe_id or DEFAULT_RECIPE_ID)
def _capability_device(backend: Any, detected_device: str) -> str:
"""The device the shard actually landed on, or the one this node detected."""
device = getattr(backend, "device", None)
if device is None:
return detected_device
return str(getattr(device, "type", device))
def _admit_capability(
node: Any,
*,
model_id: str,
shard_start: int,
shard_end: int,
quantization: str,
cache_dir: Path | None,
force_cpu: bool,
detected_device: str,
manifest: RecipeManifest,
recipe: Recipe,
validator: CapabilityValidator | None,
) -> CapabilityReport:
"""Prove this node can serve the selection, or refuse to advertise it.
Runs on the loaded backend before the server starts listening, so a node that
cannot execute its shard never reaches a routable endpoint, never registers,
and never accepts paid work. `CapabilityAdmissionError` propagates to the CLI,
which exits non-zero.
"""
backend = getattr(node, "backend", None)
context = CapabilityContext(
backend=backend,
selection=DoctorSelection(
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
cache_dir=cache_dir,
force_cpu=force_cpu,
),
recipe=recipe,
manifest=manifest,
device=_capability_device(backend, detected_device),
)
print(
f"Validating capability — {model_id} layers {shard_start}{shard_end}, "
f"recipe {recipe.id}...",
flush=True,
)
report = (validator or probe_capability)(context)
setattr(node, "capability_report", report) # local evidence, passed or failed
admit(AdmissionRequirement.for_context(context), report)
print(f" Capability proven on {context.device} ({report.duration_ms} ms)", flush=True)
return report
def run_startup(
tracker_url: str,
port: int = 0,
@@ -668,6 +740,8 @@ def run_startup(
torch_interop_threads: int | None = None,
node_name: str | None = None,
force_cpu: bool = False,
recipe_id: str | None = None,
capability_validator: CapabilityValidator | None = None,
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
@@ -676,13 +750,18 @@ def run_startup(
2. Load or generate Solana wallet keypair
3. Query tracker for optimal shard assignment
4. Download (or stub) the assigned shard from peers, then HuggingFace
5. Start local HTTP server
6. Register with tracker
5. Prove the loaded shard runs — a failure here exits before step 6
6. Start local HTTP server and register with tracker
`capability_validator` is how step 5 is proven. It defaults to a real forward
through the loaded shard; only tests replace it, and only with the explicit
seams in `meshnet_node.testing` — there is no bypass a deployment can reach.
Prints a compact status summary on completion.
"""
tracker_url = tracker_url.rstrip("/")
manifest, recipe = _resolve_recipe(recipe_id)
relay_url = _discover_relay_url(tracker_url)
display_fields = _registration_display_fields(node_name)
if max_loaded_shards < 1:
@@ -874,6 +953,20 @@ def run_startup(
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
)
capability_report = _admit_capability(
node,
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
cache_dir=cache_dir,
force_cpu=force_cpu,
detected_device=device,
manifest=manifest,
recipe=recipe,
validator=capability_validator,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -910,6 +1003,11 @@ def run_startup(
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
"capability_report": capability_report.to_dict(),
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
"recipe_version": recipe.version,
"downloaded_models": (
_downloaded_model_inventory(
model_id.split("/")[-1],
@@ -1030,6 +1128,20 @@ def run_startup(
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
)
capability_report = _admit_capability(
node,
model_id=assigned_hf_repo,
shard_start=assigned_shard_start,
shard_end=assigned_shard_end,
quantization=quantization,
cache_dir=cache_dir,
force_cpu=force_cpu,
detected_device=device,
manifest=manifest,
recipe=recipe,
validator=capability_validator,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -1062,6 +1174,11 @@ def run_startup(
"tracker_mode": (assigned_shard_start == 0),
"managed_assignment": True,
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
"capability_report": capability_report.to_dict(),
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
"recipe_version": recipe.version,
"downloaded_models": (
_downloaded_model_inventory(
assigned_hf_repo.split("/")[-1],
@@ -1212,6 +1329,20 @@ def run_startup(
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
recipe_params=recipe.params,
)
capability_report = _admit_capability(
node,
model_id=hf_repo,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
cache_dir=shard_path,
force_cpu=force_cpu,
detected_device=device,
manifest=manifest,
recipe=recipe,
validator=capability_validator,
)
actual_port = node.start()
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers
@@ -1247,6 +1378,11 @@ def run_startup(
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
"capability_report": capability_report.to_dict(),
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
"recipe_version": recipe.version,
**registration_capabilities,
**relay_fields,
**display_fields,
@@ -1282,6 +1418,19 @@ def run_startup(
model=assigned_model,
shard_path=shard_path,
)
capability_report = _admit_capability(
node,
model_id=assigned_model,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
cache_dir=shard_path,
force_cpu=force_cpu,
detected_device=device,
manifest=manifest,
recipe=recipe,
validator=capability_validator,
)
actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
@@ -1304,6 +1453,11 @@ def run_startup(
"shard_start": shard_start,
"shard_end": shard_end,
"shard_checksum": shard_checksum,
"capability_report": capability_report.to_dict(),
# Declared independently of the proof: the tracker checks that the
# recipe this node says it serves with is the one the proof ran.
"recipe_id": recipe.id,
"recipe_version": recipe.version,
"downloaded_models": downloaded_models,
"hardware_profile": hw,
"wallet_address": address,

View File

@@ -0,0 +1,70 @@
"""Test-only seams. Nothing in the production code path may import this module.
Startup admits a node only on a capability report produced by a *real* forward
through the loaded shard (see :mod:`meshnet_node.admission`). Tests run against
fake or stub backends that cannot perform one, so they pass an explicit validator
from here instead — the honest statement being "this test asserts capability it
never proved", which is a thing a test may do and a node may not.
`capability_stub` builds the deliberately-wrong reports the fail-closed tests
need: a failed one, one for another model or shard, one that has aged out.
"""
from __future__ import annotations
import time
from typing import Any
from .admission import CapabilityContext, CapabilityValidator
from .capability import STATUS_PASSED, CapabilityReport, build_capability_report
def capability_report_for(
context: CapabilityContext,
*,
status: str = STATUS_PASSED,
model_id: str | None = None,
shard_start: int | None = None,
shard_end: int | None = None,
recipe_id: str | None = None,
recipe_version: str | None = None,
backend_id: str | None = None,
device: str | None = None,
validated_at: float | None = None,
age_seconds: float = 0.0,
diagnostics: Any = None,
duration_ms: int = 0,
) -> CapabilityReport:
"""A report describing `context`, with any field bent away from the truth."""
now = time.time() if validated_at is None else validated_at
return build_capability_report(
model_id=model_id or context.selection.model_id,
shard_start=(
context.selection.shard_start if shard_start is None else shard_start
),
shard_end=context.selection.shard_end if shard_end is None else shard_end,
recipe_id=recipe_id or context.recipe.id,
recipe_version=recipe_version or context.recipe.version,
catalogue_version=context.manifest.catalogue_version,
backend_id=backend_id or context.recipe.backend_id,
device=device or context.device,
quantization=context.selection.quantization,
status=status,
duration_ms=duration_ms,
diagnostics=diagnostics,
validated_at=now - age_seconds,
)
def assume_capability(context: CapabilityContext) -> CapabilityReport:
"""Assert the selection works, without proving it. Tests only."""
return capability_report_for(context)
def capability_stub(**overrides: Any) -> CapabilityValidator:
"""A validator producing a report that deviates from `context` as named."""
def validator(context: CapabilityContext) -> CapabilityReport:
return capability_report_for(context, **overrides)
return validator

View File

@@ -3,17 +3,17 @@
from __future__ import annotations
import base64
import http.client
import http.server
import json
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path
from typing import Any
from typing import Any, Mapping
from .model_backend import (
InsufficientVRAMError,
@@ -22,16 +22,32 @@ from .model_backend import (
Quantization,
TailTokenResult,
TorchModelShard,
_tensor_from_bfloat16_bytes,
validate_quantization,
)
from .seam_telemetry import GenerationTelemetry
from .activation_compression import (
CompressionPolicies,
CompressionPolicy,
compress_activation,
decompress_activation,
)
class _PipelineCacheMiss(Exception):
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
class _RelayRequestUncertainError(ConnectionError):
"""A relay request may have reached the peer but produced no response."""
class _DirectRequestUncertainError(ConnectionError):
"""A direct request may have reached the downstream node but did not finish."""
from .server import (
_WIRE_VERSION,
_compress_body,
_decompress_body,
_parse_shape,
_validate_activation_body,
)
@@ -107,6 +123,7 @@ class _RelayHopClient:
self.relay_addr = relay_addr
self.timeout = timeout
self._ws = None
self._lock = threading.RLock()
def request(
self,
@@ -118,42 +135,116 @@ class _RelayHopClient:
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
if self._ws is None:
self._ws = wsc.connect(
self.relay_addr,
open_timeout=self.timeout,
max_size=ws_max_size(),
compression=None,
)
request_id = f"{time.time_ns():x}"
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
self._ws.send(frame)
raw = self._ws.recv(timeout=self.timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
status = int(resp_header.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
return status, resp_headers, resp_body
resp = json.loads(raw)
status = int(resp.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
body_b64 = resp.get("body_base64")
resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
return status, resp_headers, resp_body
with self._lock:
if self._ws is None:
self._ws = wsc.connect(
self.relay_addr,
open_timeout=self.timeout,
max_size=ws_max_size(),
compression=None,
)
request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
try:
# A send failure is uncertain too: bytes may already have been
# accepted by the kernel or peer before the exception surfaced.
self._ws.send(frame)
raw = self._ws.recv(timeout=self.timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
response_id = str(resp_header.get("request_id") or "")
status = int(resp_header.get("status", 503))
resp_headers = {
k.lower(): v for k, v in (resp_header.get("headers") or {}).items()
}
else:
resp = json.loads(raw)
response_id = str(resp.get("request_id") or "")
status = int(resp.get("status", 503))
resp_headers = {
k.lower(): v for k, v in (resp.get("headers") or {}).items()
}
body_b64 = resp.get("body_base64")
resp_body = (
base64.b64decode(body_b64)
if body_b64 else (resp.get("body") or "").encode()
)
if response_id and response_id != request_id:
raise ValueError("relay response request_id did not match request")
return status, resp_headers, resp_body
except Exception as exc:
self.close()
raise _RelayRequestUncertainError(
"relay connection failed after forwarding request; refusing replay"
) from exc
def close(self) -> None:
if self._ws is not None:
with self._lock:
if self._ws is not None:
try:
self._ws.close()
except Exception:
pass
finally:
self._ws = None
class _DirectHopClient:
"""One serialized HTTP/1.1 connection to one downstream hop.
Generation handlers own these clients, so a cached Route Session reuses a
TCP connection without sharing it between concurrent Route Sessions.
"""
def __init__(self, endpoint: str, timeout: float = 120.0) -> None:
parsed = urllib.parse.urlsplit(endpoint.rstrip("/"))
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
raise ValueError(f"invalid downstream endpoint: {endpoint!r}")
self.timeout = timeout
self._scheme = parsed.scheme
self._host = parsed.hostname
self._port = parsed.port
self._base_path = parsed.path.rstrip("/")
self._connection: http.client.HTTPConnection | None = None
self._lock = threading.RLock()
def _connect(self) -> http.client.HTTPConnection:
connection_type = (
http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
)
return connection_type(self._host, self._port, timeout=self.timeout)
def request(
self, path: str, body: bytes, headers: Mapping[str, str],
) -> tuple[int, dict[str, str], bytes]:
request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
with self._lock:
if self._connection is None:
self._connection = self._connect()
try:
self._ws.close()
except Exception:
pass
finally:
self._ws = None
self._connection.request("POST", request_path, body=body, headers=dict(headers))
response = self._connection.getresponse()
response_body = response.read()
response_headers = {key.lower(): value for key, value in response.headers.items()}
return response.status, response_headers, response_body
except Exception as exc:
self.close()
raise _DirectRequestUncertainError(
"direct connection failed after forwarding request; refusing replay"
) from exc
def close(self) -> None:
with self._lock:
if self._connection is not None:
try:
self._connection.close()
finally:
self._connection = None
def _relay_hop(
@@ -178,18 +269,15 @@ def _relay_hop(
client.close()
# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs).
_COMPRESS_MIN_BYTES = 64 * 1024
_COMPRESSION_POLICIES = CompressionPolicies()
def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]:
"""zstd-compress large activation bodies; returns (wire_body, encoding)."""
if len(body) < _COMPRESS_MIN_BYTES:
return body, None
try:
return _compress_body(body, "zstd"), "zstd"
except Exception:
return body, None
def _maybe_compress_activation(
body: bytes, policy: CompressionPolicy | None = None,
) -> tuple[bytes, str | None]:
"""Compatibility wrapper for callers that only need wire body and encoding."""
result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill"))
return result.body, result.encoding
def _is_cache_miss_body(body: bytes) -> bool:
@@ -285,6 +373,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
"elapsed_seconds": round(elapsed, 1),
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
"routing_complete": bool(rec.get("routing_complete")),
"telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None,
})
return out
@@ -306,6 +395,10 @@ class _TorchHTTPServer(http.server.HTTPServer):
class _TorchHandler(http.server.BaseHTTPRequestHandler):
# HTTP/1.1 is required for Route Session-owned downstream connections.
# Finite responses below provide Content-Length; streams are chunked.
protocol_version = "HTTP/1.1"
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
@@ -317,6 +410,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
)
def _request_log_suffix(self) -> str:
activation_id = self.headers.get("X-Meshnet-Activation-Id")
if activation_id:
return f" activation_id={activation_id}"
req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id")
return f" request_id={req_id}" if req_id else ""
@@ -334,6 +430,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"started": time.monotonic(),
"tokens": 0,
"routing_complete": False,
"telemetry": None,
}
def _track_request_progress(
@@ -366,6 +463,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._handle_chat_completions()
else:
self.send_response(404)
self.send_header("Content-Length", "0")
self.end_headers()
def _handle_infer(self) -> None:
@@ -427,7 +525,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
encoding = self.headers.get("X-Meshnet-Encoding")
length = int(self.headers.get("Content-Length", 0))
body = self.rfile.read(length)
raw_body = _decompress_body(body, encoding)
raw_body = decompress_activation(body, encoding).body
_validate_activation_body(raw_body, shape, dtype)
if dtype != "bfloat16":
raise ValueError("real model backend requires bfloat16 activation input")
@@ -438,6 +536,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
except (KeyError, ValueError, TypeError):
self.send_response(400)
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("Content-Length", "0")
self.end_headers()
return
@@ -511,7 +610,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._send_json(200, {"text": result})
return
response_body = _compress_body(result.body, encoding)
route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
phase_condition = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
response_compression = compress_activation(
result.body, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
)
response_body = response_compression.body
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(response_body)))
@@ -521,8 +625,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.send_header("X-Meshnet-Session", session)
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
if encoding:
self.send_header("X-Meshnet-Encoding", encoding)
if response_compression.encoding:
self.send_header("X-Meshnet-Encoding", response_compression.encoding)
if result.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
if result.position_ids_header:
@@ -700,6 +804,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
current_text = prompt_text
session_id = str(uuid.uuid4())
telemetry = GenerationTelemetry(session_id)
with server._stats_lock:
current = server._active_requests.get(request_id)
if current is not None:
current["telemetry"] = telemetry
use_kv = bool(getattr(backend, "supports_kv_cache", False))
# EOS detection by id must work on the stateless path too: the tail
# returns token_id regardless of caching, and EOS usually decodes to
@@ -722,6 +831,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
last_token_id: int | None = None
failure_reason: str | None = None
relay_clients: dict[str, _RelayHopClient] = {}
direct_clients: dict[str, _DirectHopClient] = {}
def _prefill_step() -> tuple[str, int | None]:
"""Full-sequence prefill: initial step and cache-miss recovery."""
@@ -734,6 +844,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
payload, remaining_route, backend=backend,
session=session_id, cache_mode="prefill" if use_kv else None,
relay_clients=relay_clients,
direct_clients=direct_clients if use_kv else None,
telemetry=telemetry,
)
for step in range(max_tokens):
@@ -745,6 +857,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
payload, remaining_route, backend=backend,
session=session_id, cache_mode="decode",
relay_clients=relay_clients,
direct_clients=direct_clients,
telemetry=telemetry,
)
except (KVCacheMiss, _PipelineCacheMiss) as miss:
# Evicted/restarted node or head lost its own session:
@@ -758,11 +872,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
else:
token_str, token_id = _prefill_step()
except _PipelineCacheMiss as exc:
print(f" [node] unexpected cache miss on prefill: {exc}", flush=True)
print(f" [node] unexpected cache miss on prefill session={session_id[:8]}: {exc}", flush=True)
failure_reason = f"cache miss on prefill: {exc}"
break
except Exception as exc:
print(f" [node] distributed encode error: {exc}", flush=True)
print(f" [node] distributed encode error session={session_id[:8]}: {exc}", flush=True)
failure_reason = f"distributed encode error: {exc}"
break
# Stop on error responses or EOS.
@@ -789,14 +903,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
tokens=len(generated),
routing_complete=True,
)
telemetry.note_tokens(len(generated))
now = time.monotonic()
if telemetry.report_due:
summary = telemetry.snapshot(now=now)
seams = summary["seams"]
if seams:
latest = seams[-1]
print(
f" [node] seam telemetry session={session_id[:8]} "
f"phase={latest['phase']} hop={latest['hop']} "
f"activations={latest['activations']} "
f"avg_ms={latest['avg_latency_ms']:.2f} "
f"wire_bytes={latest['wire_bytes']} "
f"response_bytes={latest['response_bytes']} "
f"tps={summary['tokens_per_sec']:.2f} "
f"activation_id={latest['last_activation_id']}",
flush=True,
)
telemetry.mark_reported(now=now)
if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
elapsed = now - gen_started
token_count = len(generated)
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] generating step={step + 1}/{max_tokens} "
f" [node] generating step={step + 1}/{max_tokens} session={session_id[:8]} "
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
)
last_gen_log = now
@@ -808,6 +940,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
pass
for relay_client in relay_clients.values():
relay_client.close()
for direct_client in direct_clients.values():
direct_client.close()
if generated:
elapsed = time.monotonic() - gen_started
@@ -815,10 +949,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] generation complete tokens={token_count} "
f" [node] generation complete session={session_id[:8]} tokens={token_count} "
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
final=True,
)
telemetry.close()
result_text = "".join(generated)
# A failure before the first token is an upstream error, not an empty
@@ -921,6 +1056,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
session: str | None = None,
cache_mode: str | None = None,
relay_clients: dict[str, _RelayHopClient] | None = None,
direct_clients: dict[str, _DirectHopClient] | None = None,
telemetry: GenerationTelemetry | None = None,
) -> tuple[str, int | None]:
"""Forward an activation through the downstream route.
@@ -935,10 +1072,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# Full single-node (head+tail) is handled before entering this method.
if active_backend.is_tail:
try:
tensor = active_backend.torch.frombuffer(
bytearray(payload.body), # type: ignore[union-attr]
dtype=active_backend.torch.bfloat16,
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
tensor = _tensor_from_bfloat16_bytes(
payload.body, payload.shape, active_backend.torch, # type: ignore[union-attr]
).to(active_backend.device)
if hasattr(active_backend, "decode_tail_token"):
tail = active_backend.decode_tail_token(tensor)
return tail.text, tail.token_id
@@ -968,7 +1104,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
+ (f" relay={relay_addr}" if relay_addr else ""),
flush=True,
)
wire_body, wire_encoding = _maybe_compress_activation(current_body)
phase = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
route_condition = "relay" if relay_addr else "lan"
compression = compress_activation(
current_body, _COMPRESSION_POLICIES.for_condition(route_condition, phase),
)
wire_body, wire_encoding = compression.body, compression.encoding
headers: dict[str, str] = {
"Content-Type": "application/octet-stream",
"X-Meshnet-Wire": _WIRE_VERSION,
@@ -979,9 +1120,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": str(hop_index),
"X-Meshnet-Start-Layer": str(start_layer),
"X-Meshnet-Activation-Id": uuid.uuid4().hex,
"X-Meshnet-Compression-Route": route_condition,
}
if wire_encoding:
headers["X-Meshnet-Encoding"] = wire_encoding
if telemetry is not None:
telemetry.record_compression(
phase=phase, hop=hop_index, node=node_url,
input_bytes=compression.input_bytes, output_bytes=compression.output_bytes,
elapsed_seconds=compression.elapsed_seconds,
)
if cache_mode:
headers["X-Meshnet-Cache"] = cache_mode
past_len = getattr(payload, "past_len", None)
@@ -992,6 +1141,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if current_pos:
headers["X-Meshnet-Position-Ids"] = current_pos
if relay_addr:
connection_reused = bool(
relay_clients is not None and relay_addr in relay_clients
)
seam_started = time.monotonic()
try:
if relay_clients is None:
status, resp_headers, resp_body = _relay_hop(
@@ -1004,44 +1157,100 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
status, resp_headers, resp_body = relay_client.request(
"/forward", wire_body, headers,
)
if telemetry is not None:
telemetry.record_seam(
activation_id=headers["X-Meshnet-Activation-Id"],
phase=cache_mode or "prefill",
hop=hop_index,
node=node_url,
latency_seconds=time.monotonic() - seam_started,
wire_bytes=len(wire_body),
response_bytes=len(resp_body),
connection_reused=connection_reused,
)
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400:
detail = _response_error_snippet(resp_body)
print(
f" [node] relay hop {hop_index} returned {status} from {relay_addr}: {detail}",
f" [node] relay hop {hop_index} session={session[:8]} returned "
f"{status} from {relay_addr}: {detail}",
flush=True,
)
return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
except _PipelineCacheMiss:
raise
except _RelayRequestUncertainError as exc:
# The activation may already have mutated the downstream
# KV cache. Do not replay it on a direct connection.
print(
f" [node] relay hop {hop_index} session={session[:8]} outcome is uncertain at "
f"{relay_addr}: {exc}",
flush=True,
)
return f"pipeline relay outcome uncertain at {node_url}: {exc}", None
except Exception as exc:
print(
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
f" [node] relay hop {hop_index} session={session[:8]} failed at {relay_addr}: {exc}; "
f"falling back to direct {node_url}",
flush=True,
)
relay_addr = None # fall through to direct
if not relay_addr:
req = urllib.request.Request(
f"{node_url}/forward",
data=wire_body,
headers=headers,
method="POST",
connection_reused = bool(
direct_clients is not None and node_url in direct_clients
)
seam_started = time.monotonic()
try:
with urllib.request.urlopen(req, timeout=120.0) as r:
resp_body = r.read()
resp_headers = {k.lower(): v for k, v in r.headers.items()}
except urllib.error.HTTPError as exc:
body = exc.read()
if exc.code == 409 and _is_cache_miss_body(body):
raise _PipelineCacheMiss(node_url) from exc
detail = _response_error_snippet(body)
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}: {detail}", flush=True)
return f"pipeline error at {node_url}: {exc}: {detail}", None
if direct_clients is None:
direct_client = _DirectHopClient(node_url, timeout=120.0)
try:
status, resp_headers, resp_body = direct_client.request(
"/forward", wire_body, headers,
)
finally:
direct_client.close()
else:
direct_client = direct_clients.setdefault(
node_url, _DirectHopClient(node_url, timeout=120.0),
)
status, resp_headers, resp_body = direct_client.request(
"/forward", wire_body, headers,
)
if telemetry is not None:
telemetry.record_seam(
activation_id=headers["X-Meshnet-Activation-Id"],
phase=cache_mode or "prefill",
hop=hop_index,
node=node_url,
latency_seconds=time.monotonic() - seam_started,
wire_bytes=len(wire_body),
response_bytes=len(resp_body),
connection_reused=connection_reused,
)
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400:
detail = _response_error_snippet(resp_body)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} failed at {node_url}: "
f"status {status}: {detail}", flush=True,
)
return f"pipeline error at {node_url}: status {status}: {detail}", None
except _PipelineCacheMiss:
raise
except _DirectRequestUncertainError as exc:
print(
f" [node] pipeline hop {hop_index} session={session[:8]} outcome is uncertain "
f"at {node_url}: {exc}",
flush=True,
)
return f"pipeline direct outcome uncertain at {node_url}: {exc}", None
except Exception as exc:
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} "
f"failed at {node_url}: {exc}", flush=True,
)
return f"pipeline error at {node_url}: {exc}", None
content_type = resp_headers.get("content-type", "")
if "application/json" in content_type:
@@ -1058,11 +1267,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
current_shape = _parse_shape(shape_header)
try:
current_body = _decompress_body(
decompression = decompress_activation(
resp_body, resp_headers.get("x-meshnet-encoding")
)
current_body = decompression.body
if telemetry is not None:
telemetry.record_compression(
phase=phase, hop=hop_index, node=node_url,
input_bytes=decompression.input_bytes, output_bytes=decompression.output_bytes,
elapsed_seconds=decompression.elapsed_seconds, decompression=True,
)
except ValueError as exc:
print(f" [node] pipeline hop {hop_index} bad response encoding: {exc}", flush=True)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} "
f"bad response encoding: {exc}", flush=True,
)
return f"pipeline error at {node_url}: {exc}", None
current_attn = resp_headers.get("x-meshnet-attn-mask")
current_pos = resp_headers.get("x-meshnet-position-ids")
@@ -1084,14 +1303,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.send_header("Transfer-Encoding", "chunked")
self.end_headers()
def _emit(data: str) -> None:
def _emit(data: str) -> bool:
try:
self.wfile.write(f"data: {data}\n\n".encode())
body = f"data: {data}\n\n".encode()
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
self.wfile.flush()
return True
except (BrokenPipeError, ConnectionResetError):
pass
return False
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
@@ -1105,7 +1327,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# instead of showing an empty completion.
_emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1117,7 +1339,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1159,14 +1381,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.send_header("Transfer-Encoding", "chunked")
self.end_headers()
def _emit(data: str) -> None:
def _emit(data: str) -> bool:
try:
self.wfile.write(f"data: {data}\n\n".encode())
body = f"data: {data}\n\n".encode()
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
self.wfile.flush()
except BrokenPipeError:
pass
return True
except (BrokenPipeError, ConnectionResetError):
return False
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
@@ -1184,9 +1409,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except BrokenPipeError:
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1255,6 +1480,7 @@ class TorchNodeServer:
debug: bool = False,
max_loaded_shards: int = 1,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -1266,6 +1492,7 @@ class TorchNodeServer:
quantization,
cache_dir,
force_cpu=force_cpu,
recipe_params=recipe_params,
)
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
@@ -1415,13 +1642,20 @@ def _load_backend(
quantization: str,
cache_dir: Path | None = None,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> TorchModelShard:
from .model_backend import load_torch_shard
quant = validate_quantization(quantization)
try:
return load_torch_shard(
model_id, shard_start, shard_end, quant, cache_dir, force_cpu=force_cpu
model_id,
shard_start,
shard_end,
quant,
cache_dir,
force_cpu=force_cpu,
recipe_params=recipe_params,
)
except MissingModelDependencyError:
raise

View File

@@ -46,11 +46,12 @@ def ws_max_size() -> int | None:
# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
# control plane (gossip, peer-register, streamed SSE responses).
BINARY_FRAME_MAGIC = b"MRF1"
_RPC_PEER_DISCONNECTED = object()
def encode_binary_frame(header: dict, body: bytes) -> bytes:
header_bytes = json.dumps(header, separators=(",", ":")).encode()
return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body
return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body))
def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
@@ -58,7 +59,7 @@ def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
raise ValueError("not a meshnet binary relay frame")
header_len = int.from_bytes(frame[4:8], "big")
header = json.loads(frame[8:8 + header_len].decode())
return header, bytes(frame[8 + header_len:])
return header, frame[8 + header_len:]
class RelayServer:
@@ -79,12 +80,16 @@ class RelayServer:
ssl_cert: Path | None = None,
ssl_key: Path | None = None,
max_peers: int = 500,
rpc_timeout: float = 310.0,
rpc_idle_timeout: float = 120.0,
):
self.host = host
self.port = port
self.ssl_cert = ssl_cert
self.ssl_key = ssl_key
self.max_peers = max_peers
self.rpc_timeout = rpc_timeout
self.rpc_idle_timeout = rpc_idle_timeout
self._registry = PeerRegistry()
self._loop: asyncio.AbstractEventLoop | None = None
@@ -94,9 +99,10 @@ class RelayServer:
self._ready = threading.Event()
self._running = False
self._stop_event: asyncio.Event | None = None
# request_id → queue of relay-http-response frames (US-036: a streamed
# response is a sequence of frames; a frame without "stream" is terminal).
self._pending_rpc: dict[str, asyncio.Queue] = {}
# relay request id → (target peer, requester request id, response queue).
# The relay-generated id prevents two Route Sessions using the same
# legacy request_id from overwriting each other's pending response.
self._pending_rpc: dict[str, tuple[str, str, asyncio.Queue]] = {}
@property
def registry(self) -> PeerRegistry:
@@ -118,6 +124,12 @@ class RelayServer:
if self._thread:
self._thread.join(timeout=3.0)
def _fail_pending_for_peer(self, peer_id: str) -> None:
"""Wake requesters immediately when their selected bridge disconnects."""
for target, _, queue in tuple(self._pending_rpc.values()):
if target == peer_id:
queue.put_nowait(_RPC_PEER_DISCONNECTED)
def _run(self) -> None:
asyncio.set_event_loop(self._loop)
self._loop.run_until_complete(self._serve())
@@ -192,9 +204,9 @@ class RelayServer:
header, _ = decode_binary_frame(bytes(raw))
except (ValueError, json.JSONDecodeError):
continue
queue = self._pending_rpc.get(header.get("request_id"))
if queue is not None:
queue.put_nowait(bytes(raw))
pending = self._pending_rpc.get(header.get("request_id"))
if pending is not None:
pending[2].put_nowait(bytes(raw))
continue
try:
envelope = json.loads(raw)
@@ -226,9 +238,9 @@ class RelayServer:
if topic == "relay-http-response":
payload = envelope.get("payload", {})
request_id = payload.get("request_id")
queue = self._pending_rpc.get(request_id)
if queue is not None:
queue.put_nowait(payload)
pending = self._pending_rpc.get(request_id)
if pending is not None:
pending[2].put_nowait(payload)
continue
# Fan out to all other registered peers
@@ -242,6 +254,9 @@ class RelayServer:
finally:
if peer_id:
self._registry.unregister(peer_id)
# Do not leave a requester waiting for its full timeout after
# the selected bridge goes away.
self._fail_pending_for_peer(peer_id)
log.debug("Peer unregistered: %s", peer_id)
async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
@@ -290,18 +305,19 @@ class RelayServer:
except Exception:
return
request_id: str | None = None
requester_request_id: str | None = None
relay_request_id = uuid.uuid4().hex
try:
if isinstance(raw, (bytes, bytearray)):
header, body = decode_binary_frame(bytes(raw))
request_id = str(header.get("request_id") or uuid.uuid4())
header["request_id"] = request_id
requester_request_id = str(header.get("request_id") or uuid.uuid4())
header["request_id"] = relay_request_id
header["target_peer"] = target_peer_id
outbound: str | bytes = encode_binary_frame(header, body)
else:
payload = json.loads(raw)
request_id = str(payload.get("request_id") or uuid.uuid4())
payload["request_id"] = request_id
requester_request_id = str(payload.get("request_id") or uuid.uuid4())
payload["request_id"] = relay_request_id
payload["target_peer"] = target_peer_id
outbound = json.dumps({
"topic": "relay-http-request",
@@ -314,16 +330,18 @@ class RelayServer:
return
queue: asyncio.Queue = asyncio.Queue()
self._pending_rpc[request_id] = queue
overall_timeout = 310.0
idle_timeout = 120.0
self._pending_rpc[relay_request_id] = (
target_peer_id, requester_request_id, queue,
)
overall_timeout = self.rpc_timeout
idle_timeout = self.rpc_idle_timeout
loop = asyncio.get_running_loop()
deadline = loop.time() + overall_timeout
target = self._registry.get(target_peer_id)
try:
if target is None:
await ws_requester.send(json.dumps({
"request_id": request_id,
"request_id": requester_request_id,
"status": 503,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
@@ -339,21 +357,34 @@ class RelayServer:
frame = await asyncio.wait_for(
queue.get(), timeout=min(idle_timeout, remaining)
)
if frame is _RPC_PEER_DISCONNECTED:
raise ConnectionError(f"peer {target_peer_id!r} disconnected")
if isinstance(frame, (bytes, bytearray)):
await ws_requester.send(frame)
header, body = decode_binary_frame(bytes(frame))
header["request_id"] = requester_request_id
await ws_requester.send(encode_binary_frame(header, body))
break
await ws_requester.send(json.dumps(frame))
if not frame.get("stream") or frame.get("done"):
response = dict(frame)
response["request_id"] = requester_request_id
await ws_requester.send(json.dumps(response))
if not response.get("stream") or response.get("done"):
break
except asyncio.TimeoutError:
await ws_requester.send(json.dumps({
"request_id": request_id,
"request_id": requester_request_id,
"status": 504,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({"error": "relay rpc timed out"}),
}))
except ConnectionError as exc:
await ws_requester.send(json.dumps({
"request_id": requester_request_id,
"status": 503,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({"error": str(exc)}),
}))
finally:
self._pending_rpc.pop(request_id, None)
self._pending_rpc.pop(relay_request_id, None)
async def _broadcast(raw: str | bytes, peers: list) -> None:

View File

@@ -0,0 +1,415 @@
"""Tracker-side validation of the capability report a Node presents at registration.
A Node proves locally that it can execute one exact combination — model artifact,
shard range, recipe, backend/device — and ships that proof with its registration
(ADR-0023, NCA-001/002/003). The tracker does not re-run the forward; it decides
whether the presented proof *covers what the node is advertising*, records the
verdict as a small sanitized enum, and routes only to nodes whose verdict is
`admitted`.
Two properties this module deliberately keeps:
* **No model knowledge.** Model ids, recipe ids, backend ids and device names are
opaque labels. They are compared, never interpreted; no vendor string is a
code path here.
* **Evidence, not assertion.** A report is treated as a claim about identity, and
the tracker only ever *narrows* what a node may serve with it. Nothing in a
report can widen a node's eligibility or its routing weight — throughput
routing stays measurement-driven (ADR-0013/0021).
Older nodes that predate the capability protocol present no report at all. They
are handled by an explicit policy (`POLICY_COMPAT` vs `POLICY_ENFORCE`), never by
silently treating "no proof" as "proven" — see `docs/adr/0023-…` for the rollout.
"""
from __future__ import annotations
import os
import re
import time
from dataclasses import dataclass, replace
from typing import Any, Callable, Mapping
# The capability report layout this tracker reads (meshnet_node.capability).
SUPPORTED_SCHEMA_VERSION = 1
# The oldest recipe catalogue whose recipe semantics this tracker still trusts.
# A node carrying an older catalogue may be running a recipe whose id has since
# been redefined, so its proof cannot be matched to a name reliably.
MIN_CATALOGUE_VERSION = "2026.07.1"
# How old a proof may be *at the moment it is presented*. Freshness after that is
# carried by liveness: a registration is re-asserted on tracker restart and the
# node is purged once heartbeats stop.
DEFAULT_MAX_REPORT_AGE_SECONDS = 900.0
# A proof timestamped further ahead than this is not fresh, it is wrong.
MAX_CLOCK_SKEW_SECONDS = 60.0
STATUS_PASSED = "passed"
# --- Admission verdicts. `admitted` is the only routable one under `enforce`. ---
STATE_ADMITTED = "admitted"
STATE_ABSENT = "absent"
STATE_INVALID = "invalid"
STATE_FAILED = "failed"
STATE_STALE = "stale"
STATE_MODEL_MISMATCH = "model-mismatch"
STATE_SHARD_MISMATCH = "shard-mismatch"
STATE_RECIPE_MISMATCH = "recipe-mismatch"
STATE_CATALOGUE_INCOMPATIBLE = "catalogue-incompatible"
ALL_STATES = (
STATE_ADMITTED,
STATE_ABSENT,
STATE_INVALID,
STATE_FAILED,
STATE_STALE,
STATE_MODEL_MISMATCH,
STATE_SHARD_MISMATCH,
STATE_RECIPE_MISMATCH,
STATE_CATALOGUE_INCOMPATIBLE,
)
# --- Compatibility policy for nodes that predate the capability protocol. ---
# `compat` — a node presenting *no* proof still routes (legacy behaviour), but a
# node presenting a *bad* proof never does. Presenting a broken or
# mismatched proof is a stronger signal than presenting none.
# `enforce` — only `admitted` routes. Absent proof is not routable.
POLICY_COMPAT = "compat"
POLICY_ENFORCE = "enforce"
ALL_POLICIES = (POLICY_COMPAT, POLICY_ENFORCE)
DEFAULT_POLICY = POLICY_COMPAT
POLICY_ENV_VAR = "MESHNET_TRACKER_CAPABILITY_POLICY"
# Operator-facing detail strings are short and never carry a raw exception.
_MAX_DETAIL_CHARS = 240
_MAX_DIAGNOSTICS = 3
_CREDENTIAL_PATTERNS = (
re.compile(r"\b[A-Za-z0-9_]{2,6}_[A-Za-z0-9]{16,}\b"), # hf_…, ghp_…, sk_live_…
re.compile(r"\bsk-[A-Za-z0-9_-]{16,}\b"),
re.compile(r"(?i)\bbearer\s+\S+"),
re.compile(r"(?i)\b(?:token|api[_-]?key|password|secret)\s*[=:]\s*\S+"),
)
_REDACTED = "[redacted]"
def normalize_policy(value: Any) -> str:
"""Return a known policy name, falling back to the default for anything else."""
if isinstance(value, str) and value.strip().lower() in ALL_POLICIES:
return value.strip().lower()
return DEFAULT_POLICY
def policy_from_env(environ: Mapping[str, str] | None = None) -> str:
env = os.environ if environ is None else environ
return normalize_policy(env.get(POLICY_ENV_VAR))
def sanitize_detail(text: Any) -> str:
"""Collapse, redact and clip a string bound for an operator view."""
cleaned = " ".join(str(text).split())
for pattern in _CREDENTIAL_PATTERNS:
cleaned = pattern.sub(_REDACTED, cleaned)
if len(cleaned) > _MAX_DETAIL_CHARS:
cleaned = cleaned[: _MAX_DETAIL_CHARS - 1].rstrip() + ""
return cleaned
def catalogue_is_compatible(version: Any) -> bool:
"""True when `version` is at least `MIN_CATALOGUE_VERSION`.
Versions are dotted integer sequences (`2026.07.1`). Anything that does not
parse is incompatible — an unparseable catalogue version cannot be shown to
be new enough.
"""
parsed = _parse_version(version)
if parsed is None:
return False
return parsed >= _parse_version(MIN_CATALOGUE_VERSION) # type: ignore[operator]
def _parse_version(value: Any) -> tuple[int, ...] | None:
if not isinstance(value, str) or not value.strip():
return None
parts = value.strip().split(".")
try:
return tuple(int(part) for part in parts)
except ValueError:
return None
@dataclass(frozen=True)
class CapabilityState:
"""The tracker's sanitized verdict on one node's presented proof.
This is what the network map exposes and what route selection consults. It
holds identity labels and a verdict — never a raw exception, a file path, or
a credential.
"""
state: str
detail: str = ""
model_id: str | None = None
shard_start: int | None = None
shard_end: int | None = None
recipe_id: str | None = None
recipe_version: str | None = None
catalogue_version: str | None = None
backend_id: str | None = None
device: str | None = None
quantization: str | None = None
validated_at: float | None = None
recorded_at: float = 0.0
schema_version: int | None = None
diagnostics: tuple[str, ...] = ()
@property
def proven(self) -> bool:
"""The presented proof covers exactly what the node advertised."""
return self.state == STATE_ADMITTED
def routable_under(self, policy: str) -> bool:
if self.proven:
return True
return self.state == STATE_ABSENT and normalize_policy(policy) == POLICY_COMPAT
def with_state(self, state: str, detail: str) -> CapabilityState:
"""Re-verdict a recorded proof against what the node advertises *now*."""
return replace(self, state=state, detail=sanitize_detail(detail))
def to_dict(self) -> dict:
return {
"state": self.state,
"detail": self.detail,
"model_id": self.model_id,
"shard_start": self.shard_start,
"shard_end": self.shard_end,
"recipe_id": self.recipe_id,
"recipe_version": self.recipe_version,
"catalogue_version": self.catalogue_version,
"backend_id": self.backend_id,
"device": self.device,
"quantization": self.quantization,
"validated_at": self.validated_at,
"recorded_at": self.recorded_at,
"schema_version": self.schema_version,
"diagnostics": list(self.diagnostics),
}
def absent_state(detail: str = "", *, now: float | None = None) -> CapabilityState:
"""The verdict for a node that presented no proof at all (legacy node)."""
return CapabilityState(
state=STATE_ABSENT,
detail=sanitize_detail(
detail
or "node registered without a capability report; it predates the "
"capability protocol or ran with admission disabled"
),
recorded_at=time.time() if now is None else now,
)
def evaluate_report(
report: Any,
*,
model_matches: Callable[[str], bool],
advertised_model: str | None,
shard_start: int | None,
shard_end: int | None,
declared_recipe_id: str | None = None,
declared_recipe_version: str | None = None,
now: float | None = None,
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS,
) -> CapabilityState:
"""Judge the proof a node presented against what that node is advertising.
`model_matches` is the tracker's own alias-aware comparison against the
node's registered model / hf_repo, so an opaque model id never has to be
parsed here.
Returns a verdict for *every* input, including malformed ones: a bad proof is
recorded and shown to the operator rather than dropped, so "why is my node not
routing" has an answer in the network map.
"""
now = time.time() if now is None else now
if report is None:
return absent_state(now=now)
if not isinstance(report, Mapping):
return CapabilityState(
state=STATE_INVALID,
detail=sanitize_detail(
f"capability_report must be a JSON object, got "
f"{type(report).__name__}"
),
recorded_at=now,
)
try:
parsed = _parse_report(report)
except _ReportError as exc:
return CapabilityState(
state=STATE_INVALID,
detail=sanitize_detail(str(exc)),
recorded_at=now,
schema_version=_maybe_int(report.get("schema_version")),
)
status = parsed.pop("_status")
base = CapabilityState(state=STATE_ADMITTED, recorded_at=now, **parsed)
if base.schema_version != SUPPORTED_SCHEMA_VERSION:
return base.with_state(
STATE_INVALID,
f"capability report declares schema version {base.schema_version}; "
f"this tracker reads version {SUPPORTED_SCHEMA_VERSION}",
)
if not catalogue_is_compatible(base.catalogue_version):
return base.with_state(
STATE_CATALOGUE_INCOMPATIBLE,
f"recipe catalogue {base.catalogue_version!r} is older than the "
f"minimum this tracker trusts ({MIN_CATALOGUE_VERSION}); upgrade the node",
)
if not model_matches(base.model_id or ""):
return base.with_state(
STATE_MODEL_MISMATCH,
f"proof is for model {base.model_id!r}, but the node registered "
f"{advertised_model!r}",
)
if shard_start is not None and shard_end is not None:
if (base.shard_start, base.shard_end) != (shard_start, shard_end):
return base.with_state(
STATE_SHARD_MISMATCH,
f"proof is for layers {base.shard_start}{base.shard_end}, but the "
f"node registered layers {shard_start}{shard_end}",
)
if declared_recipe_id is not None and base.recipe_id != declared_recipe_id:
return base.with_state(
STATE_RECIPE_MISMATCH,
f"proof is for recipe {base.recipe_id!r}, but the node declared it "
f"serves with {declared_recipe_id!r}",
)
if (
declared_recipe_version is not None
and base.recipe_version != declared_recipe_version
):
return base.with_state(
STATE_RECIPE_MISMATCH,
f"proof is for recipe {base.recipe_id!r} v{base.recipe_version}, but "
f"the node declared v{declared_recipe_version}",
)
if status != STATUS_PASSED:
return base.with_state(
STATE_FAILED,
f"capability validation {status} on the node"
+ (f"{' '.join(base.diagnostics)}" if base.diagnostics else ""),
)
age = now - (base.validated_at or 0.0)
if age > max_age_seconds:
return base.with_state(
STATE_STALE,
f"proof is {age / 60:.0f} min old (limit {max_age_seconds / 60:.0f} min); "
"the node must re-validate before it can be routed",
)
if age < -MAX_CLOCK_SKEW_SECONDS:
return base.with_state(
STATE_STALE,
f"proof is timestamped {-age:.0f}s in the future; check the node's clock",
)
return base.with_state(
STATE_ADMITTED,
f"{base.model_id} layers {base.shard_start}{base.shard_end} proven on "
f"{base.device} with recipe {base.recipe_id} (v{base.recipe_version})",
)
class _ReportError(ValueError):
"""Malformed report input. Messages name the field, never echo a payload."""
def _parse_report(doc: Mapping[str, Any]) -> dict:
model = _object(doc.get("model"), "model")
shard = _object(doc.get("shard"), "shard")
recipe = _object(doc.get("recipe"), "recipe")
backend = _object(doc.get("backend"), "backend")
validated_at = doc.get("validated_at")
if isinstance(validated_at, bool) or not isinstance(validated_at, (int, float)):
raise _ReportError("'validated_at' must be a Unix timestamp")
schema_version = doc.get("schema_version")
if isinstance(schema_version, bool) or not isinstance(schema_version, int):
raise _ReportError("'schema_version' must be an integer")
return {
"model_id": _text(model.get("model_id"), "model.model_id"),
"shard_start": _index(shard.get("start"), "shard.start"),
"shard_end": _index(shard.get("end"), "shard.end"),
"recipe_id": _text(recipe.get("recipe_id"), "recipe.recipe_id"),
"recipe_version": _text(recipe.get("recipe_version"), "recipe.recipe_version"),
"catalogue_version": _text(
recipe.get("catalogue_version"), "recipe.catalogue_version"
),
"backend_id": _text(backend.get("backend_id"), "backend.backend_id"),
"device": _text(backend.get("device"), "backend.device"),
"quantization": _optional_text(
backend.get("quantization"), "backend.quantization"
),
"validated_at": float(validated_at),
"schema_version": schema_version,
"diagnostics": _diagnostics(doc.get("diagnostics")),
"_status": _text(doc.get("status"), "status"),
}
def _object(value: Any, field_name: str) -> Mapping[str, Any]:
if not isinstance(value, Mapping):
raise _ReportError(f"{field_name!r} must be a JSON object")
return value
def _text(value: Any, field_name: str) -> str:
if not isinstance(value, str) or not value.strip():
raise _ReportError(f"{field_name!r} must be a non-empty string")
return value
def _optional_text(value: Any, field_name: str) -> str | None:
if value is None:
return None
return _text(value, field_name)
def _index(value: Any, field_name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise _ReportError(f"{field_name!r} must be a non-negative integer")
return value
def _maybe_int(value: Any) -> int | None:
if isinstance(value, bool) or not isinstance(value, int):
return None
return value
def _diagnostics(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
out = [
sanitize_detail(item)
for item in value[:_MAX_DIAGNOSTICS]
if isinstance(item, str) and item.strip()
]
return tuple(out)

View File

@@ -9,6 +9,7 @@ from pathlib import Path
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
from .billing import DEFAULT_BILLING_DB_PATH
from .capability import ALL_POLICIES as ALL_CAPABILITY_POLICIES
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
from .logging_setup import (
DEFAULT_LOG_BACKUP_COUNT,
@@ -105,6 +106,18 @@ def main() -> None:
metavar="PATH",
help="SQLite database path for persistent model usage statistics",
)
common.add_argument(
"--capability-policy",
choices=list(ALL_CAPABILITY_POLICIES),
default=None,
help=(
"How to treat nodes that present no capability proof (ADR-0023): "
"'compat' (default) still routes pre-capability nodes; 'enforce' routes "
"only nodes whose proof covers the model and shard they advertise. "
"A broken or mismatched proof is never routed under either policy. "
"Falls back to $MESHNET_TRACKER_CAPABILITY_POLICY when omitted."
),
)
common.add_argument(
"--relay-url",
default=None,
@@ -416,6 +429,7 @@ def main() -> None:
cluster_peers=cluster_peers or None,
cluster_self_url=args.self_url,
stats_db=getattr(args, "stats_db", None),
capability_policy=getattr(args, "capability_policy", None),
relay_url=relay_url,
embedded_relay=args.embedded_relay,
embedded_relay_host=args.relay_host,

View File

@@ -6,10 +6,14 @@ HTTP API contract:
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
"model": str optional, "shard_checksum": str optional,
"hardware_profile": object, "wallet_address": str optional,
"score": number optional
"score": number optional,
"capability_report": object optional, # ADR-0023 proof of this exact shard
"recipe_id": str optional, "recipe_version": str optional
}
Response 200: {"node_id": str}
Response 200: {"node_id": str, "capability": {"state": str, "routable": bool, ...}}
Response 400: {"error": str}
A node whose proof does not admit what it advertises still registers (so the
operator can see why it is dark) but is not routable -- see `capability.py`.
- POST /v1/nodes/{node_id}/heartbeat
Response 200: {}
Response 404: {"error": "node not found"}
@@ -48,6 +52,20 @@ from typing import Any
from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore
from .auth import is_validator_token, sign_hive_request, verify_hive_request
from .capability import (
DEFAULT_POLICY as DEFAULT_CAPABILITY_POLICY,
POLICY_COMPAT,
POLICY_ENFORCE,
STATE_ABSENT,
STATE_ADMITTED,
STATE_MODEL_MISMATCH,
STATE_SHARD_MISMATCH,
CapabilityState,
absent_state,
evaluate_report,
normalize_policy,
policy_from_env,
)
from .wallet_proof import binding_message, verify_wallet_signature
from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger
from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
@@ -587,6 +605,8 @@ class _NodeEntry:
"heartbeats_expected", "heartbeats_received",
# dynamic reassignment queued by the tracker
"pending_new_assignment",
# the tracker's verdict on the capability proof this node presented
"capability",
)
def __init__(
@@ -616,6 +636,7 @@ class _NodeEntry:
cert_fingerprint: str | None = None,
peer_id: str | None = None,
friendly_name: str | None = None,
capability: "CapabilityState | None" = None,
) -> None:
self.node_id = node_id
self.endpoint = endpoint
@@ -643,6 +664,9 @@ class _NodeEntry:
self.cert_fingerprint = cert_fingerprint
self.peer_id = peer_id
self.friendly_name = friendly_name
# No proof presented is `absent`, never `admitted` — a node can only earn
# `admitted` by presenting a report that covers what it advertises.
self.capability: CapabilityState = capability or absent_state()
self.pending_directives: list[dict] = []
self.last_heartbeat: float = time.monotonic()
self.total_requests: int = 0
@@ -731,12 +755,88 @@ def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float:
return 0.5 + 0.5 * reputation
def _node_admission(node: "_NodeEntry") -> CapabilityState:
"""The node's recorded verdict, re-checked against what it advertises *now*.
A proof covers one model/shard combination. The tracker can move a node to a
different range after it registered (rebalance, `pending_new_assignment`), and
the proof does not travel with it: until the node re-registers with a proof
for the range it now advertises, it is shard-mismatched and not routable.
"""
state = node.capability
if not state.proven:
return state
if state.model_id and not _node_matches_model(node, state.model_id):
return state.with_state(
STATE_MODEL_MISMATCH,
f"proof is for model {state.model_id!r}, but the node now serves "
f"{node.hf_repo or node.model!r}",
)
if (
node.shard_start is not None
and node.shard_end is not None
and (state.shard_start, state.shard_end) != (node.shard_start, node.shard_end)
):
return state.with_state(
STATE_SHARD_MISMATCH,
f"proof is for layers {state.shard_start}{state.shard_end}, but the "
f"node now serves layers {node.shard_start}{node.shard_end}",
)
return state
def _capability_from_registration(
payload: dict,
*,
model: str | None,
hf_repo: str | None,
shard_start: int | None,
shard_end: int | None,
) -> CapabilityState:
"""The tracker's verdict on the proof carried by one registration payload.
Used by the register handler and by the Raft follower path, so a replicated
registration lands with the same verdict as the one the leader recorded.
"""
aliases = _model_aliases(model) | _model_aliases(hf_repo)
recipe_id = payload.get("recipe_id")
recipe_version = payload.get("recipe_version")
return evaluate_report(
payload.get("capability_report"),
model_matches=lambda reported: bool(_model_aliases(reported) & aliases),
advertised_model=hf_repo or model,
shard_start=shard_start,
shard_end=shard_end,
declared_recipe_id=recipe_id if isinstance(recipe_id, str) else None,
declared_recipe_version=(
recipe_version if isinstance(recipe_version, str) else None
),
)
def _capability_routable(node: "_NodeEntry", policy: str) -> bool:
"""May this node carry traffic under the tracker's capability policy?"""
return _node_admission(node).routable_under(policy)
def _admitted_nodes(nodes: list["_NodeEntry"], policy: str | None) -> list["_NodeEntry"]:
"""Drop every candidate whose capability proof does not admit it (ADR-0023).
This is the single gate every route path goes through. It removes candidates;
it never reorders or reweights them, so coverage-first selection and
throughput-weighted preference among the survivors are unchanged.
"""
effective = normalize_policy(policy)
return [node for node in nodes if _capability_routable(node, effective)]
def _select_route(
nodes: list[_NodeEntry],
required_start: int,
required_end: int,
model: str | None = None,
contracts: Any | None = None,
policy: str | None = None,
) -> tuple[list[_NodeEntry], str]:
"""Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes.
@@ -747,7 +847,10 @@ def _select_route(
Tiebreak: higher shard_end (fewer hops).
"""
candidates = sorted(
[node for node in nodes if node.shard_start is not None and node.shard_end is not None],
[
node for node in _admitted_nodes(nodes, policy)
if node.shard_start is not None and node.shard_end is not None
],
key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator]
)
route: list[_NodeEntry] = []
@@ -783,6 +886,7 @@ def _enumerate_routes(
model: str | None = None,
contracts: Any | None = None,
max_candidates: int = 8,
policy: str | None = None,
) -> list["RouteCandidate"]:
"""Enumerate viable route candidates for bandit selection (ADR-0021).
@@ -791,9 +895,13 @@ def _enumerate_routes(
with the longest-advancing hops. The route's prior throughput estimate is
its bottleneck hop's queue-adjusted effective throughput — used only until
observed route samples exist.
Candidates that are not admitted by their capability proof are dropped before
enumeration, so an unproven node cannot appear in any route candidate — not
even as a scouted one.
"""
sharded = [
n for n in nodes
n for n in _admitted_nodes(nodes, policy)
if n.shard_start is not None and n.shard_end is not None
]
# Heads must start the pipeline at the first required layer (they tokenize
@@ -2675,9 +2783,13 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
route_stats: "RouteStatsStore | None" = None,
relay_status: dict | None = None,
test_runner: "TestRunManager | None" = None,
capability_policy: str | None = None,
) -> None:
super().__init__(addr, handler)
self.registry = registry
self.capability_policy = normalize_policy(
capability_policy if capability_policy is not None else policy_from_env()
)
self.lock = lock
self.heartbeat_timeout = heartbeat_timeout
self.model_presets = model_presets
@@ -3177,9 +3289,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
def model_supply_for(node: _NodeEntry) -> dict:
return _model_health_summary(server, node.model, node.hf_repo)
def capability_for(node: _NodeEntry) -> dict:
# Re-verdicted against what the node advertises now, so the operator
# view and the routing gate can never disagree.
state = _node_admission(node)
return {
**state.to_dict(),
"routable": state.routable_under(server.capability_policy),
}
self._send_json(200, {
"relay_url": server.relay_url,
"relay": dict(server.relay_status),
"capability_policy": server.capability_policy,
"pool": _pool_summary(nodes),
"memory_pool": memory_pool,
"recommended_models": [
@@ -3213,6 +3335,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"capacity": capacity_for(node),
"model_supply": model_supply_for(node),
"throughput": throughput_for(node),
"capability": capability_for(node),
"stats": _node_health(node, server.heartbeat_timeout),
}
for node in nodes
@@ -3365,6 +3488,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if n.tracker_mode
and _node_matches_model(n, model)
and _quantization_satisfies(n.quantization, requested_quantization)
and _capability_routable(n, server.capability_policy)
]
if not candidates:
@@ -3375,6 +3499,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if n.shard_start == 0
and _node_matches_model(n, model)
and _quantization_satisfies(n.quantization, requested_quantization)
and _capability_routable(n, server.capability_policy)
]
if not candidates:
@@ -3480,6 +3605,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model=route_model,
contracts=server.contracts,
max_candidates=server.route_stats.config.max_candidate_routes,
policy=server.capability_policy,
)
picked, routing_decision = choose_route(
route_candidates, server.route_stats, route_model, rng=server.route_rng,
@@ -3489,7 +3615,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
else:
# No head-anchored candidate — legacy greedy cover as fallback
# (also produces the layer-gap error message).
route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts)
route_nodes, route_error = _select_route(
all_nodes, rs, re,
model=route_model,
contracts=server.contracts,
policy=server.capability_policy,
)
routing_decision = {"mode": "greedy-fallback"}
if route_error:
_tracker_log(
@@ -4344,6 +4475,25 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._send_json(400, {"error": str(exc)})
return
recipe_id = body.get("recipe_id")
recipe_version = body.get("recipe_version")
if (recipe_id is not None and not isinstance(recipe_id, str)) or (
recipe_version is not None and not isinstance(recipe_version, str)
):
self._send_json(400, {"error": "recipe_id and recipe_version must be strings"})
return
# The capability proof (ADR-0023). A bad proof does not fail registration --
# the node is recorded with its verdict so the operator can see *why* it is
# not routing -- but only an `admitted` verdict makes it routable.
capability = _capability_from_registration(
body,
model=model,
hf_repo=hf_repo,
shard_start=shard_start,
shard_end=shard_end,
)
node_id = _node_id_for_registration(
endpoint,
model,
@@ -4378,6 +4528,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
cert_fingerprint=cert_fingerprint,
peer_id=peer_id,
friendly_name=friendly_name,
capability=capability,
)
with server.lock:
self._purge_expired_nodes()
@@ -4433,6 +4584,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
old_model_health=_model_health_summary(server, old.model, old.hf_repo),
model_health=model_health,
)
routable = _capability_routable(entry, server.capability_policy)
_tracker_log(
server,
"info",
@@ -4444,7 +4596,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
shard=f"{entry.shard_start}-{entry.shard_end}",
tracker_mode=entry.tracker_mode,
model_health=model_health,
capability=entry.capability.state,
routable=routable,
)
if not routable:
_tracker_log(
server,
"warn",
"node registered but is not routable",
node_id=node_id,
endpoint=entry.endpoint,
capability=entry.capability.state,
detail=entry.capability.detail,
)
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
repo_info = f" [{hf_repo}]" if hf_repo else ""
@@ -4456,7 +4620,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
flush=True,
)
payload = {"node_id": node_id}
payload = {
"node_id": node_id,
# Tell the node what the tracker made of its proof: a node that is
# registered but not routable must be able to see that it is dark.
"capability": {
"state": entry.capability.state,
"detail": entry.capability.detail,
"routable": routable,
"policy": server.capability_policy,
},
}
if assignment_directive is not None:
payload["directive"] = assignment_directive
self._send_json(200, payload)
@@ -5989,6 +6163,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model=route_model,
contracts=server.contracts,
max_candidates=server.route_stats.config.max_candidate_routes,
policy=server.capability_policy,
)
if candidates:
# Prefer a distributed multi-hop route when available. Greedy
@@ -5998,7 +6173,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
error = ""
else:
route, error = _select_route(
alive, required_start, required_end, contracts=server.contracts,
alive, required_start, required_end,
contracts=server.contracts,
policy=server.capability_policy,
)
if error:
_tracker_log(
@@ -6084,6 +6261,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model=stats_key,
contracts=server.contracts,
max_candidates=cfg.max_candidate_routes,
policy=server.capability_policy,
)
out[stats_key] = {
"epoch": server.route_stats.epoch(stats_key),
@@ -6177,7 +6355,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
routes = []
remaining = list(candidates)
for _ in range(redundancy):
route, error = _select_route(remaining, required_start, required_end, contracts=server.contracts)
route, error = _select_route(
remaining, required_start, required_end,
contracts=server.contracts,
policy=server.capability_policy,
)
if error:
self._send_json(503, {"error": error})
return
@@ -6262,7 +6444,11 @@ class TrackerServer:
routing_config: RoutingConfig | None = None,
enable_test_runner: bool = False,
test_runner: TestRunManager | None = None,
capability_policy: str | None = None,
) -> None:
self._capability_policy = normalize_policy(
capability_policy if capability_policy is not None else policy_from_env()
)
self._host = host
self._requested_port = port
self._heartbeat_timeout = heartbeat_timeout
@@ -6475,6 +6661,7 @@ class TrackerServer:
route_stats=self._route_stats,
relay_status=http_relay_status,
test_runner=self._test_runner,
capability_policy=self._capability_policy,
)
self.port = self._server.server_address[1]
@@ -6708,6 +6895,15 @@ class TrackerServer:
else None
),
friendly_name=_normalize_friendly_name(payload.get("friendly_name")),
# A replicated registration carries its proof: without this, a proven
# node would be routable on the leader and dark on every follower.
capability=_capability_from_registration(
payload,
model=payload.get("model", "stub-model"),
hf_repo=payload.get("hf_repo"),
shard_start=shard_start,
shard_end=shard_end,
),
)
with self._lock:
self._registry[node_id] = entry