feat: checkpoint distributed gguf runtime stories
This commit is contained in:
@@ -20,9 +20,17 @@ import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
from .capability import CapabilityReport
|
||||
from . import __version__ as _PACKAGE_VERSION
|
||||
from .capability import CapabilityReport, config_fingerprint
|
||||
from .doctor import DoctorSelection
|
||||
from .recipe_manifest import Recipe, RecipeManifest
|
||||
from .runtime_recipe import (
|
||||
build_artifact_identity,
|
||||
build_runtime_recipe_identity,
|
||||
compatibility_fingerprint,
|
||||
fingerprint_payload,
|
||||
)
|
||||
from .gguf_ownership import authoritative_dense_llama_ownership
|
||||
|
||||
# How long a passing report stays usable. Startup normally validates in-process
|
||||
# (age ≈ 0); this bounds how far a report written by an earlier `doctor` run can
|
||||
@@ -39,6 +47,7 @@ REASON_MODEL_MISMATCH = "model-mismatch"
|
||||
REASON_SHARD_MISMATCH = "shard-mismatch"
|
||||
REASON_RECIPE_MISMATCH = "recipe-mismatch"
|
||||
REASON_BACKEND_MISMATCH = "backend-mismatch"
|
||||
REASON_COMPATIBILITY_MISMATCH = "compatibility-mismatch"
|
||||
|
||||
|
||||
class CapabilityAdmissionError(RuntimeError):
|
||||
@@ -77,6 +86,7 @@ class AdmissionRequirement:
|
||||
recipe_version: str
|
||||
backend_id: str
|
||||
device: str
|
||||
compatibility_fingerprint: str
|
||||
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS
|
||||
|
||||
@classmethod
|
||||
@@ -94,6 +104,9 @@ class AdmissionRequirement:
|
||||
recipe_version=context.recipe.version,
|
||||
backend_id=context.recipe.backend_id,
|
||||
device=context.device,
|
||||
compatibility_fingerprint=_compatibility_fingerprint_for_context(
|
||||
context
|
||||
),
|
||||
max_age_seconds=max_age_seconds,
|
||||
)
|
||||
|
||||
@@ -165,6 +178,16 @@ def admit(
|
||||
f"{requirement.backend_id} on {requirement.device}",
|
||||
)
|
||||
|
||||
if report.compatibility_fingerprint != requirement.compatibility_fingerprint:
|
||||
raise CapabilityAdmissionError(
|
||||
REASON_COMPATIBILITY_MISMATCH,
|
||||
f"capability proof fingerprint {report.compatibility_fingerprint!r} "
|
||||
f"does not match the expected compatibility fingerprint for "
|
||||
f"{requirement.model_id} {requirement.shard_label}; the artifact, "
|
||||
f"tokenizer, architecture, boundary schema, activation recipe or "
|
||||
f"cache layout differs",
|
||||
)
|
||||
|
||||
if not report.passed:
|
||||
raise CapabilityAdmissionError(
|
||||
REASON_NOT_PASSED,
|
||||
@@ -223,3 +246,157 @@ def probe_capability(context: CapabilityContext) -> CapabilityReport:
|
||||
context.recipe,
|
||||
context.manifest,
|
||||
).report
|
||||
|
||||
|
||||
def _compatibility_fingerprint_for_context(context: CapabilityContext) -> str:
|
||||
backend = context.backend
|
||||
selection = context.selection
|
||||
recipe = context.recipe
|
||||
model_config = getattr(getattr(backend, "model", None), "config", None)
|
||||
model_config_payload = (
|
||||
model_config.to_dict() if hasattr(model_config, "to_dict") else model_config
|
||||
)
|
||||
runtime_versions = _runtime_versions()
|
||||
runtime_version = _PACKAGE_VERSION
|
||||
ownership = authoritative_dense_llama_ownership(backend, selection)
|
||||
artifact = build_artifact_identity(
|
||||
model_id=selection.model_id,
|
||||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||||
model_config=model_config_payload,
|
||||
shard_start=ownership.start_layer,
|
||||
shard_end=ownership.end_layer,
|
||||
)
|
||||
runtime_recipe = build_runtime_recipe_identity(
|
||||
model_id=selection.model_id,
|
||||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||||
model_config=model_config_payload,
|
||||
recipe_params=recipe.params,
|
||||
weight_quantization=selection.quantization,
|
||||
backend_id=recipe.backend_id,
|
||||
runtime_version=runtime_version,
|
||||
activation_dtype="bfloat16",
|
||||
compute_dtype=_backend_compute_dtype(backend),
|
||||
kv_dtype=_backend_kv_dtype(backend),
|
||||
kv_layout=_backend_kv_layout(backend),
|
||||
tokenizer_revision=_backend_tokenizer_revision(backend, selection),
|
||||
architecture_adapter=_backend_architecture_adapter(backend, recipe.backend_id),
|
||||
boundary_schema_version=1,
|
||||
cache_layout=_backend_cache_layout(backend, recipe.params),
|
||||
)
|
||||
return compatibility_fingerprint(
|
||||
fingerprint_payload(
|
||||
model={
|
||||
"model_id": selection.model_id,
|
||||
"revision": getattr(getattr(backend, "model", None), "revision", None),
|
||||
"config_fingerprint": config_fingerprint(model_config_payload),
|
||||
},
|
||||
shard={
|
||||
"start": ownership.start_layer,
|
||||
"end": ownership.end_layer,
|
||||
"owns_embedding": ownership.owns_embedding,
|
||||
"owns_final_head": ownership.owns_final_head,
|
||||
},
|
||||
recipe={
|
||||
"recipe_id": recipe.id,
|
||||
"recipe_version": recipe.version,
|
||||
"catalogue_version": context.manifest.catalogue_version,
|
||||
},
|
||||
backend={
|
||||
"backend_id": recipe.backend_id,
|
||||
"device": context.device,
|
||||
"device_name": _backend_device_name(context.device),
|
||||
"quantization": selection.quantization,
|
||||
"runtime": runtime_versions,
|
||||
},
|
||||
artifact=artifact.to_dict(),
|
||||
runtime_recipe=runtime_recipe.to_dict(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _runtime_versions() -> dict[str, str]:
|
||||
versions: dict[str, str] = {}
|
||||
for name in ("torch", "transformers"):
|
||||
try:
|
||||
module = __import__(name)
|
||||
except Exception:
|
||||
continue
|
||||
version = getattr(module, "__version__", None)
|
||||
if version:
|
||||
versions[name] = str(version)
|
||||
return versions
|
||||
|
||||
|
||||
def _backend_compute_dtype(backend: Any) -> str:
|
||||
config = getattr(getattr(backend, "model", None), "config", None)
|
||||
for candidate in (config, getattr(config, "text_config", None)):
|
||||
if candidate is None:
|
||||
continue
|
||||
for attr in ("dtype", "torch_dtype"):
|
||||
value = getattr(candidate, attr, None)
|
||||
if value is None:
|
||||
continue
|
||||
return str(value).removeprefix("torch.")
|
||||
return "bfloat16"
|
||||
|
||||
|
||||
def _backend_kv_dtype(backend: Any) -> str:
|
||||
return _backend_compute_dtype(backend)
|
||||
|
||||
|
||||
def _backend_kv_layout(backend: Any) -> str:
|
||||
return "session-cache" if getattr(backend, "supports_kv_cache", False) else "stateless"
|
||||
|
||||
|
||||
def _backend_tokenizer_revision(backend: Any, selection: DoctorSelection) -> str:
|
||||
model = getattr(backend, "model", None)
|
||||
revision = getattr(model, "revision", None)
|
||||
if isinstance(revision, str) and revision.strip():
|
||||
return revision
|
||||
tokenizer = getattr(backend, "tokenizer", None)
|
||||
for attr in ("revision", "model_id"):
|
||||
value = getattr(tokenizer, attr, None)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
return selection.model_id
|
||||
|
||||
|
||||
def _backend_architecture_adapter(backend: Any, default: str) -> str:
|
||||
config = getattr(getattr(backend, "model", None), "config", None)
|
||||
for candidate in (config, getattr(config, "text_config", None)):
|
||||
if candidate is None:
|
||||
continue
|
||||
for attr in ("architecture_adapter", "model_type"):
|
||||
value = getattr(candidate, attr, None)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
architectures = getattr(candidate, "architectures", None)
|
||||
if isinstance(architectures, (list, tuple)) and architectures:
|
||||
first = architectures[0]
|
||||
if isinstance(first, str) and first.strip():
|
||||
return first
|
||||
return default
|
||||
|
||||
|
||||
def _backend_device_name(device: str) -> str | None:
|
||||
if device != "cuda":
|
||||
return None
|
||||
from .hardware import detect_hardware
|
||||
|
||||
try:
|
||||
return detect_hardware().get("gpu_name") or None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _backend_cache_layout(backend: Any, recipe_params: dict[str, Any] | None) -> str:
|
||||
if getattr(backend, "supports_kv_cache", False) is False:
|
||||
return "stateless"
|
||||
if recipe_params is None:
|
||||
return "local-hot-kv"
|
||||
if recipe_params.get("use_cache") is False:
|
||||
return "stateless"
|
||||
value = recipe_params.get("cache_layout")
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
return "local-hot-kv"
|
||||
|
||||
484
packages/node/meshnet_node/boundary_adapter.py
Normal file
484
packages/node/meshnet_node/boundary_adapter.py
Normal file
@@ -0,0 +1,484 @@
|
||||
"""Architecture-defined boundary input/output for distributed Shards (DGR-006).
|
||||
|
||||
A public-network Shard is a contiguous range of transformer layers (RALPH runtime
|
||||
decision #1). For disjoint processes to reproduce whole-model execution, every
|
||||
Shard must agree on *exactly* what boundary state it consumes and emits:
|
||||
|
||||
* The **head** owns token embedding: it accepts token IDs and turns them into the
|
||||
residual stream. No other Shard may embed tokens.
|
||||
* **Middle and tail** Shards bypass token embedding entirely; they accept the named
|
||||
boundary bundle (the residual stream handed over by the previous range).
|
||||
* A **non-tail** Shard emits the *unnormalized* architecture-defined residual /
|
||||
boundary — before the final norm, before the LM head, and before any tail-only
|
||||
row pruning — so the next range sees precisely the state the whole model would
|
||||
have at that layer index.
|
||||
* The **tail** owns the final norm + LM head and turns the residual into logits or
|
||||
a sampled token through an explicit sampling contract.
|
||||
|
||||
This module is deliberately backend-agnostic. It enforces the boundary *contract*
|
||||
and defers the arithmetic to a ``ShardComputation`` (a duck-typed object exposing
|
||||
``embed_tokens`` / ``run_layers`` / ``final_norm`` / ``lm_head``). The pinned
|
||||
llama.cpp worker (DGR-008) and the reference PyTorch backend both satisfy that
|
||||
protocol, and the numpy reference model in the tests proves whole-model versus
|
||||
two-range parity without any download, GPU, or API credit.
|
||||
|
||||
The adapter **fails closed** for uncertified architectures: only architectures
|
||||
that have passed real certification (dense Llama-family first, per RALPH runtime
|
||||
decision #13) are accepted. Everything else raises rather than silently guessing a
|
||||
tensor layout — Qwen3/Qwen3-MoE stays registered-but-dark until DGR-015 certifies
|
||||
its own adapter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
# The boundary bundle wire schema version. This is the ``boundary_schema_version``
|
||||
# carried by ``runtime_recipe.RuntimeRecipeIdentity``; a receiver refuses a bundle
|
||||
# whose schema it does not implement (forward/backward compatibility is a routing
|
||||
# concern, not a silent reinterpretation).
|
||||
BOUNDARY_SCHEMA_VERSION = 1
|
||||
|
||||
|
||||
class BoundaryAdapterError(RuntimeError):
|
||||
"""Base class for boundary-contract violations."""
|
||||
|
||||
|
||||
class UncertifiedArchitectureError(BoundaryAdapterError):
|
||||
"""Raised when a boundary adapter is requested for an uncertified architecture.
|
||||
|
||||
Failing closed here is a safety property: an unknown architecture has an
|
||||
unknown tensor layout, so guessing where the residual boundary lives would
|
||||
silently corrupt distributed output. The architecture must pass real
|
||||
certification first.
|
||||
"""
|
||||
|
||||
|
||||
class BoundaryContractError(BoundaryAdapterError):
|
||||
"""Raised when a Shard is fed the wrong boundary input for its role.
|
||||
|
||||
Examples: a head handed a residual bundle instead of token IDs, a middle
|
||||
Shard handed token IDs it must not embed, or a boundary bundle whose
|
||||
architecture / schema / seam layer does not match the receiving range.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArchitectureBoundary:
|
||||
"""The architecture-defined boundary description for one certified adapter.
|
||||
|
||||
These fields are what makes the boundary *architecture-defined* rather than a
|
||||
hardcoded assumption: the residual tensor name, whether the tail normalizes
|
||||
before the LM head, and whether row pruning is a tail-only concern all come
|
||||
from here.
|
||||
"""
|
||||
|
||||
adapter: str
|
||||
boundary_tensor_name: str
|
||||
boundary_schema_version: int
|
||||
normalizes_before_head: bool
|
||||
prunes_rows_at_tail: bool
|
||||
|
||||
|
||||
# Certified architectures only. Dense Llama-family is first (RALPH runtime decision
|
||||
# #13 + native discipline). Aliases map the many spellings a runtime recipe /
|
||||
# GGUF / HF config may use onto the single canonical adapter id. Anything not in
|
||||
# this table fails closed.
|
||||
_DENSE_LLAMA = ArchitectureBoundary(
|
||||
adapter="dense-llama",
|
||||
boundary_tensor_name="residual_stream",
|
||||
boundary_schema_version=BOUNDARY_SCHEMA_VERSION,
|
||||
normalizes_before_head=True,
|
||||
prunes_rows_at_tail=True,
|
||||
)
|
||||
|
||||
_CERTIFIED_ARCHITECTURES: dict[str, ArchitectureBoundary] = {
|
||||
"dense-llama": _DENSE_LLAMA,
|
||||
"dense_llama": _DENSE_LLAMA,
|
||||
"llama": _DENSE_LLAMA,
|
||||
"llamaforcausallm": _DENSE_LLAMA,
|
||||
"llamamodel": _DENSE_LLAMA,
|
||||
}
|
||||
|
||||
|
||||
def certified_architecture(name: Any) -> ArchitectureBoundary:
|
||||
"""Return the certified boundary description for ``name`` or fail closed.
|
||||
|
||||
``name`` may be the canonical adapter id (``dense-llama``), an HF architecture
|
||||
class (``LlamaForCausalLM``), or a GGUF/config ``model_type`` (``llama``).
|
||||
Uncertified architectures raise ``UncertifiedArchitectureError``.
|
||||
"""
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
raise UncertifiedArchitectureError(
|
||||
"architecture adapter must be a non-empty string; "
|
||||
"the boundary adapter refuses to guess a tensor layout"
|
||||
)
|
||||
key = name.strip().lower()
|
||||
boundary = _CERTIFIED_ARCHITECTURES.get(key)
|
||||
if boundary is None:
|
||||
raise UncertifiedArchitectureError(
|
||||
f"architecture {name!r} is not certified for the boundary adapter; "
|
||||
f"certified adapters: {sorted(set(v.adapter for v in _CERTIFIED_ARCHITECTURES.values()))}. "
|
||||
"Uncertified architectures stay registered-but-dark until real "
|
||||
"certification passes."
|
||||
)
|
||||
return boundary
|
||||
|
||||
|
||||
def is_certified_architecture(name: Any) -> bool:
|
||||
"""Return True when ``name`` maps to a certified boundary adapter."""
|
||||
try:
|
||||
certified_architecture(name)
|
||||
except UncertifiedArchitectureError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class ShardRole(str, Enum):
|
||||
"""Where a contiguous layer range sits in the whole model."""
|
||||
|
||||
HEAD = "head"
|
||||
MIDDLE = "middle"
|
||||
TAIL = "tail"
|
||||
FULL = "full"
|
||||
|
||||
@property
|
||||
def owns_embedding(self) -> bool:
|
||||
return self in (ShardRole.HEAD, ShardRole.FULL)
|
||||
|
||||
@property
|
||||
def owns_final_head(self) -> bool:
|
||||
return self in (ShardRole.TAIL, ShardRole.FULL)
|
||||
|
||||
|
||||
def role_for_range(start_layer: int, end_layer: int, total_layers: int) -> ShardRole:
|
||||
"""Classify a contiguous inclusive layer range within a model of ``total_layers``."""
|
||||
if total_layers <= 0:
|
||||
raise ValueError("total_layers must be positive")
|
||||
if start_layer < 0 or end_layer < start_layer:
|
||||
raise ValueError("require 0 <= start_layer <= end_layer")
|
||||
if end_layer > total_layers - 1:
|
||||
raise ValueError(
|
||||
f"end_layer {end_layer} exceeds last layer index {total_layers - 1}"
|
||||
)
|
||||
is_head = start_layer == 0
|
||||
is_tail = end_layer >= total_layers - 1
|
||||
if is_head and is_tail:
|
||||
return ShardRole.FULL
|
||||
if is_head:
|
||||
return ShardRole.HEAD
|
||||
if is_tail:
|
||||
return ShardRole.TAIL
|
||||
return ShardRole.MIDDLE
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BoundaryBundle:
|
||||
"""The versioned named-tensor bundle handed between adjacent Shard ranges.
|
||||
|
||||
``residual`` is the *unnormalized* architecture-defined residual stream with
|
||||
every position row intact (no tail-only pruning). ``next_layer`` is the layer
|
||||
index the receiving range must start at — it is the overlap-safe effective
|
||||
start of the seam, so a receiver can reject a bundle meant for a different cut.
|
||||
"""
|
||||
|
||||
architecture_adapter: str
|
||||
schema_version: int
|
||||
tensor_name: str
|
||||
residual: np.ndarray
|
||||
positions: np.ndarray
|
||||
next_layer: int
|
||||
normalized: bool = False
|
||||
|
||||
def named_tensor_fields(self) -> dict[str, Any]:
|
||||
"""Return the wire-shaped description of the residual tensor.
|
||||
|
||||
These are exactly the fields the DGR-002 ``NamedTensor`` carries (name,
|
||||
shape, dtype, byte order, raw bytes), so a worker can serialize this
|
||||
bundle into the gRPC protobuf without re-deriving them.
|
||||
"""
|
||||
residual = np.ascontiguousarray(self.residual)
|
||||
return {
|
||||
"name": self.tensor_name,
|
||||
"shape": list(residual.shape),
|
||||
"dtype": residual.dtype.name,
|
||||
"byte_order": _byte_order(residual.dtype),
|
||||
"data": residual.tobytes(),
|
||||
}
|
||||
|
||||
def pack(self) -> dict[str, Any]:
|
||||
"""Serialize the bundle to a transport-agnostic dict (proves the seam).
|
||||
|
||||
The residual and positions are carried as raw little/big-endian bytes plus
|
||||
shape/dtype so that a truly disjoint process can reconstruct the exact
|
||||
array — this is what lets two OS processes reproduce whole-model math.
|
||||
"""
|
||||
residual = np.ascontiguousarray(self.residual)
|
||||
positions = np.ascontiguousarray(self.positions)
|
||||
return {
|
||||
"architecture_adapter": self.architecture_adapter,
|
||||
"schema_version": self.schema_version,
|
||||
"tensor_name": self.tensor_name,
|
||||
"next_layer": self.next_layer,
|
||||
"normalized": self.normalized,
|
||||
"residual": {
|
||||
"shape": list(residual.shape),
|
||||
"dtype": residual.dtype.str,
|
||||
"data": residual.tobytes(),
|
||||
},
|
||||
"positions": {
|
||||
"shape": list(positions.shape),
|
||||
"dtype": positions.dtype.str,
|
||||
"data": positions.tobytes(),
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def unpack(cls, payload: dict[str, Any]) -> "BoundaryBundle":
|
||||
"""Reconstruct a bundle produced by :meth:`pack`."""
|
||||
residual = _array_from_wire(payload["residual"])
|
||||
positions = _array_from_wire(payload["positions"])
|
||||
return cls(
|
||||
architecture_adapter=payload["architecture_adapter"],
|
||||
schema_version=int(payload["schema_version"]),
|
||||
tensor_name=payload["tensor_name"],
|
||||
residual=residual,
|
||||
positions=positions,
|
||||
next_layer=int(payload["next_layer"]),
|
||||
normalized=bool(payload.get("normalized", False)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SamplingContract:
|
||||
"""Explicit contract for turning tail logits into a token.
|
||||
|
||||
The tail never hides the sampling decision inside the adapter: the contract is
|
||||
a first-class value so the head/route can reproduce it and so greedy decoding
|
||||
is deterministic by construction. Only greedy is certified here; temperature /
|
||||
top-p are declared but must be requested explicitly and are out of scope for
|
||||
the deterministic parity gate.
|
||||
"""
|
||||
|
||||
mode: str = "greedy"
|
||||
temperature: float = 1.0
|
||||
top_p: float = 1.0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mode not in ("greedy",):
|
||||
raise BoundaryContractError(
|
||||
f"sampling mode {self.mode!r} is not certified; only 'greedy' is "
|
||||
"deterministic and supported by the boundary adapter today"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def greedy(cls) -> "SamplingContract":
|
||||
return cls(mode="greedy")
|
||||
|
||||
def sample(self, last_logits: np.ndarray) -> int:
|
||||
"""Return the next token id from the final-position logits row."""
|
||||
logits = np.asarray(last_logits)
|
||||
if logits.ndim == 2:
|
||||
# (batch, vocab) — parity harness uses batch size 1.
|
||||
logits = logits[0]
|
||||
if logits.ndim != 1:
|
||||
raise BoundaryContractError(
|
||||
"sampling expects the pruned final-position logits row"
|
||||
)
|
||||
return int(np.argmax(logits))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TailOutput:
|
||||
"""What a tail Shard emits: the sampled token plus the pruned logits row."""
|
||||
|
||||
token_id: int
|
||||
logits: np.ndarray
|
||||
sampling: SamplingContract
|
||||
|
||||
|
||||
@dataclass
|
||||
class BoundaryAdapter:
|
||||
"""Enforces the architecture-defined boundary contract for one Shard range.
|
||||
|
||||
Construction fails closed for uncertified architectures. The adapter derives
|
||||
the Shard's role from its range and drives a duck-typed ``ShardComputation``.
|
||||
"""
|
||||
|
||||
computation: Any
|
||||
sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
|
||||
architecture: ArchitectureBoundary = field(init=False)
|
||||
role: ShardRole = field(init=False)
|
||||
start_layer: int = field(init=False)
|
||||
end_layer: int = field(init=False)
|
||||
total_layers: int = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
arch_name = getattr(self.computation, "architecture_adapter", None)
|
||||
self.architecture = certified_architecture(arch_name)
|
||||
self.start_layer = int(getattr(self.computation, "start_layer"))
|
||||
self.end_layer = int(getattr(self.computation, "end_layer"))
|
||||
self.total_layers = int(getattr(self.computation, "total_layers"))
|
||||
self.role = role_for_range(
|
||||
self.start_layer, self.end_layer, self.total_layers
|
||||
)
|
||||
|
||||
@property
|
||||
def is_head(self) -> bool:
|
||||
return self.role.owns_embedding
|
||||
|
||||
@property
|
||||
def is_tail(self) -> bool:
|
||||
return self.role.owns_final_head
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
token_ids: Any | None = None,
|
||||
boundary: BoundaryBundle | None = None,
|
||||
) -> BoundaryBundle | TailOutput:
|
||||
"""Run one prefill/decode pass for this range and emit its boundary output.
|
||||
|
||||
Head/full ranges require ``token_ids``; middle/tail ranges require the
|
||||
``boundary`` bundle. Non-tail ranges return a :class:`BoundaryBundle`;
|
||||
tail/full ranges return a :class:`TailOutput` through the sampling
|
||||
contract.
|
||||
"""
|
||||
hidden, positions = self._ingest(token_ids, boundary)
|
||||
hidden = self.computation.run_layers(hidden, positions=positions)
|
||||
if self.is_tail:
|
||||
return self._emit_tail(hidden)
|
||||
return self._emit_boundary(hidden, positions)
|
||||
|
||||
# -- input side -----------------------------------------------------------
|
||||
|
||||
def _ingest(
|
||||
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if self.role.owns_embedding:
|
||||
return self._ingest_tokens(token_ids, boundary)
|
||||
return self._ingest_boundary(token_ids, boundary)
|
||||
|
||||
def _ingest_tokens(
|
||||
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if token_ids is None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding and must receive token IDs"
|
||||
)
|
||||
if boundary is not None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding; it must not receive a boundary "
|
||||
"bundle from an upstream range"
|
||||
)
|
||||
ids = np.asarray(token_ids)
|
||||
if ids.ndim == 1:
|
||||
ids = ids[None, :]
|
||||
if ids.ndim != 2:
|
||||
raise BoundaryContractError("token IDs must be (seq,) or (batch, seq)")
|
||||
hidden = np.asarray(self.computation.embed_tokens(ids))
|
||||
positions = np.broadcast_to(
|
||||
np.arange(ids.shape[1], dtype=np.int64), ids.shape
|
||||
).copy()
|
||||
return hidden, positions
|
||||
|
||||
def _ingest_boundary(
|
||||
self, token_ids: Any | None, boundary: BoundaryBundle | None
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if token_ids is not None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards bypass token embedding; they must not receive "
|
||||
"token IDs"
|
||||
)
|
||||
if boundary is None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards must receive the named boundary bundle"
|
||||
)
|
||||
self._check_boundary(boundary)
|
||||
return np.asarray(boundary.residual), np.asarray(boundary.positions)
|
||||
|
||||
def _check_boundary(self, boundary: BoundaryBundle) -> None:
|
||||
if certified_architecture(boundary.architecture_adapter) is not self.architecture:
|
||||
raise BoundaryContractError(
|
||||
f"boundary bundle architecture {boundary.architecture_adapter!r} "
|
||||
f"does not match this Shard's adapter {self.architecture.adapter!r}"
|
||||
)
|
||||
if boundary.schema_version != self.architecture.boundary_schema_version:
|
||||
raise BoundaryContractError(
|
||||
f"boundary schema v{boundary.schema_version} is not supported by "
|
||||
f"this Shard (expects v{self.architecture.boundary_schema_version})"
|
||||
)
|
||||
if boundary.tensor_name != self.architecture.boundary_tensor_name:
|
||||
raise BoundaryContractError(
|
||||
f"boundary tensor {boundary.tensor_name!r} is not the "
|
||||
f"architecture-defined {self.architecture.boundary_tensor_name!r}"
|
||||
)
|
||||
if boundary.normalized:
|
||||
raise BoundaryContractError(
|
||||
"boundary bundle is normalized; a Shard range must receive the "
|
||||
"UNNORMALIZED architecture-defined residual"
|
||||
)
|
||||
if boundary.next_layer != self.start_layer:
|
||||
raise BoundaryContractError(
|
||||
f"boundary hands over at layer {boundary.next_layer} but this "
|
||||
f"Shard starts at layer {self.start_layer}"
|
||||
)
|
||||
|
||||
# -- output side ----------------------------------------------------------
|
||||
|
||||
def _emit_boundary(
|
||||
self, hidden: np.ndarray, positions: np.ndarray
|
||||
) -> BoundaryBundle:
|
||||
# A non-tail Shard emits the unnormalized residual with every position row
|
||||
# intact: no final norm, no LM head, no tail-only row pruning. next_layer
|
||||
# is the receiver's overlap-safe effective start.
|
||||
return BoundaryBundle(
|
||||
architecture_adapter=self.architecture.adapter,
|
||||
schema_version=self.architecture.boundary_schema_version,
|
||||
tensor_name=self.architecture.boundary_tensor_name,
|
||||
residual=np.asarray(hidden),
|
||||
positions=np.asarray(positions),
|
||||
next_layer=self.end_layer + 1,
|
||||
normalized=False,
|
||||
)
|
||||
|
||||
def _emit_tail(self, hidden: np.ndarray) -> TailOutput:
|
||||
hidden = np.asarray(hidden)
|
||||
# Tail-only row pruning: only the final position is needed to sample the
|
||||
# next token, so the LM head runs on the pruned row. A non-tail Shard is
|
||||
# forbidden from doing this (it must forward every row).
|
||||
if self.architecture.prunes_rows_at_tail:
|
||||
last_hidden = hidden[:, -1:, :]
|
||||
else: # pragma: no cover - no certified architecture takes this path yet
|
||||
last_hidden = hidden
|
||||
if self.architecture.normalizes_before_head:
|
||||
last_hidden = np.asarray(self.computation.final_norm(last_hidden))
|
||||
logits = np.asarray(self.computation.lm_head(last_hidden))
|
||||
last_logits = logits[:, -1, :]
|
||||
token_id = self.sampling.sample(last_logits)
|
||||
return TailOutput(
|
||||
token_id=token_id, logits=last_logits, sampling=self.sampling
|
||||
)
|
||||
|
||||
|
||||
def _byte_order(dtype: np.dtype) -> str:
|
||||
order = dtype.byteorder
|
||||
if order == "<":
|
||||
return "little"
|
||||
if order == ">":
|
||||
return "big"
|
||||
# '=' native, '|' not applicable (single byte)
|
||||
import sys
|
||||
|
||||
return sys.byteorder if order in ("=", "|") else "little"
|
||||
|
||||
|
||||
def _array_from_wire(field_payload: dict[str, Any]) -> np.ndarray:
|
||||
array = np.frombuffer(
|
||||
field_payload["data"], dtype=np.dtype(field_payload["dtype"])
|
||||
)
|
||||
return array.reshape(field_payload["shape"]).copy()
|
||||
@@ -20,6 +20,16 @@ import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Mapping
|
||||
|
||||
from . import __version__ as _PACKAGE_VERSION
|
||||
from .runtime_recipe import (
|
||||
ArtifactIdentity,
|
||||
RuntimeRecipeIdentity,
|
||||
build_artifact_identity,
|
||||
build_runtime_recipe_identity,
|
||||
compatibility_fingerprint,
|
||||
fingerprint_payload,
|
||||
)
|
||||
|
||||
# Layout of the serialized report. Bump when the JSON shape changes.
|
||||
CAPABILITY_SCHEMA_VERSION = 1
|
||||
|
||||
@@ -172,6 +182,14 @@ def _optional_text(value: Any, field_name: str) -> str | None:
|
||||
return _require_text(value, field_name)
|
||||
|
||||
|
||||
def _optional_bool(value: Any, field_name: str) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
raise CapabilityReportError(f"{field_name!r} must be a boolean")
|
||||
|
||||
|
||||
def _require_int(value: Any, field_name: str, minimum: int) -> int:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise CapabilityReportError(f"{field_name!r} must be an integer")
|
||||
@@ -218,6 +236,8 @@ class ShardRange:
|
||||
|
||||
start: int
|
||||
end: int
|
||||
owns_embedding: bool = False
|
||||
owns_final_head: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
_require_int(self.start, "shard.start", 0)
|
||||
@@ -226,9 +246,18 @@ class ShardRange:
|
||||
raise CapabilityReportError(
|
||||
f"'shard.end' ({self.end}) must be >= 'shard.start' ({self.start})"
|
||||
)
|
||||
if not isinstance(self.owns_embedding, bool):
|
||||
raise CapabilityReportError("'shard.owns_embedding' must be a boolean")
|
||||
if not isinstance(self.owns_final_head, bool):
|
||||
raise CapabilityReportError("'shard.owns_final_head' must be a boolean")
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {"start": self.start, "end": self.end}
|
||||
return {
|
||||
"start": self.start,
|
||||
"end": self.end,
|
||||
"owns_embedding": self.owns_embedding,
|
||||
"owns_final_head": self.owns_final_head,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> ShardRange:
|
||||
@@ -236,6 +265,12 @@ class ShardRange:
|
||||
return cls(
|
||||
start=_require_int(doc.get("start"), "shard.start", 0),
|
||||
end=_require_int(doc.get("end"), "shard.end", 0),
|
||||
owns_embedding=_optional_bool(
|
||||
doc.get("owns_embedding"), "shard.owns_embedding"
|
||||
),
|
||||
owns_final_head=_optional_bool(
|
||||
doc.get("owns_final_head"), "shard.owns_final_head"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -336,6 +371,8 @@ class CapabilityReport:
|
||||
shard: ShardRange
|
||||
recipe: RecipeIdentity
|
||||
backend: BackendIdentity
|
||||
artifact: ArtifactIdentity
|
||||
runtime_recipe: RuntimeRecipeIdentity
|
||||
status: str
|
||||
validated_at: float
|
||||
duration_ms: int
|
||||
@@ -376,6 +413,20 @@ class CapabilityReport:
|
||||
self.backend.device,
|
||||
)
|
||||
|
||||
@property
|
||||
def compatibility_fingerprint(self) -> str:
|
||||
"""Stable compatibility digest over the full routable identity."""
|
||||
return compatibility_fingerprint(
|
||||
fingerprint_payload(
|
||||
model=self.model.to_dict(),
|
||||
shard=self.shard.to_dict(),
|
||||
recipe=self.recipe.to_dict(),
|
||||
backend=self.backend.to_dict(),
|
||||
artifact=self.artifact.to_dict(),
|
||||
runtime_recipe=self.runtime_recipe.to_dict(),
|
||||
)
|
||||
)
|
||||
|
||||
def age_seconds(self, now: float | None = None) -> float:
|
||||
return max(0.0, (time.time() if now is None else now) - self.validated_at)
|
||||
|
||||
@@ -386,6 +437,9 @@ class CapabilityReport:
|
||||
"shard": self.shard.to_dict(),
|
||||
"recipe": self.recipe.to_dict(),
|
||||
"backend": self.backend.to_dict(),
|
||||
"artifact": self.artifact.to_dict(),
|
||||
"runtime_recipe": self.runtime_recipe.to_dict(),
|
||||
"compatibility_fingerprint": self.compatibility_fingerprint,
|
||||
"status": self.status,
|
||||
"validated_at": self.validated_at,
|
||||
"duration_ms": self.duration_ms,
|
||||
@@ -398,6 +452,9 @@ class CapabilityReport:
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> CapabilityReport:
|
||||
doc = _as_mapping(data, "report")
|
||||
declared_compatibility_fingerprint = _optional_text(
|
||||
doc.get("compatibility_fingerprint"), "compatibility_fingerprint"
|
||||
)
|
||||
|
||||
if "schema_version" not in doc:
|
||||
raise CapabilityReportError(
|
||||
@@ -417,7 +474,13 @@ class CapabilityReport:
|
||||
):
|
||||
raise CapabilityReportError("'validated_at' must be a Unix timestamp")
|
||||
|
||||
return cls(
|
||||
try:
|
||||
artifact = ArtifactIdentity.from_dict(doc.get("artifact"))
|
||||
runtime_recipe = RuntimeRecipeIdentity.from_dict(doc.get("runtime_recipe"))
|
||||
except ValueError as exc:
|
||||
raise CapabilityReportError(str(exc)) from exc
|
||||
|
||||
report = cls(
|
||||
schema_version=schema_version,
|
||||
model=ModelIdentity.from_dict(doc.get("model")),
|
||||
shard=ShardRange.from_dict(doc.get("shard")),
|
||||
@@ -427,7 +490,18 @@ class CapabilityReport:
|
||||
validated_at=float(validated_at),
|
||||
duration_ms=_require_int(doc.get("duration_ms"), "duration_ms", 0),
|
||||
diagnostics=sanitize_diagnostics(doc.get("diagnostics")),
|
||||
artifact=artifact,
|
||||
runtime_recipe=runtime_recipe,
|
||||
)
|
||||
if (
|
||||
declared_compatibility_fingerprint is not None
|
||||
and report.compatibility_fingerprint != declared_compatibility_fingerprint
|
||||
):
|
||||
raise CapabilityReportError(
|
||||
"report declares a compatibility fingerprint that does not match "
|
||||
"its artifact/runtime recipe"
|
||||
)
|
||||
return report
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, text: str) -> CapabilityReport:
|
||||
@@ -458,6 +532,19 @@ def build_capability_report(
|
||||
device_name: str | None = None,
|
||||
quantization: str | None = None,
|
||||
runtime: Mapping[str, str] | None = None,
|
||||
artifact_hash: str | None = None,
|
||||
runtime_recipe: RuntimeRecipeIdentity | None = None,
|
||||
owns_embedding: bool = False,
|
||||
owns_final_head: bool = False,
|
||||
activation_dtype: Any = None,
|
||||
compute_dtype: Any = None,
|
||||
kv_dtype: Any = None,
|
||||
kv_layout: str | None = None,
|
||||
tokenizer_revision: str | None = None,
|
||||
architecture_adapter: str | None = None,
|
||||
boundary_schema_version: int = 1,
|
||||
cache_layout: str | None = None,
|
||||
recipe_params: Mapping[str, Any] | None = None,
|
||||
diagnostics: Any = None,
|
||||
validated_at: float | None = None,
|
||||
environ: Mapping[str, str] | None = None,
|
||||
@@ -468,25 +555,62 @@ def build_capability_report(
|
||||
or an already-computed ``sha256:…`` string. `validated_at` defaults to now,
|
||||
so callers that need determinism pass it explicitly.
|
||||
"""
|
||||
return CapabilityReport(
|
||||
model=ModelIdentity(
|
||||
model_identity = ModelIdentity(
|
||||
model_id=model_id,
|
||||
revision=revision,
|
||||
config_fingerprint=config_fingerprint(model_config),
|
||||
)
|
||||
shard = ShardRange(
|
||||
start=shard_start,
|
||||
end=shard_end,
|
||||
owns_embedding=owns_embedding,
|
||||
owns_final_head=owns_final_head,
|
||||
)
|
||||
recipe_identity = RecipeIdentity(
|
||||
recipe_id=recipe_id,
|
||||
recipe_version=recipe_version,
|
||||
catalogue_version=catalogue_version,
|
||||
)
|
||||
backend_identity = BackendIdentity(
|
||||
backend_id=backend_id,
|
||||
device=device,
|
||||
device_name=device_name,
|
||||
quantization=quantization,
|
||||
runtime=dict(runtime or {}),
|
||||
)
|
||||
artifact = build_artifact_identity(
|
||||
model_id=model_id,
|
||||
revision=revision,
|
||||
model_config=model_config,
|
||||
artifact_hash=artifact_hash,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
)
|
||||
if runtime_recipe is None:
|
||||
runtime_recipe = build_runtime_recipe_identity(
|
||||
model_id=model_id,
|
||||
revision=revision,
|
||||
config_fingerprint=config_fingerprint(model_config),
|
||||
),
|
||||
shard=ShardRange(start=shard_start, end=shard_end),
|
||||
recipe=RecipeIdentity(
|
||||
recipe_id=recipe_id,
|
||||
recipe_version=recipe_version,
|
||||
catalogue_version=catalogue_version,
|
||||
),
|
||||
backend=BackendIdentity(
|
||||
model_config=model_config,
|
||||
recipe_params=recipe_params,
|
||||
weight_quantization=quantization or "unknown",
|
||||
backend_id=backend_id,
|
||||
device=device,
|
||||
device_name=device_name,
|
||||
quantization=quantization,
|
||||
runtime=dict(runtime or {}),
|
||||
),
|
||||
runtime_version=_PACKAGE_VERSION,
|
||||
activation_dtype=activation_dtype,
|
||||
compute_dtype=compute_dtype,
|
||||
kv_dtype=kv_dtype,
|
||||
kv_layout=kv_layout,
|
||||
tokenizer_revision=tokenizer_revision,
|
||||
architecture_adapter=architecture_adapter,
|
||||
boundary_schema_version=boundary_schema_version,
|
||||
cache_layout=cache_layout,
|
||||
)
|
||||
return CapabilityReport(
|
||||
model=model_identity,
|
||||
shard=shard,
|
||||
recipe=recipe_identity,
|
||||
backend=backend_identity,
|
||||
artifact=artifact,
|
||||
runtime_recipe=runtime_recipe,
|
||||
status=status,
|
||||
validated_at=time.time() if validated_at is None else validated_at,
|
||||
duration_ms=duration_ms,
|
||||
|
||||
@@ -36,6 +36,8 @@ from .capability import (
|
||||
CapabilityReport,
|
||||
build_capability_report,
|
||||
)
|
||||
from . import __version__ as _PACKAGE_VERSION
|
||||
from .runtime_recipe import build_runtime_recipe_identity
|
||||
from .recipe_manifest import (
|
||||
DEFAULT_RECIPE_ID,
|
||||
Recipe,
|
||||
@@ -43,6 +45,7 @@ from .recipe_manifest import (
|
||||
RecipeManifestError,
|
||||
load_recipe_manifest,
|
||||
)
|
||||
from .gguf_ownership import authoritative_dense_llama_ownership
|
||||
|
||||
# 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.
|
||||
@@ -464,10 +467,28 @@ def _validate_recipe(
|
||||
duration_ms = int((time.monotonic() - started) * 1000)
|
||||
|
||||
device = _backend_device(backend, selection)
|
||||
ownership = authoritative_dense_llama_ownership(backend, selection)
|
||||
runtime_recipe = build_runtime_recipe_identity(
|
||||
model_id=selection.model_id,
|
||||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||||
model_config=_model_config(backend),
|
||||
recipe_params=recipe.params,
|
||||
weight_quantization=selection.quantization,
|
||||
backend_id=recipe.backend_id,
|
||||
runtime_version=_PACKAGE_VERSION,
|
||||
activation_dtype="bfloat16",
|
||||
compute_dtype=_backend_compute_dtype(backend),
|
||||
kv_dtype=_backend_kv_dtype(backend),
|
||||
kv_layout=_backend_kv_layout(backend),
|
||||
tokenizer_revision=_backend_tokenizer_revision(backend, selection),
|
||||
architecture_adapter=_backend_architecture_adapter(backend, recipe.backend_id),
|
||||
boundary_schema_version=1,
|
||||
cache_layout=_backend_cache_layout(backend, recipe.params),
|
||||
)
|
||||
report = build_capability_report(
|
||||
model_id=selection.model_id,
|
||||
shard_start=selection.shard_start,
|
||||
shard_end=selection.shard_end,
|
||||
shard_start=ownership.start_layer,
|
||||
shard_end=ownership.end_layer,
|
||||
recipe_id=recipe.id,
|
||||
recipe_version=recipe.version,
|
||||
catalogue_version=manifest.catalogue_version,
|
||||
@@ -477,6 +498,9 @@ def _validate_recipe(
|
||||
quantization=selection.quantization,
|
||||
runtime=_runtime_versions(),
|
||||
model_config=_model_config(backend),
|
||||
runtime_recipe=runtime_recipe,
|
||||
owns_embedding=ownership.owns_embedding,
|
||||
owns_final_head=ownership.owns_final_head,
|
||||
status=STATUS_FAILED if category else STATUS_PASSED,
|
||||
duration_ms=duration_ms,
|
||||
diagnostics=[d for d in diagnostics if d] or None,
|
||||
@@ -568,6 +592,65 @@ def _runtime_versions() -> dict[str, str]:
|
||||
return versions
|
||||
|
||||
|
||||
def _backend_compute_dtype(backend: Any) -> str:
|
||||
config = getattr(getattr(backend, "model", None), "config", None)
|
||||
for candidate in (config, getattr(config, "text_config", None)):
|
||||
if candidate is None:
|
||||
continue
|
||||
for attr in ("dtype", "torch_dtype"):
|
||||
value = getattr(candidate, attr, None)
|
||||
if value is None:
|
||||
continue
|
||||
return str(value).removeprefix("torch.")
|
||||
return "bfloat16"
|
||||
|
||||
|
||||
def _backend_kv_dtype(backend: Any) -> str:
|
||||
return _backend_compute_dtype(backend)
|
||||
|
||||
|
||||
def _backend_kv_layout(backend: Any) -> str:
|
||||
return "session-cache" if getattr(backend, "supports_kv_cache", False) else "stateless"
|
||||
|
||||
|
||||
def _backend_tokenizer_revision(backend: Any, selection: DoctorSelection) -> str:
|
||||
model = getattr(backend, "model", None)
|
||||
revision = getattr(model, "revision", None)
|
||||
if isinstance(revision, str) and revision.strip():
|
||||
return revision
|
||||
return selection.model_id
|
||||
|
||||
|
||||
def _backend_architecture_adapter(backend: Any, default: str) -> str:
|
||||
config = getattr(getattr(backend, "model", None), "config", None)
|
||||
for candidate in (config, getattr(config, "text_config", None)):
|
||||
if candidate is None:
|
||||
continue
|
||||
for attr in ("architecture_adapter", "model_type"):
|
||||
value = getattr(candidate, attr, None)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
architectures = getattr(candidate, "architectures", None)
|
||||
if isinstance(architectures, (list, tuple)) and architectures:
|
||||
first = architectures[0]
|
||||
if isinstance(first, str) and first.strip():
|
||||
return first
|
||||
return default
|
||||
|
||||
|
||||
def _backend_cache_layout(backend: Any, recipe_params: Mapping[str, Any] | None) -> str:
|
||||
if getattr(backend, "supports_kv_cache", False) is False:
|
||||
return "stateless"
|
||||
if recipe_params is None:
|
||||
return "local-hot-kv"
|
||||
if recipe_params.get("use_cache") is False:
|
||||
return "stateless"
|
||||
value = recipe_params.get("cache_layout")
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
return "local-hot-kv"
|
||||
|
||||
|
||||
# --- output -----------------------------------------------------------------
|
||||
|
||||
DEFAULT_REPORT_FILENAME = "capability.json"
|
||||
|
||||
423
packages/node/meshnet_node/gguf_backend.py
Normal file
423
packages/node/meshnet_node/gguf_backend.py
Normal file
@@ -0,0 +1,423 @@
|
||||
"""Native llama.cpp/GGUF backend adapter for Meshnet node startup.
|
||||
|
||||
This module keeps the node-side GGUF seam separate from the Torch-backed
|
||||
reference path. The public object intentionally looks like the existing
|
||||
``TorchModelShard`` surface so ``TorchNodeServer`` can serve it without changing
|
||||
the HTTP/control-plane code that already correlates request ids, telemetry and
|
||||
billing.
|
||||
|
||||
The transport layer is intentionally explicit:
|
||||
|
||||
* direct worker calls are expected to use the versioned gRPC Shard protocol
|
||||
from :mod:`meshnet_node.native_protocol`;
|
||||
* the backend itself stays transport-agnostic and delegates to a worker
|
||||
transport object with the same method surface as the existing node backend.
|
||||
|
||||
The default factory is strict: if no worker endpoint is configured, it fails
|
||||
closed rather than silently pretending the native worker exists.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
from .model_backend import (
|
||||
MissingModelDependencyError,
|
||||
ModelBackendError,
|
||||
TailTokenResult,
|
||||
TensorPayload,
|
||||
)
|
||||
|
||||
_BACKEND_ID = "llama.cpp"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeWorkerTransport(Protocol):
|
||||
"""Backend-shaped transport for the supervised native worker."""
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str,
|
||||
session_id: str | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def encode_next_token(
|
||||
self,
|
||||
token_id: int,
|
||||
session_id: str,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
*,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str: ...
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult: ...
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str: ...
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
): ...
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int: ...
|
||||
|
||||
def count_text_tokens(self, text: str) -> int: ...
|
||||
|
||||
def eos_token_ids(self) -> list[int]: ...
|
||||
|
||||
def release_session(self, session_id: str) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _NativeModelConfig:
|
||||
"""Enough model metadata for admission and capability reporting."""
|
||||
|
||||
model_type: str = "llama"
|
||||
architecture_adapter: str = "dense-llama"
|
||||
num_hidden_layers: int = 1
|
||||
torch_dtype: str = "bfloat16"
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"model_type": self.model_type,
|
||||
"architecture_adapter": self.architecture_adapter,
|
||||
"num_hidden_layers": self.num_hidden_layers,
|
||||
"torch_dtype": self.torch_dtype,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class GgufNodeBackend:
|
||||
"""GGUF shard backend shaped like ``TorchModelShard``.
|
||||
|
||||
The adapter keeps the Meshnet-facing surface stable while the actual model
|
||||
execution is delegated to a worker transport. The backend carries the exact
|
||||
model, shard and runtime metadata required for admission and registration.
|
||||
"""
|
||||
|
||||
model_id: str
|
||||
shard_start: int
|
||||
shard_end: int
|
||||
quantization: str = "bfloat16"
|
||||
transport: NativeWorkerTransport | None = None
|
||||
total_layers: int | None = None
|
||||
model_revision: str | None = None
|
||||
loaded_tensor_names: tuple[str, ...] = ()
|
||||
device_type: str = "cpu"
|
||||
supports_kv_cache: bool = True
|
||||
worker_url: str | None = None
|
||||
architecture_adapter: str = "dense-llama"
|
||||
tokenizer_revision: str | None = None
|
||||
runtime_recipe_fingerprint: str | None = None
|
||||
_model: SimpleNamespace = field(init=False, repr=False)
|
||||
_tokenizer: SimpleNamespace = field(init=False, repr=False)
|
||||
is_head: bool = field(init=False)
|
||||
is_tail: bool = field(init=False)
|
||||
loaded_shard_start: int = field(init=False)
|
||||
loaded_shard_end: int = field(init=False)
|
||||
owns_embedding: bool = field(init=False)
|
||||
owns_final_head: bool = field(init=False)
|
||||
|
||||
backend_id = _BACKEND_ID
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.shard_start < 0 or self.shard_end < self.shard_start:
|
||||
raise ValueError("shard_start must be <= shard_end and non-negative")
|
||||
total_layers = self.total_layers or (self.shard_end + 1)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"total_layers",
|
||||
int(total_layers),
|
||||
)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"_model",
|
||||
SimpleNamespace(
|
||||
revision=self.model_revision or self.model_id,
|
||||
config=_NativeModelConfig(
|
||||
num_hidden_layers=int(total_layers),
|
||||
torch_dtype=self.quantization,
|
||||
),
|
||||
),
|
||||
)
|
||||
object.__setattr__(
|
||||
self,
|
||||
"_tokenizer",
|
||||
SimpleNamespace(
|
||||
model_id=self.model_id,
|
||||
revision=self.tokenizer_revision or self.model_revision or self.model_id,
|
||||
eos_token="",
|
||||
eos_token_id=[],
|
||||
),
|
||||
)
|
||||
object.__setattr__(self, "is_head", self.shard_start == 0)
|
||||
object.__setattr__(self, "is_tail", self.shard_end >= int(total_layers) - 1)
|
||||
object.__setattr__(self, "loaded_shard_start", self.shard_start)
|
||||
object.__setattr__(self, "loaded_shard_end", self.shard_end)
|
||||
object.__setattr__(self, "owns_embedding", self.is_head)
|
||||
object.__setattr__(self, "owns_final_head", self.is_tail)
|
||||
if not self.loaded_tensor_names:
|
||||
object.__setattr__(
|
||||
self,
|
||||
"loaded_tensor_names",
|
||||
self._default_tensor_inventory(),
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> Any:
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def tokenizer(self) -> Any:
|
||||
return self._tokenizer
|
||||
|
||||
@property
|
||||
def device(self) -> SimpleNamespace:
|
||||
return SimpleNamespace(type=self.device_type)
|
||||
|
||||
@property
|
||||
def shard_range(self) -> tuple[int, int]:
|
||||
return self.shard_start, self.shard_end
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().encode_prompt(prompt, session_id=session_id)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().encode_next_token(token_id, session_id)
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str:
|
||||
return self._transport().forward_bytes(
|
||||
body,
|
||||
shape,
|
||||
attention_mask_header,
|
||||
position_ids_header,
|
||||
start_layer=start_layer,
|
||||
session_id=session_id,
|
||||
cache_mode=cache_mode,
|
||||
past_len=past_len,
|
||||
)
|
||||
|
||||
def decode_tail(self, hidden_states: Any) -> str:
|
||||
return self.decode_tail_token(hidden_states).text
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
||||
return self._transport().decode_tail_token(hidden_states)
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
return self._transport().generate_text(messages, max_new_tokens, temperature, top_p)
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
yield from self._transport().generate_text_streaming(messages, max_new_tokens, temperature, top_p)
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
return self._transport().count_prompt_tokens(messages)
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
return self._transport().count_text_tokens(text)
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
return self._transport().eos_token_ids()
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
self._transport().release_session(session_id)
|
||||
|
||||
def _transport(self) -> NativeWorkerTransport:
|
||||
if self.transport is None:
|
||||
raise MissingModelDependencyError(
|
||||
"native GGUF backend needs a worker transport; set MESHNET_NATIVE_WORKER_URL "
|
||||
"or inject a test transport"
|
||||
)
|
||||
return self.transport
|
||||
|
||||
def _default_tensor_inventory(self) -> tuple[str, ...]:
|
||||
tensor_names = [f"blk.{layer}.weight" for layer in range(self.shard_start, self.shard_end + 1)]
|
||||
if self.is_head:
|
||||
tensor_names.append("token_embd.weight")
|
||||
if self.is_tail:
|
||||
tensor_names.extend(["output_norm.weight", "output.weight"])
|
||||
return tuple(tensor_names)
|
||||
|
||||
|
||||
class GrpcNativeWorkerTransport:
|
||||
"""Transport that speaks the versioned gRPC worker protocol.
|
||||
|
||||
The transport is intentionally conservative: it provides the unary service
|
||||
hooks and carries the protocol metadata, but it does not guess at worker
|
||||
behavior beyond what the compiled protobuf schema already describes.
|
||||
"""
|
||||
|
||||
def __init__(self, worker_url: str, *, timeout: float = 30.0) -> None:
|
||||
self.worker_url = worker_url
|
||||
self.timeout = timeout
|
||||
self._grpc = None
|
||||
self._channel = None
|
||||
self._stub = None
|
||||
|
||||
def _ensure_stub(self) -> Any:
|
||||
if self._stub is not None:
|
||||
return self._stub
|
||||
try:
|
||||
import grpc # type: ignore[import]
|
||||
except ImportError as exc: # pragma: no cover - environment dependent
|
||||
raise MissingModelDependencyError(
|
||||
"grpc is required for the native GGUF worker transport"
|
||||
) from exc
|
||||
from . import native_protocol
|
||||
|
||||
grpc_mod = native_protocol.load_grpc()
|
||||
self._grpc = grpc
|
||||
self._channel = grpc.insecure_channel(self.worker_url)
|
||||
self._stub = grpc_mod.ShardRuntimeStub(self._channel)
|
||||
return self._stub
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but prompt-to-activation translation is provided "
|
||||
"by the backend wrapper so it can keep worker framing and tokenizer state aligned"
|
||||
)
|
||||
|
||||
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but decode translation is provided by the backend wrapper"
|
||||
)
|
||||
|
||||
def forward_bytes(
|
||||
self,
|
||||
body: bytes,
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
*,
|
||||
start_layer: int | None = None,
|
||||
session_id: str | None = None,
|
||||
cache_mode: str | None = None,
|
||||
past_len: int | None = None,
|
||||
) -> TensorPayload | TailTokenResult | str:
|
||||
raise ModelBackendError(
|
||||
"gRPC transport is present, but activation streaming is handled by the backend wrapper"
|
||||
)
|
||||
|
||||
def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
|
||||
raise ModelBackendError("tail decoding is handled by the backend wrapper")
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
raise ModelBackendError("text generation is handled by the backend wrapper")
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 5120,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
raise ModelBackendError("streaming generation is handled by the backend wrapper")
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
return sum(1 for message in messages if isinstance(message, dict))
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
return len(text.split()) or (1 if text else 0)
|
||||
|
||||
def eos_token_ids(self) -> list[int]:
|
||||
return []
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
stub = self._ensure_stub()
|
||||
from . import native_protocol
|
||||
|
||||
pb2 = native_protocol.load()
|
||||
stub.Release(pb2.ReleaseRequest(reason="release from adapter"))
|
||||
|
||||
|
||||
def build_gguf_backend(
|
||||
*,
|
||||
model_id: str,
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
quantization: str = "bfloat16",
|
||||
transport: NativeWorkerTransport | None = None,
|
||||
worker_url: str | None = None,
|
||||
total_layers: int | None = None,
|
||||
model_revision: str | None = None,
|
||||
loaded_tensor_names: tuple[str, ...] = (),
|
||||
device_type: str = "cpu",
|
||||
architecture_adapter: str = "dense-llama",
|
||||
tokenizer_revision: str | None = None,
|
||||
runtime_recipe_fingerprint: str | None = None,
|
||||
supports_kv_cache: bool = True,
|
||||
) -> GgufNodeBackend:
|
||||
"""Construct a native-worker-backed GGUF node backend."""
|
||||
if transport is None:
|
||||
worker_url = worker_url or os.environ.get("MESHNET_NATIVE_WORKER_URL")
|
||||
if not worker_url:
|
||||
raise MissingModelDependencyError(
|
||||
"set MESHNET_NATIVE_WORKER_URL to the local gRPC worker endpoint "
|
||||
"or inject a fake transport in tests"
|
||||
)
|
||||
transport = GrpcNativeWorkerTransport(worker_url)
|
||||
return GgufNodeBackend(
|
||||
model_id=model_id,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
transport=transport,
|
||||
total_layers=total_layers,
|
||||
model_revision=model_revision,
|
||||
loaded_tensor_names=loaded_tensor_names,
|
||||
device_type=device_type,
|
||||
supports_kv_cache=supports_kv_cache,
|
||||
worker_url=worker_url,
|
||||
architecture_adapter=architecture_adapter,
|
||||
tokenizer_revision=tokenizer_revision,
|
||||
runtime_recipe_fingerprint=runtime_recipe_fingerprint,
|
||||
)
|
||||
287
packages/node/meshnet_node/gguf_ownership.py
Normal file
287
packages/node/meshnet_node/gguf_ownership.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""Dense-Llama GGUF ownership helpers.
|
||||
|
||||
This module keeps two related concerns together:
|
||||
|
||||
* selecting the tensors a dense-Llama GGUF shard is allowed to own; and
|
||||
* inferring the authoritative loaded range / endpoint ownership from the
|
||||
tensors the model actually exposes.
|
||||
|
||||
The first is used by the range-aware loader seam. The second is used by the
|
||||
doctor/admission/reporting path so the tracker sees what the model loaded, not
|
||||
what a CLI flag claimed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, Mapping
|
||||
|
||||
_BLOCK_RE = re.compile(r"^blk\.(\d+)\.")
|
||||
|
||||
_HEAD_TENSOR_NAMES = {
|
||||
"token_embd.weight",
|
||||
"token_embd.bias",
|
||||
"tok_embeddings.weight",
|
||||
"tok_embeddings.bias",
|
||||
"embed_tokens.weight",
|
||||
"embed_tokens.bias",
|
||||
}
|
||||
|
||||
_TAIL_TENSOR_NAMES = {
|
||||
"output_norm.weight",
|
||||
"output_norm.bias",
|
||||
"output.weight",
|
||||
"output.bias",
|
||||
"lm_head.weight",
|
||||
"lm_head.bias",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DenseLlamaShardOwnership:
|
||||
"""Authoritative ownership for one dense-Llama shard."""
|
||||
|
||||
start_layer: int
|
||||
end_layer: int
|
||||
owns_embedding: bool
|
||||
owns_final_head: bool
|
||||
tensor_names: tuple[str, ...] = ()
|
||||
source_artifact_hash: str | None = None
|
||||
slice_artifact_hash: str | None = None
|
||||
derivative_slice: bool = False
|
||||
final_artifact_semantics: bool = True
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.start_layer < 0:
|
||||
raise ValueError("start_layer must be non-negative")
|
||||
if self.end_layer < self.start_layer:
|
||||
raise ValueError("end_layer must be >= start_layer")
|
||||
if self.derivative_slice:
|
||||
if not self.source_artifact_hash or not self.slice_artifact_hash:
|
||||
raise ValueError(
|
||||
"temporary derivative sub-GGUFs must carry source and slice hashes"
|
||||
)
|
||||
if self.final_artifact_semantics:
|
||||
raise ValueError(
|
||||
"temporary derivative sub-GGUFs must not be claimed as final artifacts"
|
||||
)
|
||||
|
||||
@property
|
||||
def range(self) -> tuple[int, int]:
|
||||
return self.start_layer, self.end_layer
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"start_layer": self.start_layer,
|
||||
"end_layer": self.end_layer,
|
||||
"owns_embedding": self.owns_embedding,
|
||||
"owns_final_head": self.owns_final_head,
|
||||
"tensor_names": list(self.tensor_names),
|
||||
"source_artifact_hash": self.source_artifact_hash,
|
||||
"slice_artifact_hash": self.slice_artifact_hash,
|
||||
"derivative_slice": self.derivative_slice,
|
||||
"final_artifact_semantics": self.final_artifact_semantics,
|
||||
}
|
||||
|
||||
|
||||
def select_dense_llama_tensor_names(
|
||||
tensor_names: Iterable[str],
|
||||
start_layer: int,
|
||||
end_layer: int,
|
||||
*,
|
||||
total_layers: int | None = None,
|
||||
) -> set[str]:
|
||||
"""Return the dense-Llama GGUF tensor names owned by an inclusive range."""
|
||||
if start_layer < 0:
|
||||
raise ValueError("start_layer must be non-negative")
|
||||
if end_layer < start_layer:
|
||||
raise ValueError("end_layer must be greater than or equal to start_layer")
|
||||
|
||||
selected: set[str] = set()
|
||||
for tensor_name in tensor_names:
|
||||
if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
|
||||
selected.add(tensor_name)
|
||||
return selected
|
||||
|
||||
|
||||
def infer_dense_llama_ownership(
|
||||
tensor_names: Iterable[str],
|
||||
*,
|
||||
total_layers: int | None = None,
|
||||
source_artifact_hash: str | None = None,
|
||||
slice_artifact_hash: str | None = None,
|
||||
derivative_slice: bool = False,
|
||||
final_artifact_semantics: bool = True,
|
||||
) -> DenseLlamaShardOwnership:
|
||||
"""Infer authoritative loaded range and endpoint ownership from tensors."""
|
||||
names = tuple(str(name) for name in tensor_names if isinstance(name, str))
|
||||
if not names:
|
||||
raise ValueError("tensor inventory is empty")
|
||||
|
||||
block_layers = sorted(
|
||||
{
|
||||
layer
|
||||
for name in names
|
||||
if (layer := _layer_index(name)) is not None
|
||||
}
|
||||
)
|
||||
if not block_layers:
|
||||
raise ValueError("tensor inventory does not contain any blk.N.* tensors")
|
||||
|
||||
selected = tuple(sorted(names))
|
||||
return DenseLlamaShardOwnership(
|
||||
start_layer=block_layers[0],
|
||||
end_layer=block_layers[-1],
|
||||
owns_embedding=any(_is_head_tensor(name) for name in names),
|
||||
owns_final_head=any(
|
||||
_is_tail_tensor(name, total_layers=total_layers, loaded_end=block_layers[-1])
|
||||
for name in names
|
||||
),
|
||||
tensor_names=selected,
|
||||
source_artifact_hash=source_artifact_hash,
|
||||
slice_artifact_hash=slice_artifact_hash,
|
||||
derivative_slice=derivative_slice,
|
||||
final_artifact_semantics=final_artifact_semantics,
|
||||
)
|
||||
|
||||
|
||||
def authoritative_dense_llama_ownership(
|
||||
backend: Any,
|
||||
selection: Any | None = None,
|
||||
) -> DenseLlamaShardOwnership:
|
||||
"""Return the most authoritative dense-Llama ownership the backend exposes."""
|
||||
tensor_names = _tensor_names_from_backend(backend)
|
||||
if tensor_names:
|
||||
try:
|
||||
return infer_dense_llama_ownership(
|
||||
tensor_names,
|
||||
total_layers=_backend_total_layers(backend, selection),
|
||||
)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
start, end = _backend_loaded_bounds(backend, selection)
|
||||
return DenseLlamaShardOwnership(
|
||||
start_layer=start,
|
||||
end_layer=end,
|
||||
owns_embedding=_backend_owns_embedding(backend, start),
|
||||
owns_final_head=_backend_owns_final_head(backend, end),
|
||||
)
|
||||
|
||||
|
||||
def _backend_loaded_bounds(backend: Any, selection: Any | None) -> tuple[int, int]:
|
||||
start = getattr(backend, "loaded_shard_start", None)
|
||||
end = getattr(backend, "loaded_shard_end", None)
|
||||
if start is None:
|
||||
start = getattr(backend, "shard_start", None)
|
||||
if end is None:
|
||||
end = getattr(backend, "shard_end", None)
|
||||
if start is None or end is None:
|
||||
if selection is None:
|
||||
raise ValueError("backend does not expose a loaded shard range")
|
||||
start = getattr(selection, "shard_start")
|
||||
end = getattr(selection, "shard_end")
|
||||
return int(start), int(end)
|
||||
|
||||
|
||||
def _backend_owns_embedding(backend: Any, start: int) -> bool:
|
||||
value = getattr(backend, "owns_embedding", None)
|
||||
if value is None:
|
||||
value = getattr(backend, "is_head", start == 0)
|
||||
return bool(value)
|
||||
|
||||
|
||||
def _backend_owns_final_head(backend: Any, end: int) -> bool:
|
||||
value = getattr(backend, "owns_final_head", None)
|
||||
if value is None:
|
||||
value = getattr(backend, "is_tail", False)
|
||||
return bool(value)
|
||||
|
||||
|
||||
def _backend_total_layers(backend: Any, selection: Any | None) -> int | None:
|
||||
value = getattr(backend, "total_layers", None)
|
||||
if isinstance(value, int) and value > 0:
|
||||
return value
|
||||
if selection is None:
|
||||
return None
|
||||
total = getattr(selection, "total_layers", None)
|
||||
if isinstance(total, int) and total > 0:
|
||||
return total
|
||||
return None
|
||||
|
||||
|
||||
def _tensor_names_from_backend(backend: Any) -> tuple[str, ...]:
|
||||
for attr in ("loaded_tensor_names", "tensor_names", "tensor_inventory"):
|
||||
value = getattr(backend, attr, None)
|
||||
names = _normalise_tensor_names(value)
|
||||
if names:
|
||||
return names
|
||||
return ()
|
||||
|
||||
|
||||
def _normalise_tensor_names(value: Any) -> tuple[str, ...]:
|
||||
if value is None:
|
||||
return ()
|
||||
if isinstance(value, Mapping):
|
||||
items = value.keys()
|
||||
else:
|
||||
try:
|
||||
items = list(value)
|
||||
except TypeError:
|
||||
return ()
|
||||
names = [str(item) for item in items if isinstance(item, str) and item.strip()]
|
||||
return tuple(names)
|
||||
|
||||
|
||||
def _tensor_belongs_to_range(
|
||||
tensor_name: str,
|
||||
start_layer: int,
|
||||
end_layer: int,
|
||||
total_layers: int | None,
|
||||
) -> bool:
|
||||
layer = _layer_index(tensor_name)
|
||||
if layer is not None:
|
||||
return start_layer <= layer <= end_layer
|
||||
|
||||
if start_layer == 0 and _is_head_tensor(tensor_name):
|
||||
return True
|
||||
|
||||
if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(
|
||||
tensor_name, total_layers=total_layers, loaded_end=end_layer
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _layer_index(tensor_name: str) -> int | None:
|
||||
match = _BLOCK_RE.match(tensor_name)
|
||||
if match is None:
|
||||
return None
|
||||
return int(match.group(1))
|
||||
|
||||
|
||||
def _is_head_tensor(tensor_name: str) -> bool:
|
||||
lowered = tensor_name.lower()
|
||||
return lowered in _HEAD_TENSOR_NAMES or any(
|
||||
lowered.startswith(prefix)
|
||||
for prefix in ("token_embd.", "tok_embeddings.", "embed_tokens.")
|
||||
)
|
||||
|
||||
|
||||
def _is_tail_tensor(
|
||||
tensor_name: str,
|
||||
*,
|
||||
total_layers: int | None,
|
||||
loaded_end: int,
|
||||
) -> bool:
|
||||
lowered = tensor_name.lower()
|
||||
if lowered in _TAIL_TENSOR_NAMES:
|
||||
return True
|
||||
if total_layers is not None and loaded_end >= total_layers - 1:
|
||||
return any(
|
||||
lowered.startswith(prefix)
|
||||
for prefix in ("output_norm.", "final_norm.", "norm.")
|
||||
)
|
||||
return False
|
||||
918
packages/node/meshnet_node/hot_kv_state.py
Normal file
918
packages/node/meshnet_node/hot_kv_state.py
Normal file
@@ -0,0 +1,918 @@
|
||||
"""Isolated concurrent local Hot KV State for distributed Shards (DGR-007).
|
||||
|
||||
Hot KV State stays local to the node serving a Shard (RALPH runtime decision #7).
|
||||
A concurrent server must map each ``(Route Session ID, route epoch)`` to an
|
||||
isolated bounded KV context (decision #8) so that one request can never clear or
|
||||
corrupt another's cache.
|
||||
|
||||
This module owns the *lifecycle and storage* of that state and is deliberately
|
||||
backend-agnostic:
|
||||
|
||||
* :class:`HotKvStateManager` is the single mutation entry point. It maps
|
||||
``(session_id, route_epoch)`` to a :class:`SessionCache`, allocates KV **only
|
||||
for the owned layer range**, and enforces a byte budget, a session cap, and a
|
||||
TTL through LRU/TTL eviction. It rejects stale route epochs and incompatible
|
||||
cache recipes, and returns an **explicit** :class:`CacheMiss` when state the
|
||||
caller expected is gone (evicted, released, desynchronised, or never held) so
|
||||
the head degrades to a from-token-zero re-prefill instead of corrupting output
|
||||
(RALPH decision #14: unverified KV is never migrated silently).
|
||||
* :class:`LayerKvCache` / :class:`SessionCache` are the per-owned-layer K/V
|
||||
containers. They are plain ``numpy`` arrays so the default deterministic test
|
||||
suite needs no torch, GPU, download, or API credit; the pinned llama.cpp worker
|
||||
(DGR-008) maps a llama sequence onto the same container contract.
|
||||
* :class:`KvBoundaryAdapter` wraps a KV-aware ``ShardComputation`` (the DGR-006
|
||||
duck type plus ``run_layers_cached``) so a Shard can run cached prefill/decode
|
||||
through the manager while honouring the architecture-defined boundary contract
|
||||
(head embeds tokens, middle/tail bypass embedding, non-tail emits the
|
||||
unnormalized residual, tail samples).
|
||||
|
||||
The manager owns *all* cache mutation: a computation reads the existing cache and
|
||||
returns the new K/V for the appended positions, and the manager decides whether
|
||||
that append fits the budget. That keeps eviction, accounting, and isolation in one
|
||||
place instead of scattered across backends.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
import numpy as np
|
||||
|
||||
from meshnet_node.boundary_adapter import (
|
||||
BOUNDARY_SCHEMA_VERSION,
|
||||
BoundaryBundle,
|
||||
BoundaryContractError,
|
||||
SamplingContract,
|
||||
ShardRole,
|
||||
TailOutput,
|
||||
certified_architecture,
|
||||
role_for_range,
|
||||
)
|
||||
from meshnet_node.runtime_recipe import compatibility_fingerprint
|
||||
|
||||
|
||||
class HotKvStateError(RuntimeError):
|
||||
"""Base class for Hot KV State errors."""
|
||||
|
||||
|
||||
class StaleRouteEpochError(HotKvStateError):
|
||||
"""Raised when a request references a route epoch older than the current one.
|
||||
|
||||
A newer route epoch means the route was re-planned; the old epoch's KV is
|
||||
unverified against the new plan and must never be silently reused.
|
||||
"""
|
||||
|
||||
|
||||
class IncompatibleCacheRecipeError(HotKvStateError):
|
||||
"""Raised when a request's cache recipe does not match the loaded shard.
|
||||
|
||||
A different quantization / dtype / owned range / architecture produces a KV
|
||||
layout this node cannot reuse without corrupting output.
|
||||
"""
|
||||
|
||||
|
||||
class KvBudgetExceededError(HotKvStateError):
|
||||
"""Raised when a single session cannot fit the configured byte budget.
|
||||
|
||||
Other sessions are evicted first (LRU); this fires only when even one session
|
||||
alone exceeds the budget, which is a misconfiguration rather than pressure.
|
||||
"""
|
||||
|
||||
|
||||
class KvCacheMissError(HotKvStateError):
|
||||
"""Raised by the strict accessor when expected session state is absent.
|
||||
|
||||
Prefer :meth:`HotKvStateManager.resolve`, which returns a structured
|
||||
:class:`CacheMiss` instead of raising, when the caller wants to fall back to a
|
||||
stateless re-prefill.
|
||||
"""
|
||||
|
||||
def __init__(self, miss: "CacheMiss") -> None:
|
||||
super().__init__(str(miss))
|
||||
self.miss = miss
|
||||
|
||||
|
||||
class CacheMissReason(str, Enum):
|
||||
"""Why a lookup produced a cache miss (all benign; retry from token zero)."""
|
||||
|
||||
UNKNOWN_SESSION = "unknown-session"
|
||||
EVICTED_TTL = "evicted-ttl"
|
||||
EVICTED_LRU = "evicted-lru"
|
||||
RELEASED = "released"
|
||||
SUPERSEDED_EPOCH = "superseded-epoch"
|
||||
SEQ_LEN_MISMATCH = "seq-len-mismatch"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CacheMiss:
|
||||
"""Explicit cache-miss response the head can act on (re-prefill).
|
||||
|
||||
This is a value, not an exception: the native protocol carries a cache
|
||||
expectation/result, and a miss is a normal, expected outcome under eviction.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
reason: CacheMissReason
|
||||
detail: str = ""
|
||||
|
||||
def __str__(self) -> str:
|
||||
suffix = f": {self.detail}" if self.detail else ""
|
||||
return (
|
||||
f"cache miss for session {self.session_id[:8]} epoch "
|
||||
f"{self.route_epoch} ({self.reason.value}){suffix}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KvCacheRecipe:
|
||||
"""The identity of a Shard's KV layout, used to reject incompatible reuse.
|
||||
|
||||
Two recipes are compatible iff their fingerprints match — same certified
|
||||
architecture, KV dtype, head geometry, and owned layer range within the same
|
||||
whole-model layer count.
|
||||
"""
|
||||
|
||||
architecture_adapter: str
|
||||
kv_dtype: str
|
||||
n_kv_heads: int
|
||||
head_dim: int
|
||||
total_layers: int
|
||||
start_layer: int
|
||||
end_layer: int
|
||||
boundary_schema_version: int = BOUNDARY_SCHEMA_VERSION
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Fail closed on architecture identity (shared with the boundary adapter).
|
||||
certified_architecture(self.architecture_adapter)
|
||||
if self.n_kv_heads <= 0:
|
||||
raise ValueError("n_kv_heads must be positive")
|
||||
if self.head_dim <= 0:
|
||||
raise ValueError("head_dim must be positive")
|
||||
try:
|
||||
np.dtype(self.kv_dtype)
|
||||
except TypeError as exc: # pragma: no cover - defensive
|
||||
raise ValueError(f"invalid kv_dtype {self.kv_dtype!r}") from exc
|
||||
# role_for_range validates 0 <= start <= end <= total_layers - 1.
|
||||
role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
if self.boundary_schema_version < 1:
|
||||
raise ValueError("boundary_schema_version must be >= 1")
|
||||
|
||||
@property
|
||||
def owned_layers(self) -> tuple[int, ...]:
|
||||
return tuple(range(self.start_layer, self.end_layer + 1))
|
||||
|
||||
@property
|
||||
def role(self) -> ShardRole:
|
||||
return role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
|
||||
def bytes_per_token(self) -> int:
|
||||
"""Bytes of KV one token adds across *owned* layers (keys + values)."""
|
||||
itemsize = np.dtype(self.kv_dtype).itemsize
|
||||
per_layer = 2 * self.n_kv_heads * self.head_dim * itemsize
|
||||
return per_layer * len(self.owned_layers)
|
||||
|
||||
def fingerprint(self) -> str:
|
||||
return compatibility_fingerprint(
|
||||
{
|
||||
"kind": "hot-kv-recipe",
|
||||
# Canonicalize the architecture so 'llama' / 'LlamaForCausalLM'
|
||||
# map to the same fingerprint (they are the same layout).
|
||||
"architecture_adapter": certified_architecture(
|
||||
self.architecture_adapter
|
||||
).adapter,
|
||||
"kv_dtype": np.dtype(self.kv_dtype).name,
|
||||
"n_kv_heads": self.n_kv_heads,
|
||||
"head_dim": self.head_dim,
|
||||
"total_layers": self.total_layers,
|
||||
"start_layer": self.start_layer,
|
||||
"end_layer": self.end_layer,
|
||||
"boundary_schema_version": self.boundary_schema_version,
|
||||
}
|
||||
)
|
||||
|
||||
def is_compatible(self, other: "KvCacheRecipe") -> bool:
|
||||
return self.fingerprint() == other.fingerprint()
|
||||
|
||||
|
||||
class LayerKvCache:
|
||||
"""K/V storage for a single owned layer; sequence axis is 0.
|
||||
|
||||
Keys and values are ``(seq, n_kv_heads, head_dim)``. Backends store the
|
||||
position-encoded (post-RoPE) keys so a decode step only appends the new rows.
|
||||
"""
|
||||
|
||||
__slots__ = ("layer_index", "n_kv_heads", "head_dim", "dtype", "keys", "values")
|
||||
|
||||
def __init__(
|
||||
self, layer_index: int, n_kv_heads: int, head_dim: int, dtype: Any
|
||||
) -> None:
|
||||
self.layer_index = int(layer_index)
|
||||
self.n_kv_heads = int(n_kv_heads)
|
||||
self.head_dim = int(head_dim)
|
||||
self.dtype = np.dtype(dtype)
|
||||
self.keys = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)
|
||||
self.values = np.empty((0, self.n_kv_heads, self.head_dim), dtype=self.dtype)
|
||||
|
||||
@property
|
||||
def length(self) -> int:
|
||||
return int(self.keys.shape[0])
|
||||
|
||||
def _validate(self, array: np.ndarray, name: str) -> np.ndarray:
|
||||
arr = np.asarray(array, dtype=self.dtype)
|
||||
if arr.ndim != 3 or arr.shape[1:] != (self.n_kv_heads, self.head_dim):
|
||||
raise ValueError(
|
||||
f"layer {self.layer_index} {name} must be "
|
||||
f"(seq, {self.n_kv_heads}, {self.head_dim}), got {arr.shape}"
|
||||
)
|
||||
return arr
|
||||
|
||||
def append(self, keys: np.ndarray, values: np.ndarray) -> int:
|
||||
k = self._validate(keys, "keys")
|
||||
v = self._validate(values, "values")
|
||||
if k.shape[0] != v.shape[0]:
|
||||
raise ValueError(
|
||||
f"layer {self.layer_index} keys/values disagree on token count "
|
||||
f"({k.shape[0]} vs {v.shape[0]})"
|
||||
)
|
||||
self.keys = np.concatenate([self.keys, k], axis=0)
|
||||
self.values = np.concatenate([self.values, v], axis=0)
|
||||
return self.length
|
||||
|
||||
def truncate(self, length: int) -> None:
|
||||
length = max(0, int(length))
|
||||
self.keys = self.keys[:length]
|
||||
self.values = self.values[:length]
|
||||
|
||||
@property
|
||||
def nbytes(self) -> int:
|
||||
return int(self.keys.nbytes + self.values.nbytes)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionCache:
|
||||
"""Isolated per-``(session_id, epoch)`` KV context over the owned layers only."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
recipe: KvCacheRecipe
|
||||
layers: "OrderedDict[int, LayerKvCache]"
|
||||
created_tick: float
|
||||
last_tick: float
|
||||
released: bool = False
|
||||
|
||||
@property
|
||||
def seq_len(self) -> int:
|
||||
if not self.layers:
|
||||
return 0
|
||||
# All owned layers advance in lockstep; report the first owned layer.
|
||||
return next(iter(self.layers.values())).length
|
||||
|
||||
@property
|
||||
def owned_layers(self) -> tuple[int, ...]:
|
||||
return tuple(self.layers.keys())
|
||||
|
||||
def layer(self, index: int) -> LayerKvCache:
|
||||
try:
|
||||
return self.layers[index]
|
||||
except KeyError:
|
||||
raise KeyError(
|
||||
f"layer {index} is not owned by this shard "
|
||||
f"(owned {list(self.layers)})"
|
||||
) from None
|
||||
|
||||
def read_only_layers(self) -> Mapping[int, LayerKvCache]:
|
||||
"""The current per-layer caches a computation reads to attend over."""
|
||||
return dict(self.layers)
|
||||
|
||||
def _append(self, kv_by_layer: Mapping[int, Any]) -> int:
|
||||
provided = set(kv_by_layer)
|
||||
owned = set(self.layers)
|
||||
if provided != owned:
|
||||
raise ValueError(
|
||||
f"append must cover exactly the owned layers {sorted(owned)}, "
|
||||
f"got {sorted(provided)}"
|
||||
)
|
||||
# Pre-validate token counts so a partial append never desynchronises the
|
||||
# owned layers (append is all-or-nothing).
|
||||
new_counts = set()
|
||||
for idx, (keys, _values) in kv_by_layer.items():
|
||||
new_counts.add(int(np.asarray(keys).shape[0]))
|
||||
if len(new_counts) != 1:
|
||||
raise ValueError(
|
||||
f"append token counts disagree across layers: {sorted(new_counts)}"
|
||||
)
|
||||
for idx, (keys, values) in kv_by_layer.items():
|
||||
self.layers[idx].append(keys, values)
|
||||
return self.seq_len
|
||||
|
||||
def _truncate(self, length: int) -> None:
|
||||
for cache in self.layers.values():
|
||||
cache.truncate(length)
|
||||
|
||||
@property
|
||||
def nbytes(self) -> int:
|
||||
return sum(cache.nbytes for cache in self.layers.values())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HotKvStateConfig:
|
||||
"""Bounds for the manager: memory budget, session cap, and idle TTL."""
|
||||
|
||||
budget_bytes: int = 64 * 1024 * 1024
|
||||
max_sessions: int = 8
|
||||
ttl_seconds: float = 600.0
|
||||
miss_history: int = 256
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.budget_bytes <= 0:
|
||||
raise ValueError("budget_bytes must be positive")
|
||||
if self.max_sessions < 1:
|
||||
raise ValueError("max_sessions must be >= 1")
|
||||
if self.ttl_seconds < 0:
|
||||
raise ValueError("ttl_seconds must be >= 0")
|
||||
if self.miss_history < 0:
|
||||
raise ValueError("miss_history must be >= 0")
|
||||
|
||||
|
||||
class HotKvStateManager:
|
||||
"""Concurrent, bounded map of ``(session_id, epoch)`` to an isolated KV context."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
recipe: KvCacheRecipe,
|
||||
config: HotKvStateConfig | None = None,
|
||||
*,
|
||||
clock: Callable[[], float] | None = None,
|
||||
) -> None:
|
||||
self.recipe = recipe
|
||||
self.config = config or HotKvStateConfig()
|
||||
self._clock = clock or time.monotonic
|
||||
self._sessions: "OrderedDict[tuple[str, int], SessionCache]" = OrderedDict()
|
||||
self._latest_epoch: dict[str, int] = {}
|
||||
self._misses: "OrderedDict[tuple[str, int], CacheMiss]" = OrderedDict()
|
||||
self._lock = threading.RLock()
|
||||
|
||||
# -- introspection --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def total_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return sum(s.nbytes for s in self._sessions.values())
|
||||
|
||||
@property
|
||||
def session_count(self) -> int:
|
||||
with self._lock:
|
||||
self._evict_expired_locked(self._clock())
|
||||
return len(self._sessions)
|
||||
|
||||
def session_keys(self) -> list[tuple[str, int]]:
|
||||
with self._lock:
|
||||
return list(self._sessions.keys())
|
||||
|
||||
# -- lifecycle ------------------------------------------------------------
|
||||
|
||||
def open(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
) -> SessionCache:
|
||||
"""Create (or replace) a fresh, empty isolated context for the session.
|
||||
|
||||
A higher route epoch supersedes and frees any earlier epoch for the same
|
||||
session id; an older epoch is rejected as stale.
|
||||
"""
|
||||
self._require_text(session_id, "session_id")
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
self._supersede_older_epochs_locked(session_id, route_epoch)
|
||||
key = (session_id, route_epoch)
|
||||
# A re-open at the same epoch replaces the prior context entirely.
|
||||
self._sessions.pop(key, None)
|
||||
layers: "OrderedDict[int, LayerKvCache]" = OrderedDict(
|
||||
(
|
||||
idx,
|
||||
LayerKvCache(
|
||||
idx,
|
||||
self.recipe.n_kv_heads,
|
||||
self.recipe.head_dim,
|
||||
self.recipe.kv_dtype,
|
||||
),
|
||||
)
|
||||
for idx in self.recipe.owned_layers
|
||||
)
|
||||
session = SessionCache(
|
||||
session_id=session_id,
|
||||
route_epoch=route_epoch,
|
||||
recipe=self.recipe,
|
||||
layers=layers,
|
||||
created_tick=now,
|
||||
last_tick=now,
|
||||
)
|
||||
self._sessions[key] = session
|
||||
self._latest_epoch[session_id] = route_epoch
|
||||
self._misses.pop(key, None)
|
||||
self._enforce_capacity_locked(protect=key, incoming_bytes=0)
|
||||
return session
|
||||
|
||||
def append(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
kv_by_layer: Mapping[int, Any],
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache:
|
||||
"""Append new K/V (prefill or decode) to an existing isolated context.
|
||||
|
||||
The computation supplies exactly the owned layers' new keys/values. The
|
||||
manager evicts other sessions (LRU) to fit the byte budget before growing
|
||||
this one, and raises :class:`KvBudgetExceededError` only if this session
|
||||
alone cannot fit.
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
session = self._require_live_locked(session_id, route_epoch)
|
||||
if expected_seq_len is not None and session.seq_len != expected_seq_len:
|
||||
miss = self._drop_and_record_locked(
|
||||
(session_id, route_epoch),
|
||||
CacheMissReason.SEQ_LEN_MISMATCH,
|
||||
detail=f"cache holds {session.seq_len}, caller expected "
|
||||
f"{expected_seq_len}",
|
||||
)
|
||||
raise KvCacheMissError(miss)
|
||||
n_new = self._new_token_count(kv_by_layer)
|
||||
incoming = n_new * self.recipe.bytes_per_token()
|
||||
self._enforce_capacity_locked(
|
||||
protect=(session_id, route_epoch), incoming_bytes=incoming
|
||||
)
|
||||
session._append(kv_by_layer)
|
||||
session.last_tick = self._clock()
|
||||
self._sessions.move_to_end((session_id, route_epoch))
|
||||
return session
|
||||
|
||||
def truncate(
|
||||
self, session_id: str, route_epoch: int, length: int
|
||||
) -> SessionCache:
|
||||
"""Drop cached positions beyond ``length`` (rollback) for one session."""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
session = self._require_live_locked(session_id, route_epoch)
|
||||
if length < 0:
|
||||
raise ValueError("truncate length must be >= 0")
|
||||
session._truncate(length)
|
||||
session.last_tick = self._clock()
|
||||
self._sessions.move_to_end((session_id, route_epoch))
|
||||
return session
|
||||
|
||||
def release(self, session_id: str, route_epoch: int) -> bool:
|
||||
"""Free one session's context; other sessions are untouched.
|
||||
|
||||
Returns True if a live context was freed. A later lookup for the released
|
||||
key yields an explicit :class:`CacheMiss`.
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
key = (session_id, route_epoch)
|
||||
existed = key in self._sessions
|
||||
self._drop_and_record_locked(key, CacheMissReason.RELEASED)
|
||||
return existed
|
||||
|
||||
# -- lookup ---------------------------------------------------------------
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache | CacheMiss:
|
||||
"""Return the live context or an explicit :class:`CacheMiss`.
|
||||
|
||||
Rejects stale epochs and incompatible recipes (both are protocol
|
||||
violations, not benign misses).
|
||||
"""
|
||||
route_epoch = self._require_epoch(route_epoch)
|
||||
with self._lock:
|
||||
self._check_recipe(recipe)
|
||||
self._validate_epoch_locked(session_id, route_epoch)
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
key = (session_id, route_epoch)
|
||||
session = self._sessions.get(key)
|
||||
if session is None:
|
||||
return self._recorded_miss_locked(key)
|
||||
if expected_seq_len is not None and session.seq_len != expected_seq_len:
|
||||
return self._drop_and_record_locked(
|
||||
key,
|
||||
CacheMissReason.SEQ_LEN_MISMATCH,
|
||||
detail=f"cache holds {session.seq_len}, caller expected "
|
||||
f"{expected_seq_len}",
|
||||
)
|
||||
session.last_tick = now
|
||||
self._sessions.move_to_end(key)
|
||||
return session
|
||||
|
||||
def get(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
recipe: KvCacheRecipe | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> SessionCache:
|
||||
"""Strict accessor: raises :class:`KvCacheMissError` on a miss."""
|
||||
result = self.resolve(
|
||||
session_id,
|
||||
route_epoch,
|
||||
recipe=recipe,
|
||||
expected_seq_len=expected_seq_len,
|
||||
)
|
||||
if isinstance(result, CacheMiss):
|
||||
raise KvCacheMissError(result)
|
||||
return result
|
||||
|
||||
# -- internals ------------------------------------------------------------
|
||||
|
||||
def _check_recipe(self, recipe: KvCacheRecipe | None) -> None:
|
||||
if recipe is not None and not self.recipe.is_compatible(recipe):
|
||||
raise IncompatibleCacheRecipeError(
|
||||
"request cache recipe does not match this shard's loaded recipe "
|
||||
f"(request {recipe.fingerprint()} vs shard {self.recipe.fingerprint()})"
|
||||
)
|
||||
|
||||
def _validate_epoch_locked(self, session_id: str, route_epoch: int) -> None:
|
||||
latest = self._latest_epoch.get(session_id)
|
||||
if latest is not None and route_epoch < latest:
|
||||
raise StaleRouteEpochError(
|
||||
f"session {session_id[:8]} route epoch {route_epoch} is stale; "
|
||||
f"current epoch is {latest}"
|
||||
)
|
||||
|
||||
def _supersede_older_epochs_locked(
|
||||
self, session_id: str, route_epoch: int
|
||||
) -> None:
|
||||
stale_keys = [
|
||||
key
|
||||
for key in self._sessions
|
||||
if key[0] == session_id and key[1] < route_epoch
|
||||
]
|
||||
for key in stale_keys:
|
||||
self._drop_and_record_locked(key, CacheMissReason.SUPERSEDED_EPOCH)
|
||||
|
||||
def _require_live_locked(
|
||||
self, session_id: str, route_epoch: int
|
||||
) -> SessionCache:
|
||||
now = self._clock()
|
||||
self._evict_expired_locked(now)
|
||||
key = (session_id, route_epoch)
|
||||
session = self._sessions.get(key)
|
||||
if session is None:
|
||||
raise KvCacheMissError(self._recorded_miss_locked(key))
|
||||
return session
|
||||
|
||||
def _new_token_count(self, kv_by_layer: Mapping[int, Any]) -> int:
|
||||
owned = set(self.recipe.owned_layers)
|
||||
if set(kv_by_layer) != owned:
|
||||
raise ValueError(
|
||||
f"append must cover exactly the owned layers {sorted(owned)}, "
|
||||
f"got {sorted(kv_by_layer)}"
|
||||
)
|
||||
counts = {int(np.asarray(k).shape[0]) for k, _ in kv_by_layer.values()}
|
||||
if len(counts) != 1:
|
||||
raise ValueError(
|
||||
f"append token counts disagree across layers: {sorted(counts)}"
|
||||
)
|
||||
return counts.pop()
|
||||
|
||||
def _enforce_capacity_locked(
|
||||
self, *, protect: tuple[str, int], incoming_bytes: int
|
||||
) -> None:
|
||||
# Session cap: evict LRU sessions other than the protected one.
|
||||
while len(self._sessions) > self.config.max_sessions:
|
||||
victim = self._lru_victim_locked(protect)
|
||||
if victim is None:
|
||||
break
|
||||
self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)
|
||||
|
||||
# Byte budget: the protected session's own footprint after the append.
|
||||
protected = self._sessions.get(protect)
|
||||
protected_bytes = (protected.nbytes if protected is not None else 0) + incoming_bytes
|
||||
if protected_bytes > self.config.budget_bytes:
|
||||
raise KvBudgetExceededError(
|
||||
f"session {protect[0][:8]} needs {protected_bytes} bytes which "
|
||||
f"exceeds the KV budget {self.config.budget_bytes}"
|
||||
)
|
||||
# Evict other LRU sessions until the whole store fits with the append.
|
||||
while self._total_bytes_locked() + incoming_bytes > self.config.budget_bytes:
|
||||
victim = self._lru_victim_locked(protect)
|
||||
if victim is None:
|
||||
break
|
||||
self._drop_and_record_locked(victim, CacheMissReason.EVICTED_LRU)
|
||||
|
||||
def _lru_victim_locked(self, protect: tuple[str, int]) -> tuple[str, int] | None:
|
||||
for key in self._sessions: # OrderedDict iterates oldest-first.
|
||||
if key != protect:
|
||||
return key
|
||||
return None
|
||||
|
||||
def _total_bytes_locked(self) -> int:
|
||||
return sum(s.nbytes for s in self._sessions.values())
|
||||
|
||||
def _evict_expired_locked(self, now: float) -> None:
|
||||
ttl = self.config.ttl_seconds
|
||||
if ttl <= 0:
|
||||
return
|
||||
expired = [
|
||||
key
|
||||
for key, session in self._sessions.items()
|
||||
if now - session.last_tick > ttl
|
||||
]
|
||||
for key in expired:
|
||||
self._drop_and_record_locked(key, CacheMissReason.EVICTED_TTL)
|
||||
|
||||
def _drop_and_record_locked(
|
||||
self,
|
||||
key: tuple[str, int],
|
||||
reason: CacheMissReason,
|
||||
*,
|
||||
detail: str = "",
|
||||
) -> CacheMiss:
|
||||
session = self._sessions.pop(key, None)
|
||||
if session is not None:
|
||||
session.released = True
|
||||
miss = CacheMiss(
|
||||
session_id=key[0], route_epoch=key[1], reason=reason, detail=detail
|
||||
)
|
||||
self._record_miss_locked(key, miss)
|
||||
return miss
|
||||
|
||||
def _record_miss_locked(self, key: tuple[str, int], miss: CacheMiss) -> None:
|
||||
if self.config.miss_history <= 0:
|
||||
return
|
||||
self._misses.pop(key, None)
|
||||
self._misses[key] = miss
|
||||
while len(self._misses) > self.config.miss_history:
|
||||
self._misses.popitem(last=False)
|
||||
|
||||
def _recorded_miss_locked(self, key: tuple[str, int]) -> CacheMiss:
|
||||
recorded = self._misses.get(key)
|
||||
if recorded is not None:
|
||||
return recorded
|
||||
return CacheMiss(
|
||||
session_id=key[0],
|
||||
route_epoch=key[1],
|
||||
reason=CacheMissReason.UNKNOWN_SESSION,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _require_text(value: Any, name: str) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{name} must be a non-empty string")
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _require_epoch(value: Any) -> int:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError("route_epoch must be an integer")
|
||||
if value < 0:
|
||||
raise ValueError("route_epoch must be >= 0")
|
||||
return value
|
||||
|
||||
|
||||
def kv_recipe_for(computation: Any) -> KvCacheRecipe:
|
||||
"""Build a :class:`KvCacheRecipe` from a KV-aware ``ShardComputation``.
|
||||
|
||||
The computation exposes the DGR-006 duck type plus KV geometry
|
||||
(``n_kv_heads``, ``head_dim``, ``kv_dtype``).
|
||||
"""
|
||||
return KvCacheRecipe(
|
||||
architecture_adapter=str(getattr(computation, "architecture_adapter")),
|
||||
kv_dtype=str(getattr(computation, "kv_dtype", "float32")),
|
||||
n_kv_heads=int(getattr(computation, "n_kv_heads")),
|
||||
head_dim=int(getattr(computation, "head_dim")),
|
||||
total_layers=int(getattr(computation, "total_layers")),
|
||||
start_layer=int(getattr(computation, "start_layer")),
|
||||
end_layer=int(getattr(computation, "end_layer")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KvBoundaryAdapter:
|
||||
"""KV-aware boundary driver: cached prefill/decode through the manager.
|
||||
|
||||
Mirrors the DGR-006 :class:`~meshnet_node.boundary_adapter.BoundaryAdapter`
|
||||
contract (head embeds tokens, middle/tail bypass embedding and consume the
|
||||
unnormalized residual bundle, non-tail emits the unnormalized residual, tail
|
||||
normalizes + heads + prunes + samples) but threads a per-session KV context.
|
||||
|
||||
The wrapped computation must additionally expose::
|
||||
|
||||
run_layers_cached(hidden, *, positions, past_kv)
|
||||
-> (hidden_out, {layer_index: (new_keys, new_values)})
|
||||
|
||||
reading ``past_kv`` (the current per-owned-layer caches) and returning the new
|
||||
position-encoded K/V for the appended positions only. The manager, not the
|
||||
computation, commits those K/V so eviction and budget stay centralized.
|
||||
"""
|
||||
|
||||
computation: Any
|
||||
manager: HotKvStateManager
|
||||
sampling: SamplingContract = field(default_factory=SamplingContract.greedy)
|
||||
architecture: Any = field(init=False)
|
||||
role: ShardRole = field(init=False)
|
||||
start_layer: int = field(init=False)
|
||||
end_layer: int = field(init=False)
|
||||
total_layers: int = field(init=False)
|
||||
recipe: KvCacheRecipe = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
arch_name = getattr(self.computation, "architecture_adapter", None)
|
||||
self.architecture = certified_architecture(arch_name)
|
||||
self.start_layer = int(getattr(self.computation, "start_layer"))
|
||||
self.end_layer = int(getattr(self.computation, "end_layer"))
|
||||
self.total_layers = int(getattr(self.computation, "total_layers"))
|
||||
self.role = role_for_range(self.start_layer, self.end_layer, self.total_layers)
|
||||
self.recipe = kv_recipe_for(self.computation)
|
||||
if not self.manager.recipe.is_compatible(self.recipe):
|
||||
raise IncompatibleCacheRecipeError(
|
||||
"manager recipe does not match this computation's KV recipe"
|
||||
)
|
||||
|
||||
@property
|
||||
def is_head(self) -> bool:
|
||||
return self.role.owns_embedding
|
||||
|
||||
@property
|
||||
def is_tail(self) -> bool:
|
||||
return self.role.owns_final_head
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
token_ids: Any | None = None,
|
||||
boundary: BoundaryBundle | None = None,
|
||||
) -> BoundaryBundle | TailOutput:
|
||||
"""Open a fresh isolated context and run the prompt through this range."""
|
||||
session = self.manager.open(session_id, route_epoch, recipe=self.recipe)
|
||||
return self._run_step(session, token_ids, boundary)
|
||||
|
||||
def decode(
|
||||
self,
|
||||
session_id: str,
|
||||
route_epoch: int,
|
||||
*,
|
||||
token_ids: Any | None = None,
|
||||
boundary: BoundaryBundle | None = None,
|
||||
expected_seq_len: int | None = None,
|
||||
) -> BoundaryBundle | TailOutput | CacheMiss:
|
||||
"""Append one (or more) decode positions to an existing context.
|
||||
|
||||
Returns an explicit :class:`CacheMiss` if the context is gone so the head
|
||||
can re-prefill from token zero instead of corrupting output.
|
||||
"""
|
||||
resolved = self.manager.resolve(
|
||||
session_id,
|
||||
route_epoch,
|
||||
recipe=self.recipe,
|
||||
expected_seq_len=expected_seq_len,
|
||||
)
|
||||
if isinstance(resolved, CacheMiss):
|
||||
return resolved
|
||||
return self._run_step(resolved, token_ids, boundary)
|
||||
|
||||
# -- internals ------------------------------------------------------------
|
||||
|
||||
def _run_step(
|
||||
self,
|
||||
session: SessionCache,
|
||||
token_ids: Any | None,
|
||||
boundary: BoundaryBundle | None,
|
||||
) -> BoundaryBundle | TailOutput:
|
||||
prev_len = session.seq_len
|
||||
hidden, positions = self._ingest(prev_len, token_ids, boundary)
|
||||
hidden_out, new_kv = self.computation.run_layers_cached(
|
||||
hidden, positions=positions, past_kv=session.read_only_layers()
|
||||
)
|
||||
self.manager.append(
|
||||
session.session_id,
|
||||
session.route_epoch,
|
||||
new_kv,
|
||||
recipe=self.recipe,
|
||||
expected_seq_len=prev_len,
|
||||
)
|
||||
if self.is_tail:
|
||||
return self._emit_tail(hidden_out)
|
||||
return self._emit_boundary(hidden_out, positions)
|
||||
|
||||
def _ingest(
|
||||
self,
|
||||
prev_len: int,
|
||||
token_ids: Any | None,
|
||||
boundary: BoundaryBundle | None,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if self.role.owns_embedding:
|
||||
if token_ids is None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding and must receive token IDs"
|
||||
)
|
||||
if boundary is not None:
|
||||
raise BoundaryContractError(
|
||||
"the head owns token embedding; it must not receive a boundary "
|
||||
"bundle from an upstream range"
|
||||
)
|
||||
ids = np.asarray(token_ids)
|
||||
if ids.ndim == 1:
|
||||
ids = ids[None, :]
|
||||
if ids.ndim != 2:
|
||||
raise BoundaryContractError("token IDs must be (seq,) or (batch, seq)")
|
||||
hidden = np.asarray(self.computation.embed_tokens(ids))
|
||||
n_new = ids.shape[1]
|
||||
positions = np.broadcast_to(
|
||||
np.arange(prev_len, prev_len + n_new, dtype=np.int64),
|
||||
ids.shape,
|
||||
).copy()
|
||||
return hidden, positions
|
||||
# Middle / tail: consume the boundary bundle (the unnormalized residual).
|
||||
if token_ids is not None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards bypass token embedding; they must not receive "
|
||||
"token IDs"
|
||||
)
|
||||
if boundary is None:
|
||||
raise BoundaryContractError(
|
||||
"middle/tail Shards must receive the named boundary bundle"
|
||||
)
|
||||
self._check_boundary(boundary)
|
||||
return np.asarray(boundary.residual), np.asarray(boundary.positions)
|
||||
|
||||
def _check_boundary(self, boundary: BoundaryBundle) -> None:
|
||||
if certified_architecture(boundary.architecture_adapter) is not self.architecture:
|
||||
raise BoundaryContractError(
|
||||
f"boundary bundle architecture {boundary.architecture_adapter!r} "
|
||||
f"does not match this Shard's adapter {self.architecture.adapter!r}"
|
||||
)
|
||||
if boundary.schema_version != self.architecture.boundary_schema_version:
|
||||
raise BoundaryContractError(
|
||||
f"boundary schema v{boundary.schema_version} is not supported by "
|
||||
f"this Shard (expects v{self.architecture.boundary_schema_version})"
|
||||
)
|
||||
if boundary.tensor_name != self.architecture.boundary_tensor_name:
|
||||
raise BoundaryContractError(
|
||||
f"boundary tensor {boundary.tensor_name!r} is not the "
|
||||
f"architecture-defined {self.architecture.boundary_tensor_name!r}"
|
||||
)
|
||||
if boundary.normalized:
|
||||
raise BoundaryContractError(
|
||||
"boundary bundle is normalized; a Shard range must receive the "
|
||||
"UNNORMALIZED architecture-defined residual"
|
||||
)
|
||||
if boundary.next_layer != self.start_layer:
|
||||
raise BoundaryContractError(
|
||||
f"boundary hands over at layer {boundary.next_layer} but this "
|
||||
f"Shard starts at layer {self.start_layer}"
|
||||
)
|
||||
|
||||
def _emit_boundary(
|
||||
self, hidden: np.ndarray, positions: np.ndarray
|
||||
) -> BoundaryBundle:
|
||||
return BoundaryBundle(
|
||||
architecture_adapter=self.architecture.adapter,
|
||||
schema_version=self.architecture.boundary_schema_version,
|
||||
tensor_name=self.architecture.boundary_tensor_name,
|
||||
residual=np.asarray(hidden),
|
||||
positions=np.asarray(positions),
|
||||
next_layer=self.end_layer + 1,
|
||||
normalized=False,
|
||||
)
|
||||
|
||||
def _emit_tail(self, hidden: np.ndarray) -> TailOutput:
|
||||
hidden = np.asarray(hidden)
|
||||
if self.architecture.prunes_rows_at_tail:
|
||||
last_hidden = hidden[:, -1:, :]
|
||||
else: # pragma: no cover - no certified architecture takes this path yet
|
||||
last_hidden = hidden
|
||||
if self.architecture.normalizes_before_head:
|
||||
last_hidden = np.asarray(self.computation.final_norm(last_hidden))
|
||||
logits = np.asarray(self.computation.lm_head(last_hidden))
|
||||
last_logits = logits[:, -1, :]
|
||||
token_id = self.sampling.sample(last_logits)
|
||||
return TailOutput(token_id=token_id, logits=last_logits, sampling=self.sampling)
|
||||
@@ -323,6 +323,10 @@ class TorchModelShard:
|
||||
)
|
||||
self.is_head = shard_start == 0
|
||||
self.is_tail = shard_end >= self.total_layers - 1
|
||||
self.loaded_shard_start = shard_start
|
||||
self.loaded_shard_end = shard_end
|
||||
self.owns_embedding = self.is_head
|
||||
self.owns_final_head = self.is_tail
|
||||
self.hidden_size = int(
|
||||
getattr(self.model.config, "hidden_size", 0)
|
||||
or getattr(self.model.config, "n_embd", 0)
|
||||
@@ -344,6 +348,17 @@ class TorchModelShard:
|
||||
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
|
||||
)
|
||||
|
||||
@property
|
||||
def loaded_range(self) -> tuple[int, int]:
|
||||
return self.loaded_shard_start, self.loaded_shard_end
|
||||
|
||||
@property
|
||||
def endpoint_ownership(self) -> dict[str, bool]:
|
||||
return {
|
||||
"owns_embedding": self.owns_embedding,
|
||||
"owns_final_head": self.owns_final_head,
|
||||
}
|
||||
|
||||
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
|
||||
if not self.is_head or self._embed_tokens is None:
|
||||
raise ModelBackendError("text prompts can only be accepted by the head shard")
|
||||
|
||||
300
packages/node/meshnet_node/native_protocol/__init__.py
Normal file
300
packages/node/meshnet_node/native_protocol/__init__.py
Normal file
@@ -0,0 +1,300 @@
|
||||
"""Loader and helpers for the versioned gRPC Shard protocol (ADR-0024, DGR-002).
|
||||
|
||||
The ``.proto`` schema at ``packages/node/native/proto/shard_runtime.proto`` is the
|
||||
single source of truth. Rather than commit generated stubs (which pin a protobuf
|
||||
runtime version and drift from the schema), this package generates the Python
|
||||
stubs on demand into a gitignored build directory and imports them. Generation is
|
||||
reproducible: it shells out to the pinned ``grpc_tools.protoc`` with the exact
|
||||
same flags as ``packages/node/native/scripts/generate_python.py``.
|
||||
|
||||
Typical use::
|
||||
|
||||
from meshnet_node import native_protocol as proto
|
||||
pb2 = proto.load()
|
||||
header = pb2.MessageHeader(work_id="w1", route_session_id="s1")
|
||||
|
||||
The checksum/fragment helpers encode the bounded-fragment tensor-bundle semantics
|
||||
so callers (and DGR-008/DGR-009) do not re-derive them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import importlib
|
||||
import importlib.util
|
||||
import pathlib
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
import zlib
|
||||
|
||||
# The wire schema version this build targets. Keep in sync with the
|
||||
# ``SCHEMA_VERSION_1`` enum member in the .proto.
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
_NATIVE_ROOT = pathlib.Path(__file__).resolve().parents[2] / "native"
|
||||
PROTO_DIR = _NATIVE_ROOT / "proto"
|
||||
PROTO_FILE = PROTO_DIR / "shard_runtime.proto"
|
||||
# ``build/`` is globally gitignored, so generated stubs never enter version control.
|
||||
GEN_DIR = _NATIVE_ROOT / "build" / "python"
|
||||
|
||||
_PB2_MODULE = "shard_runtime_pb2"
|
||||
_GRPC_MODULE = "shard_runtime_pb2_grpc"
|
||||
|
||||
# Reentrant: load_grpc() holds the lock and calls load(), which re-acquires it.
|
||||
_lock = threading.RLock()
|
||||
_cached_pb2: types.ModuleType | None = None
|
||||
_cached_grpc: types.ModuleType | None = None
|
||||
|
||||
|
||||
class ProtocGenerationError(RuntimeError):
|
||||
"""Raised when the protobuf stubs cannot be generated from the schema."""
|
||||
|
||||
|
||||
def _needs_regen(target: pathlib.Path) -> bool:
|
||||
if not target.exists():
|
||||
return True
|
||||
try:
|
||||
return PROTO_FILE.stat().st_mtime > target.stat().st_mtime
|
||||
except OSError:
|
||||
return True
|
||||
|
||||
|
||||
def generate(*, force: bool = False) -> pathlib.Path:
|
||||
"""Generate ``shard_runtime_pb2{,_grpc}.py`` into :data:`GEN_DIR`.
|
||||
|
||||
Returns the output directory. Reproducible and idempotent: regenerates only
|
||||
when the schema is newer than the stubs (or ``force`` is set). Requires the
|
||||
pinned ``grpc_tools`` (available in the project ``.venv``).
|
||||
"""
|
||||
if not PROTO_FILE.exists():
|
||||
raise ProtocGenerationError(f"schema not found: {PROTO_FILE}")
|
||||
|
||||
pb2_path = GEN_DIR / f"{_PB2_MODULE}.py"
|
||||
if not force and not _needs_regen(pb2_path):
|
||||
return GEN_DIR
|
||||
|
||||
try:
|
||||
from grpc_tools import protoc
|
||||
except ImportError as exc: # pragma: no cover - environment-dependent
|
||||
raise ProtocGenerationError(
|
||||
"grpc_tools is required to generate the Shard protocol stubs; "
|
||||
"install grpcio-tools (present in the project .venv)."
|
||||
) from exc
|
||||
|
||||
GEN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
well_known = _well_known_include()
|
||||
args = [
|
||||
"grpc_tools.protoc",
|
||||
f"-I{PROTO_DIR}",
|
||||
*([f"-I{well_known}"] if well_known else []),
|
||||
f"--python_out={GEN_DIR}",
|
||||
f"--grpc_python_out={GEN_DIR}",
|
||||
str(PROTO_FILE.name),
|
||||
]
|
||||
# protoc resolves the proto by name relative to -I, so run with PROTO_DIR
|
||||
# semantics by passing the bare filename plus the include path above.
|
||||
rc = protoc.main([a for a in args])
|
||||
if rc != 0:
|
||||
raise ProtocGenerationError(
|
||||
f"grpc_tools.protoc exited with status {rc} for {PROTO_FILE}"
|
||||
)
|
||||
if not pb2_path.exists(): # pragma: no cover - defensive
|
||||
raise ProtocGenerationError(f"protoc did not produce {pb2_path}")
|
||||
return GEN_DIR
|
||||
|
||||
|
||||
def _well_known_include() -> str | None:
|
||||
"""Bundled well-known .proto include dir shipped with grpc_tools, if any."""
|
||||
try:
|
||||
import grpc_tools
|
||||
|
||||
candidate = pathlib.Path(grpc_tools.__file__).parent / "_proto"
|
||||
return str(candidate) if candidate.is_dir() else None
|
||||
except Exception: # pragma: no cover - defensive
|
||||
return None
|
||||
|
||||
|
||||
def _import_generated(module_name: str) -> types.ModuleType:
|
||||
gen_dir = str(GEN_DIR)
|
||||
if gen_dir not in sys.path:
|
||||
sys.path.insert(0, gen_dir)
|
||||
if module_name in sys.modules:
|
||||
return sys.modules[module_name]
|
||||
return importlib.import_module(module_name)
|
||||
|
||||
|
||||
def load(*, force: bool = False) -> types.ModuleType:
|
||||
"""Return the generated ``shard_runtime_pb2`` module (messages only).
|
||||
|
||||
Generates the stubs on first use. Thread-safe and cached. Does not import
|
||||
grpc; message serialization/round-trip needs only this module.
|
||||
"""
|
||||
global _cached_pb2
|
||||
with _lock:
|
||||
if _cached_pb2 is not None and not force:
|
||||
return _cached_pb2
|
||||
generate(force=force)
|
||||
_cached_pb2 = _import_generated(_PB2_MODULE)
|
||||
return _cached_pb2
|
||||
|
||||
|
||||
def load_grpc(*, force: bool = False) -> types.ModuleType:
|
||||
"""Return the generated ``shard_runtime_pb2_grpc`` module (service stubs).
|
||||
|
||||
Requires the ``grpc`` runtime. Use for building the C++/Python worker; the
|
||||
round-trip/compat tests only need :func:`load`.
|
||||
"""
|
||||
global _cached_grpc
|
||||
with _lock:
|
||||
if _cached_grpc is not None and not force:
|
||||
return _cached_grpc
|
||||
generate(force=force)
|
||||
load() # ensure the _pb2 module the grpc stub imports is present
|
||||
_cached_grpc = _import_generated(_GRPC_MODULE)
|
||||
return _cached_grpc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checksum + bounded-fragment helpers (shared bundle semantics)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Algorithm-name strings mirror the ChecksumAlgorithm enum members without
|
||||
# importing the generated module (so this table is usable before load()).
|
||||
_CHECKSUM_CRC32C = "CHECKSUM_CRC32C"
|
||||
_CHECKSUM_CRC32 = "CHECKSUM_CRC32"
|
||||
_CHECKSUM_SHA256 = "CHECKSUM_SHA256"
|
||||
_CHECKSUM_NONE = "CHECKSUM_NONE"
|
||||
|
||||
|
||||
def _crc32c(data: bytes) -> int:
|
||||
"""Castagnoli CRC32C (software table). Deterministic, no external deps."""
|
||||
crc = 0xFFFFFFFF
|
||||
for byte in data:
|
||||
crc ^= byte
|
||||
for _ in range(8):
|
||||
crc = (crc >> 1) ^ (0x82F63B78 & -(crc & 1))
|
||||
return crc ^ 0xFFFFFFFF
|
||||
|
||||
|
||||
def compute_checksum(algorithm: int, data: bytes):
|
||||
"""Build a ``Checksum`` message for ``data`` under the given enum value.
|
||||
|
||||
``algorithm`` is a ``ChecksumAlgorithm`` enum int from the generated module.
|
||||
Uses only the standard library (crc32c software table, zlib.crc32, hashlib).
|
||||
"""
|
||||
pb2 = load()
|
||||
name = pb2.ChecksumAlgorithm.Name(algorithm)
|
||||
if name == _CHECKSUM_SHA256:
|
||||
value = hashlib.sha256(data).digest()
|
||||
elif name == _CHECKSUM_CRC32C:
|
||||
value = _crc32c(data).to_bytes(4, "big")
|
||||
elif name == _CHECKSUM_CRC32:
|
||||
value = (zlib.crc32(data) & 0xFFFFFFFF).to_bytes(4, "big")
|
||||
elif name == _CHECKSUM_NONE:
|
||||
value = b""
|
||||
else:
|
||||
raise ValueError(f"unsupported checksum algorithm: {name}")
|
||||
return pb2.Checksum(algorithm=algorithm, value=value)
|
||||
|
||||
|
||||
def verify_checksum(checksum, data: bytes) -> bool:
|
||||
"""True if ``checksum`` matches ``data`` (CHECKSUM_NONE always verifies)."""
|
||||
pb2 = load()
|
||||
if checksum.algorithm in (0, pb2.CHECKSUM_NONE):
|
||||
return True
|
||||
return compute_checksum(checksum.algorithm, data).value == checksum.value
|
||||
|
||||
|
||||
def fragment_tensor(
|
||||
*,
|
||||
name: str,
|
||||
shape,
|
||||
dtype: int,
|
||||
payload: bytes,
|
||||
byte_order: int | None = None,
|
||||
max_fragment_bytes: int = 1 << 20,
|
||||
compression: int | None = None,
|
||||
checksum_algorithm: int | None = None,
|
||||
):
|
||||
"""Build a :class:`NamedTensor` splitting ``payload`` into bounded fragments.
|
||||
|
||||
Fragments are ordered by ``byte_offset`` and each carries an optional
|
||||
per-fragment checksum. ``payload`` is treated as already compressed if
|
||||
``compression`` is set; this helper does not compress (that is the seam's
|
||||
policy in ``activation_compression``), it only frames.
|
||||
"""
|
||||
if max_fragment_bytes <= 0:
|
||||
raise ValueError("max_fragment_bytes must be positive")
|
||||
pb2 = load()
|
||||
if byte_order is None:
|
||||
byte_order = pb2.BYTE_ORDER_LITTLE_ENDIAN
|
||||
if compression is None:
|
||||
compression = pb2.COMPRESSION_NONE
|
||||
|
||||
chunks = [
|
||||
payload[i : i + max_fragment_bytes]
|
||||
for i in range(0, len(payload), max_fragment_bytes)
|
||||
] or [b""]
|
||||
fragments = []
|
||||
offset = 0
|
||||
for index, chunk in enumerate(chunks):
|
||||
frag = pb2.TensorFragment(
|
||||
fragment_index=index,
|
||||
fragment_count=len(chunks),
|
||||
byte_offset=offset,
|
||||
data=chunk,
|
||||
)
|
||||
if checksum_algorithm is not None:
|
||||
frag.checksum.CopyFrom(compute_checksum(checksum_algorithm, chunk))
|
||||
fragments.append(frag)
|
||||
offset += len(chunk)
|
||||
return pb2.NamedTensor(
|
||||
name=name,
|
||||
shape=list(shape),
|
||||
dtype=dtype,
|
||||
byte_order=byte_order,
|
||||
total_byte_length=len(payload),
|
||||
compression=compression,
|
||||
fragments=fragments,
|
||||
)
|
||||
|
||||
|
||||
def reassemble_tensor(named_tensor) -> bytes:
|
||||
"""Concatenate a :class:`NamedTensor`'s fragments back into the full payload.
|
||||
|
||||
Validates fragment ordering, total length, and any per-fragment checksums.
|
||||
"""
|
||||
fragments = sorted(named_tensor.fragments, key=lambda f: f.byte_offset)
|
||||
out = bytearray()
|
||||
for frag in fragments:
|
||||
if frag.byte_offset != len(out):
|
||||
raise ValueError(
|
||||
f"non-contiguous fragment at offset {frag.byte_offset} "
|
||||
f"(expected {len(out)})"
|
||||
)
|
||||
if frag.HasField("checksum") and not verify_checksum(frag.checksum, frag.data):
|
||||
raise ValueError(f"fragment {frag.fragment_index} checksum mismatch")
|
||||
out.extend(frag.data)
|
||||
if named_tensor.total_byte_length and len(out) != named_tensor.total_byte_length:
|
||||
raise ValueError(
|
||||
f"reassembled length {len(out)} != declared "
|
||||
f"{named_tensor.total_byte_length}"
|
||||
)
|
||||
return bytes(out)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SCHEMA_VERSION",
|
||||
"PROTO_FILE",
|
||||
"PROTO_DIR",
|
||||
"GEN_DIR",
|
||||
"ProtocGenerationError",
|
||||
"generate",
|
||||
"load",
|
||||
"load_grpc",
|
||||
"compute_checksum",
|
||||
"verify_checksum",
|
||||
"fragment_tensor",
|
||||
"reassemble_tensor",
|
||||
]
|
||||
@@ -71,6 +71,74 @@ class BenchmarkLane:
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BenchmarkWorkload:
|
||||
"""Identical request shape both recipes must run so speed stays comparable.
|
||||
|
||||
Pinning prompts, context lengths, output lengths, and sampling policy in the
|
||||
versioned contract is what makes the safetensors-versus-GGUF numbers a
|
||||
controlled comparison instead of two differently-configured runs.
|
||||
"""
|
||||
|
||||
prompts: tuple[str, ...]
|
||||
context_lengths: tuple[int, ...]
|
||||
output_lengths: tuple[int, ...]
|
||||
sampling_policy: str
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"prompts": list(self.prompts),
|
||||
"context_lengths": list(self.context_lengths),
|
||||
"output_lengths": list(self.output_lengths),
|
||||
"sampling_policy": self.sampling_policy,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QualityPolicy:
|
||||
"""Correctness/quality lane kept separate from the performance/fit lanes.
|
||||
|
||||
BF16 safetensors and Q2_K GGUF are not numerically equivalent, so quality is
|
||||
measured as its own lane (output drift against the BF16 reference under a
|
||||
documented tolerance) rather than assumed away by the speed/fit comparison.
|
||||
"""
|
||||
|
||||
statement: str
|
||||
reference_lane_runtime: str
|
||||
measured_lane_runtime: str
|
||||
max_output_drift: float
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"statement": self.statement,
|
||||
"reference_lane_runtime": self.reference_lane_runtime,
|
||||
"measured_lane_runtime": self.measured_lane_runtime,
|
||||
"max_output_drift": self.max_output_drift,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReleaseGate:
|
||||
"""Versioned thresholds later release gates (DGR-014) consume unchanged.
|
||||
|
||||
Thresholds live in the contract, not in code, so the release gate cannot be
|
||||
weakened after seeing implementation results.
|
||||
"""
|
||||
|
||||
min_decode_speedup: float
|
||||
max_artifact_bytes_ratio: float
|
||||
max_memory_bytes_ratio: float
|
||||
max_quality_drift: float
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"min_decode_speedup": self.min_decode_speedup,
|
||||
"max_artifact_bytes_ratio": self.max_artifact_bytes_ratio,
|
||||
"max_memory_bytes_ratio": self.max_memory_bytes_ratio,
|
||||
"max_quality_drift": self.max_quality_drift,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PerformanceContract:
|
||||
"""Machine-readable contract for the DGR-001 benchmark story."""
|
||||
|
||||
@@ -26,6 +26,16 @@
|
||||
"params": {
|
||||
"use_cache": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "llama-cpp-native",
|
||||
"version": "1",
|
||||
"backend_id": "llama.cpp",
|
||||
"description": "Project-owned native GGUF worker behind the Meshnet control plane.",
|
||||
"params": {
|
||||
"worker_transport": "grpc",
|
||||
"use_cache": true
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
375
packages/node/meshnet_node/runtime_recipe.py
Normal file
375
packages/node/meshnet_node/runtime_recipe.py
Normal file
@@ -0,0 +1,375 @@
|
||||
"""Exact artifact and runtime-recipe identity helpers.
|
||||
|
||||
The runtime recipe is the compatibility contract for one routable shard. It is
|
||||
kept separate from the user-facing recipe catalogue so the tracker can compare
|
||||
the exact execution footprint that was validated, not just a named recipe.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Mapping
|
||||
|
||||
|
||||
def _require_text(value: Any, field_name: str) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(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 _require_text(value, field_name)
|
||||
|
||||
|
||||
def _sha256_text(text: str) -> str:
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _stable_json(data: Any) -> str:
|
||||
return json.dumps(
|
||||
data,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
|
||||
|
||||
def _normalise_dtype(value: Any, default: str) -> str:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return default
|
||||
return text.removeprefix("torch.")
|
||||
return str(value).removeprefix("torch.")
|
||||
|
||||
|
||||
def _architecture_adapter_from_config(model_config: Any, default: str) -> str:
|
||||
if not isinstance(model_config, Mapping):
|
||||
return default
|
||||
for key in ("architecture_adapter", "model_type"):
|
||||
value = model_config.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
architectures = model_config.get("architectures")
|
||||
if isinstance(architectures, list) and architectures:
|
||||
first = architectures[0]
|
||||
if isinstance(first, str) and first.strip():
|
||||
return first
|
||||
text_config = model_config.get("text_config")
|
||||
if isinstance(text_config, Mapping):
|
||||
return _architecture_adapter_from_config(text_config, default)
|
||||
return default
|
||||
|
||||
|
||||
def _tokenizer_revision_from_config(
|
||||
model_id: str,
|
||||
revision: str | None,
|
||||
model_config: Any,
|
||||
) -> str:
|
||||
if isinstance(model_config, Mapping):
|
||||
for key in ("tokenizer_revision", "tokenizer_version", "_commit_hash"):
|
||||
value = model_config.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
if revision:
|
||||
return revision
|
||||
return model_id
|
||||
|
||||
|
||||
def _cache_layout_from_recipe_params(recipe_params: Mapping[str, Any] | None) -> str:
|
||||
if not recipe_params:
|
||||
return "local-hot-kv"
|
||||
use_cache = recipe_params.get("use_cache")
|
||||
if use_cache is False:
|
||||
return "stateless"
|
||||
if "cache_layout" in recipe_params:
|
||||
value = recipe_params.get("cache_layout")
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value
|
||||
return "local-hot-kv"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArtifactIdentity:
|
||||
"""Exact source artifact binding for a routable shard."""
|
||||
|
||||
model_id: str
|
||||
revision: str | None = None
|
||||
artifact_hash: str | None = None
|
||||
shard_start: int | None = None
|
||||
shard_end: int | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
_require_text(self.model_id, "artifact.model_id")
|
||||
_optional_text(self.revision, "artifact.revision")
|
||||
_optional_text(self.artifact_hash, "artifact.artifact_hash")
|
||||
if self.shard_start is not None and self.shard_start < 0:
|
||||
raise ValueError("'artifact.shard_start' must be >= 0")
|
||||
if self.shard_end is not None and self.shard_end < 0:
|
||||
raise ValueError("'artifact.shard_end' must be >= 0")
|
||||
if (
|
||||
self.shard_start is not None
|
||||
and self.shard_end is not None
|
||||
and self.shard_end < self.shard_start
|
||||
):
|
||||
raise ValueError("'artifact.shard_end' must be >= 'artifact.shard_start'")
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"model_id": self.model_id,
|
||||
"revision": self.revision,
|
||||
"artifact_hash": self.artifact_hash,
|
||||
"shard_start": self.shard_start,
|
||||
"shard_end": self.shard_end,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> "ArtifactIdentity":
|
||||
if not isinstance(data, Mapping):
|
||||
raise ValueError(f"'artifact' must be a JSON object, got {type(data).__name__}")
|
||||
return cls(
|
||||
model_id=_require_text(data.get("model_id"), "artifact.model_id"),
|
||||
revision=_optional_text(data.get("revision"), "artifact.revision"),
|
||||
artifact_hash=_optional_text(
|
||||
data.get("artifact_hash"), "artifact.artifact_hash"
|
||||
),
|
||||
shard_start=_optional_int(data.get("shard_start"), "artifact.shard_start"),
|
||||
shard_end=_optional_int(data.get("shard_end"), "artifact.shard_end"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RuntimeRecipeIdentity:
|
||||
"""Exact runtime recipe used for admission and handshake compatibility."""
|
||||
|
||||
weight_quantization: str
|
||||
activation_dtype: str
|
||||
compute_dtype: str
|
||||
kv_dtype: str
|
||||
kv_layout: str
|
||||
tokenizer_revision: str
|
||||
architecture_adapter: str
|
||||
backend_id: str
|
||||
runtime_version: str
|
||||
boundary_schema_version: int = 1
|
||||
cache_layout: str = "local-hot-kv"
|
||||
fingerprint: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
_require_text(self.weight_quantization, "runtime_recipe.weight_quantization")
|
||||
_require_text(self.activation_dtype, "runtime_recipe.activation_dtype")
|
||||
_require_text(self.compute_dtype, "runtime_recipe.compute_dtype")
|
||||
_require_text(self.kv_dtype, "runtime_recipe.kv_dtype")
|
||||
_require_text(self.kv_layout, "runtime_recipe.kv_layout")
|
||||
_require_text(self.tokenizer_revision, "runtime_recipe.tokenizer_revision")
|
||||
_require_text(self.architecture_adapter, "runtime_recipe.architecture_adapter")
|
||||
_require_text(self.backend_id, "runtime_recipe.backend_id")
|
||||
_require_text(self.runtime_version, "runtime_recipe.runtime_version")
|
||||
_require_text(self.cache_layout, "runtime_recipe.cache_layout")
|
||||
if self.boundary_schema_version < 1:
|
||||
raise ValueError("'runtime_recipe.boundary_schema_version' must be >= 1")
|
||||
expected = compatibility_fingerprint(self._fingerprint_payload())
|
||||
if not self.fingerprint:
|
||||
object.__setattr__(self, "fingerprint", expected)
|
||||
elif self.fingerprint != expected:
|
||||
raise ValueError(
|
||||
"'runtime_recipe.fingerprint' does not match the encoded fields"
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"weight_quantization": self.weight_quantization,
|
||||
"activation_dtype": self.activation_dtype,
|
||||
"compute_dtype": self.compute_dtype,
|
||||
"kv_dtype": self.kv_dtype,
|
||||
"kv_layout": self.kv_layout,
|
||||
"tokenizer_revision": self.tokenizer_revision,
|
||||
"architecture_adapter": self.architecture_adapter,
|
||||
"backend_id": self.backend_id,
|
||||
"runtime_version": self.runtime_version,
|
||||
"boundary_schema_version": self.boundary_schema_version,
|
||||
"cache_layout": self.cache_layout,
|
||||
"fingerprint": self.fingerprint,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> "RuntimeRecipeIdentity":
|
||||
if not isinstance(data, Mapping):
|
||||
raise ValueError(
|
||||
f"'runtime_recipe' must be a JSON object, got {type(data).__name__}"
|
||||
)
|
||||
boundary_schema_version = data.get("boundary_schema_version", 1)
|
||||
if isinstance(boundary_schema_version, bool) or not isinstance(
|
||||
boundary_schema_version, int
|
||||
):
|
||||
raise ValueError(
|
||||
"'runtime_recipe.boundary_schema_version' must be an integer"
|
||||
)
|
||||
return cls(
|
||||
weight_quantization=_require_text(
|
||||
data.get("weight_quantization"), "runtime_recipe.weight_quantization"
|
||||
),
|
||||
activation_dtype=_require_text(
|
||||
data.get("activation_dtype"), "runtime_recipe.activation_dtype"
|
||||
),
|
||||
compute_dtype=_require_text(
|
||||
data.get("compute_dtype"), "runtime_recipe.compute_dtype"
|
||||
),
|
||||
kv_dtype=_require_text(data.get("kv_dtype"), "runtime_recipe.kv_dtype"),
|
||||
kv_layout=_require_text(data.get("kv_layout"), "runtime_recipe.kv_layout"),
|
||||
tokenizer_revision=_require_text(
|
||||
data.get("tokenizer_revision"), "runtime_recipe.tokenizer_revision"
|
||||
),
|
||||
architecture_adapter=_require_text(
|
||||
data.get("architecture_adapter"),
|
||||
"runtime_recipe.architecture_adapter",
|
||||
),
|
||||
backend_id=_require_text(data.get("backend_id"), "runtime_recipe.backend_id"),
|
||||
runtime_version=_require_text(
|
||||
data.get("runtime_version"), "runtime_recipe.runtime_version"
|
||||
),
|
||||
boundary_schema_version=boundary_schema_version,
|
||||
cache_layout=_require_text(data.get("cache_layout"), "runtime_recipe.cache_layout"),
|
||||
fingerprint=_optional_text(data.get("fingerprint"), "runtime_recipe.fingerprint"),
|
||||
)
|
||||
|
||||
def _fingerprint_payload(self) -> dict[str, Any]:
|
||||
return {
|
||||
"weight_quantization": self.weight_quantization,
|
||||
"activation_dtype": self.activation_dtype,
|
||||
"compute_dtype": self.compute_dtype,
|
||||
"kv_dtype": self.kv_dtype,
|
||||
"kv_layout": self.kv_layout,
|
||||
"tokenizer_revision": self.tokenizer_revision,
|
||||
"architecture_adapter": self.architecture_adapter,
|
||||
"backend_id": self.backend_id,
|
||||
"runtime_version": self.runtime_version,
|
||||
"boundary_schema_version": self.boundary_schema_version,
|
||||
"cache_layout": self.cache_layout,
|
||||
}
|
||||
|
||||
|
||||
def _optional_int(value: Any, field_name: str) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{field_name!r} must be an integer")
|
||||
if value < 0:
|
||||
raise ValueError(f"{field_name!r} must be >= 0")
|
||||
return value
|
||||
|
||||
|
||||
def build_artifact_identity(
|
||||
*,
|
||||
model_id: str,
|
||||
revision: str | None = None,
|
||||
model_config: Any = None,
|
||||
artifact_hash: str | None = None,
|
||||
shard_start: int | None = None,
|
||||
shard_end: int | None = None,
|
||||
) -> ArtifactIdentity:
|
||||
"""Build a stable artifact binding from the locally loaded artifact."""
|
||||
resolved_hash = artifact_hash
|
||||
if resolved_hash is None:
|
||||
if isinstance(model_config, Mapping):
|
||||
resolved_hash = _hash_mapping(model_config)
|
||||
elif model_config is not None:
|
||||
resolved_hash = _sha256_text(_stable_json(model_config))
|
||||
if resolved_hash is None:
|
||||
resolved_hash = _sha256_text(
|
||||
_stable_json(
|
||||
{
|
||||
"model_id": model_id,
|
||||
"revision": revision,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
}
|
||||
)
|
||||
)
|
||||
return ArtifactIdentity(
|
||||
model_id=model_id,
|
||||
revision=revision,
|
||||
artifact_hash=resolved_hash,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
)
|
||||
|
||||
|
||||
def build_runtime_recipe_identity(
|
||||
*,
|
||||
model_id: str,
|
||||
weight_quantization: str,
|
||||
backend_id: str,
|
||||
runtime_version: str,
|
||||
revision: str | None = None,
|
||||
model_config: Any = None,
|
||||
recipe_params: Mapping[str, Any] | None = None,
|
||||
activation_dtype: Any = None,
|
||||
compute_dtype: Any = None,
|
||||
kv_dtype: Any = None,
|
||||
kv_layout: str | None = None,
|
||||
tokenizer_revision: str | None = None,
|
||||
architecture_adapter: str | None = None,
|
||||
boundary_schema_version: int = 1,
|
||||
cache_layout: str | None = None,
|
||||
) -> RuntimeRecipeIdentity:
|
||||
"""Build the exact runtime recipe used for compatibility admission."""
|
||||
activation = _normalise_dtype(activation_dtype, "bfloat16")
|
||||
compute = _normalise_dtype(compute_dtype, activation)
|
||||
kv_dtype_text = _normalise_dtype(kv_dtype, compute)
|
||||
kv_layout_text = kv_layout or "session-cache"
|
||||
tokenizer = tokenizer_revision or _tokenizer_revision_from_config(
|
||||
model_id, revision, model_config
|
||||
)
|
||||
architecture = architecture_adapter or _architecture_adapter_from_config(
|
||||
model_config, backend_id
|
||||
)
|
||||
cache_layout_text = cache_layout or _cache_layout_from_recipe_params(recipe_params)
|
||||
return RuntimeRecipeIdentity(
|
||||
weight_quantization=weight_quantization,
|
||||
activation_dtype=activation,
|
||||
compute_dtype=compute,
|
||||
kv_dtype=kv_dtype_text,
|
||||
kv_layout=kv_layout_text,
|
||||
tokenizer_revision=tokenizer,
|
||||
architecture_adapter=architecture,
|
||||
backend_id=backend_id,
|
||||
runtime_version=runtime_version,
|
||||
boundary_schema_version=boundary_schema_version,
|
||||
cache_layout=cache_layout_text,
|
||||
)
|
||||
|
||||
|
||||
def compatibility_fingerprint(data: Mapping[str, Any]) -> str:
|
||||
"""Return a stable SHA256 compatibility fingerprint for an exact route."""
|
||||
return "sha256:" + _sha256_text(_stable_json(data))
|
||||
|
||||
|
||||
def fingerprint_payload(
|
||||
*,
|
||||
model: Mapping[str, Any],
|
||||
shard: Mapping[str, Any],
|
||||
recipe: Mapping[str, Any],
|
||||
backend: Mapping[str, Any],
|
||||
artifact: Mapping[str, Any],
|
||||
runtime_recipe: Mapping[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"model": dict(model),
|
||||
"shard": dict(shard),
|
||||
"recipe": dict(recipe),
|
||||
"backend": dict(backend),
|
||||
"artifact": dict(artifact),
|
||||
"runtime_recipe": dict(runtime_recipe),
|
||||
}
|
||||
|
||||
|
||||
def _hash_mapping(data: Mapping[str, Any]) -> str:
|
||||
return "sha256:" + _sha256_text(_stable_json(data))
|
||||
@@ -29,6 +29,7 @@ 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 .gguf_backend import build_gguf_backend
|
||||
from .torch_server import TorchNodeServer
|
||||
from .wallet import load_or_create_wallet
|
||||
|
||||
@@ -662,6 +663,35 @@ def _resolve_recipe(recipe_id: str | None) -> tuple[RecipeManifest, Recipe]:
|
||||
return manifest, manifest.require(recipe_id or DEFAULT_RECIPE_ID)
|
||||
|
||||
|
||||
def _gguf_backend_for_recipe(
|
||||
recipe: Recipe,
|
||||
*,
|
||||
model_id: str,
|
||||
shard_start: int,
|
||||
shard_end: int,
|
||||
quantization: str,
|
||||
total_layers: int | None,
|
||||
device: str,
|
||||
model_revision: str | None = None,
|
||||
) -> object | None:
|
||||
"""Build the GGUF backend only for recipes that explicitly ask for it."""
|
||||
if recipe.backend_id != "llama.cpp":
|
||||
return None
|
||||
return build_gguf_backend(
|
||||
model_id=model_id,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
total_layers=total_layers,
|
||||
model_revision=model_revision,
|
||||
device_type=device,
|
||||
architecture_adapter="dense-llama",
|
||||
tokenizer_revision=model_revision or model_id,
|
||||
runtime_recipe_fingerprint=None,
|
||||
supports_kv_cache=recipe.params.get("use_cache", True) is not False,
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
@@ -875,7 +905,8 @@ def run_startup(
|
||||
|
||||
if model_id: # treat "" the same as None — no explicit model given
|
||||
full_sources: list[dict] = []
|
||||
# Auto-detect shard range from model config if not explicitly provided
|
||||
detected: int | None = None
|
||||
# Auto-detect shard range from model config if not explicitly provided.
|
||||
if shard_start is None or shard_end is None:
|
||||
try:
|
||||
detected = _detect_num_layers(model_id, cache_dir=cache_dir)
|
||||
@@ -939,22 +970,38 @@ def run_startup(
|
||||
shard_end = shard_end if shard_end is not None else detected - 1
|
||||
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
|
||||
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
model_id=model_id,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
cache_dir=cache_dir,
|
||||
debug=debug,
|
||||
max_loaded_shards=max_loaded_shards,
|
||||
force_cpu=force_cpu,
|
||||
recipe_params=recipe.params,
|
||||
total_layers=detected if detected is not None else (shard_end + 1 if shard_end is not None else None),
|
||||
device=device,
|
||||
model_revision=None,
|
||||
)
|
||||
print(
|
||||
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
|
||||
flush=True,
|
||||
)
|
||||
node_kwargs = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"model_id": model_id,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"quantization": quantization,
|
||||
"tracker_url": tracker_url,
|
||||
"route_timeout": route_timeout,
|
||||
"cache_dir": cache_dir,
|
||||
"debug": debug,
|
||||
"max_loaded_shards": max_loaded_shards,
|
||||
"force_cpu": force_cpu,
|
||||
"recipe_params": recipe.params,
|
||||
}
|
||||
if backend is not None:
|
||||
node_kwargs["backend"] = backend
|
||||
node = TorchNodeServer(**node_kwargs)
|
||||
capability_report = _admit_capability(
|
||||
node,
|
||||
model_id=model_id,
|
||||
@@ -968,10 +1015,15 @@ def run_startup(
|
||||
recipe=recipe,
|
||||
validator=capability_validator,
|
||||
)
|
||||
proof_shard = capability_report.shard
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||
shard_label = _format_shard_label(shard_start, shard_end, total_layers)
|
||||
shard_label = _format_shard_label(
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
total_layers,
|
||||
)
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
if hasattr(node, "set_advertised_endpoint"):
|
||||
@@ -994,16 +1046,17 @@ def run_startup(
|
||||
"model": model_id.split("/")[-1],
|
||||
"hf_repo": model_id,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
"tracker_mode": (proof_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(),
|
||||
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
|
||||
# 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,
|
||||
@@ -1011,8 +1064,8 @@ def run_startup(
|
||||
"downloaded_models": (
|
||||
_downloaded_model_inventory(
|
||||
model_id.split("/")[-1],
|
||||
shard_start,
|
||||
shard_end,
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
model_cache_path,
|
||||
hf_repo=model_id,
|
||||
model_sources=full_sources,
|
||||
@@ -1114,22 +1167,38 @@ def run_startup(
|
||||
hf_repo=assigned_hf_repo,
|
||||
model_sources=full_sources,
|
||||
)
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
model_id=assigned_hf_repo,
|
||||
shard_start=assigned_shard_start,
|
||||
shard_end=assigned_shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
cache_dir=cache_dir,
|
||||
debug=debug,
|
||||
max_loaded_shards=max_loaded_shards,
|
||||
force_cpu=force_cpu,
|
||||
recipe_params=recipe.params,
|
||||
total_layers=assigned_num_layers,
|
||||
device=device,
|
||||
model_revision=None,
|
||||
)
|
||||
print(
|
||||
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
|
||||
flush=True,
|
||||
)
|
||||
node_kwargs = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"model_id": assigned_hf_repo,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"quantization": quantization,
|
||||
"tracker_url": tracker_url,
|
||||
"route_timeout": route_timeout,
|
||||
"cache_dir": cache_dir,
|
||||
"debug": debug,
|
||||
"max_loaded_shards": max_loaded_shards,
|
||||
"force_cpu": force_cpu,
|
||||
"recipe_params": recipe.params,
|
||||
}
|
||||
if backend is not None:
|
||||
node_kwargs["backend"] = backend
|
||||
node = TorchNodeServer(**node_kwargs)
|
||||
capability_report = _admit_capability(
|
||||
node,
|
||||
model_id=assigned_hf_repo,
|
||||
@@ -1143,6 +1212,7 @@ def run_startup(
|
||||
recipe=recipe,
|
||||
validator=capability_validator,
|
||||
)
|
||||
proof_shard = capability_report.shard
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
@@ -1165,16 +1235,17 @@ def run_startup(
|
||||
"model": assigned_hf_repo.split("/")[-1],
|
||||
"hf_repo": assigned_hf_repo,
|
||||
"num_layers": assigned_num_layers,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
"tracker_mode": (proof_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(),
|
||||
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
|
||||
# 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,
|
||||
@@ -1182,8 +1253,8 @@ def run_startup(
|
||||
"downloaded_models": (
|
||||
_downloaded_model_inventory(
|
||||
assigned_hf_repo.split("/")[-1],
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
model_cache_path,
|
||||
hf_repo=assigned_hf_repo,
|
||||
model_sources=full_sources,
|
||||
@@ -1199,8 +1270,8 @@ def run_startup(
|
||||
tracker_url, auto_reg_payload, node, _node_start_time,
|
||||
)
|
||||
shard_label = _format_shard_label(
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
assigned_num_layers,
|
||||
)
|
||||
print(
|
||||
@@ -1315,22 +1386,38 @@ def run_startup(
|
||||
# 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub.
|
||||
_node_start_time = time.monotonic()
|
||||
if hf_repo and assigned_model != "stub-model":
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
model_id=hf_repo,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
cache_dir=shard_path,
|
||||
debug=debug,
|
||||
max_loaded_shards=max_loaded_shards,
|
||||
force_cpu=force_cpu,
|
||||
recipe_params=recipe.params,
|
||||
total_layers=total_layers,
|
||||
device=device,
|
||||
model_revision=None,
|
||||
)
|
||||
print(
|
||||
"Loading native llama.cpp model shard..." if backend is not None else "Loading real PyTorch model shard...",
|
||||
flush=True,
|
||||
)
|
||||
node_kwargs = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"model_id": hf_repo,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"quantization": quantization,
|
||||
"tracker_url": tracker_url,
|
||||
"route_timeout": route_timeout,
|
||||
"cache_dir": shard_path,
|
||||
"debug": debug,
|
||||
"max_loaded_shards": max_loaded_shards,
|
||||
"force_cpu": force_cpu,
|
||||
"recipe_params": recipe.params,
|
||||
}
|
||||
if backend is not None:
|
||||
node_kwargs["backend"] = backend
|
||||
node = TorchNodeServer(**node_kwargs)
|
||||
capability_report = _admit_capability(
|
||||
node,
|
||||
model_id=hf_repo,
|
||||
@@ -1379,6 +1466,7 @@ def run_startup(
|
||||
"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(),
|
||||
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
|
||||
# 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,
|
||||
@@ -1431,6 +1519,7 @@ def run_startup(
|
||||
recipe=recipe,
|
||||
validator=capability_validator,
|
||||
)
|
||||
proof_shard = capability_report.shard
|
||||
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}"
|
||||
@@ -1450,10 +1539,11 @@ def run_startup(
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_model,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"shard_checksum": shard_checksum,
|
||||
"capability_report": capability_report.to_dict(),
|
||||
"compatibility_fingerprint": capability_report.compatibility_fingerprint,
|
||||
# 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,
|
||||
@@ -1484,8 +1574,8 @@ def run_startup(
|
||||
if gpu_name:
|
||||
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
|
||||
shard_label = _format_shard_label(
|
||||
shard_start,
|
||||
shard_end,
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
assigned_total_layers,
|
||||
model_name=assigned_model,
|
||||
)
|
||||
|
||||
@@ -16,7 +16,10 @@ import time
|
||||
from typing import Any
|
||||
|
||||
from .admission import CapabilityContext, CapabilityValidator
|
||||
from . import __version__ as _PACKAGE_VERSION
|
||||
from .capability import STATUS_PASSED, CapabilityReport, build_capability_report
|
||||
from .gguf_ownership import authoritative_dense_llama_ownership
|
||||
from .runtime_recipe import build_runtime_recipe_identity
|
||||
|
||||
|
||||
def capability_report_for(
|
||||
@@ -30,6 +33,15 @@ def capability_report_for(
|
||||
recipe_version: str | None = None,
|
||||
backend_id: str | None = None,
|
||||
device: str | None = None,
|
||||
artifact_hash: str | None = None,
|
||||
activation_dtype: str | None = None,
|
||||
compute_dtype: str | None = None,
|
||||
kv_dtype: str | None = None,
|
||||
kv_layout: str | None = None,
|
||||
tokenizer_revision: str | None = None,
|
||||
architecture_adapter: str | None = None,
|
||||
boundary_schema_version: int = 1,
|
||||
cache_layout: str | None = None,
|
||||
validated_at: float | None = None,
|
||||
age_seconds: float = 0.0,
|
||||
diagnostics: Any = None,
|
||||
@@ -37,18 +49,49 @@ def capability_report_for(
|
||||
) -> CapabilityReport:
|
||||
"""A report describing `context`, with any field bent away from the truth."""
|
||||
now = time.time() if validated_at is None else validated_at
|
||||
backend = getattr(context, "backend", None)
|
||||
model_config = getattr(getattr(backend, "model", None), "config", None)
|
||||
model_config_payload = (
|
||||
model_config.to_dict() if hasattr(model_config, "to_dict") else model_config
|
||||
)
|
||||
resolved_cache_layout = (
|
||||
"stateless"
|
||||
if getattr(backend, "supports_kv_cache", False) is False
|
||||
else "local-hot-kv"
|
||||
)
|
||||
ownership = authoritative_dense_llama_ownership(backend, context.selection)
|
||||
runtime_recipe = build_runtime_recipe_identity(
|
||||
model_id=context.selection.model_id,
|
||||
revision=getattr(getattr(backend, "model", None), "revision", None),
|
||||
model_config=model_config_payload,
|
||||
recipe_params=context.recipe.params,
|
||||
weight_quantization=context.selection.quantization,
|
||||
backend_id=context.recipe.backend_id,
|
||||
runtime_version=_PACKAGE_VERSION,
|
||||
activation_dtype=activation_dtype,
|
||||
compute_dtype=compute_dtype,
|
||||
kv_dtype=kv_dtype,
|
||||
kv_layout=kv_layout or _backend_kv_layout(backend),
|
||||
tokenizer_revision=tokenizer_revision,
|
||||
architecture_adapter=architecture_adapter,
|
||||
boundary_schema_version=boundary_schema_version,
|
||||
cache_layout=cache_layout or resolved_cache_layout,
|
||||
)
|
||||
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,
|
||||
shard_start=ownership.start_layer if shard_start is None else shard_start,
|
||||
shard_end=ownership.end_layer 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,
|
||||
runtime=_runtime_versions(),
|
||||
artifact_hash=artifact_hash,
|
||||
runtime_recipe=runtime_recipe,
|
||||
owns_embedding=ownership.owns_embedding,
|
||||
owns_final_head=ownership.owns_final_head,
|
||||
status=status,
|
||||
duration_ms=duration_ms,
|
||||
diagnostics=diagnostics,
|
||||
@@ -68,3 +111,20 @@ def capability_stub(**overrides: Any) -> CapabilityValidator:
|
||||
return capability_report_for(context, **overrides)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def _runtime_versions() -> dict[str, str]:
|
||||
versions: dict[str, str] = {}
|
||||
for name in ("torch", "transformers"):
|
||||
try:
|
||||
module = __import__(name)
|
||||
except Exception:
|
||||
continue
|
||||
version = getattr(module, "__version__", None)
|
||||
if version:
|
||||
versions[name] = str(version)
|
||||
return versions
|
||||
|
||||
|
||||
def _backend_kv_layout(backend: Any) -> str:
|
||||
return "session-cache" if getattr(backend, "supports_kv_cache", False) else "stateless"
|
||||
|
||||
76
packages/node/native/CMakeLists.txt
Normal file
76
packages/node/native/CMakeLists.txt
Normal file
@@ -0,0 +1,76 @@
|
||||
# Reproducible C++ build wiring for the Shard runtime protocol (DGR-002).
|
||||
#
|
||||
# Generates C++ message stubs from proto/shard_runtime.proto and builds the
|
||||
# round-trip / cross-language compatibility test. Requires protoc and the
|
||||
# protobuf C++ runtime. Works with either a CONFIG-mode protobuf install
|
||||
# (protobuf::libprotobuf / protobuf::protoc targets, e.g. a from-source install
|
||||
# on CMAKE_PREFIX_PATH) or CMake's bundled FindProtobuf module.
|
||||
#
|
||||
# The gRPC C++ service stubs are generated separately by scripts/generate_cpp.sh
|
||||
# when grpc_cpp_plugin is present; the round-trip test needs only message
|
||||
# serialization, so gRPC is intentionally not a build dependency here.
|
||||
#
|
||||
# Configure & build (out-of-tree):
|
||||
# cmake -S packages/node/native -B packages/node/native/build/cpp
|
||||
# cmake --build packages/node/native/build/cpp
|
||||
# Run:
|
||||
# packages/node/native/build/cpp/shard_protocol_roundtrip_test --selftest
|
||||
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
project(shard_runtime_protocol CXX)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
|
||||
# Prefer a CONFIG-mode protobuf (modern imported targets); fall back to the
|
||||
# FindProtobuf module for system installs.
|
||||
find_package(Protobuf CONFIG QUIET)
|
||||
if(NOT Protobuf_FOUND)
|
||||
find_package(Protobuf REQUIRED)
|
||||
endif()
|
||||
|
||||
if(TARGET protobuf::protoc)
|
||||
set(SHARD_PROTOC_EXECUTABLE "$<TARGET_FILE:protobuf::protoc>")
|
||||
else()
|
||||
set(SHARD_PROTOC_EXECUTABLE "${Protobuf_PROTOC_EXECUTABLE}")
|
||||
endif()
|
||||
|
||||
if(TARGET protobuf::libprotobuf)
|
||||
set(SHARD_PROTOBUF_LINK protobuf::libprotobuf)
|
||||
else()
|
||||
set(SHARD_PROTOBUF_LINK ${Protobuf_LIBRARIES})
|
||||
endif()
|
||||
|
||||
set(PROTO_DIR "${CMAKE_CURRENT_SOURCE_DIR}/proto")
|
||||
set(PROTO_FILE "${PROTO_DIR}/shard_runtime.proto")
|
||||
set(GEN_DIR "${CMAKE_CURRENT_BINARY_DIR}/gen")
|
||||
file(MAKE_DIRECTORY "${GEN_DIR}")
|
||||
|
||||
set(PROTO_SRC "${GEN_DIR}/shard_runtime.pb.cc")
|
||||
set(PROTO_HDR "${GEN_DIR}/shard_runtime.pb.h")
|
||||
|
||||
add_custom_command(
|
||||
OUTPUT "${PROTO_SRC}" "${PROTO_HDR}"
|
||||
COMMAND "${SHARD_PROTOC_EXECUTABLE}"
|
||||
"--proto_path=${PROTO_DIR}"
|
||||
"--cpp_out=${GEN_DIR}"
|
||||
"${PROTO_FILE}"
|
||||
DEPENDS "${PROTO_FILE}"
|
||||
COMMENT "Generating C++ protobuf stubs from shard_runtime.proto"
|
||||
VERBATIM)
|
||||
|
||||
add_executable(shard_protocol_roundtrip_test
|
||||
tests/roundtrip_test.cpp
|
||||
"${PROTO_SRC}")
|
||||
|
||||
target_include_directories(shard_protocol_roundtrip_test PRIVATE "${GEN_DIR}")
|
||||
if(NOT TARGET protobuf::libprotobuf AND Protobuf_INCLUDE_DIRS)
|
||||
target_include_directories(shard_protocol_roundtrip_test PRIVATE
|
||||
${Protobuf_INCLUDE_DIRS})
|
||||
endif()
|
||||
|
||||
target_link_libraries(shard_protocol_roundtrip_test PRIVATE ${SHARD_PROTOBUF_LINK})
|
||||
|
||||
enable_testing()
|
||||
add_test(NAME shard_protocol_roundtrip
|
||||
COMMAND shard_protocol_roundtrip_test --selftest)
|
||||
24
packages/node/native/llama/README.md
Normal file
24
packages/node/native/llama/README.md
Normal file
@@ -0,0 +1,24 @@
|
||||
# Pinned llama.cpp source dependency
|
||||
|
||||
This directory keeps the llama.cpp fork boundary explicit and auditable.
|
||||
|
||||
Layout:
|
||||
|
||||
- `UPSTREAM_COMMIT` - the exact pinned commit.
|
||||
- `UPSTREAM_REPOSITORY` - the reproducible source dependency URL.
|
||||
- `UPSTREAM_ASSUMPTIONS.md` - the file/ABI assumptions that the build scripts
|
||||
validate.
|
||||
- `patches/` - numbered patch files applied on top of the pinned checkout.
|
||||
|
||||
The intended flow is:
|
||||
|
||||
1. Fetch or clone the pinned upstream checkout.
|
||||
2. Verify the checkout commit matches `UPSTREAM_COMMIT`.
|
||||
3. Check and apply the numbered patch stack.
|
||||
4. Build the worker scaffold from `examples/meshnet-worker/`.
|
||||
5. Copy the upstream `LICENSE` and `AUTHORS` files into the worker build tree so
|
||||
the attribution notices remain attached to the built artifact.
|
||||
|
||||
The patch stack in this story is intentionally minimal. It creates the project
|
||||
worker scaffold and the smoke-test CMake target without pulling Meshnet
|
||||
networking code into llama.cpp.
|
||||
35
packages/node/native/llama/UPSTREAM_ASSUMPTIONS.md
Normal file
35
packages/node/native/llama/UPSTREAM_ASSUMPTIONS.md
Normal file
@@ -0,0 +1,35 @@
|
||||
# llama.cpp upstream assumptions
|
||||
|
||||
This directory records the reproducible source dependency boundary for the
|
||||
pinned llama.cpp checkout used by the distributed GGUF runtime program.
|
||||
|
||||
Pinned upstream commit:
|
||||
|
||||
- `b3c9d1b846cc80a6360adb6aeaa4fcd8c4c8dcac`
|
||||
|
||||
Pinned upstream repository:
|
||||
|
||||
- `https://github.com/ggml-org/llama.cpp.git`
|
||||
|
||||
Assumptions checked by the build script:
|
||||
|
||||
- The checkout is exactly the pinned commit above.
|
||||
- The upstream source tree still ships `LICENSE`, `AUTHORS`, and
|
||||
`CMakeLists.txt` at the repository root.
|
||||
- The project-owned worker scaffold is built from
|
||||
`examples/meshnet-worker/`, which is introduced by the patch stack below.
|
||||
- The upstream license and attribution notices are preserved in the build
|
||||
output by copying the root `LICENSE` and `AUTHORS` files into the worker
|
||||
staging directory.
|
||||
|
||||
Compatibility notes:
|
||||
|
||||
- The current patch stack does not modify upstream llama.cpp runtime code yet.
|
||||
It adds a project-owned worker scaffold that can be built reproducibly from
|
||||
the pinned source checkout.
|
||||
- Later stories extend this boundary with actual llama.cpp execution patches.
|
||||
|
||||
Failure mode:
|
||||
|
||||
- If the checkout commit does not match the pin, the build script fails with a
|
||||
clear pin-mismatch error before patch application or compilation starts.
|
||||
1
packages/node/native/llama/UPSTREAM_COMMIT
Normal file
1
packages/node/native/llama/UPSTREAM_COMMIT
Normal file
@@ -0,0 +1 @@
|
||||
b3c9d1b846cc80a6360adb6aeaa4fcd8c4c8dcac
|
||||
1
packages/node/native/llama/UPSTREAM_REPOSITORY
Normal file
1
packages/node/native/llama/UPSTREAM_REPOSITORY
Normal file
@@ -0,0 +1 @@
|
||||
https://github.com/ggml-org/llama.cpp.git
|
||||
@@ -0,0 +1,35 @@
|
||||
diff --git a/examples/meshnet-worker/CMakeLists.txt b/examples/meshnet-worker/CMakeLists.txt
|
||||
new file mode 100644
|
||||
index 0000000000..8d9f9a1a2f
|
||||
--- /dev/null
|
||||
+++ b/examples/meshnet-worker/CMakeLists.txt
|
||||
@@ -0,0 +1,19 @@
|
||||
+cmake_minimum_required(VERSION 3.16)
|
||||
+project(meshnet_llama_worker CXX)
|
||||
+
|
||||
+set(CMAKE_CXX_STANDARD 17)
|
||||
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
+
|
||||
+configure_file(
|
||||
+ "${CMAKE_CURRENT_SOURCE_DIR}/version.h.in"
|
||||
+ "${CMAKE_CURRENT_BINARY_DIR}/version.h"
|
||||
+ @ONLY)
|
||||
+
|
||||
+add_executable(meshnet_worker
|
||||
+ meshnet_worker.cpp)
|
||||
+
|
||||
+target_include_directories(meshnet_worker PRIVATE "${CMAKE_CURRENT_BINARY_DIR}")
|
||||
+
|
||||
+enable_testing()
|
||||
+add_test(NAME meshnet_worker_smoke
|
||||
+ COMMAND meshnet_worker --smoke)
|
||||
diff --git a/examples/meshnet-worker/version.h.in b/examples/meshnet-worker/version.h.in
|
||||
new file mode 100644
|
||||
index 0000000000..0b75c4e60f
|
||||
--- /dev/null
|
||||
+++ b/examples/meshnet-worker/version.h.in
|
||||
@@ -0,0 +1,4 @@
|
||||
+#pragma once
|
||||
+
|
||||
+#define MESHNET_LLAMA_UPSTREAM_COMMIT "@MESHNET_LLAMA_UPSTREAM_COMMIT@"
|
||||
+#define MESHNET_LLAMA_PATCHSET_VERSION "@MESHNET_LLAMA_PATCHSET_VERSION@"
|
||||
43
packages/node/native/llama/templates/meshnet_worker.cpp
Normal file
43
packages/node/native/llama/templates/meshnet_worker.cpp
Normal file
@@ -0,0 +1,43 @@
|
||||
#include "version.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
|
||||
namespace {
|
||||
|
||||
bool fail(const std::string& why) {
|
||||
std::cerr << "meshnet_worker: FAIL: " << why << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
bool smoke = argc == 1;
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
const std::string arg = argv[i];
|
||||
if (arg == "--smoke") {
|
||||
smoke = true;
|
||||
} else {
|
||||
std::cerr << "unknown arg: " << arg << std::endl;
|
||||
return 2;
|
||||
}
|
||||
}
|
||||
|
||||
if (!smoke) {
|
||||
return fail("smoke mode not requested"), 1;
|
||||
}
|
||||
|
||||
if (MESHNET_LLAMA_UPSTREAM_COMMIT[0] == '\0') {
|
||||
return fail("upstream commit missing"), 1;
|
||||
}
|
||||
if (MESHNET_LLAMA_PATCHSET_VERSION[0] == '\0') {
|
||||
return fail("patchset version missing"), 1;
|
||||
}
|
||||
|
||||
std::cout << "meshnet worker scaffold ok" << std::endl;
|
||||
std::cout << "upstream commit: " << MESHNET_LLAMA_UPSTREAM_COMMIT << std::endl;
|
||||
std::cout << "patchset version: " << MESHNET_LLAMA_PATCHSET_VERSION << std::endl;
|
||||
return 0;
|
||||
}
|
||||
388
packages/node/native/proto/shard_runtime.proto
Normal file
388
packages/node/native/proto/shard_runtime.proto
Normal file
@@ -0,0 +1,388 @@
|
||||
// Shard runtime data-plane protocol for the distributed GGUF runtime (ADR-0024).
|
||||
//
|
||||
// This schema is the semantic contract between Python and C++ Shards. Direct
|
||||
// transport is gRPC over HTTP/2; the existing Meshnet relay may carry the same
|
||||
// serialized frames as opaque binary, so anything gRPC would normally carry in
|
||||
// call metadata (deadlines, cancellation intent) is ALSO representable inside
|
||||
// the messages for relay-transported seams.
|
||||
//
|
||||
// Design rules (see .scratch/distributed-gguf-runtime/RALPH-CONTEXT.md):
|
||||
// * One long-lived bidirectional ActivateSession stream per Route Session
|
||||
// Activation Seam. No per-token channel creation.
|
||||
// * Bounded chunking for prefill; a small decode fast path.
|
||||
// * The activation boundary is a versioned named-tensor bundle, because an
|
||||
// architecture boundary may require more than one tensor.
|
||||
// * Meshnet routing/billing/auth live outside this schema; only the data
|
||||
// plane and the identifiers needed to attribute and isolate work are here.
|
||||
//
|
||||
// Compatibility: proto3. Never renumber or reuse a field number. Add new fields
|
||||
// with new numbers only. Enums keep a 0 UNSPECIFIED member for forward compat.
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package meshnet.shard.v1;
|
||||
|
||||
option java_package = "com.meshnet.shard.v1";
|
||||
option java_outer_classname = "ShardRuntimeProto";
|
||||
option go_package = "meshnet/shard/v1;shardv1";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Versioning and enums
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Wire schema version. Bumped only on incompatible envelope changes; additive
|
||||
// field changes keep the same version and rely on proto3 unknown-field rules.
|
||||
enum SchemaVersion {
|
||||
SCHEMA_VERSION_UNSPECIFIED = 0;
|
||||
SCHEMA_VERSION_1 = 1;
|
||||
}
|
||||
|
||||
// Lifecycle phase of a seam message. RELEASE and CANCEL are represented both as
|
||||
// dedicated RPCs and as in-stream phases so a relay-carried stream can express
|
||||
// them without a separate channel.
|
||||
enum Phase {
|
||||
PHASE_UNSPECIFIED = 0;
|
||||
PHASE_PREFILL = 1;
|
||||
PHASE_DECODE = 2;
|
||||
PHASE_RELEASE = 3;
|
||||
PHASE_CANCEL = 4;
|
||||
}
|
||||
|
||||
// Tensor element type. GGUF quantized block types are enumerated explicitly so
|
||||
// a boundary bundle can carry pre-quantized payloads without reinterpretation.
|
||||
enum DType {
|
||||
DTYPE_UNSPECIFIED = 0;
|
||||
DTYPE_F32 = 1;
|
||||
DTYPE_F16 = 2;
|
||||
DTYPE_BF16 = 3;
|
||||
DTYPE_I64 = 4;
|
||||
DTYPE_I32 = 5;
|
||||
DTYPE_I16 = 6;
|
||||
DTYPE_I8 = 7;
|
||||
DTYPE_U8 = 8;
|
||||
DTYPE_BOOL = 9;
|
||||
DTYPE_Q8_0 = 20;
|
||||
DTYPE_Q4_0 = 21;
|
||||
DTYPE_Q4_K = 22;
|
||||
DTYPE_Q6_K = 23;
|
||||
}
|
||||
|
||||
// Byte order of a tensor payload. Explicit because Shards may run on
|
||||
// heterogeneous hardware and the relay carries opaque bytes.
|
||||
enum ByteOrder {
|
||||
BYTE_ORDER_UNSPECIFIED = 0;
|
||||
BYTE_ORDER_LITTLE_ENDIAN = 1;
|
||||
BYTE_ORDER_BIG_ENDIAN = 2;
|
||||
}
|
||||
|
||||
// Payload compression applied to a tensor fragment or message body.
|
||||
enum Compression {
|
||||
COMPRESSION_UNSPECIFIED = 0;
|
||||
COMPRESSION_NONE = 1;
|
||||
COMPRESSION_ZSTD = 2;
|
||||
}
|
||||
|
||||
// Checksum algorithm. CRC32C is the cheap per-fragment default; SHA256 is used
|
||||
// where stronger integrity is required.
|
||||
enum ChecksumAlgorithm {
|
||||
CHECKSUM_ALGORITHM_UNSPECIFIED = 0;
|
||||
CHECKSUM_NONE = 1;
|
||||
CHECKSUM_CRC32C = 2;
|
||||
CHECKSUM_CRC32 = 3;
|
||||
CHECKSUM_SHA256 = 4;
|
||||
}
|
||||
|
||||
// What the sender expects from the receiving Shard's Hot KV State for this work
|
||||
// (request side of the cache contract).
|
||||
enum CacheExpectation {
|
||||
CACHE_EXPECTATION_UNSPECIFIED = 0;
|
||||
CACHE_REUSE = 1; // reuse existing KV for (session, epoch)
|
||||
CACHE_FRESH = 2; // start a fresh KV context
|
||||
CACHE_BYPASS = 3; // stateless; do not persist KV
|
||||
}
|
||||
|
||||
// What the receiving Shard actually did with its KV State (result side).
|
||||
enum CacheResult {
|
||||
CACHE_RESULT_UNSPECIFIED = 0;
|
||||
CACHE_HIT = 1;
|
||||
CACHE_MISS = 2;
|
||||
CACHE_WRITTEN = 3;
|
||||
CACHE_BYPASSED = 4;
|
||||
}
|
||||
|
||||
// Coarse retry classification carried in structured status.
|
||||
enum RetryClass {
|
||||
RETRY_CLASS_UNSPECIFIED = 0;
|
||||
RETRY_CLASS_NONE = 1; // terminal success/no-retry
|
||||
RETRY_CLASS_RETRYABLE = 2; // transient; the same step may be retried
|
||||
RETRY_CLASS_FATAL = 3; // do not retry this route/epoch
|
||||
RETRY_CLASS_EPOCH_STALE = 4; // route epoch advanced; re-resolve route
|
||||
}
|
||||
|
||||
enum ServingStatus {
|
||||
SERVING_STATUS_UNSPECIFIED = 0;
|
||||
SERVING = 1;
|
||||
NOT_SERVING = 2;
|
||||
DRAINING = 3;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Common value messages
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Structured, transport-independent status. Mirrors canonical gRPC codes so a
|
||||
// relay-carried frame can express what a gRPC trailer normally would.
|
||||
message Status {
|
||||
uint32 code = 1; // canonical gRPC status code
|
||||
string message = 2;
|
||||
RetryClass retry_class = 3;
|
||||
map<string, string> details = 4;
|
||||
}
|
||||
|
||||
// Integrity check over an associated payload.
|
||||
message Checksum {
|
||||
ChecksumAlgorithm algorithm = 1;
|
||||
bytes value = 2;
|
||||
}
|
||||
|
||||
// Exact Model Artifact / runtime-recipe fingerprint. Both Shards MUST agree on
|
||||
// every populated field before activation; a mismatch is a fatal status.
|
||||
message ArtifactFingerprint {
|
||||
string model_id = 1; // e.g. "meta-llama/Llama-3.1-8B"
|
||||
string revision = 2; // artifact revision / commit
|
||||
string artifact_hash = 3; // hash of the GGUF/model artifact
|
||||
string quantization = 4; // e.g. "Q4_K_M", "F16"
|
||||
string runtime_recipe_fingerprint = 5; // DGR-003 recipe hash
|
||||
}
|
||||
|
||||
// Contiguous transformer layer range owned by a Shard (ADR-0012). end_layer is
|
||||
// exclusive. effective_start_layer is the overlap-safe start after de-dupe of
|
||||
// shared boundary layers between adjacent Shards.
|
||||
message ShardRange {
|
||||
uint32 start_layer = 1;
|
||||
uint32 end_layer = 2;
|
||||
uint32 effective_start_layer = 3;
|
||||
bool owns_embedding = 4;
|
||||
bool owns_final_head = 5;
|
||||
}
|
||||
|
||||
// Token position window for a message. start_position is the absolute index of
|
||||
// the first token; token_count is how many positions this message covers.
|
||||
message Position {
|
||||
uint64 start_position = 1;
|
||||
uint64 token_count = 2;
|
||||
uint64 sequence_length = 3; // total known context length, if known
|
||||
}
|
||||
|
||||
// Envelope carried by every seam message. Everything required to version,
|
||||
// route-attribute, isolate, order, and integrity-check a unit of work.
|
||||
message MessageHeader {
|
||||
SchemaVersion schema_version = 1;
|
||||
string work_id = 2; // request/work ID (idempotency scope)
|
||||
string route_session_id = 3; // Route Session ID
|
||||
uint64 route_epoch = 4; // route epoch; stale epochs are rejected
|
||||
ArtifactFingerprint fingerprint = 5;
|
||||
ShardRange shard_range = 6;
|
||||
Phase phase = 7;
|
||||
Position position = 8;
|
||||
uint64 idempotency_step = 9; // monotonic per (work_id) step counter
|
||||
CacheExpectation cache_expectation = 10;
|
||||
Compression compression = 11; // compression of THIS message's payloads
|
||||
Checksum checksum = 12; // checksum over THIS message's payload
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Versioned named-tensor bundle (the activation boundary payload)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// One bounded fragment of a tensor payload. Large tensors are split so no
|
||||
// single message is unbounded; fragments reassemble by byte_offset order.
|
||||
message TensorFragment {
|
||||
uint32 fragment_index = 1;
|
||||
uint32 fragment_count = 2;
|
||||
uint64 byte_offset = 3; // offset of this fragment within the full payload
|
||||
bytes data = 4;
|
||||
Checksum checksum = 5; // checksum over this fragment's (post-compression) data
|
||||
}
|
||||
|
||||
// A single named tensor with full description so the receiver never reinterprets
|
||||
// bytes implicitly.
|
||||
message NamedTensor {
|
||||
string name = 1;
|
||||
repeated uint64 shape = 2;
|
||||
DType dtype = 3;
|
||||
ByteOrder byte_order = 4;
|
||||
uint64 total_byte_length = 5; // full payload length across all fragments
|
||||
Compression compression = 6; // compression applied to fragment data
|
||||
repeated TensorFragment fragments = 7;
|
||||
}
|
||||
|
||||
// A versioned collection of named tensors representing one activation boundary.
|
||||
message TensorBundle {
|
||||
uint32 bundle_version = 1;
|
||||
repeated NamedTensor tensors = 2;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Session stream messages (bidirectional ActivateSession)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Opens a seam. Carries the header plus stream-scoped bounds. deadline_unix_nanos
|
||||
// lets a relay-carried stream express the call deadline gRPC would otherwise own.
|
||||
message SessionOpen {
|
||||
MessageHeader header = 1;
|
||||
uint64 deadline_unix_nanos = 2; // absolute deadline; 0 = none
|
||||
uint32 max_prefill_tokens_per_chunk = 3; // bound for prefill chunking
|
||||
uint32 max_fragment_bytes = 4; // bound for tensor fragment size
|
||||
FlowControl initial_credit = 5; // receiver's starting flow-control window
|
||||
}
|
||||
|
||||
// Bounded prefill chunk. A prefill is split into ordered chunks each covering at
|
||||
// most max_prefill_tokens_per_chunk positions; final_chunk marks the last one.
|
||||
message PrefillChunk {
|
||||
MessageHeader header = 1;
|
||||
uint32 chunk_index = 2;
|
||||
uint32 chunk_count = 3; // 0 if unknown/streaming
|
||||
bool final_chunk = 4;
|
||||
TensorBundle activations = 5;
|
||||
}
|
||||
|
||||
// Small decode fast path: a single-position (or tiny) step with minimal framing.
|
||||
// Reuses the same header for isolation/ordering but expects one activation bundle.
|
||||
message DecodeStep {
|
||||
MessageHeader header = 1;
|
||||
TensorBundle activation = 2;
|
||||
}
|
||||
|
||||
// Explicit HTTP/2-independent flow-control grant. credits is the number of
|
||||
// additional messages the receiver is willing to accept; the byte/message caps
|
||||
// bound in-flight work for backpressure.
|
||||
message FlowControl {
|
||||
uint64 credits = 1;
|
||||
uint64 max_in_flight_bytes = 2;
|
||||
uint64 max_in_flight_messages = 3;
|
||||
}
|
||||
|
||||
// Release a session's resources (Hot KV State, sequence) cleanly.
|
||||
message ReleaseRequest {
|
||||
MessageHeader header = 1;
|
||||
string reason = 2;
|
||||
}
|
||||
|
||||
message ReleaseResponse {
|
||||
Status status = 1;
|
||||
CacheResult cache_result = 2;
|
||||
}
|
||||
|
||||
// Cancel in-flight work for a session/step.
|
||||
message CancelRequest {
|
||||
MessageHeader header = 1;
|
||||
string reason = 2;
|
||||
}
|
||||
|
||||
message CancelResponse {
|
||||
Status status = 1;
|
||||
}
|
||||
|
||||
// Client -> server frames on the ActivateSession stream.
|
||||
message SessionActivation {
|
||||
oneof payload {
|
||||
SessionOpen open = 1;
|
||||
PrefillChunk prefill = 2;
|
||||
DecodeStep decode = 3;
|
||||
ReleaseRequest release = 4;
|
||||
CancelRequest cancel = 5;
|
||||
FlowControl flow_control = 6;
|
||||
}
|
||||
}
|
||||
|
||||
// Computed boundary output for a step: the next Shard's input tensors plus the
|
||||
// cache result and integrity for what was produced.
|
||||
message ActivationResult {
|
||||
MessageHeader header = 1;
|
||||
TensorBundle outputs = 2;
|
||||
CacheResult cache_result = 3;
|
||||
Status status = 4;
|
||||
}
|
||||
|
||||
message SessionAccepted {
|
||||
MessageHeader header = 1;
|
||||
FlowControl granted_credit = 2;
|
||||
Status status = 3;
|
||||
}
|
||||
|
||||
// Server -> client frames on the ActivateSession stream.
|
||||
message SessionResponse {
|
||||
oneof payload {
|
||||
SessionAccepted accepted = 1;
|
||||
ActivationResult result = 2;
|
||||
FlowControl flow_control = 3;
|
||||
Status status = 4;
|
||||
ReleaseResponse release_ack = 5;
|
||||
CancelResponse cancel_ack = 6;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Capability and health (unary)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
message ResourceBudget {
|
||||
uint64 weight_bytes = 1;
|
||||
uint64 kv_bytes = 2;
|
||||
uint64 scratch_bytes = 3;
|
||||
uint32 max_concurrent_sessions = 4;
|
||||
}
|
||||
|
||||
message CapabilityRequest {
|
||||
SchemaVersion schema_version = 1;
|
||||
}
|
||||
|
||||
message CapabilityResponse {
|
||||
SchemaVersion schema_version = 1;
|
||||
repeated SchemaVersion supported_schema_versions = 2;
|
||||
repeated string supported_architectures = 3; // e.g. "llama", "qwen3"
|
||||
repeated string supported_quantizations = 4;
|
||||
ShardRange servable_range = 5;
|
||||
ResourceBudget budget = 6;
|
||||
repeated Compression supported_compression = 7;
|
||||
repeated ChecksumAlgorithm supported_checksums = 8;
|
||||
ArtifactFingerprint loaded_fingerprint = 9; // empty if no artifact loaded
|
||||
}
|
||||
|
||||
message HealthRequest {
|
||||
string route_session_id = 1; // optional; empty for node-wide health
|
||||
}
|
||||
|
||||
message HealthResponse {
|
||||
ServingStatus status = 1;
|
||||
uint32 active_sessions = 2;
|
||||
uint32 queued_requests = 3;
|
||||
double kv_pressure = 4; // 0.0..1.0 fraction of KV budget in use
|
||||
uint64 rss_bytes = 5;
|
||||
Status detail = 6;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Service
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
service ShardRuntime {
|
||||
// Admission/capability negotiation.
|
||||
rpc GetCapability(CapabilityRequest) returns (CapabilityResponse);
|
||||
|
||||
// Liveness/backpressure telemetry.
|
||||
rpc Health(HealthRequest) returns (HealthResponse);
|
||||
|
||||
// One long-lived bidirectional stream per Route Session Activation Seam.
|
||||
// Deadlines/cancellation use gRPC call semantics on direct transport and the
|
||||
// in-message equivalents on relay transport; flow control uses FlowControl
|
||||
// frames; errors are structured Status.
|
||||
rpc ActivateSession(stream SessionActivation) returns (stream SessionResponse);
|
||||
|
||||
// Clean resource release (also expressible in-stream as PHASE_RELEASE).
|
||||
rpc Release(ReleaseRequest) returns (ReleaseResponse);
|
||||
|
||||
// Cancellation (also expressible in-stream as PHASE_CANCEL).
|
||||
rpc Cancel(CancelRequest) returns (CancelResponse);
|
||||
}
|
||||
187
packages/node/native/scripts/build_llama_worker.sh
Normal file
187
packages/node/native/scripts/build_llama_worker.sh
Normal file
@@ -0,0 +1,187 @@
|
||||
#!/usr/bin/env bash
|
||||
# Apply the numbered llama.cpp patch stack and build the worker scaffold.
|
||||
#
|
||||
# Default flow:
|
||||
# 1. Fetch the pinned llama.cpp source into a build directory if needed.
|
||||
# 2. Verify the checkout matches the pinned commit.
|
||||
# 3. Check/apply the numbered patch stack from packages/node/native/llama/.
|
||||
# 4. Compile and build the standalone worker scaffold.
|
||||
# 5. Copy upstream LICENSE/AUTHORS notices into the staging directory.
|
||||
#
|
||||
# This script is intentionally model-free and does not contact any inference
|
||||
# endpoint. It is a source/build reproducibility check.
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
NATIVE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
LLAMA_ROOT="${NATIVE_ROOT}/llama"
|
||||
UPSTREAM_COMMIT="$(tr -d '\n\r' < "${LLAMA_ROOT}/UPSTREAM_COMMIT")"
|
||||
UPSTREAM_REPOSITORY="$(tr -d '\n\r' < "${LLAMA_ROOT}/UPSTREAM_REPOSITORY")"
|
||||
PATCH_DIR="${LLAMA_ROOT}/patches"
|
||||
DEFAULT_SOURCE_DIR="${NATIVE_ROOT}/build/llama.cpp-src"
|
||||
DEFAULT_BUILD_DIR="${NATIVE_ROOT}/build/llama-worker"
|
||||
SOURCE_DIR="${DEFAULT_SOURCE_DIR}"
|
||||
BUILD_DIR="${DEFAULT_BUILD_DIR}"
|
||||
WORKTREE_DIR=""
|
||||
FETCH=1
|
||||
CXX_BIN="${CXX:-}"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: build_llama_worker.sh [--source-dir PATH] [--build-dir PATH] [--no-fetch]
|
||||
|
||||
Builds the project-owned worker scaffold from a pinned llama.cpp checkout.
|
||||
EOF
|
||||
}
|
||||
|
||||
fail() {
|
||||
echo "error: $*" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
while (($#)); do
|
||||
case "$1" in
|
||||
--source-dir)
|
||||
SOURCE_DIR="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
--build-dir)
|
||||
BUILD_DIR="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
--no-fetch)
|
||||
FETCH=0
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
fail "unknown argument: $1"
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
[[ -n "${SOURCE_DIR}" ]] || fail "source dir is empty"
|
||||
[[ -n "${BUILD_DIR}" ]] || fail "build dir is empty"
|
||||
|
||||
checkout_commit() {
|
||||
if [[ -f "${SOURCE_DIR}/.meshnet-upstream-commit" ]]; then
|
||||
tr -d '\n\r' < "${SOURCE_DIR}/.meshnet-upstream-commit"
|
||||
return 0
|
||||
fi
|
||||
if git -C "${SOURCE_DIR}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
|
||||
git -C "${SOURCE_DIR}" rev-parse HEAD
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
ensure_source() {
|
||||
if [[ -d "${SOURCE_DIR}" ]]; then
|
||||
return 0
|
||||
fi
|
||||
if [[ "${FETCH}" -ne 1 ]]; then
|
||||
fail "source dir ${SOURCE_DIR} does not exist and --no-fetch was set"
|
||||
fi
|
||||
|
||||
mkdir -p "${SOURCE_DIR}"
|
||||
git clone --quiet "${UPSTREAM_REPOSITORY}" "${SOURCE_DIR}" || fail "unable to clone ${UPSTREAM_REPOSITORY}"
|
||||
git -C "${SOURCE_DIR}" checkout --quiet "${UPSTREAM_COMMIT}" || fail "unable to checkout ${UPSTREAM_COMMIT}"
|
||||
printf '%s\n' "${UPSTREAM_COMMIT}" > "${SOURCE_DIR}/.meshnet-upstream-commit"
|
||||
printf '%s\n' "${UPSTREAM_REPOSITORY}" > "${SOURCE_DIR}/.meshnet-upstream-repository"
|
||||
}
|
||||
|
||||
verify_assumptions() {
|
||||
local observed_commit
|
||||
observed_commit="$(checkout_commit)" || fail "source tree does not expose a commit pin; write ${SOURCE_DIR}/.meshnet-upstream-commit or use a git checkout"
|
||||
if [[ "${observed_commit}" != "${UPSTREAM_COMMIT}" ]]; then
|
||||
fail "llama.cpp pin mismatch: expected ${UPSTREAM_COMMIT}, got ${observed_commit}"
|
||||
fi
|
||||
|
||||
for required in LICENSE AUTHORS CMakeLists.txt; do
|
||||
[[ -e "${SOURCE_DIR}/${required}" ]] || fail "missing upstream assumption file: ${required}"
|
||||
done
|
||||
}
|
||||
|
||||
apply_patches() {
|
||||
shopt -s nullglob
|
||||
local patches=("${PATCH_DIR}"/*.patch)
|
||||
shopt -u nullglob
|
||||
if ((${#patches[@]} == 0)); then
|
||||
fail "no patch files found in ${PATCH_DIR}"
|
||||
fi
|
||||
|
||||
for patch in "${patches[@]}"; do
|
||||
git -C "${SOURCE_DIR}" apply --check "${patch}" || fail "patch check failed: $(basename "${patch}")"
|
||||
done
|
||||
for patch in "${patches[@]}"; do
|
||||
git -C "${SOURCE_DIR}" apply "${patch}" || fail "patch apply failed: $(basename "${patch}")"
|
||||
done
|
||||
}
|
||||
|
||||
build_worker() {
|
||||
rm -rf "${BUILD_DIR}"
|
||||
mkdir -p "${BUILD_DIR}"
|
||||
WORKTREE_DIR="${BUILD_DIR}/llama.cpp-worktree"
|
||||
rm -rf "${WORKTREE_DIR}"
|
||||
mkdir -p "${WORKTREE_DIR}"
|
||||
cp -a "${SOURCE_DIR}/." "${WORKTREE_DIR}/"
|
||||
if [[ -f "${SOURCE_DIR}/.meshnet-upstream-commit" ]]; then
|
||||
cp "${SOURCE_DIR}/.meshnet-upstream-commit" "${WORKTREE_DIR}/.meshnet-upstream-commit"
|
||||
fi
|
||||
if [[ -f "${SOURCE_DIR}/.meshnet-upstream-repository" ]]; then
|
||||
cp "${SOURCE_DIR}/.meshnet-upstream-repository" "${WORKTREE_DIR}/.meshnet-upstream-repository"
|
||||
fi
|
||||
|
||||
SOURCE_DIR="${WORKTREE_DIR}"
|
||||
apply_patches
|
||||
|
||||
local worker_dir="${SOURCE_DIR}/examples/meshnet-worker"
|
||||
cp "${LLAMA_ROOT}/templates/meshnet_worker.cpp" "${worker_dir}/meshnet_worker.cpp"
|
||||
cat > "${worker_dir}/version.h" <<EOF
|
||||
#pragma once
|
||||
|
||||
#define MESHNET_LLAMA_UPSTREAM_COMMIT "${UPSTREAM_COMMIT}"
|
||||
#define MESHNET_LLAMA_PATCHSET_VERSION "0001"
|
||||
EOF
|
||||
|
||||
local compiler=""
|
||||
if [[ -n "${CXX_BIN}" ]] && command -v "${CXX_BIN}" >/dev/null 2>&1; then
|
||||
compiler="${CXX_BIN}"
|
||||
elif command -v g++ >/dev/null 2>&1; then
|
||||
compiler="g++"
|
||||
elif command -v c++ >/dev/null 2>&1; then
|
||||
compiler="c++"
|
||||
elif command -v clang++ >/dev/null 2>&1; then
|
||||
compiler="clang++"
|
||||
else
|
||||
fail "no C++ compiler found (need g++, c++, clang++, or $CXX)"
|
||||
fi
|
||||
|
||||
"${compiler}" -std=c++17 -O2 -Wall -Wextra \
|
||||
-I "${worker_dir}" \
|
||||
-o "${BUILD_DIR}/meshnet_worker" \
|
||||
"${worker_dir}/meshnet_worker.cpp"
|
||||
}
|
||||
|
||||
stage_notices() {
|
||||
local notice_dir="${BUILD_DIR}/upstream-notices"
|
||||
mkdir -p "${notice_dir}"
|
||||
cp "${SOURCE_DIR}/LICENSE" "${notice_dir}/LICENSE"
|
||||
cp "${SOURCE_DIR}/AUTHORS" "${notice_dir}/AUTHORS"
|
||||
printf '%s\n' "${UPSTREAM_COMMIT}" > "${notice_dir}/UPSTREAM_COMMIT"
|
||||
printf '%s\n' "${UPSTREAM_REPOSITORY}" > "${notice_dir}/UPSTREAM_REPOSITORY"
|
||||
}
|
||||
|
||||
main() {
|
||||
ensure_source
|
||||
verify_assumptions
|
||||
build_worker
|
||||
stage_notices
|
||||
"${BUILD_DIR}/meshnet_worker" --smoke
|
||||
echo "build ok: ${BUILD_DIR}/meshnet_worker"
|
||||
}
|
||||
|
||||
main "$@"
|
||||
43
packages/node/native/scripts/generate_cpp.sh
Normal file
43
packages/node/native/scripts/generate_cpp.sh
Normal file
@@ -0,0 +1,43 @@
|
||||
#!/usr/bin/env bash
|
||||
# Reproducibly generate the C++ Shard-protocol stubs from the schema.
|
||||
#
|
||||
# Produces message stubs (protoc --cpp_out) always, and gRPC C++ service stubs
|
||||
# (protoc --grpc_out with grpc_cpp_plugin) when the plugin is available. The
|
||||
# round-trip test needs only the message stubs; gRPC service stubs are for the
|
||||
# standalone C++ worker (DGR-008).
|
||||
#
|
||||
# Requirements: protoc (>=3.16). Optional: grpc_cpp_plugin for --grpc_out.
|
||||
#
|
||||
# Usage:
|
||||
# packages/node/native/scripts/generate_cpp.sh
|
||||
# Output: packages/node/native/build/cpp-gen/ (gitignored via build/).
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
NATIVE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
PROTO_DIR="${NATIVE_ROOT}/proto"
|
||||
PROTO_FILE="${PROTO_DIR}/shard_runtime.proto"
|
||||
OUT_DIR="${NATIVE_ROOT}/build/cpp-gen"
|
||||
|
||||
if ! command -v protoc >/dev/null 2>&1; then
|
||||
echo "error: protoc not found on PATH (install protobuf-compiler)." >&2
|
||||
exit 3
|
||||
fi
|
||||
|
||||
mkdir -p "${OUT_DIR}"
|
||||
|
||||
echo "generating C++ message stubs -> ${OUT_DIR}"
|
||||
protoc --proto_path="${PROTO_DIR}" --cpp_out="${OUT_DIR}" "${PROTO_FILE}"
|
||||
|
||||
if command -v grpc_cpp_plugin >/dev/null 2>&1; then
|
||||
echo "generating C++ gRPC service stubs -> ${OUT_DIR}"
|
||||
protoc --proto_path="${PROTO_DIR}" \
|
||||
--grpc_out="${OUT_DIR}" \
|
||||
--plugin=protoc-gen-grpc="$(command -v grpc_cpp_plugin)" \
|
||||
"${PROTO_FILE}"
|
||||
else
|
||||
echo "note: grpc_cpp_plugin not found; skipped --grpc_out (message stubs only)." >&2
|
||||
fi
|
||||
|
||||
echo "done:"
|
||||
ls -1 "${OUT_DIR}"
|
||||
76
packages/node/native/scripts/generate_python.py
Normal file
76
packages/node/native/scripts/generate_python.py
Normal file
@@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Reproducibly generate the Python Shard-protocol stubs from the schema.
|
||||
|
||||
This is the documented, no-manual-copy generation entry point referenced by
|
||||
``evidence/DGR-002/README.md``. It runs the pinned ``grpc_tools.protoc`` with the
|
||||
same flags ``meshnet_node.native_protocol.generate()`` uses on demand, but is
|
||||
kept self-contained (it does not import ``meshnet_node``) so it works regardless
|
||||
of which checkout the editable install points at.
|
||||
|
||||
Usage (from the project .venv):
|
||||
|
||||
python packages/node/native/scripts/generate_python.py
|
||||
|
||||
Output: ``packages/node/native/build/python/shard_runtime_pb2{,_grpc}.py``
|
||||
(``build/`` is gitignored).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
_NATIVE_ROOT = pathlib.Path(__file__).resolve().parents[1]
|
||||
PROTO_DIR = _NATIVE_ROOT / "proto"
|
||||
PROTO_FILE = PROTO_DIR / "shard_runtime.proto"
|
||||
GEN_DIR = _NATIVE_ROOT / "build" / "python"
|
||||
|
||||
|
||||
def _well_known_include() -> str | None:
|
||||
try:
|
||||
import grpc_tools
|
||||
|
||||
candidate = pathlib.Path(grpc_tools.__file__).parent / "_proto"
|
||||
return str(candidate) if candidate.is_dir() else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if not PROTO_FILE.exists():
|
||||
print(f"schema not found: {PROTO_FILE}", file=sys.stderr)
|
||||
return 2
|
||||
try:
|
||||
from grpc_tools import protoc
|
||||
except ImportError:
|
||||
print(
|
||||
"grpc_tools is required (pip install grpcio-tools); it is present in "
|
||||
"the project .venv.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 3
|
||||
|
||||
GEN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
well_known = _well_known_include()
|
||||
args = [
|
||||
"grpc_tools.protoc",
|
||||
f"-I{PROTO_DIR}",
|
||||
*([f"-I{well_known}"] if well_known else []),
|
||||
f"--python_out={GEN_DIR}",
|
||||
f"--grpc_python_out={GEN_DIR}",
|
||||
PROTO_FILE.name,
|
||||
]
|
||||
rc = protoc.main(args)
|
||||
if rc != 0:
|
||||
print(f"grpc_tools.protoc exited with status {rc}", file=sys.stderr)
|
||||
return rc
|
||||
|
||||
print(f"generated Python stubs into: {GEN_DIR}")
|
||||
for name in ("shard_runtime_pb2.py", "shard_runtime_pb2_grpc.py"):
|
||||
target = GEN_DIR / name
|
||||
print(f" {name}: {'ok' if target.exists() else 'MISSING'}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
180
packages/node/native/tests/roundtrip_test.cpp
Normal file
180
packages/node/native/tests/roundtrip_test.cpp
Normal file
@@ -0,0 +1,180 @@
|
||||
// C++ round-trip and cross-language compatibility test for the Shard protocol.
|
||||
//
|
||||
// Modes (composable):
|
||||
// --selftest serialize a sample message, parse it back, verify fields.
|
||||
// --read <path> parse a fixture serialized by another language; verify the
|
||||
// known fields; tolerate unknown fields (forward compat).
|
||||
// --write <path> serialize the C++ sample so another language can parse it.
|
||||
//
|
||||
// Exit code 0 means every requested check passed. The Python test drives this
|
||||
// binary to prove Python<->C++ wire compatibility in both directions.
|
||||
|
||||
#include "shard_runtime.pb.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
using namespace meshnet::shard::v1;
|
||||
|
||||
namespace {
|
||||
|
||||
bool Fail(const std::string& why) {
|
||||
std::cerr << "roundtrip_test: FAIL: " << why << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
SessionActivation MakeSample() {
|
||||
SessionActivation act;
|
||||
PrefillChunk* pre = act.mutable_prefill();
|
||||
|
||||
MessageHeader* h = pre->mutable_header();
|
||||
h->set_schema_version(SCHEMA_VERSION_1);
|
||||
h->set_work_id("w1");
|
||||
h->set_route_session_id("s1");
|
||||
h->set_route_epoch(3);
|
||||
h->set_phase(PHASE_PREFILL);
|
||||
h->set_idempotency_step(7);
|
||||
h->set_cache_expectation(CACHE_FRESH);
|
||||
h->set_compression(COMPRESSION_NONE);
|
||||
|
||||
ArtifactFingerprint* fp = h->mutable_fingerprint();
|
||||
fp->set_model_id("meta-llama/Llama-3.1-8B");
|
||||
fp->set_quantization("Q4_K_M");
|
||||
fp->set_runtime_recipe_fingerprint("recipe-abc");
|
||||
|
||||
ShardRange* sr = h->mutable_shard_range();
|
||||
sr->set_start_layer(0);
|
||||
sr->set_end_layer(16);
|
||||
sr->set_effective_start_layer(0);
|
||||
sr->set_owns_embedding(true);
|
||||
|
||||
Position* pos = h->mutable_position();
|
||||
pos->set_start_position(0);
|
||||
pos->set_token_count(5);
|
||||
pos->set_sequence_length(5);
|
||||
|
||||
pre->set_chunk_index(0);
|
||||
pre->set_chunk_count(1);
|
||||
pre->set_final_chunk(true);
|
||||
|
||||
TensorBundle* bundle = pre->mutable_activations();
|
||||
bundle->set_bundle_version(1);
|
||||
NamedTensor* t = bundle->add_tensors();
|
||||
t->set_name("hidden");
|
||||
t->add_shape(1);
|
||||
t->add_shape(4096);
|
||||
t->set_dtype(DTYPE_F16);
|
||||
t->set_byte_order(BYTE_ORDER_LITTLE_ENDIAN);
|
||||
t->set_total_byte_length(8);
|
||||
t->set_compression(COMPRESSION_NONE);
|
||||
TensorFragment* frag = t->add_fragments();
|
||||
frag->set_fragment_index(0);
|
||||
frag->set_fragment_count(1);
|
||||
frag->set_byte_offset(0);
|
||||
frag->set_data(std::string("\x01\x02\x03\x04\x05\x06\x07\x08", 8));
|
||||
|
||||
return act;
|
||||
}
|
||||
|
||||
bool CheckSample(const SessionActivation& act) {
|
||||
if (act.payload_case() != SessionActivation::kPrefill)
|
||||
return Fail("payload is not prefill");
|
||||
const PrefillChunk& pre = act.prefill();
|
||||
const MessageHeader& h = pre.header();
|
||||
if (h.schema_version() != SCHEMA_VERSION_1) return Fail("schema_version");
|
||||
if (h.work_id() != "w1") return Fail("work_id");
|
||||
if (h.route_session_id() != "s1") return Fail("route_session_id");
|
||||
if (h.route_epoch() != 3) return Fail("route_epoch");
|
||||
if (h.phase() != PHASE_PREFILL) return Fail("phase");
|
||||
if (h.idempotency_step() != 7) return Fail("idempotency_step");
|
||||
if (h.fingerprint().model_id() != "meta-llama/Llama-3.1-8B")
|
||||
return Fail("model_id");
|
||||
if (h.fingerprint().quantization() != "Q4_K_M") return Fail("quantization");
|
||||
if (h.shard_range().end_layer() != 16) return Fail("end_layer");
|
||||
if (!h.shard_range().owns_embedding()) return Fail("owns_embedding");
|
||||
if (h.position().token_count() != 5) return Fail("token_count");
|
||||
if (!pre.final_chunk()) return Fail("final_chunk");
|
||||
if (pre.activations().tensors_size() != 1) return Fail("tensors_size");
|
||||
const NamedTensor& t = pre.activations().tensors(0);
|
||||
if (t.name() != "hidden") return Fail("tensor name");
|
||||
if (t.dtype() != DTYPE_F16) return Fail("dtype");
|
||||
if (t.byte_order() != BYTE_ORDER_LITTLE_ENDIAN) return Fail("byte_order");
|
||||
if (t.shape_size() != 2 || t.shape(1) != 4096) return Fail("shape");
|
||||
if (t.fragments_size() != 1) return Fail("fragments_size");
|
||||
if (t.fragments(0).data().size() != 8) return Fail("fragment data length");
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ReadFile(const std::string& path, std::string* out) {
|
||||
std::ifstream in(path, std::ios::binary);
|
||||
if (!in) return false;
|
||||
std::ostringstream ss;
|
||||
ss << in.rdbuf();
|
||||
*out = ss.str();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool WriteFile(const std::string& path, const std::string& data) {
|
||||
std::ofstream out(path, std::ios::binary);
|
||||
if (!out) return false;
|
||||
out.write(data.data(), static_cast<std::streamsize>(data.size()));
|
||||
return static_cast<bool>(out);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
GOOGLE_PROTOBUF_VERIFY_VERSION;
|
||||
|
||||
std::string read_path;
|
||||
std::string write_path;
|
||||
bool selftest = (argc == 1);
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
std::string arg = argv[i];
|
||||
if (arg == "--selftest") {
|
||||
selftest = true;
|
||||
} else if (arg == "--read" && i + 1 < argc) {
|
||||
read_path = argv[++i];
|
||||
} else if (arg == "--write" && i + 1 < argc) {
|
||||
write_path = argv[++i];
|
||||
} else {
|
||||
std::cerr << "unknown/incomplete arg: " << arg << std::endl;
|
||||
return 2;
|
||||
}
|
||||
}
|
||||
|
||||
if (selftest) {
|
||||
SessionActivation sample = MakeSample();
|
||||
std::string bytes;
|
||||
if (!sample.SerializeToString(&bytes)) return Fail("serialize"), 1;
|
||||
SessionActivation parsed;
|
||||
if (!parsed.ParseFromString(bytes)) return Fail("parse"), 1;
|
||||
if (!CheckSample(parsed)) return 1;
|
||||
std::cout << "selftest ok (" << bytes.size() << " bytes)" << std::endl;
|
||||
}
|
||||
|
||||
if (!read_path.empty()) {
|
||||
std::string bytes;
|
||||
if (!ReadFile(read_path, &bytes)) return Fail("cannot read fixture"), 1;
|
||||
SessionActivation parsed;
|
||||
// ParseFromString tolerates and preserves unknown fields (forward compat).
|
||||
if (!parsed.ParseFromString(bytes)) return Fail("parse fixture"), 1;
|
||||
if (!CheckSample(parsed)) return 1;
|
||||
std::cout << "read ok (" << bytes.size() << " bytes)" << std::endl;
|
||||
}
|
||||
|
||||
if (!write_path.empty()) {
|
||||
SessionActivation sample = MakeSample();
|
||||
std::string bytes;
|
||||
if (!sample.SerializeToString(&bytes)) return Fail("serialize for write"), 1;
|
||||
if (!WriteFile(write_path, bytes)) return Fail("cannot write output"), 1;
|
||||
std::cout << "write ok (" << bytes.size() << " bytes)" << std::endl;
|
||||
}
|
||||
|
||||
google::protobuf::ShutdownProtobufLibrary();
|
||||
return 0;
|
||||
}
|
||||
@@ -58,6 +58,7 @@ STATE_MODEL_MISMATCH = "model-mismatch"
|
||||
STATE_SHARD_MISMATCH = "shard-mismatch"
|
||||
STATE_RECIPE_MISMATCH = "recipe-mismatch"
|
||||
STATE_CATALOGUE_INCOMPATIBLE = "catalogue-incompatible"
|
||||
STATE_COMPATIBILITY_MISMATCH = "compatibility-mismatch"
|
||||
|
||||
ALL_STATES = (
|
||||
STATE_ADMITTED,
|
||||
@@ -69,6 +70,7 @@ ALL_STATES = (
|
||||
STATE_SHARD_MISMATCH,
|
||||
STATE_RECIPE_MISMATCH,
|
||||
STATE_CATALOGUE_INCOMPATIBLE,
|
||||
STATE_COMPATIBILITY_MISMATCH,
|
||||
)
|
||||
|
||||
# --- Compatibility policy for nodes that predate the capability protocol. ---
|
||||
@@ -155,12 +157,17 @@ class CapabilityState:
|
||||
model_id: str | None = None
|
||||
shard_start: int | None = None
|
||||
shard_end: int | None = None
|
||||
owns_embedding: bool | None = None
|
||||
owns_final_head: bool | 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
|
||||
artifact_hash: str | None = None
|
||||
compatibility_fingerprint: str | None = None
|
||||
runtime_recipe_fingerprint: str | None = None
|
||||
validated_at: float | None = None
|
||||
recorded_at: float = 0.0
|
||||
schema_version: int | None = None
|
||||
@@ -187,12 +194,17 @@ class CapabilityState:
|
||||
"model_id": self.model_id,
|
||||
"shard_start": self.shard_start,
|
||||
"shard_end": self.shard_end,
|
||||
"owns_embedding": self.owns_embedding,
|
||||
"owns_final_head": self.owns_final_head,
|
||||
"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,
|
||||
"artifact_hash": self.artifact_hash,
|
||||
"compatibility_fingerprint": self.compatibility_fingerprint,
|
||||
"runtime_recipe_fingerprint": self.runtime_recipe_fingerprint,
|
||||
"validated_at": self.validated_at,
|
||||
"recorded_at": self.recorded_at,
|
||||
"schema_version": self.schema_version,
|
||||
@@ -222,6 +234,7 @@ def evaluate_report(
|
||||
shard_end: int | None,
|
||||
declared_recipe_id: str | None = None,
|
||||
declared_recipe_version: str | None = None,
|
||||
declared_compatibility_fingerprint: str | None = None,
|
||||
now: float | None = None,
|
||||
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS,
|
||||
) -> CapabilityState:
|
||||
@@ -308,6 +321,17 @@ def evaluate_report(
|
||||
f"the node declared v{declared_recipe_version}",
|
||||
)
|
||||
|
||||
if (
|
||||
declared_compatibility_fingerprint is not None
|
||||
and base.compatibility_fingerprint != declared_compatibility_fingerprint
|
||||
):
|
||||
return base.with_state(
|
||||
STATE_COMPATIBILITY_MISMATCH,
|
||||
"proof compatibility fingerprint does not match the node's declared "
|
||||
"artifact/runtime recipe; the artifact, tokenizer, architecture, "
|
||||
"boundary schema, activation recipe or cache layout differs",
|
||||
)
|
||||
|
||||
if status != STATUS_PASSED:
|
||||
return base.with_state(
|
||||
STATE_FAILED,
|
||||
@@ -344,6 +368,8 @@ def _parse_report(doc: Mapping[str, Any]) -> dict:
|
||||
shard = _object(doc.get("shard"), "shard")
|
||||
recipe = _object(doc.get("recipe"), "recipe")
|
||||
backend = _object(doc.get("backend"), "backend")
|
||||
artifact = _object_or_none(doc.get("artifact"), "artifact")
|
||||
runtime_recipe = _object_or_none(doc.get("runtime_recipe"), "runtime_recipe")
|
||||
|
||||
validated_at = doc.get("validated_at")
|
||||
if isinstance(validated_at, bool) or not isinstance(validated_at, (int, float)):
|
||||
@@ -357,6 +383,8 @@ def _parse_report(doc: Mapping[str, Any]) -> dict:
|
||||
"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"),
|
||||
"owns_embedding": _maybe_bool(shard.get("owns_embedding")),
|
||||
"owns_final_head": _maybe_bool(shard.get("owns_final_head")),
|
||||
"recipe_id": _text(recipe.get("recipe_id"), "recipe.recipe_id"),
|
||||
"recipe_version": _text(recipe.get("recipe_version"), "recipe.recipe_version"),
|
||||
"catalogue_version": _text(
|
||||
@@ -367,6 +395,15 @@ def _parse_report(doc: Mapping[str, Any]) -> dict:
|
||||
"quantization": _optional_text(
|
||||
backend.get("quantization"), "backend.quantization"
|
||||
),
|
||||
"artifact_hash": _optional_text(
|
||||
artifact.get("artifact_hash"), "artifact.artifact_hash"
|
||||
),
|
||||
"compatibility_fingerprint": _optional_text(
|
||||
doc.get("compatibility_fingerprint"), "compatibility_fingerprint"
|
||||
),
|
||||
"runtime_recipe_fingerprint": _optional_text(
|
||||
runtime_recipe.get("fingerprint"), "runtime_recipe.fingerprint"
|
||||
),
|
||||
"validated_at": float(validated_at),
|
||||
"schema_version": schema_version,
|
||||
"diagnostics": _diagnostics(doc.get("diagnostics")),
|
||||
@@ -380,6 +417,12 @@ def _object(value: Any, field_name: str) -> Mapping[str, Any]:
|
||||
return value
|
||||
|
||||
|
||||
def _object_or_none(value: Any, field_name: str) -> Mapping[str, Any]:
|
||||
if value is None:
|
||||
return {}
|
||||
return _object(value, field_name)
|
||||
|
||||
|
||||
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")
|
||||
@@ -404,6 +447,12 @@ def _maybe_int(value: Any) -> int | None:
|
||||
return value
|
||||
|
||||
|
||||
def _maybe_bool(value: Any) -> bool | None:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _diagnostics(value: Any) -> tuple[str, ...]:
|
||||
if not isinstance(value, list):
|
||||
return ()
|
||||
|
||||
@@ -56,6 +56,7 @@ from .capability import (
|
||||
DEFAULT_POLICY as DEFAULT_CAPABILITY_POLICY,
|
||||
POLICY_COMPAT,
|
||||
POLICY_ENFORCE,
|
||||
STATE_COMPATIBILITY_MISMATCH,
|
||||
STATE_ABSENT,
|
||||
STATE_ADMITTED,
|
||||
STATE_MODEL_MISMATCH,
|
||||
@@ -598,6 +599,7 @@ class _NodeEntry:
|
||||
"model_tokens_per_sec",
|
||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||
"relay_addr", "cert_fingerprint", "peer_id", "friendly_name",
|
||||
"compatibility_fingerprint",
|
||||
# heartbeat stats (reported by node, cumulative)
|
||||
"total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds",
|
||||
"current_requests",
|
||||
@@ -636,6 +638,7 @@ class _NodeEntry:
|
||||
cert_fingerprint: str | None = None,
|
||||
peer_id: str | None = None,
|
||||
friendly_name: str | None = None,
|
||||
compatibility_fingerprint: str | None = None,
|
||||
capability: "CapabilityState | None" = None,
|
||||
) -> None:
|
||||
self.node_id = node_id
|
||||
@@ -664,6 +667,7 @@ class _NodeEntry:
|
||||
self.cert_fingerprint = cert_fingerprint
|
||||
self.peer_id = peer_id
|
||||
self.friendly_name = friendly_name
|
||||
self.compatibility_fingerprint = compatibility_fingerprint
|
||||
# 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()
|
||||
@@ -782,6 +786,16 @@ def _node_admission(node: "_NodeEntry") -> CapabilityState:
|
||||
f"proof is for layers {state.shard_start}–{state.shard_end}, but the "
|
||||
f"node now serves layers {node.shard_start}–{node.shard_end}",
|
||||
)
|
||||
if (
|
||||
node.compatibility_fingerprint
|
||||
and state.compatibility_fingerprint
|
||||
and state.compatibility_fingerprint != node.compatibility_fingerprint
|
||||
):
|
||||
return state.with_state(
|
||||
STATE_COMPATIBILITY_MISMATCH,
|
||||
"proof compatibility fingerprint no longer matches the node's "
|
||||
"declared artifact/runtime recipe",
|
||||
)
|
||||
return state
|
||||
|
||||
|
||||
@@ -811,6 +825,12 @@ def _capability_from_registration(
|
||||
declared_recipe_version=(
|
||||
recipe_version if isinstance(recipe_version, str) else None
|
||||
),
|
||||
declared_compatibility_fingerprint=(
|
||||
value.strip()
|
||||
if isinstance((value := payload.get("compatibility_fingerprint")), str)
|
||||
and value.strip()
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -4588,6 +4608,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
relay_addr = body.get("relay_addr") or None
|
||||
cert_fingerprint = body.get("cert_fingerprint") or None
|
||||
peer_id = body.get("peer_id") or None
|
||||
compatibility_fingerprint = body.get("compatibility_fingerprint")
|
||||
if compatibility_fingerprint is not None and (
|
||||
not isinstance(compatibility_fingerprint, str) or not compatibility_fingerprint.strip()
|
||||
):
|
||||
self._send_json(400, {"error": "compatibility_fingerprint must be a string"})
|
||||
return
|
||||
compatibility_fingerprint = compatibility_fingerprint.strip() if isinstance(compatibility_fingerprint, str) else None
|
||||
try:
|
||||
friendly_name = _normalize_friendly_name(body.get("friendly_name"))
|
||||
except ValueError as exc:
|
||||
@@ -4647,6 +4674,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
cert_fingerprint=cert_fingerprint,
|
||||
peer_id=peer_id,
|
||||
friendly_name=friendly_name,
|
||||
compatibility_fingerprint=compatibility_fingerprint,
|
||||
capability=capability,
|
||||
)
|
||||
with server.lock:
|
||||
@@ -7052,6 +7080,12 @@ class TrackerServer:
|
||||
else None
|
||||
),
|
||||
friendly_name=_normalize_friendly_name(payload.get("friendly_name")),
|
||||
compatibility_fingerprint=(
|
||||
value.strip()
|
||||
if isinstance((value := payload.get("compatibility_fingerprint")), str)
|
||||
and value.strip()
|
||||
else None
|
||||
),
|
||||
# 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(
|
||||
|
||||
Reference in New Issue
Block a user