219 lines
8.9 KiB
Python
219 lines
8.9 KiB
Python
"""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,
|
|
)
|