131 lines
3.8 KiB
Python
131 lines
3.8 KiB
Python
"""DGR-021: versioned activation envelope for shard traffic."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hashlib
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.model_backend import TensorPayload
|
|
from meshnet_node.protocol import ActivationEnvelope, NamedTensor, TensorFragment
|
|
|
|
|
|
def _payload(body: bytes = b"\x01\x02\x03\x04", *, shape=(1, 2, 1)) -> TensorPayload:
|
|
return TensorPayload(
|
|
body=body,
|
|
shape=list(shape),
|
|
attention_mask_header="1,2:AAAA",
|
|
position_ids_header="1,2:BBBB",
|
|
past_len=7,
|
|
)
|
|
|
|
|
|
def test_envelope_roundtrips_deterministically():
|
|
payload = _payload()
|
|
env_a = payload.to_envelope(
|
|
name="activations",
|
|
request_id="req-1",
|
|
work_id="work-1",
|
|
route_session="route-1",
|
|
route_epoch=9,
|
|
shard_start=12,
|
|
effective_start=18,
|
|
phase="prefill",
|
|
position=0,
|
|
idempotency_step=0,
|
|
extensions={"zeta": 3, "alpha": "x"},
|
|
)
|
|
env_b = ActivationEnvelope.from_bytes(env_a.to_bytes())
|
|
|
|
assert env_a.to_bytes() == env_b.to_bytes()
|
|
assert env_b.request_id == "req-1"
|
|
assert env_b.work_id == "work-1"
|
|
assert env_b.route_session == "route-1"
|
|
assert env_b.shard_start == 12
|
|
assert env_b.effective_start == 18
|
|
assert env_b.phase == "prefill"
|
|
assert env_b.tensors[0].body() == payload.body
|
|
assert env_b.tensors[0].extensions["past_len"] == 7
|
|
assert env_b.extensions == {"alpha": "x", "zeta": 3}
|
|
|
|
|
|
def test_fragmentation_and_checksums_are_bounded():
|
|
body = b"0123456789abcdef"
|
|
tensor = NamedTensor.from_bytes(
|
|
name="activations",
|
|
body=body,
|
|
shape=[1, 8, 1],
|
|
dtype="bfloat16",
|
|
max_fragment_bytes=5,
|
|
)
|
|
|
|
assert len(tensor.fragments) == 4
|
|
assert all(isinstance(fragment, TensorFragment) for fragment in tensor.fragments)
|
|
assert all(len(fragment.body) <= 5 for fragment in tensor.fragments)
|
|
assert tensor.body() == body
|
|
assert tensor.checksum == hashlib.sha256(body).hexdigest()
|
|
|
|
|
|
def test_unknown_fields_are_preserved_on_roundtrip():
|
|
env = _payload().to_envelope(
|
|
name="activations",
|
|
request_id="req-2",
|
|
work_id="work-2",
|
|
route_session="route-2",
|
|
route_epoch=1,
|
|
shard_start=0,
|
|
effective_start=0,
|
|
phase="decode",
|
|
position=3,
|
|
idempotency_step=2,
|
|
)
|
|
raw = json.loads(env.to_bytes())
|
|
raw["future_top_level"] = {"x": 1}
|
|
raw["tensors"][0]["future_tensor_field"] = "ok"
|
|
roundtrip = ActivationEnvelope.from_bytes(json.dumps(raw).encode())
|
|
|
|
assert roundtrip.extensions["future_top_level"] == {"x": 1}
|
|
assert roundtrip.tensors[0].extensions["future_tensor_field"] == "ok"
|
|
|
|
|
|
def test_size_limit_rejects_large_envelopes():
|
|
env = _payload(body=b"x" * 64).to_envelope(
|
|
name="activations",
|
|
request_id="req-3",
|
|
work_id="work-3",
|
|
route_session="route-3",
|
|
route_epoch=1,
|
|
shard_start=0,
|
|
effective_start=0,
|
|
phase="decode",
|
|
position=5,
|
|
idempotency_step=4,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="size limit"):
|
|
env.to_bytes(max_bytes=32)
|
|
|
|
|
|
def test_tensor_payload_roundtrips_through_envelope():
|
|
payload = _payload(body=b"abcde", shape=(1, 5, 1))
|
|
env = payload.to_envelope(
|
|
name="activations",
|
|
request_id="req-4",
|
|
work_id="work-4",
|
|
route_session="route-4",
|
|
route_epoch=2,
|
|
shard_start=6,
|
|
effective_start=11,
|
|
phase="decode",
|
|
position=8,
|
|
idempotency_step=9,
|
|
)
|
|
|
|
restored = TensorPayload.from_envelope(env)
|
|
assert restored.body == payload.body
|
|
assert restored.shape == payload.shape
|
|
assert restored.attention_mask_header == payload.attention_mask_header
|
|
assert restored.position_ids_header == payload.position_ids_header
|
|
assert restored.past_len == payload.past_len
|