story: DGR-034 Implement dense-Llama range-aware GGUF ownership
This commit is contained in:
Dobromir Popov
218
packages/node/meshnet_node/range_report.py
Normal file
218
packages/node/meshnet_node/range_report.py
Normal file
@@ -0,0 +1,218 @@
|
||||
"""Authoritative dense-Llama owned-range reports from the loaded engine state.
|
||||
|
||||
DGR-034 loads only the tensors a shard range owns through the Meshnet
|
||||
owned-range loader (``llama_model_params::meshnet_owned_layer_start/end`` in
|
||||
the pinned llama.cpp patch stack). The project-owned ``meshnet-range-report``
|
||||
native tool runs that load and prints a JSON document derived from the loaded
|
||||
model state — the registered tensor set and the backend buffers — never from
|
||||
caller-asserted values. This module is the strict consumer of that document:
|
||||
it parses it into :class:`OwnedRangeReport` and fails closed on any
|
||||
inconsistency, so a range or endpoint claim that the loaded engine state does
|
||||
not back is rejected before it can reach identity, admission, or routing.
|
||||
|
||||
Ownership contract enforced here (dense Llama only):
|
||||
|
||||
- every registered ``blk.N.*`` tensor lies inside the half-open owned range
|
||||
``[start, end)``, and every layer in that range is present — a gapped or
|
||||
out-of-range registration is rejected;
|
||||
- ``token_embd.weight`` is registered only by the head shard (``start == 0``),
|
||||
or by a tail shard whose model ties the output head to the embedding
|
||||
(``end == n_layer`` and no separate ``output.weight``);
|
||||
- ``output_norm.weight`` and ``output.weight`` are registered only by the
|
||||
tail shard (``end == n_layer``);
|
||||
- any other registered tensor name is unexpected and rejected;
|
||||
- byte counts are consistent: an mmap load maps a file span at least the
|
||||
registered tensor bytes and at most the artifact size; a non-mmap load
|
||||
reports a resident allocation at least the registered tensor bytes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Mapping
|
||||
|
||||
|
||||
class RangeReportError(ValueError):
|
||||
"""A range report is malformed, or the loaded state breaks ownership."""
|
||||
|
||||
|
||||
_DENSE_ARCHITECTURE = "llama"
|
||||
|
||||
_INT_FIELDS = (
|
||||
"n_layer",
|
||||
"file_bytes",
|
||||
"mapped_bytes",
|
||||
"resident_bytes",
|
||||
"registered_tensors",
|
||||
"registered_bytes",
|
||||
)
|
||||
|
||||
_BOOL_FIELDS = (
|
||||
"mmap",
|
||||
"touched",
|
||||
"has_token_embeddings",
|
||||
"has_output_head",
|
||||
"tied_output_head",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OwnedRangeReport:
|
||||
"""One validated owned-range load, derived from loaded engine state.
|
||||
|
||||
``start_layer``/``end_layer`` are the authoritative half-open owned range
|
||||
the engine actually registered (the tool already refused a report whose
|
||||
loaded bounds differ from the requested ones). ``has_token_embeddings`` is
|
||||
true for the head shard, and also for a tail shard on a tied-output model
|
||||
(the embedding tensor *is* its output head); ``tied_output_head``
|
||||
disambiguates those two cases. ``mapped_bytes``/``resident_bytes`` come
|
||||
from the backend buffers: with mmap they are the mapped file span holding
|
||||
the owned tensors, without mmap the resident allocation holding them.
|
||||
"""
|
||||
|
||||
architecture: str
|
||||
n_layer: int
|
||||
start_layer: int
|
||||
end_layer: int
|
||||
has_token_embeddings: bool
|
||||
has_output_head: bool
|
||||
tied_output_head: bool
|
||||
mapped_bytes: int
|
||||
resident_bytes: int
|
||||
registered_tensors: int
|
||||
registered_bytes: int
|
||||
file_bytes: int
|
||||
mmap: bool
|
||||
touched: bool
|
||||
vm_size_bytes: int | None
|
||||
vm_rss_bytes: int | None
|
||||
vm_hwm_bytes: int | None
|
||||
|
||||
@property
|
||||
def is_head(self) -> bool:
|
||||
return self.start_layer == 0
|
||||
|
||||
@property
|
||||
def is_tail(self) -> bool:
|
||||
return self.end_layer == self.n_layer
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.architecture != _DENSE_ARCHITECTURE:
|
||||
raise RangeReportError(
|
||||
f"owned-range loading supports dense Llama only, got {self.architecture!r}"
|
||||
)
|
||||
if isinstance(self.n_layer, bool) or self.n_layer < 1:
|
||||
raise RangeReportError("report must record a positive GGUF block count")
|
||||
for name in _INT_FIELDS:
|
||||
value = getattr(self, name)
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
raise RangeReportError(f"report field {name!r} must be a non-negative integer")
|
||||
for name in _BOOL_FIELDS:
|
||||
if not isinstance(getattr(self, name), bool):
|
||||
raise RangeReportError(f"report field {name!r} must be a boolean")
|
||||
if not 0 <= self.start_layer < self.end_layer <= self.n_layer:
|
||||
raise RangeReportError(
|
||||
f"owned range [{self.start_layer}, {self.end_layer}) is empty or "
|
||||
f"outside the model's {self.n_layer} layers"
|
||||
)
|
||||
if self.tied_output_head and not self.is_tail:
|
||||
raise RangeReportError("a tied output head can only belong to the tail shard")
|
||||
expected_embeddings = self.is_head or self.tied_output_head
|
||||
if self.has_token_embeddings != expected_embeddings:
|
||||
raise RangeReportError(
|
||||
"token-embedding registration disagrees with endpoint ownership: "
|
||||
"embeddings belong to the head shard (or to a tied-output tail)"
|
||||
)
|
||||
if self.has_output_head != self.is_tail:
|
||||
raise RangeReportError(
|
||||
"output-head registration disagrees with endpoint ownership: "
|
||||
"the final norm and output head belong to the tail shard"
|
||||
)
|
||||
if self.registered_tensors < 1 or self.registered_bytes < 1:
|
||||
raise RangeReportError("the owned range registered no tensors")
|
||||
if self.file_bytes < 1:
|
||||
raise RangeReportError("report must record the artifact size")
|
||||
if self.mmap:
|
||||
if self.mapped_bytes < self.registered_bytes:
|
||||
raise RangeReportError(
|
||||
"mapped span undercounts the registered owned tensors"
|
||||
)
|
||||
if self.mapped_bytes > self.file_bytes:
|
||||
raise RangeReportError("mapped span exceeds the artifact size")
|
||||
else:
|
||||
if self.mapped_bytes != 0:
|
||||
raise RangeReportError("a non-mmap load must not claim a mapped span")
|
||||
if self.resident_bytes < self.registered_bytes:
|
||||
raise RangeReportError(
|
||||
"resident allocation undercounts the registered owned tensors"
|
||||
)
|
||||
for name in ("vm_size_bytes", "vm_rss_bytes", "vm_hwm_bytes"):
|
||||
value = getattr(self, name)
|
||||
if value is not None and (
|
||||
isinstance(value, bool) or not isinstance(value, int) or value < 0
|
||||
):
|
||||
raise RangeReportError(f"report field {name!r} must be a non-negative integer or null")
|
||||
|
||||
|
||||
def _require_range(doc: Mapping[str, Any], key: str) -> tuple[int, int]:
|
||||
value = doc.get(key)
|
||||
if (
|
||||
not isinstance(value, (list, tuple))
|
||||
or len(value) != 2
|
||||
or any(isinstance(v, bool) or not isinstance(v, int) for v in value)
|
||||
):
|
||||
raise RangeReportError(f"report field {key!r} must be a [start, end] integer pair")
|
||||
return value[0], value[1]
|
||||
|
||||
|
||||
def parse_owned_range_report(doc: Mapping[str, Any]) -> OwnedRangeReport:
|
||||
"""Parse and validate one ``meshnet-range-report`` JSON document.
|
||||
|
||||
Fails closed: a load the tool rejected (``ok: false``), a requested range
|
||||
the loaded state did not match, a gapped or out-of-range registration, an
|
||||
unexpected registered tensor, and any byte-count inconsistency all raise
|
||||
:class:`RangeReportError` instead of producing a report.
|
||||
"""
|
||||
if not isinstance(doc, Mapping):
|
||||
raise RangeReportError("range report must be a JSON object")
|
||||
if doc.get("ok") is not True:
|
||||
error = doc.get("error")
|
||||
detail = f": {error}" if isinstance(error, str) and error else ""
|
||||
raise RangeReportError(f"the owned-range load was rejected{detail}")
|
||||
|
||||
requested = _require_range(doc, "requested_range")
|
||||
reported = _require_range(doc, "reported_range")
|
||||
if requested != reported:
|
||||
raise RangeReportError(
|
||||
f"reported range {reported} does not match the requested range {requested}; "
|
||||
"ownership must be derived from the loaded engine state"
|
||||
)
|
||||
|
||||
for key in ("unexpected_registered_tensors", "missing_owned_layers"):
|
||||
value = doc.get(key)
|
||||
if not isinstance(value, list):
|
||||
raise RangeReportError(f"report field {key!r} must be a list")
|
||||
if value:
|
||||
raise RangeReportError(
|
||||
f"ownership audit failed: {key} is {value!r}; the registered "
|
||||
"tensor set must exactly cover the owned range and its endpoints"
|
||||
)
|
||||
|
||||
architecture = doc.get("architecture")
|
||||
if not isinstance(architecture, str):
|
||||
raise RangeReportError("report field 'architecture' must be a string")
|
||||
|
||||
fields: dict[str, Any] = {}
|
||||
for name in _INT_FIELDS + _BOOL_FIELDS:
|
||||
if name not in doc:
|
||||
raise RangeReportError(f"range report is missing field {name!r}")
|
||||
fields[name] = doc[name]
|
||||
for name in ("vm_size_bytes", "vm_rss_bytes", "vm_hwm_bytes"):
|
||||
fields[name] = doc.get(name)
|
||||
|
||||
return OwnedRangeReport(
|
||||
architecture=architecture,
|
||||
start_layer=reported[0],
|
||||
end_layer=reported[1],
|
||||
**fields,
|
||||
)
|
||||
Reference in New Issue
Block a user