Files
neuron-tai/tests/test_range_report.py

274 lines
11 KiB
Python

"""DGR-034: strict consumption of owned-range reports from loaded engine state.
The ``meshnet-range-report`` native tool loads one dense-Llama GGUF through
the Meshnet owned-range loader and prints a JSON document derived from the
loaded model state. ``meshnet_node.range_report`` is the strict consumer:
it must accept exactly the documents that encode the dense-Llama ownership
contract and fail closed on everything else — invalid, empty, or
out-of-model ranges, endpoint registrations that disagree with the loaded
state, gapped or unexpected tensor registrations, and inconsistent byte
counts.
"""
from __future__ import annotations
from typing import Any
import pytest
from meshnet_node.range_report import (
OwnedRangeReport,
RangeReportError,
parse_owned_range_report,
)
N_LAYER = 40
LAYER_BYTES = 300 * 2**20
EMBD_BYTES = 360 * 2**20
OUT_BYTES = 525 * 2**20
FILE_BYTES = 13669 * 2**20
def _doc(**overrides: Any) -> dict[str, Any]:
"""A valid middle-range [10, 20) mmap report the consumer must accept."""
doc: dict[str, Any] = {
"ok": True,
"model": "/models/dense.gguf",
"architecture": "llama",
"n_layer": N_LAYER,
"file_bytes": FILE_BYTES,
"requested_range": [10, 20],
"reported_range": [10, 20],
"mmap": True,
"touched": False,
"use_extra_bufts": True,
"has_token_embeddings": False,
"has_output_head": False,
"tied_output_head": False,
"mapped_bytes": 10 * LAYER_BYTES,
"resident_bytes": 10 * LAYER_BYTES,
"registered_tensors": 90,
"registered_bytes": 10 * LAYER_BYTES,
"unexpected_registered_tensors": [],
"missing_owned_layers": [],
"vm_size_bytes": FILE_BYTES + 2**28,
"vm_rss_bytes": 2**28,
"vm_hwm_bytes": 2**28,
}
doc.update(overrides)
return doc
def _head_doc(**overrides: Any) -> dict[str, Any]:
base = _doc(
requested_range=[0, 10],
reported_range=[0, 10],
has_token_embeddings=True,
mapped_bytes=10 * LAYER_BYTES + EMBD_BYTES,
resident_bytes=10 * LAYER_BYTES + EMBD_BYTES,
registered_tensors=91,
registered_bytes=10 * LAYER_BYTES + EMBD_BYTES,
)
base.update(overrides)
return base
def _tail_doc(**overrides: Any) -> dict[str, Any]:
base = _doc(
requested_range=[30, 40],
reported_range=[30, 40],
has_output_head=True,
mapped_bytes=10 * LAYER_BYTES + OUT_BYTES,
resident_bytes=10 * LAYER_BYTES + OUT_BYTES,
registered_tensors=92,
registered_bytes=10 * LAYER_BYTES + OUT_BYTES,
)
base.update(overrides)
return base
class TestAcceptance:
def test_middle_range_registers_only_per_layer_tensors(self) -> None:
report = parse_owned_range_report(_doc())
assert (report.start_layer, report.end_layer) == (10, 20)
assert not report.is_head and not report.is_tail
assert not report.has_token_embeddings and not report.has_output_head
def test_head_range_owns_embeddings_only_at_the_head(self) -> None:
report = parse_owned_range_report(_head_doc())
assert report.is_head and not report.is_tail
assert report.has_token_embeddings and not report.has_output_head
def test_tail_range_owns_norm_and_output_only_at_the_tail(self) -> None:
report = parse_owned_range_report(_tail_doc())
assert report.is_tail and not report.is_head
assert report.has_output_head and not report.has_token_embeddings
def test_whole_model_range_owns_both_endpoints(self) -> None:
report = parse_owned_range_report(
_head_doc(
requested_range=[0, 40],
reported_range=[0, 40],
has_output_head=True,
mapped_bytes=FILE_BYTES,
resident_bytes=FILE_BYTES,
registered_tensors=363,
registered_bytes=N_LAYER * LAYER_BYTES + EMBD_BYTES + OUT_BYTES,
)
)
assert report.is_head and report.is_tail
assert report.has_token_embeddings and report.has_output_head
def test_tied_output_tail_registers_the_embedding_as_its_output_head(self) -> None:
report = parse_owned_range_report(
_tail_doc(
has_token_embeddings=True,
tied_output_head=True,
registered_tensors=91,
registered_bytes=10 * LAYER_BYTES + EMBD_BYTES,
mapped_bytes=10 * LAYER_BYTES + EMBD_BYTES,
resident_bytes=10 * LAYER_BYTES + EMBD_BYTES,
)
)
assert report.tied_output_head and report.has_output_head
def test_non_mmap_load_reports_resident_allocation_only(self) -> None:
report = parse_owned_range_report(
_doc(mmap=False, mapped_bytes=0, resident_bytes=10 * LAYER_BYTES)
)
assert report.mapped_bytes == 0
assert report.resident_bytes == 10 * LAYER_BYTES
def test_process_counters_may_be_absent_off_linux(self) -> None:
report = parse_owned_range_report(
_doc(vm_size_bytes=None, vm_rss_bytes=None, vm_hwm_bytes=None)
)
assert report.vm_hwm_bytes is None
class TestRangeRejection:
def test_rejected_load_fails_closed_with_the_tool_error(self) -> None:
with pytest.raises(RangeReportError, match="dense Llama only"):
parse_owned_range_report(
{"ok": False, "error": "owned-range load rejected the artifact or range: dense Llama only"}
)
def test_reported_range_must_match_the_requested_range(self) -> None:
with pytest.raises(RangeReportError, match="loaded engine state"):
parse_owned_range_report(_doc(reported_range=[10, 21]))
def test_out_of_model_range_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="outside the model"):
parse_owned_range_report(
_doc(requested_range=[30, 41], reported_range=[30, 41], has_output_head=True)
)
def test_empty_range_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="empty or"):
parse_owned_range_report(_doc(requested_range=[10, 10], reported_range=[10, 10]))
def test_inverted_range_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="empty or"):
parse_owned_range_report(_doc(requested_range=[20, 10], reported_range=[20, 10]))
def test_boolean_range_bounds_are_rejected(self) -> None:
with pytest.raises(RangeReportError, match="integer pair"):
parse_owned_range_report(_doc(reported_range=[True, 20]))
class TestEndpointRejection:
def test_embeddings_registered_below_the_head_are_rejected(self) -> None:
with pytest.raises(RangeReportError, match="embeddings belong to the head"):
parse_owned_range_report(_doc(has_token_embeddings=True))
def test_output_head_registered_above_the_tail_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="output head belong to the tail"):
parse_owned_range_report(_tail_doc(requested_range=[20, 30], reported_range=[20, 30]))
def test_tail_without_an_output_head_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="output head belong to the tail"):
parse_owned_range_report(_tail_doc(has_output_head=False))
def test_tied_output_below_the_tail_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="only belong to the tail"):
parse_owned_range_report(_doc(tied_output_head=True))
def test_unexpected_registered_tensors_are_rejected(self) -> None:
with pytest.raises(RangeReportError, match="unexpected_registered_tensors"):
parse_owned_range_report(
_doc(unexpected_registered_tensors=["blk.10.attn_q.weight.extra"])
)
def test_missing_owned_layers_are_rejected_as_gaps(self) -> None:
with pytest.raises(RangeReportError, match="missing_owned_layers"):
parse_owned_range_report(_doc(missing_owned_layers=[12]))
class TestByteCountRejection:
def test_mapped_span_must_cover_the_registered_tensors(self) -> None:
with pytest.raises(RangeReportError, match="undercounts"):
parse_owned_range_report(_doc(mapped_bytes=LAYER_BYTES))
def test_mapped_span_must_not_exceed_the_artifact(self) -> None:
with pytest.raises(RangeReportError, match="exceeds the artifact"):
parse_owned_range_report(
_tail_doc(mapped_bytes=FILE_BYTES + 1, resident_bytes=FILE_BYTES + 1)
)
def test_non_mmap_load_must_not_claim_a_mapped_span(self) -> None:
with pytest.raises(RangeReportError, match="must not claim"):
parse_owned_range_report(_doc(mmap=False, mapped_bytes=LAYER_BYTES))
def test_resident_allocation_must_cover_the_registered_tensors(self) -> None:
with pytest.raises(RangeReportError, match="undercounts"):
parse_owned_range_report(
_doc(mmap=False, mapped_bytes=0, resident_bytes=LAYER_BYTES)
)
def test_an_empty_registration_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="no tensors"):
parse_owned_range_report(_doc(registered_tensors=0, registered_bytes=0))
class TestSchemaRejection:
def test_wrong_architecture_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="dense Llama only"):
parse_owned_range_report(_doc(architecture="qwen2"))
def test_missing_field_is_rejected(self) -> None:
doc = _doc()
del doc["mapped_bytes"]
with pytest.raises(RangeReportError, match="missing field"):
parse_owned_range_report(doc)
def test_boolean_bytes_are_rejected(self) -> None:
with pytest.raises(RangeReportError, match="non-negative integer"):
parse_owned_range_report(_doc(mapped_bytes=True))
def test_non_mapping_document_is_rejected(self) -> None:
with pytest.raises(RangeReportError, match="JSON object"):
parse_owned_range_report(["not", "a", "report"]) # type: ignore[arg-type]
def test_owned_range_report_rejects_direct_construction_outside_the_contract() -> None:
with pytest.raises(RangeReportError, match="dense Llama only"):
OwnedRangeReport(
architecture="qwen2",
n_layer=N_LAYER,
start_layer=10,
end_layer=20,
has_token_embeddings=False,
has_output_head=False,
tied_output_head=False,
mapped_bytes=10 * LAYER_BYTES,
resident_bytes=10 * LAYER_BYTES,
registered_tensors=90,
registered_bytes=10 * LAYER_BYTES,
file_bytes=FILE_BYTES,
mmap=True,
touched=False,
vm_size_bytes=None,
vm_rss_bytes=None,
vm_hwm_bytes=None,
)