[verified] feat: complete Ralph task workstreams
This commit is contained in:
130
packages/gateway/meshnet_gateway/prefill_backpressure.py
Normal file
130
packages/gateway/meshnet_gateway/prefill_backpressure.py
Normal 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
|
||||
@@ -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:
|
||||
|
||||
142
packages/node/meshnet_node/activation_compression.py
Normal file
142
packages/node/meshnet_node/activation_compression.py
Normal 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"}
|
||||
225
packages/node/meshnet_node/admission.py
Normal file
225
packages/node/meshnet_node/admission.py
Normal 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
|
||||
@@ -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:
|
||||
|
||||
633
packages/node/meshnet_node/doctor.py
Normal file
633
packages/node/meshnet_node/doctor.py
Normal 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)
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
385
packages/node/meshnet_node/route_session_benchmark.py
Normal file
385
packages/node/meshnet_node/route_session_benchmark.py
Normal 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())
|
||||
155
packages/node/meshnet_node/seam_telemetry.py
Normal file
155
packages/node/meshnet_node/seam_telemetry.py
Normal 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
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
70
packages/node/meshnet_node/testing.py
Normal file
70
packages/node/meshnet_node/testing.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
415
packages/tracker/meshnet_tracker/capability.py
Normal file
415
packages/tracker/meshnet_tracker/capability.py
Normal 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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user