chore: clean superseded GGUF scaffolding
This commit is contained in:
Dobromir Popov
@@ -20,17 +20,9 @@ import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
from . import __version__ as _PACKAGE_VERSION
|
||||
from .capability import CapabilityReport, config_fingerprint
|
||||
from .capability import CapabilityReport
|
||||
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
|
||||
@@ -47,7 +39,6 @@ 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):
|
||||
@@ -86,7 +77,6 @@ class AdmissionRequirement:
|
||||
recipe_version: str
|
||||
backend_id: str
|
||||
device: str
|
||||
compatibility_fingerprint: str
|
||||
max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS
|
||||
|
||||
@classmethod
|
||||
@@ -104,9 +94,6 @@ 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,
|
||||
)
|
||||
|
||||
@@ -178,16 +165,6 @@ 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,
|
||||
@@ -246,157 +223,3 @@ 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"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,484 +0,0 @@
|
||||
"""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,16 +20,6 @@ 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
|
||||
|
||||
@@ -182,14 +172,6 @@ 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")
|
||||
@@ -236,8 +218,6 @@ 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)
|
||||
@@ -246,18 +226,9 @@ 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,
|
||||
"owns_embedding": self.owns_embedding,
|
||||
"owns_final_head": self.owns_final_head,
|
||||
}
|
||||
return {"start": self.start, "end": self.end}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> ShardRange:
|
||||
@@ -265,12 +236,6 @@ 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"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -371,8 +336,6 @@ class CapabilityReport:
|
||||
shard: ShardRange
|
||||
recipe: RecipeIdentity
|
||||
backend: BackendIdentity
|
||||
artifact: ArtifactIdentity
|
||||
runtime_recipe: RuntimeRecipeIdentity
|
||||
status: str
|
||||
validated_at: float
|
||||
duration_ms: int
|
||||
@@ -413,20 +376,6 @@ 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)
|
||||
|
||||
@@ -437,9 +386,6 @@ 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,
|
||||
@@ -452,9 +398,6 @@ 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(
|
||||
@@ -474,13 +417,7 @@ class CapabilityReport:
|
||||
):
|
||||
raise CapabilityReportError("'validated_at' must be a Unix timestamp")
|
||||
|
||||
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(
|
||||
return cls(
|
||||
schema_version=schema_version,
|
||||
model=ModelIdentity.from_dict(doc.get("model")),
|
||||
shard=ShardRange.from_dict(doc.get("shard")),
|
||||
@@ -490,18 +427,7 @@ 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:
|
||||
@@ -532,19 +458,6 @@ 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,
|
||||
@@ -555,62 +468,25 @@ def build_capability_report(
|
||||
or an already-computed ``sha256:…`` string. `validated_at` defaults to now,
|
||||
so callers that need determinism pass it explicitly.
|
||||
"""
|
||||
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(
|
||||
return CapabilityReport(
|
||||
model=ModelIdentity(
|
||||
model_id=model_id,
|
||||
revision=revision,
|
||||
model_config=model_config,
|
||||
recipe_params=recipe_params,
|
||||
weight_quantization=quantization or "unknown",
|
||||
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(
|
||||
backend_id=backend_id,
|
||||
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,
|
||||
device=device,
|
||||
device_name=device_name,
|
||||
quantization=quantization,
|
||||
runtime=dict(runtime or {}),
|
||||
),
|
||||
status=status,
|
||||
validated_at=time.time() if validated_at is None else validated_at,
|
||||
duration_ms=duration_ms,
|
||||
|
||||
@@ -36,8 +36,6 @@ 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,
|
||||
@@ -45,7 +43,6 @@ 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.
|
||||
@@ -467,28 +464,10 @@ 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=ownership.start_layer,
|
||||
shard_end=ownership.end_layer,
|
||||
shard_start=selection.shard_start,
|
||||
shard_end=selection.shard_end,
|
||||
recipe_id=recipe.id,
|
||||
recipe_version=recipe.version,
|
||||
catalogue_version=manifest.catalogue_version,
|
||||
@@ -498,9 +477,6 @@ 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,
|
||||
@@ -592,65 +568,6 @@ 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"
|
||||
|
||||
@@ -1,893 +0,0 @@
|
||||
"""Bounded failure, cancellation, and restart semantics for Shard streams (DGR-013).
|
||||
|
||||
Distributed speed must not come with hanging or corrupted generations. This module
|
||||
hardens the per-Route-Session decode stream that runs over the DGR-007 Hot KV State
|
||||
manager (isolated ``(session, epoch)`` KV) and the DGR-012 continuous-batch
|
||||
scheduler. It is deliberately backend-agnostic and pure-Python: it drives the same
|
||||
``KvBoundaryAdapter`` the default deterministic gate uses, so the whole matrix stays
|
||||
download-free, GPU-free, and API-credit-free while exercising the *real* KV
|
||||
isolation path (the pinned llama.cpp worker, DGR-008, implements the identical
|
||||
adapter contract).
|
||||
|
||||
The guarantees, mapped to the story's acceptance criteria:
|
||||
|
||||
* **Deadlines and heartbeat/health loss terminate blocked stream operations.**
|
||||
:class:`DeadlineGuard` bounds every step against an absolute deadline and a
|
||||
heartbeat-timeout; when either is breached it raises :class:`StreamTerminated`
|
||||
so a blocked stream never hangs.
|
||||
* **Cancellation propagates across every Shard and releases local KV and queued
|
||||
buffers.** :class:`ShardCancellationGroup` fans a single cancel across every
|
||||
node-local KV manager serving a Route Session and releases queued activation
|
||||
buffers; the DGR-012 scheduler's :meth:`~meshnet_node.batch_scheduler.
|
||||
ContinuousBatchScheduler.cancel` drops queued/active work on this node.
|
||||
* **Duplicate steps are idempotent; uncertain mutations are never replayed
|
||||
silently.** :class:`IdempotencyLedger` records each committed
|
||||
``(session, epoch, step)`` and returns the recorded token for a duplicate
|
||||
delivery instead of re-running it. A step whose outcome is *uncertain* (the
|
||||
worker died mid-mutation) is marked uncertain and can never be silently
|
||||
replayed — a replay attempt raises :class:`UncertainMutationError`, forcing an
|
||||
explicit verify-or-restart.
|
||||
* **Alpha failover restarts from token zero on a newly compatible route rather
|
||||
than importing unverified KV.** :class:`RestartController` opens a *new* route
|
||||
epoch, releases every shard's prior-epoch KV, and the restart re-prefills the
|
||||
whole prompt from token zero. The old epoch becomes stale (rejected by the KV
|
||||
manager); unverified KV is never migrated (RALPH runtime decision #14).
|
||||
* **Billing/work records distinguish completed, cancelled, failed, and unverified
|
||||
work.** :class:`WorkLedger` records a typed :class:`WorkRecord` per attempt;
|
||||
only :attr:`WorkStatus.COMPLETED` records are billable, so cancelled, failed,
|
||||
and uncertain (unverified) work is accounted but never charged.
|
||||
|
||||
:class:`HardenedSessionRunner` composes these into one drivable stream: it runs a
|
||||
single session's prefill+decode through the adapter under a deadline/heartbeat
|
||||
guard and a cancellation token, records the typed work outcome, and — via
|
||||
:meth:`HardenedSessionRunner.run_with_failover` — restarts a transient failure
|
||||
from token zero on a fresh epoch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field, replace
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Mapping, Sequence
|
||||
|
||||
from meshnet_node.batch_scheduler import DoneReason, GenerationRequest
|
||||
from meshnet_node.boundary_adapter import BoundaryContractError, TailOutput
|
||||
from meshnet_node.hot_kv_state import (
|
||||
CacheMiss,
|
||||
HotKvStateManager,
|
||||
IncompatibleCacheRecipeError,
|
||||
KvBoundaryAdapter,
|
||||
KvCacheMissError,
|
||||
StaleRouteEpochError,
|
||||
)
|
||||
|
||||
|
||||
class FailureSemanticsError(RuntimeError):
|
||||
"""Base class for failure/cancellation/restart errors."""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Typed outcomes: failure kinds and billing/work statuses.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class FailureKind(str, Enum):
|
||||
"""Why a stream step failed. Stable strings for the protocol's structured status."""
|
||||
|
||||
# Bounded termination of a blocked op.
|
||||
DEADLINE_EXCEEDED = "deadline-exceeded"
|
||||
HEARTBEAT_LOST = "heartbeat-lost"
|
||||
# Transport / worker loss (transient — a restart from token zero may succeed).
|
||||
WORKER_DEATH = "worker-death"
|
||||
STREAM_RESET = "stream-reset"
|
||||
# Protocol violations (deterministic — a restart would fail identically).
|
||||
MALFORMED_BUNDLE = "malformed-bundle"
|
||||
STALE_EPOCH = "stale-epoch"
|
||||
INCOMPATIBLE_RECIPE = "incompatible-recipe"
|
||||
# KV state expected by the caller is gone; re-prefill from token zero.
|
||||
CACHE_MISS = "cache-miss"
|
||||
# Explicit client cancellation.
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
# Failure kinds that a from-token-zero restart on a fresh route may recover from.
|
||||
# A protocol violation or an explicit bound (deadline/cancel) is NOT restartable —
|
||||
# retrying it would hang or fail identically, so we surface it instead.
|
||||
_RESTARTABLE = frozenset(
|
||||
{
|
||||
FailureKind.WORKER_DEATH,
|
||||
FailureKind.STREAM_RESET,
|
||||
FailureKind.CACHE_MISS,
|
||||
}
|
||||
)
|
||||
|
||||
# Failure kinds whose mutation outcome is *uncertain* — the KV may or may not have
|
||||
# advanced, so the confirmed work is billed as UNVERIFIED and never replayed
|
||||
# silently. Only an *unexpected* error raised while a step was executing is
|
||||
# uncertain (mapped to WORKER_DEATH). A stream reset, deadline, or cache miss
|
||||
# detected at a step boundary is certain: nothing committed for that step.
|
||||
_UNCERTAIN = frozenset({FailureKind.WORKER_DEATH})
|
||||
|
||||
|
||||
class WorkStatus(str, Enum):
|
||||
"""The billing-relevant outcome class of a unit of work (AC: billing records).
|
||||
|
||||
Only :attr:`COMPLETED` work is billable. Cancelled, failed, and unverified
|
||||
work is recorded distinctly so a client is never charged for a generation that
|
||||
hung, was cancelled, or whose mutations could not be verified.
|
||||
"""
|
||||
|
||||
COMPLETED = "completed"
|
||||
CANCELLED = "cancelled"
|
||||
FAILED = "failed"
|
||||
UNVERIFIED = "unverified"
|
||||
|
||||
|
||||
def work_status_for(kind: FailureKind) -> WorkStatus:
|
||||
"""Map a terminal failure kind to its billing/work status."""
|
||||
if kind is FailureKind.CANCELLED:
|
||||
return WorkStatus.CANCELLED
|
||||
if kind in _UNCERTAIN:
|
||||
return WorkStatus.UNVERIFIED
|
||||
return WorkStatus.FAILED
|
||||
|
||||
|
||||
def classify_exception(exc: BaseException) -> FailureKind:
|
||||
"""Classify a raised error into a :class:`FailureKind`.
|
||||
|
||||
Protocol violations map to their specific kind; a :class:`StreamTerminated`
|
||||
carries its own kind; any *unexpected* error is treated as worker death
|
||||
(an uncertain, transient loss), never silently ignored.
|
||||
"""
|
||||
if isinstance(exc, StreamTerminated):
|
||||
return exc.kind
|
||||
if isinstance(exc, OperationCancelled):
|
||||
return FailureKind.CANCELLED
|
||||
if isinstance(exc, StaleRouteEpochError):
|
||||
return FailureKind.STALE_EPOCH
|
||||
if isinstance(exc, IncompatibleCacheRecipeError):
|
||||
return FailureKind.INCOMPATIBLE_RECIPE
|
||||
if isinstance(exc, BoundaryContractError):
|
||||
return FailureKind.MALFORMED_BUNDLE
|
||||
if isinstance(exc, KvCacheMissError):
|
||||
return FailureKind.CACHE_MISS
|
||||
return FailureKind.WORKER_DEATH
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Deadlines and heartbeat/health loss.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class StreamTerminated(FailureSemanticsError):
|
||||
"""A blocked stream op was terminated by a deadline or heartbeat/health loss."""
|
||||
|
||||
def __init__(self, kind: FailureKind, detail: str = "") -> None:
|
||||
self.kind = kind
|
||||
self.detail = detail
|
||||
suffix = f": {detail}" if detail else ""
|
||||
super().__init__(f"stream terminated ({kind.value}){suffix}")
|
||||
|
||||
|
||||
class OperationCancelled(FailureSemanticsError):
|
||||
"""Raised when a step observes its :class:`CancellationToken` is cancelled."""
|
||||
|
||||
def __init__(self, reason: str = "client-cancel") -> None:
|
||||
self.reason = reason
|
||||
super().__init__(f"operation cancelled: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeadlineGuard:
|
||||
"""Bounds a blocked stream op against an absolute deadline and heartbeat loss.
|
||||
|
||||
``deadline`` is an absolute time on ``clock``'s scale (``None`` disables it).
|
||||
``heartbeat_timeout`` is the maximum tolerated gap since the last observed
|
||||
heartbeat; when the peer stops sending heartbeats (its health is lost) the gap
|
||||
grows past the timeout and :meth:`check` raises rather than blocking forever.
|
||||
Both bounds are checked with an injected ``clock`` so the matrix is
|
||||
deterministic.
|
||||
"""
|
||||
|
||||
deadline: float | None = None
|
||||
heartbeat_timeout: float | None = None
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
_last_heartbeat: float = field(default=0.0, init=False)
|
||||
_started: bool = field(default=False, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.heartbeat_timeout is not None and self.heartbeat_timeout <= 0:
|
||||
raise FailureSemanticsError("heartbeat_timeout must be positive")
|
||||
|
||||
def start(self) -> None:
|
||||
self._last_heartbeat = self.clock()
|
||||
self._started = True
|
||||
|
||||
def heartbeat(self) -> None:
|
||||
"""Record that the peer is alive (resets the heartbeat gap)."""
|
||||
self._last_heartbeat = self.clock()
|
||||
|
||||
def check(self) -> None:
|
||||
"""Raise :class:`StreamTerminated` if the deadline or heartbeat lapsed."""
|
||||
if not self._started:
|
||||
self.start()
|
||||
now = self.clock()
|
||||
if self.deadline is not None and now >= self.deadline:
|
||||
raise StreamTerminated(
|
||||
FailureKind.DEADLINE_EXCEEDED,
|
||||
f"deadline {self.deadline} reached at {now}",
|
||||
)
|
||||
if self.heartbeat_timeout is not None:
|
||||
gap = now - self._last_heartbeat
|
||||
if gap > self.heartbeat_timeout:
|
||||
raise StreamTerminated(
|
||||
FailureKind.HEARTBEAT_LOST,
|
||||
f"no heartbeat for {gap} > {self.heartbeat_timeout}",
|
||||
)
|
||||
|
||||
def remaining(self) -> float | None:
|
||||
if self.deadline is None:
|
||||
return None
|
||||
return self.deadline - self.clock()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Cancellation that propagates across shards and releases KV + queued buffers.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
"""A thread-safe one-shot cancellation flag shared by a Route Session's steps."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._cancelled = False
|
||||
self._reason = ""
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def cancel(self, reason: str = "client-cancel") -> None:
|
||||
with self._lock:
|
||||
if not self._cancelled:
|
||||
self._cancelled = True
|
||||
self._reason = reason
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
with self._lock:
|
||||
return self._cancelled
|
||||
|
||||
@property
|
||||
def reason(self) -> str:
|
||||
with self._lock:
|
||||
return self._reason
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
with self._lock:
|
||||
if self._cancelled:
|
||||
raise OperationCancelled(self._reason)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CancellationOutcome:
|
||||
"""What a :meth:`ShardCancellationGroup.cancel` released (for observability)."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
shards_released: int
|
||||
buffers_released: int
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"route_epoch": self.route_epoch,
|
||||
"shards_released": self.shards_released,
|
||||
"buffers_released": self.buffers_released,
|
||||
}
|
||||
|
||||
|
||||
class ShardCancellationGroup:
|
||||
"""Fan one cancellation across every node-local Shard of a Route Session.
|
||||
|
||||
A Route Session spans a chain of Shards, each with its own local Hot KV State
|
||||
manager (KV is never migrated between nodes). Cancelling the session must free
|
||||
*all* of that state: this group releases the ``(session, epoch)`` KV on every
|
||||
registered manager and invokes every registered queued-buffer release callback
|
||||
(the pending activation bundles a node holds for the session). Release is
|
||||
idempotent, so cancelling twice is safe.
|
||||
"""
|
||||
|
||||
def __init__(self, session_id: str, route_epoch: int) -> None:
|
||||
if not isinstance(session_id, str) or not session_id.strip():
|
||||
raise FailureSemanticsError("session_id must be a non-empty string")
|
||||
self.session_id = session_id
|
||||
self.route_epoch = int(route_epoch)
|
||||
self._managers: list[HotKvStateManager] = []
|
||||
self._buffers: list[Callable[[], None]] = []
|
||||
self._lock = threading.Lock()
|
||||
self._cancelled = False
|
||||
|
||||
def add_shard(self, manager: HotKvStateManager) -> "ShardCancellationGroup":
|
||||
with self._lock:
|
||||
self._managers.append(manager)
|
||||
return self
|
||||
|
||||
def add_queued_buffer(
|
||||
self, release: Callable[[], None]
|
||||
) -> "ShardCancellationGroup":
|
||||
"""Register a queued activation buffer's release callback."""
|
||||
with self._lock:
|
||||
self._buffers.append(release)
|
||||
return self
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
with self._lock:
|
||||
return self._cancelled
|
||||
|
||||
def cancel(self) -> CancellationOutcome:
|
||||
"""Release every shard's KV and every queued buffer for this session."""
|
||||
with self._lock:
|
||||
managers = list(self._managers)
|
||||
buffers = list(self._buffers)
|
||||
self._buffers.clear()
|
||||
self._cancelled = True
|
||||
shards_released = 0
|
||||
for manager in managers:
|
||||
if manager.release(self.session_id, self.route_epoch):
|
||||
shards_released += 1
|
||||
buffers_released = 0
|
||||
for release in buffers:
|
||||
release()
|
||||
buffers_released += 1
|
||||
return CancellationOutcome(
|
||||
session_id=self.session_id,
|
||||
route_epoch=self.route_epoch,
|
||||
shards_released=shards_released,
|
||||
buffers_released=buffers_released,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Idempotency: duplicate steps are no-ops; uncertain mutations never replay.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class StepPhase(str, Enum):
|
||||
IN_FLIGHT = "in-flight"
|
||||
COMMITTED = "committed"
|
||||
UNCERTAIN = "uncertain"
|
||||
|
||||
|
||||
class UncertainMutationError(FailureSemanticsError):
|
||||
"""Raised when a caller tries to replay a step whose outcome is uncertain.
|
||||
|
||||
A step is uncertain when its mutation may or may not have been applied (worker
|
||||
death / stream reset mid-append). Replaying it silently could double-apply KV
|
||||
or bill unverified work, so the ledger refuses: the caller must verify against
|
||||
the actual KV length or restart from token zero on a fresh epoch instead.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepKey:
|
||||
"""Identity of one idempotent stream step within a route epoch."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
step_index: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepDisposition:
|
||||
"""What :meth:`IdempotencyLedger.begin` decided for a step."""
|
||||
|
||||
fresh: bool
|
||||
token: int | None = None
|
||||
|
||||
@property
|
||||
def duplicate(self) -> bool:
|
||||
return not self.fresh
|
||||
|
||||
|
||||
class IdempotencyLedger:
|
||||
"""Records committed/uncertain stream steps so duplicates never re-mutate.
|
||||
|
||||
Keyed by ``(session, epoch, step_index)`` — the protocol's idempotency step.
|
||||
|
||||
* :meth:`begin` on a *fresh* key marks it in-flight and returns "execute".
|
||||
* :meth:`begin` on a *committed* key returns the recorded token so a duplicate
|
||||
delivery is a no-op (idempotent replay).
|
||||
* :meth:`begin` on an *in-flight* or *uncertain* key raises
|
||||
:class:`UncertainMutationError` — a concurrent duplicate or a replay of an
|
||||
unverified mutation is never silently applied.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._phase: dict[StepKey, StepPhase] = {}
|
||||
self._token: dict[StepKey, int] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def begin(self, key: StepKey) -> StepDisposition:
|
||||
with self._lock:
|
||||
phase = self._phase.get(key)
|
||||
if phase is None:
|
||||
self._phase[key] = StepPhase.IN_FLIGHT
|
||||
return StepDisposition(fresh=True)
|
||||
if phase is StepPhase.COMMITTED:
|
||||
return StepDisposition(fresh=False, token=self._token[key])
|
||||
# IN_FLIGHT (concurrent duplicate) or UNCERTAIN (post-crash replay):
|
||||
# both are unverified and must not be silently re-applied.
|
||||
raise UncertainMutationError(
|
||||
f"step {key.step_index} for session {key.session_id[:8]} epoch "
|
||||
f"{key.route_epoch} is {phase.value}; refusing silent replay"
|
||||
)
|
||||
|
||||
def commit(self, key: StepKey, token: int) -> None:
|
||||
with self._lock:
|
||||
self._phase[key] = StepPhase.COMMITTED
|
||||
self._token[key] = int(token)
|
||||
|
||||
def mark_uncertain(self, key: StepKey, detail: str = "") -> None:
|
||||
with self._lock:
|
||||
# A committed step is verified; never downgrade it.
|
||||
if self._phase.get(key) is StepPhase.COMMITTED:
|
||||
return
|
||||
self._phase[key] = StepPhase.UNCERTAIN
|
||||
|
||||
def phase_of(self, key: StepKey) -> StepPhase | None:
|
||||
with self._lock:
|
||||
return self._phase.get(key)
|
||||
|
||||
def committed_token(self, key: StepKey) -> int | None:
|
||||
with self._lock:
|
||||
return self._token.get(key)
|
||||
|
||||
def has_uncertain(self) -> bool:
|
||||
with self._lock:
|
||||
return any(p is StepPhase.UNCERTAIN for p in self._phase.values())
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Restart / alpha failover: from token zero on a fresh compatible route.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class RestartController:
|
||||
"""Alpha failover that restarts from token zero, never importing prior KV.
|
||||
|
||||
RALPH runtime decision #14: when the alpha (the head owning embedding + final
|
||||
head) fails, the route retries from token zero; unverified KV is never
|
||||
migrated. :meth:`failover` opens the *next* route epoch and releases every
|
||||
node-local shard's prior-epoch KV, so the restart begins with empty caches. The
|
||||
KV manager then treats the failed epoch as stale (a later reference to it is
|
||||
rejected), which is what keeps a half-computed cache from being reused.
|
||||
"""
|
||||
|
||||
def __init__(self, managers: Sequence[HotKvStateManager]) -> None:
|
||||
self._managers = list(managers)
|
||||
|
||||
def failover(self, session_id: str, failed_epoch: int) -> int:
|
||||
"""Advance to a fresh epoch and drop the failed epoch's KV on every shard."""
|
||||
new_epoch = int(failed_epoch) + 1
|
||||
for manager in self._managers:
|
||||
manager.release(session_id, failed_epoch)
|
||||
return new_epoch
|
||||
|
||||
def assert_fresh_start(self, session_id: str, new_epoch: int) -> None:
|
||||
"""Verify no shard carries KV for the new epoch (a true token-zero restart).
|
||||
|
||||
Any residual KV under the new epoch would be unverified imported state;
|
||||
this fails closed so a restart can never silently attend over it.
|
||||
"""
|
||||
for manager in self._managers:
|
||||
result = manager.resolve(session_id, new_epoch)
|
||||
if not isinstance(result, CacheMiss):
|
||||
raise FailureSemanticsError(
|
||||
f"restart epoch {new_epoch} for session {session_id[:8]} is not "
|
||||
"empty; refusing to import unverified KV"
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Billing / work records.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkRecord:
|
||||
"""A typed unit of served work, distinguishing what may be billed.
|
||||
|
||||
``tokens`` counts only *committed* generated tokens. Only a
|
||||
:attr:`WorkStatus.COMPLETED` record is billable; cancelled/failed/unverified
|
||||
records carry their confirmed token count for observability but are excluded
|
||||
from billing so uncompleted or unverified work is never charged.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
status: WorkStatus
|
||||
tokens: int
|
||||
failure_kind: FailureKind | None = None
|
||||
detail: str = ""
|
||||
|
||||
@property
|
||||
def billable(self) -> bool:
|
||||
return self.status is WorkStatus.COMPLETED
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"route_epoch": self.route_epoch,
|
||||
"status": self.status.value,
|
||||
"tokens": self.tokens,
|
||||
"failure_kind": self.failure_kind.value if self.failure_kind else None,
|
||||
"detail": self.detail,
|
||||
"billable": self.billable,
|
||||
}
|
||||
|
||||
|
||||
class WorkLedger:
|
||||
"""Append-only ledger of :class:`WorkRecord`, split by billing status."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._records: list[WorkRecord] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def record(self, record: WorkRecord) -> WorkRecord:
|
||||
with self._lock:
|
||||
self._records.append(record)
|
||||
return record
|
||||
|
||||
def records(self) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return list(self._records)
|
||||
|
||||
def records_for(self, session_id: str) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return [r for r in self._records if r.session_id == session_id]
|
||||
|
||||
def billable_records(self) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return [r for r in self._records if r.billable]
|
||||
|
||||
def billable_tokens(self) -> int:
|
||||
"""Total tokens that may be charged (completed work only)."""
|
||||
with self._lock:
|
||||
return sum(r.tokens for r in self._records if r.billable)
|
||||
|
||||
def counts_by_status(self) -> dict[str, int]:
|
||||
counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
|
||||
with self._lock:
|
||||
for record in self._records:
|
||||
counts[record.status.value] += 1
|
||||
return counts
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
with self._lock:
|
||||
records = [r.to_dict() for r in self._records]
|
||||
counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
|
||||
for record in records:
|
||||
counts[record["status"]] += 1
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"records": records,
|
||||
"counts_by_status": counts,
|
||||
"billable_tokens": sum(r["tokens"] for r in records if r["billable"]),
|
||||
}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# The hardened single-session stream runner.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunOutcome:
|
||||
"""The typed result of one hardened generation attempt."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
status: WorkStatus
|
||||
tokens: tuple[int, ...]
|
||||
failure_kind: FailureKind | None
|
||||
detail: str
|
||||
|
||||
@property
|
||||
def completed(self) -> bool:
|
||||
return self.status is WorkStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def token_count(self) -> int:
|
||||
return len(self.tokens)
|
||||
|
||||
@property
|
||||
def restartable(self) -> bool:
|
||||
return self.failure_kind in _RESTARTABLE
|
||||
|
||||
def work_record(self) -> WorkRecord:
|
||||
return WorkRecord(
|
||||
session_id=self.session_id,
|
||||
route_epoch=self.route_epoch,
|
||||
status=self.status,
|
||||
tokens=len(self.tokens),
|
||||
failure_kind=self.failure_kind,
|
||||
detail=self.detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FailoverResult:
|
||||
"""The result of a run that may have restarted from token zero after a failure."""
|
||||
|
||||
outcome: RunOutcome
|
||||
attempts: tuple[RunOutcome, ...]
|
||||
restarts: int
|
||||
|
||||
@property
|
||||
def completed(self) -> bool:
|
||||
return self.outcome.completed
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"final_status": self.outcome.status.value,
|
||||
"final_epoch": self.outcome.route_epoch,
|
||||
"restarts": self.restarts,
|
||||
"attempts": [
|
||||
{
|
||||
"route_epoch": a.route_epoch,
|
||||
"status": a.status.value,
|
||||
"failure_kind": a.failure_kind.value if a.failure_kind else None,
|
||||
"tokens": a.token_count,
|
||||
}
|
||||
for a in self.attempts
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class HardenedSessionRunner:
|
||||
"""Drive one Route Session's decode stream with bounded failure semantics.
|
||||
|
||||
The runner owns a single full-shard :class:`KvBoundaryAdapter` (head **and**
|
||||
tail, so a step samples a token) and threads every DGR-013 guarantee through a
|
||||
step loop:
|
||||
|
||||
* every step is bounded by a :class:`DeadlineGuard` and can observe a
|
||||
:class:`CancellationToken`;
|
||||
* every step is idempotent through an :class:`IdempotencyLedger` (a duplicate
|
||||
returns the recorded token; an uncertain mutation is never replayed);
|
||||
* any failure releases this session's KV (cancellation) and is recorded as a
|
||||
typed :class:`WorkRecord` in the :class:`WorkLedger`;
|
||||
* :meth:`run_with_failover` restarts a transient failure from token zero on a
|
||||
fresh epoch via a :class:`RestartController`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
adapter: KvBoundaryAdapter,
|
||||
*,
|
||||
clock: Callable[[], float] | None = None,
|
||||
work_ledger: WorkLedger | None = None,
|
||||
idempotency: IdempotencyLedger | None = None,
|
||||
) -> None:
|
||||
if not (adapter.is_head and adapter.is_tail):
|
||||
raise FailureSemanticsError(
|
||||
"HardenedSessionRunner needs a full (head+tail) shard so decode "
|
||||
"steps sample tokens; got a partial range "
|
||||
f"(head={adapter.is_head} tail={adapter.is_tail})"
|
||||
)
|
||||
self._adapter = adapter
|
||||
self._manager: HotKvStateManager = adapter.manager
|
||||
self._clock = clock or time.monotonic
|
||||
self.work_ledger = work_ledger or WorkLedger()
|
||||
self.idempotency = idempotency or IdempotencyLedger()
|
||||
|
||||
# -- single attempt -------------------------------------------------------
|
||||
|
||||
def run(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
*,
|
||||
deadline: float | None = None,
|
||||
heartbeat_timeout: float | None = None,
|
||||
cancel_token: CancellationToken | None = None,
|
||||
heartbeat: Callable[[int], bool] | None = None,
|
||||
before_step: Callable[[int], None] | None = None,
|
||||
) -> RunOutcome:
|
||||
"""Run one attempt of ``request``; record and return a typed outcome.
|
||||
|
||||
``deadline`` (absolute, on the injected clock) and ``heartbeat_timeout``
|
||||
bound blocked steps. ``cancel_token`` lets a client cancel mid-stream.
|
||||
``heartbeat(step)`` returns ``True`` when a heartbeat was heard before that
|
||||
step (resetting the health timer); ``before_step(step)`` is a fault-
|
||||
injection / clock-advance hook run before each step and may raise
|
||||
:class:`StreamTerminated` (e.g. a stream reset) or
|
||||
:class:`OperationCancelled`.
|
||||
"""
|
||||
sid = request.session_id
|
||||
epoch = request.route_epoch
|
||||
guard = DeadlineGuard(
|
||||
deadline=deadline,
|
||||
heartbeat_timeout=heartbeat_timeout,
|
||||
clock=self._clock,
|
||||
)
|
||||
guard.start()
|
||||
tokens: list[int] = []
|
||||
current_key: StepKey | None = None
|
||||
try:
|
||||
# step 0 is the prefill (emits the first token); steps 1..N are decodes.
|
||||
for step_index in range(request.max_new_tokens):
|
||||
# before_step is the fault-injection / clock-advance hook and may
|
||||
# itself terminate the step (stream reset, cancel); run it first so
|
||||
# a fault it raises takes effect on this step, then re-check the
|
||||
# bounds it may have advanced (deadline / heartbeat / cancel).
|
||||
if before_step is not None:
|
||||
before_step(step_index)
|
||||
if cancel_token is not None:
|
||||
cancel_token.raise_if_cancelled()
|
||||
if heartbeat is not None and heartbeat(step_index):
|
||||
guard.heartbeat()
|
||||
guard.check()
|
||||
|
||||
current_key = StepKey(sid, epoch, step_index)
|
||||
disposition = self.idempotency.begin(current_key)
|
||||
if disposition.duplicate:
|
||||
# Idempotent replay: reuse the recorded token, do not re-mutate.
|
||||
assert disposition.token is not None
|
||||
tokens.append(disposition.token)
|
||||
continue
|
||||
|
||||
token = self._execute_step(request, step_index, tokens)
|
||||
if isinstance(token, CacheMiss):
|
||||
# The expected KV was gone; the append never started, so this is
|
||||
# a certain (not uncertain) miss — restartable from token zero.
|
||||
return self._finish_failure(
|
||||
request,
|
||||
tokens,
|
||||
FailureKind.CACHE_MISS,
|
||||
str(token),
|
||||
cancel_token,
|
||||
)
|
||||
self.idempotency.commit(current_key, token)
|
||||
tokens.append(token)
|
||||
except (StreamTerminated, OperationCancelled) as exc:
|
||||
return self._finish_failure(
|
||||
request, tokens, classify_exception(exc), str(exc), cancel_token
|
||||
)
|
||||
except (
|
||||
BoundaryContractError,
|
||||
StaleRouteEpochError,
|
||||
IncompatibleCacheRecipeError,
|
||||
KvCacheMissError,
|
||||
) as exc:
|
||||
# Deterministic protocol/state errors, all validated before any KV
|
||||
# append committed — certain, not uncertain.
|
||||
return self._finish_failure(
|
||||
request, tokens, classify_exception(exc), str(exc), cancel_token
|
||||
)
|
||||
except UncertainMutationError as exc:
|
||||
# A replay of an unverified step reached the ledger — never silent.
|
||||
return self._finish_failure(
|
||||
request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - unexpected == worker death
|
||||
# An unexpected error mid-step may have left the KV half-mutated; mark
|
||||
# the step uncertain so it can never be silently replayed, then fail
|
||||
# closed as unverified work.
|
||||
if current_key is not None:
|
||||
self.idempotency.mark_uncertain(current_key, str(exc))
|
||||
return self._finish_failure(
|
||||
request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
|
||||
)
|
||||
|
||||
return self._finish_completed(request, tokens)
|
||||
|
||||
def _execute_step(
|
||||
self, request: GenerationRequest, step_index: int, tokens: list[int]
|
||||
) -> int | CacheMiss:
|
||||
sid = request.session_id
|
||||
epoch = request.route_epoch
|
||||
if step_index == 0:
|
||||
out = self._adapter.prefill(
|
||||
sid, epoch, token_ids=list(request.prompt_token_ids)
|
||||
)
|
||||
else:
|
||||
# expected_seq_len defends the KV layer against a desynchronised decode:
|
||||
# prompt positions plus the tokens already committed this run.
|
||||
expected = request.prompt_len + (step_index - 1)
|
||||
out = self._adapter.decode(
|
||||
sid,
|
||||
epoch,
|
||||
token_ids=[tokens[-1]],
|
||||
expected_seq_len=expected,
|
||||
)
|
||||
if isinstance(out, CacheMiss):
|
||||
return out
|
||||
if not isinstance(out, TailOutput):
|
||||
raise FailureSemanticsError(
|
||||
"full-shard step did not yield a sampled token; got "
|
||||
f"{type(out).__name__}"
|
||||
)
|
||||
return int(out.token_id)
|
||||
|
||||
# -- failover across restarts --------------------------------------------
|
||||
|
||||
def run_with_failover(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
controller: RestartController,
|
||||
*,
|
||||
max_restarts: int = 3,
|
||||
**run_kwargs: Any,
|
||||
) -> FailoverResult:
|
||||
"""Run ``request``, restarting a transient failure from token zero.
|
||||
|
||||
On a restartable failure (worker death, stream reset, cache miss) the
|
||||
controller advances to a fresh epoch and drops the failed epoch's KV; the
|
||||
next attempt re-prefills the whole prompt from token zero. A deterministic
|
||||
failure (deadline, cancel, malformed bundle, stale epoch) is returned as-is
|
||||
— retrying it would hang or fail identically. Per-attempt fault-injection
|
||||
hooks (``before_step`` / ``heartbeat``) are only applied to the *first*
|
||||
attempt so a restart runs clean.
|
||||
"""
|
||||
if max_restarts < 0:
|
||||
raise FailureSemanticsError("max_restarts must be >= 0")
|
||||
epoch = request.route_epoch
|
||||
attempts: list[RunOutcome] = []
|
||||
first_kwargs = run_kwargs
|
||||
for attempt in range(max_restarts + 1):
|
||||
attempt_request = replace(request, route_epoch=epoch)
|
||||
kwargs = first_kwargs if attempt == 0 else {}
|
||||
outcome = self.run(attempt_request, **kwargs)
|
||||
attempts.append(outcome)
|
||||
if outcome.completed or not outcome.restartable or attempt == max_restarts:
|
||||
return FailoverResult(
|
||||
outcome=outcome, attempts=tuple(attempts), restarts=attempt
|
||||
)
|
||||
# Alpha failover: fresh epoch, drop prior-epoch KV on every shard, and
|
||||
# verify the new epoch starts empty (no unverified KV import).
|
||||
epoch = controller.failover(request.session_id, epoch)
|
||||
controller.assert_fresh_start(request.session_id, epoch)
|
||||
# Unreachable: the loop always returns, but keep the type-checker happy.
|
||||
raise FailureSemanticsError("run_with_failover exhausted without returning")
|
||||
|
||||
# -- outcome bookkeeping --------------------------------------------------
|
||||
|
||||
def _finish_completed(
|
||||
self, request: GenerationRequest, tokens: list[int]
|
||||
) -> RunOutcome:
|
||||
outcome = RunOutcome(
|
||||
session_id=request.session_id,
|
||||
route_epoch=request.route_epoch,
|
||||
status=WorkStatus.COMPLETED,
|
||||
tokens=tuple(tokens),
|
||||
failure_kind=None,
|
||||
detail="",
|
||||
)
|
||||
self.work_ledger.record(outcome.work_record())
|
||||
return outcome
|
||||
|
||||
def _finish_failure(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
tokens: list[int],
|
||||
kind: FailureKind,
|
||||
detail: str,
|
||||
cancel_token: CancellationToken | None,
|
||||
) -> RunOutcome:
|
||||
# Cancellation semantics: release this session's local KV so a failed or
|
||||
# cancelled stream never leaks its cache. release() is idempotent.
|
||||
self._manager.release(request.session_id, request.route_epoch)
|
||||
if cancel_token is not None and kind is not FailureKind.CANCELLED:
|
||||
# Ensure downstream shards sharing the token also stop.
|
||||
cancel_token.cancel(kind.value)
|
||||
outcome = RunOutcome(
|
||||
session_id=request.session_id,
|
||||
route_epoch=request.route_epoch,
|
||||
status=work_status_for(kind),
|
||||
tokens=tuple(tokens),
|
||||
failure_kind=kind,
|
||||
detail=detail,
|
||||
)
|
||||
self.work_ledger.record(outcome.work_record())
|
||||
return outcome
|
||||
@@ -1,423 +0,0 @@
|
||||
"""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,
|
||||
)
|
||||
@@ -1,287 +0,0 @@
|
||||
"""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
|
||||
@@ -1,918 +0,0 @@
|
||||
"""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,10 +323,6 @@ 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)
|
||||
@@ -348,17 +344,6 @@ 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")
|
||||
|
||||
@@ -1,300 +0,0 @@
|
||||
"""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",
|
||||
]
|
||||
@@ -1,563 +0,0 @@
|
||||
"""Versioned performance contract metadata and stub benchmark runner for DGR-001.
|
||||
|
||||
This module captures the *contract* first: the model family, architecture
|
||||
alignment, benchmark lanes, and stop condition that benchmark runs must
|
||||
satisfy. It also runs the contract's lanes through a deterministic stub
|
||||
backend so the report data shape exists end to end. It never downloads or
|
||||
executes a model; real transformers / llama.cpp backends plug in behind the
|
||||
same ``run()`` seam later.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Mapping
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
CONTRACT_ID = "DGR-001"
|
||||
DEFAULT_OUTPUT_PATH = Path(".scratch/distributed-gguf-runtime/evidence/DGR-001/performance-contract.json")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelTarget:
|
||||
"""Architecture-aligned model target for the DGR-001 benchmark contract."""
|
||||
|
||||
name: str
|
||||
architecture: str
|
||||
safetensors_repo: str
|
||||
safetensors_precision: str
|
||||
gguf_repo: str
|
||||
gguf_quant: str
|
||||
gguf_size_gb: float
|
||||
comparison_policy: str
|
||||
rationale: str
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"name": self.name,
|
||||
"architecture": self.architecture,
|
||||
"safetensors_repo": self.safetensors_repo,
|
||||
"safetensors_precision": self.safetensors_precision,
|
||||
"gguf_repo": self.gguf_repo,
|
||||
"gguf_quant": self.gguf_quant,
|
||||
"gguf_size_gb": self.gguf_size_gb,
|
||||
"comparison_policy": self.comparison_policy,
|
||||
"rationale": self.rationale,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BenchmarkLane:
|
||||
"""One side of the comparison the contract requires."""
|
||||
|
||||
id: str
|
||||
runtime: str
|
||||
device: str
|
||||
recipe: str
|
||||
concurrency_levels: tuple[int, ...]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"id": self.id,
|
||||
"runtime": self.runtime,
|
||||
"device": self.device,
|
||||
"recipe": self.recipe,
|
||||
"concurrency_levels": list(self.concurrency_levels),
|
||||
}
|
||||
|
||||
|
||||
@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."""
|
||||
|
||||
schema_version: int
|
||||
story_id: str
|
||||
model_target: ModelTarget
|
||||
benchmark_lanes: tuple[BenchmarkLane, ...]
|
||||
metrics: tuple[str, ...]
|
||||
stop_condition: str
|
||||
notes: tuple[str, ...] = ()
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"schema_version": self.schema_version,
|
||||
"story_id": self.story_id,
|
||||
"model_target": self.model_target.to_dict(),
|
||||
"benchmark_lanes": [lane.to_dict() for lane in self.benchmark_lanes],
|
||||
"metrics": list(self.metrics),
|
||||
"stop_condition": self.stop_condition,
|
||||
"notes": list(self.notes),
|
||||
}
|
||||
|
||||
def write_json(self, path: str | Path) -> Path:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
DEFAULT_CONTRACT = PerformanceContract(
|
||||
schema_version=SCHEMA_VERSION,
|
||||
story_id=CONTRACT_ID,
|
||||
model_target=ModelTarget(
|
||||
name="DeepSeek-V2-Lite-Chat",
|
||||
architecture="deepseek2",
|
||||
safetensors_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
|
||||
safetensors_precision="bfloat16",
|
||||
gguf_repo="second-state/DeepSeek-V2-Lite-Chat-GGUF",
|
||||
gguf_quant="Q2_K",
|
||||
gguf_size_gb=6.43,
|
||||
comparison_policy=(
|
||||
"same model/revision, closest practical low-footprint precision pair: "
|
||||
"BF16 safetensors versus Q2_K GGUF"
|
||||
),
|
||||
rationale=(
|
||||
"Smallest DeepSeek-family benchmark anchor that still points toward "
|
||||
"DeepSeek-V4-Flash; keeps the runtime on the DeepSeek2 path instead "
|
||||
"of falling back to a tiny but architecture-mismatched smoke model."
|
||||
),
|
||||
),
|
||||
benchmark_lanes=(
|
||||
BenchmarkLane(
|
||||
id="transformers-safetensors-cpu",
|
||||
runtime="transformers",
|
||||
device="cpu",
|
||||
recipe="current safetensors recipe",
|
||||
concurrency_levels=(1, 4),
|
||||
),
|
||||
BenchmarkLane(
|
||||
id="llama-cpp-gguf-cpu",
|
||||
runtime="llama.cpp",
|
||||
device="cpu",
|
||||
recipe="whole-model GGUF recipe",
|
||||
concurrency_levels=(1, 4),
|
||||
),
|
||||
BenchmarkLane(
|
||||
id="transformers-safetensors-gpu",
|
||||
runtime="transformers",
|
||||
device="gpu",
|
||||
recipe="current safetensors recipe",
|
||||
concurrency_levels=(1, 4),
|
||||
),
|
||||
BenchmarkLane(
|
||||
id="llama-cpp-gguf-gpu",
|
||||
runtime="llama.cpp",
|
||||
device="gpu",
|
||||
recipe="whole-model GGUF recipe",
|
||||
concurrency_levels=(1, 4),
|
||||
),
|
||||
),
|
||||
metrics=(
|
||||
"ttft_ms",
|
||||
"prefill_tok_per_sec",
|
||||
"decode_tok_per_sec",
|
||||
"p50_latency_ms",
|
||||
"p95_latency_ms",
|
||||
"aggregate_throughput_tok_per_sec",
|
||||
"rss_bytes",
|
||||
"vram_bytes",
|
||||
"artifact_bytes",
|
||||
"failure_count",
|
||||
"output_drift",
|
||||
),
|
||||
stop_condition=(
|
||||
"Stop if GGUF does not provide a meaningful speed or fit benefit over the "
|
||||
"safetensors baseline for the chosen DeepSeek-family model target."
|
||||
),
|
||||
notes=(
|
||||
"Real model execution stays opt-in and must keep model artifacts on the mounted drive.",
|
||||
"Use the tiny fallback only for loader plumbing smoke tests; it does not replace the architecture-aligned baseline.",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_default_contract() -> PerformanceContract:
|
||||
return DEFAULT_CONTRACT
|
||||
|
||||
|
||||
BENCHMARK_SCHEMA_VERSION = 1
|
||||
STUB_OUTPUT_TOKENS = ("mesh", "activation", "seam", "baseline")
|
||||
# DeepSeek-V2-Lite is ~15.7B params at 2 bytes each; metadata only, nothing downloaded.
|
||||
_SAFETENSORS_BF16_ARTIFACT_GB = 31.4
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LaneSample:
|
||||
"""Raw single-stream measurements one backend produces for a lane."""
|
||||
|
||||
ttft_ms: float
|
||||
prefill_tok_per_sec: float
|
||||
decode_tok_per_sec: float
|
||||
rss_bytes: int
|
||||
vram_bytes: int
|
||||
artifact_bytes: int
|
||||
output_tokens: tuple[str, ...]
|
||||
failure_count: int = 0
|
||||
|
||||
|
||||
def _gb(value: float) -> int:
|
||||
return int(value * 1024**3)
|
||||
|
||||
|
||||
class StubLaneBackend:
|
||||
"""Deterministic placeholder measurements until real lane execution lands.
|
||||
|
||||
The numbers are synthetic but directionally shaped — the Q2_K GGUF loads a
|
||||
far smaller artifact and decodes faster than BF16 safetensors — so the
|
||||
comparison and stop-condition plumbing can be exercised in CI.
|
||||
"""
|
||||
|
||||
source = "stub-backend"
|
||||
|
||||
# (runtime, device) -> (ttft_ms, prefill tok/s, decode tok/s, rss GB, vram GB)
|
||||
_PROFILES = {
|
||||
("transformers", "cpu"): (1800.0, 45.0, 6.0, 33.0, 0.0),
|
||||
("llama.cpp", "cpu"): (950.0, 90.0, 14.0, 7.1, 0.0),
|
||||
("transformers", "gpu"): (420.0, 850.0, 34.0, 4.0, 33.0),
|
||||
("llama.cpp", "gpu"): (260.0, 640.0, 52.0, 1.5, 7.5),
|
||||
}
|
||||
|
||||
def __init__(self, contract: PerformanceContract) -> None:
|
||||
self._contract = contract
|
||||
|
||||
def run(self, lane: BenchmarkLane) -> LaneSample:
|
||||
ttft_ms, prefill, decode, rss_gb, vram_gb = self._PROFILES[(lane.runtime, lane.device)]
|
||||
artifact_gb = (
|
||||
self._contract.model_target.gguf_size_gb
|
||||
if lane.runtime == "llama.cpp"
|
||||
else _SAFETENSORS_BF16_ARTIFACT_GB
|
||||
)
|
||||
return LaneSample(
|
||||
ttft_ms=ttft_ms,
|
||||
prefill_tok_per_sec=prefill,
|
||||
decode_tok_per_sec=decode,
|
||||
rss_bytes=_gb(rss_gb),
|
||||
vram_bytes=_gb(vram_gb),
|
||||
artifact_bytes=_gb(artifact_gb),
|
||||
output_tokens=STUB_OUTPUT_TOKENS,
|
||||
)
|
||||
|
||||
|
||||
def _output_drift(tokens: tuple[str, ...], reference: tuple[str, ...]) -> float:
|
||||
"""Fraction of positions where a lane's output diverges from its reference."""
|
||||
length = max(len(tokens), len(reference))
|
||||
if length == 0:
|
||||
return 0.0
|
||||
mismatches = sum(a != b for a, b in zip(tokens, reference)) + abs(len(tokens) - len(reference))
|
||||
return round(mismatches / length, 4)
|
||||
|
||||
|
||||
def _metrics_for(sample: LaneSample, concurrency: int, output_drift: float) -> dict:
|
||||
# Stub concurrency model: batching scales throughput at 85% efficiency and
|
||||
# stretches per-request token latency and TTFT accordingly.
|
||||
efficiency = 1.0 if concurrency == 1 else 0.85
|
||||
p50_latency_ms = round(1000.0 / (sample.decode_tok_per_sec * efficiency), 4)
|
||||
return {
|
||||
"ttft_ms": round(sample.ttft_ms * (1 + 0.1 * (concurrency - 1)), 4),
|
||||
"prefill_tok_per_sec": round(sample.prefill_tok_per_sec * efficiency, 4),
|
||||
"decode_tok_per_sec": round(sample.decode_tok_per_sec * efficiency, 4),
|
||||
"p50_latency_ms": p50_latency_ms,
|
||||
"p95_latency_ms": round(p50_latency_ms * 1.25, 4),
|
||||
"aggregate_throughput_tok_per_sec": round(sample.decode_tok_per_sec * concurrency * efficiency, 4),
|
||||
"rss_bytes": sample.rss_bytes,
|
||||
"vram_bytes": sample.vram_bytes,
|
||||
"artifact_bytes": sample.artifact_bytes,
|
||||
"failure_count": sample.failure_count,
|
||||
"output_drift": output_drift,
|
||||
}
|
||||
|
||||
|
||||
def _compare_device(lanes: list[tuple[BenchmarkLane, LaneSample]], device: str) -> dict:
|
||||
by_runtime = {lane.runtime: (lane, sample) for lane, sample in lanes if lane.device == device}
|
||||
safetensors_lane, safetensors = by_runtime["transformers"]
|
||||
gguf_lane, gguf = by_runtime["llama.cpp"]
|
||||
memory_metric = "vram_bytes" if device == "gpu" else "rss_bytes"
|
||||
decode_speedup = round(gguf.decode_tok_per_sec / safetensors.decode_tok_per_sec, 4)
|
||||
artifact_bytes_ratio = round(gguf.artifact_bytes / max(1, safetensors.artifact_bytes), 4)
|
||||
return {
|
||||
"safetensors_lane": safetensors_lane.id,
|
||||
"gguf_lane": gguf_lane.id,
|
||||
"decode_speedup": decode_speedup,
|
||||
"ttft_speedup": round(safetensors.ttft_ms / max(0.001, gguf.ttft_ms), 4),
|
||||
"artifact_bytes_ratio": artifact_bytes_ratio,
|
||||
"memory_metric": memory_metric,
|
||||
"memory_bytes_ratio": round(
|
||||
getattr(gguf, memory_metric) / max(1, getattr(safetensors, memory_metric)), 4
|
||||
),
|
||||
"output_drift": _output_drift(gguf.output_tokens, safetensors.output_tokens),
|
||||
"gguf_benefit": decode_speedup >= 1.10 or artifact_bytes_ratio <= 0.5,
|
||||
}
|
||||
|
||||
|
||||
def run_performance_benchmark(
|
||||
contract: PerformanceContract = DEFAULT_CONTRACT,
|
||||
backend: StubLaneBackend | None = None,
|
||||
) -> dict:
|
||||
"""Run every contract lane through a backend and compare GGUF to safetensors."""
|
||||
backend = backend if backend is not None else StubLaneBackend(contract)
|
||||
lanes = [(lane, backend.run(lane)) for lane in contract.benchmark_lanes]
|
||||
references = {
|
||||
lane.device: sample.output_tokens for lane, sample in lanes if lane.runtime == "transformers"
|
||||
}
|
||||
lane_reports = []
|
||||
for lane, sample in lanes:
|
||||
drift = _output_drift(sample.output_tokens, references.get(lane.device, sample.output_tokens))
|
||||
lane_reports.append({
|
||||
**lane.to_dict(),
|
||||
"output_tokens": list(sample.output_tokens),
|
||||
"results": [
|
||||
{"concurrency": level, "metrics": _metrics_for(sample, level, drift)}
|
||||
for level in lane.concurrency_levels
|
||||
],
|
||||
})
|
||||
devices = sorted({lane.device for lane, _ in lanes})
|
||||
comparisons = {device: _compare_device(lanes, device) for device in devices}
|
||||
gguf_benefit = any(comparison["gguf_benefit"] for comparison in comparisons.values())
|
||||
return {
|
||||
"schema_version": BENCHMARK_SCHEMA_VERSION,
|
||||
"story_id": contract.story_id,
|
||||
"source": getattr(backend, "source", "custom-backend"),
|
||||
"model_target": contract.model_target.to_dict(),
|
||||
"lanes": lane_reports,
|
||||
"comparisons": comparisons,
|
||||
"stop_condition": {
|
||||
"text": contract.stop_condition,
|
||||
"gguf_benefit": gguf_benefit,
|
||||
"triggered": not gguf_benefit,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def run_real_model_endpoint_benchmark(
|
||||
endpoints: Mapping[str, str],
|
||||
*,
|
||||
model: str,
|
||||
contract: PerformanceContract = DEFAULT_CONTRACT,
|
||||
timeout: float = 120.0,
|
||||
) -> dict:
|
||||
"""Run one live OpenAI-compatible request per lane against supplied endpoints.
|
||||
|
||||
The caller provides one URL per benchmark lane. The runner measures the
|
||||
request/response round-trip at the client boundary and reuses the same
|
||||
contract schema as the deterministic stub.
|
||||
"""
|
||||
|
||||
def _sample_for_lane(lane: BenchmarkLane, endpoint: str) -> LaneSample:
|
||||
prompt = " ".join(contract.model_target.rationale.split()[:6])
|
||||
body = json.dumps(
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": len(STUB_OUTPUT_TOKENS),
|
||||
"temperature": 0,
|
||||
}
|
||||
).encode("utf-8")
|
||||
request = urllib.request.Request(
|
||||
f"{endpoint.rstrip('/')}/v1/chat/completions",
|
||||
data=body,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Lane": lane.id,
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
started = time.monotonic()
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
response_body = response.read()
|
||||
session_id = response.headers.get("X-Meshnet-Session", f"{lane.id}-session")
|
||||
elapsed_ms = round((time.monotonic() - started) * 1000, 4)
|
||||
payload = json.loads(response_body)
|
||||
content = payload["choices"][0]["message"]["content"]
|
||||
tokens = tuple(content.split())
|
||||
token_count = max(1, len(tokens))
|
||||
artifact_gb = (
|
||||
contract.model_target.gguf_size_gb
|
||||
if lane.runtime == "llama.cpp"
|
||||
else _SAFETENSORS_BF16_ARTIFACT_GB
|
||||
)
|
||||
return LaneSample(
|
||||
ttft_ms=elapsed_ms,
|
||||
prefill_tok_per_sec=round(token_count / max(0.001, elapsed_ms / 1000), 4),
|
||||
decode_tok_per_sec=round(token_count / max(0.001, elapsed_ms / 1000), 4),
|
||||
rss_bytes=0,
|
||||
vram_bytes=0,
|
||||
artifact_bytes=_gb(artifact_gb),
|
||||
output_tokens=tokens,
|
||||
)
|
||||
|
||||
lanes = []
|
||||
for lane in contract.benchmark_lanes:
|
||||
if lane.id not in endpoints:
|
||||
raise KeyError(f"missing endpoint for lane {lane.id}")
|
||||
lanes.append((lane, _sample_for_lane(lane, endpoints[lane.id])))
|
||||
references = {
|
||||
lane.device: sample.output_tokens for lane, sample in lanes if lane.runtime == "transformers"
|
||||
}
|
||||
lane_reports = []
|
||||
for lane, sample in lanes:
|
||||
drift = _output_drift(sample.output_tokens, references.get(lane.device, sample.output_tokens))
|
||||
lane_reports.append({
|
||||
**lane.to_dict(),
|
||||
"output_tokens": list(sample.output_tokens),
|
||||
"results": [
|
||||
{"concurrency": level, "metrics": _metrics_for(sample, level, drift)}
|
||||
for level in lane.concurrency_levels
|
||||
],
|
||||
})
|
||||
devices = sorted({lane.device for lane, _ in lanes})
|
||||
comparisons = {device: _compare_device(lanes, device) for device in devices}
|
||||
gguf_benefit = any(comparison["gguf_benefit"] for comparison in comparisons.values())
|
||||
return {
|
||||
"schema_version": BENCHMARK_SCHEMA_VERSION,
|
||||
"story_id": contract.story_id,
|
||||
"source": "real-model-endpoints",
|
||||
"model_target": contract.model_target.to_dict(),
|
||||
"lanes": lane_reports,
|
||||
"comparisons": comparisons,
|
||||
"stop_condition": {
|
||||
"text": contract.stop_condition,
|
||||
"gguf_benefit": gguf_benefit,
|
||||
"triggered": not gguf_benefit,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _parse_lane_endpoints(pairs: list[str], parser: argparse.ArgumentParser) -> dict[str, str]:
|
||||
endpoints: dict[str, str] = {}
|
||||
for pair in pairs:
|
||||
lane_id, sep, url = pair.partition("=")
|
||||
if not sep or not lane_id or not url:
|
||||
parser.error(f"--live-endpoint expects LANE_ID=URL, got {pair!r}")
|
||||
endpoints[lane_id] = url
|
||||
return endpoints
|
||||
|
||||
|
||||
def _write_report(report: dict, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Write the DGR-001 performance contract JSON")
|
||||
parser.add_argument("--json-out", type=Path, default=DEFAULT_OUTPUT_PATH, help="output JSON path")
|
||||
parser.add_argument(
|
||||
"--benchmark-out",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="also run the deterministic stub benchmark and write its JSON report here",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--live-endpoint",
|
||||
action="append",
|
||||
default=None,
|
||||
metavar="LANE_ID=URL",
|
||||
help="lane-to-endpoint mapping for the live benchmark; repeat once per contract lane",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--live-model",
|
||||
default=None,
|
||||
help="model name sent to live endpoints (default: contract safetensors repo)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--live-benchmark-out",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="run the live endpoint benchmark against --live-endpoint lanes and write its JSON report here",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
if args.live_endpoint and args.live_benchmark_out is None:
|
||||
parser.error("--live-endpoint requires --live-benchmark-out")
|
||||
if args.live_benchmark_out is not None and not args.live_endpoint:
|
||||
parser.error("--live-benchmark-out requires at least one --live-endpoint")
|
||||
contract = build_default_contract()
|
||||
path = contract.write_json(args.json_out)
|
||||
print(path)
|
||||
if args.benchmark_out is not None:
|
||||
_write_report(run_performance_benchmark(contract), args.benchmark_out)
|
||||
print(args.benchmark_out)
|
||||
if args.live_endpoint:
|
||||
report = run_real_model_endpoint_benchmark(
|
||||
_parse_lane_endpoints(args.live_endpoint, parser),
|
||||
model=args.live_model or contract.model_target.safetensors_repo,
|
||||
contract=contract,
|
||||
)
|
||||
_write_report(report, args.live_benchmark_out)
|
||||
print(args.live_benchmark_out)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover - CLI entry point
|
||||
raise SystemExit(main())
|
||||
@@ -26,16 +26,6 @@
|
||||
"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
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -1,375 +0,0 @@
|
||||
"""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,7 +29,6 @@ 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
|
||||
|
||||
@@ -703,35 +702,6 @@ 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)
|
||||
@@ -993,8 +963,7 @@ def run_startup(
|
||||
|
||||
if model_id: # treat "" the same as None — no explicit model given
|
||||
full_sources: list[dict] = []
|
||||
detected: int | None = None
|
||||
# Auto-detect shard range from model config if not explicitly provided.
|
||||
# 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)
|
||||
@@ -1058,38 +1027,22 @@ 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)
|
||||
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=model_id,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
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,
|
||||
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,
|
||||
)
|
||||
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,
|
||||
@@ -1103,15 +1056,10 @@ 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(
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
total_layers,
|
||||
)
|
||||
shard_label = _format_shard_label(shard_start, 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"):
|
||||
@@ -1134,17 +1082,16 @@ def run_startup(
|
||||
"model": model_id.split("/")[-1],
|
||||
"hf_repo": model_id,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (proof_shard.start == 0),
|
||||
"tracker_mode": (shard_start == 0),
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
|
||||
"capability_report": capability_report.to_dict(),
|
||||
"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,
|
||||
@@ -1152,8 +1099,8 @@ def run_startup(
|
||||
"downloaded_models": (
|
||||
_downloaded_model_inventory(
|
||||
model_id.split("/")[-1],
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
shard_start,
|
||||
shard_end,
|
||||
model_cache_path,
|
||||
hf_repo=model_id,
|
||||
model_sources=full_sources,
|
||||
@@ -1264,38 +1211,22 @@ def run_startup(
|
||||
hf_repo=assigned_hf_repo,
|
||||
model_sources=full_sources,
|
||||
)
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=assigned_hf_repo,
|
||||
shard_start=assigned_shard_start,
|
||||
shard_end=assigned_shard_end,
|
||||
quantization=quantization,
|
||||
total_layers=assigned_num_layers,
|
||||
device=device,
|
||||
model_revision=None,
|
||||
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,
|
||||
)
|
||||
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,
|
||||
@@ -1309,7 +1240,6 @@ 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)
|
||||
@@ -1332,17 +1262,16 @@ def run_startup(
|
||||
"model": assigned_hf_repo.split("/")[-1],
|
||||
"hf_repo": assigned_hf_repo,
|
||||
"num_layers": assigned_num_layers,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (proof_shard.start == 0),
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
"managed_assignment": True,
|
||||
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
|
||||
"capability_report": capability_report.to_dict(),
|
||||
"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,
|
||||
@@ -1350,8 +1279,8 @@ def run_startup(
|
||||
"downloaded_models": (
|
||||
_downloaded_model_inventory(
|
||||
assigned_hf_repo.split("/")[-1],
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
model_cache_path,
|
||||
hf_repo=assigned_hf_repo,
|
||||
model_sources=full_sources,
|
||||
@@ -1376,8 +1305,8 @@ def run_startup(
|
||||
),
|
||||
)
|
||||
shard_label = _format_shard_label(
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
assigned_shard_start,
|
||||
assigned_shard_end,
|
||||
assigned_num_layers,
|
||||
)
|
||||
print(
|
||||
@@ -1492,38 +1421,22 @@ 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":
|
||||
backend = _gguf_backend_for_recipe(
|
||||
recipe,
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=hf_repo,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
total_layers=total_layers,
|
||||
device=device,
|
||||
model_revision=None,
|
||||
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,
|
||||
)
|
||||
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,
|
||||
@@ -1572,7 +1485,6 @@ 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,
|
||||
@@ -1634,7 +1546,6 @@ 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}"
|
||||
@@ -1654,11 +1565,10 @@ def run_startup(
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_model,
|
||||
"shard_start": proof_shard.start,
|
||||
"shard_end": proof_shard.end,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": 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,
|
||||
@@ -1704,8 +1614,8 @@ def run_startup(
|
||||
if gpu_name:
|
||||
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
|
||||
shard_label = _format_shard_label(
|
||||
proof_shard.start,
|
||||
proof_shard.end,
|
||||
shard_start,
|
||||
shard_end,
|
||||
assigned_total_layers,
|
||||
model_name=assigned_model,
|
||||
)
|
||||
|
||||
@@ -16,10 +16,7 @@ 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(
|
||||
@@ -33,15 +30,6 @@ 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,
|
||||
@@ -49,49 +37,18 @@ 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=ownership.start_layer if shard_start is None else shard_start,
|
||||
shard_end=ownership.end_layer if shard_end is None else shard_end,
|
||||
shard_start=(
|
||||
context.selection.shard_start if shard_start is None else shard_start
|
||||
),
|
||||
shard_end=context.selection.shard_end if shard_end is None else shard_end,
|
||||
recipe_id=recipe_id or context.recipe.id,
|
||||
recipe_version=recipe_version or context.recipe.version,
|
||||
catalogue_version=context.manifest.catalogue_version,
|
||||
backend_id=backend_id or context.recipe.backend_id,
|
||||
device=device or context.device,
|
||||
quantization=context.selection.quantization,
|
||||
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,
|
||||
@@ -111,20 +68,3 @@ 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"
|
||||
|
||||
Reference in New Issue
Block a user