428 lines
14 KiB
Python
428 lines
14 KiB
Python
"""AH-25: sharded per-node KV cache for distributed generation.
|
|
|
|
Covers the SessionCacheStore (TTL + LRU + mismatch handling), the HTTP
|
|
session protocol (stable session id, O(1) decode payloads, 409 cache-miss
|
|
fallback, legacy stateless compatibility), and an env-gated golden test that
|
|
proves cached and stateless distributed generation produce identical tokens
|
|
on a real two-shard Qwen2.5-0.5B split.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.model_backend import (
|
|
KVCacheMiss,
|
|
SessionCacheStore,
|
|
TailTokenResult,
|
|
TensorPayload,
|
|
)
|
|
from meshnet_node.torch_server import TorchNodeServer
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SessionCacheStore units
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class _Clock:
|
|
def __init__(self) -> None:
|
|
self.now = 0.0
|
|
|
|
def __call__(self) -> float:
|
|
return self.now
|
|
|
|
|
|
def test_store_lookup_roundtrip_advances_lru():
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0, clock=_Clock())
|
|
store.store("s1", cache=object(), seq_len=6, effective_start=12)
|
|
entry = store.lookup("s1", expected_seq_len=6, effective_start=12)
|
|
assert entry.seq_len == 6
|
|
entry.seq_len += 1
|
|
assert store.lookup("s1", expected_seq_len=7).seq_len == 7
|
|
|
|
|
|
def test_lookup_unknown_session_raises_cache_miss():
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("nope")
|
|
|
|
|
|
def test_seq_len_mismatch_drops_entry_and_raises():
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
|
store.store("s1", cache=object(), seq_len=6, effective_start=0)
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s1", expected_seq_len=9)
|
|
# Entry must be gone — a poisoned cache is never reused.
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s1")
|
|
|
|
|
|
def test_effective_start_mismatch_raises():
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
|
store.store("s1", cache=object(), seq_len=6, effective_start=12)
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s1", effective_start=21)
|
|
|
|
|
|
def test_ttl_expiry_evicts_stale_sessions():
|
|
clock = _Clock()
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=60.0, clock=clock)
|
|
store.store("s1", cache=object(), seq_len=6, effective_start=0)
|
|
clock.now = 61.0
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s1")
|
|
assert len(store) == 0
|
|
|
|
|
|
def test_lru_eviction_bounds_session_count():
|
|
clock = _Clock()
|
|
store = SessionCacheStore(max_sessions=2, ttl_seconds=1000.0, clock=clock)
|
|
store.store("s1", cache=object(), seq_len=1, effective_start=0)
|
|
store.store("s2", cache=object(), seq_len=1, effective_start=0)
|
|
store.lookup("s1") # s1 becomes most recent → s2 is LRU
|
|
store.store("s3", cache=object(), seq_len=1, effective_start=0)
|
|
assert len(store) == 2
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s2")
|
|
store.lookup("s1")
|
|
store.lookup("s3")
|
|
|
|
|
|
def test_drop_removes_session():
|
|
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
|
|
store.store("s1", cache=object(), seq_len=1, effective_start=0)
|
|
store.drop("s1")
|
|
with pytest.raises(KVCacheMiss):
|
|
store.lookup("s1")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HTTP session protocol with fake cached backends
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class _ChatTokenizer:
|
|
eos_token = ""
|
|
|
|
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
|
return "debug prompt"
|
|
|
|
|
|
class _CachedHeadBackend:
|
|
model_id = "fake-model"
|
|
total_layers = 12
|
|
is_head = True
|
|
is_tail = False
|
|
supports_kv_cache = True
|
|
tokenizer = _ChatTokenizer()
|
|
|
|
def __init__(self) -> None:
|
|
self.prefills: list[str | None] = []
|
|
self.decode_calls: list[tuple[int, str]] = []
|
|
self.released: list[str] = []
|
|
self._seq: dict[str, int] = {}
|
|
|
|
def eos_token_ids(self) -> list[int]:
|
|
return [99]
|
|
|
|
def release_session(self, session_id: str) -> None:
|
|
self.released.append(session_id)
|
|
|
|
def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
|
|
self.prefills.append(session_id)
|
|
if session_id:
|
|
self._seq[session_id] = 6
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 6 * 8 * 2),
|
|
shape=[1, 6, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
)
|
|
|
|
def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
|
|
self.decode_calls.append((token_id, session_id))
|
|
past = self._seq[session_id]
|
|
self._seq[session_id] = past + 1
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 1 * 8 * 2),
|
|
shape=[1, 1, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
past_len=past,
|
|
)
|
|
|
|
|
|
class _CachedTailBackend:
|
|
model_id = "fake-model"
|
|
total_layers = 12
|
|
is_head = False
|
|
is_tail = True
|
|
supports_kv_cache = True
|
|
|
|
def __init__(self, tokens, miss_on_call: int | None = None) -> None:
|
|
self._tokens = list(tokens)
|
|
self.miss_on_call = miss_on_call
|
|
self.calls: list[dict] = []
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body,
|
|
shape,
|
|
attention_mask_header,
|
|
position_ids_header,
|
|
start_layer=None,
|
|
session_id=None,
|
|
cache_mode=None,
|
|
past_len=None,
|
|
):
|
|
call_index = len(self.calls)
|
|
self.calls.append({
|
|
"session": session_id,
|
|
"mode": cache_mode,
|
|
"past_len": past_len,
|
|
"shape": list(shape),
|
|
})
|
|
if self.miss_on_call is not None and call_index == self.miss_on_call:
|
|
raise KVCacheMiss("session evicted (test)")
|
|
text, token_id = self._tokens.pop(0)
|
|
return TailTokenResult(text=text, token_id=token_id)
|
|
|
|
|
|
def _chat_once(head_port: int, tail_port: int, max_tokens: int) -> str:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": max_tokens,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 6},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
body = json.loads(resp.read())
|
|
return body["choices"][0]["message"]["content"]
|
|
|
|
|
|
def test_session_is_stable_and_decode_payloads_are_single_token():
|
|
head_backend = _CachedHeadBackend()
|
|
tail_backend = _CachedTailBackend([(" a", 1), (" b", 2), (" c", 3)])
|
|
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
content = _chat_once(head_port, tail_port, max_tokens=3)
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert content == " a b c"
|
|
assert len(tail_backend.calls) == 3
|
|
# Step 0 is a full-prompt prefill; steps 1+ carry only the new token.
|
|
assert tail_backend.calls[0]["mode"] == "prefill"
|
|
assert tail_backend.calls[0]["shape"] == [1, 6, 8]
|
|
for step, call in enumerate(tail_backend.calls[1:], start=1):
|
|
assert call["mode"] == "decode"
|
|
assert call["shape"] == [1, 1, 8]
|
|
assert call["past_len"] == 6 + (step - 1)
|
|
# One session id across every step of the generation.
|
|
sessions = {call["session"] for call in tail_backend.calls}
|
|
assert len(sessions) == 1
|
|
session_id = sessions.pop()
|
|
assert head_backend.prefills == [session_id]
|
|
assert head_backend.decode_calls == [(1, session_id), (2, session_id)]
|
|
# Head releases its own session state when the generation ends.
|
|
assert head_backend.released == [session_id]
|
|
|
|
|
|
def test_eos_token_id_stops_generation():
|
|
head_backend = _CachedHeadBackend()
|
|
tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])
|
|
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
content = _chat_once(head_port, tail_port, max_tokens=8)
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert content == " a"
|
|
assert len(tail_backend.calls) == 2
|
|
|
|
|
|
def test_stateless_fallback_stops_at_eos_token_id():
|
|
"""When kv caching is off, EOS must still stop generation by token id —
|
|
EOS decodes to "" (skip_special_tokens) so the text check never fires."""
|
|
|
|
class _StatelessHead(_CachedHeadBackend):
|
|
supports_kv_cache = False
|
|
|
|
head_backend = _StatelessHead()
|
|
tail_backend = _CachedTailBackend([(" a", 1), ("", 99), (" never", 3)])
|
|
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
content = _chat_once(head_port, tail_port, max_tokens=8)
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert content == " a"
|
|
# Stops after the EOS step instead of burning steps until max_tokens.
|
|
assert len(tail_backend.calls) == 2
|
|
assert head_backend.decode_calls == []
|
|
|
|
|
|
def test_decode_forward_logging_is_rate_limited():
|
|
"""Shard nodes log a per-session decode summary, not one line per token."""
|
|
tail_backend = _CachedTailBackend([])
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
tail.start()
|
|
try:
|
|
srv = tail._server
|
|
assert srv.note_decode_step("s1", now=0.0) == 1
|
|
assert srv.note_decode_step("s1", now=1.0) is None
|
|
assert srv.note_decode_step("s1", now=4.9) is None
|
|
assert srv.note_decode_step("s1", now=5.5) == 4
|
|
assert srv.note_decode_step("s1", now=6.0) is None
|
|
# Sessions are throttled independently.
|
|
assert srv.note_decode_step("s2", now=6.0) == 1
|
|
finally:
|
|
tail.stop()
|
|
|
|
|
|
def test_downstream_cache_miss_falls_back_to_full_reprefill():
|
|
head_backend = _CachedHeadBackend()
|
|
# Call 1 (the first decode) raises KVCacheMiss → node answers 409 →
|
|
# head re-prefills the full sequence and keeps generating.
|
|
tail_backend = _CachedTailBackend(
|
|
[(" a", 1), (" b", 2), (" c", 3)], miss_on_call=1,
|
|
)
|
|
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
content = _chat_once(head_port, tail_port, max_tokens=3)
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert content == " a b c"
|
|
modes = [call["mode"] for call in tail_backend.calls]
|
|
assert modes == ["prefill", "decode", "prefill", "decode"]
|
|
# Head re-prefilled once, with the same stable session id.
|
|
assert len(head_backend.prefills) == 2
|
|
assert len(set(head_backend.prefills)) == 1
|
|
|
|
|
|
def test_kv_head_with_legacy_tail_reprefills_every_step():
|
|
"""Mixed fleet: tail predates the protocol and returns no token_id."""
|
|
|
|
class _LegacyTailBackend:
|
|
model_id = "fake-model"
|
|
total_layers = 12
|
|
is_head = False
|
|
is_tail = True
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def forward_bytes(self, body, shape, attention_mask_header,
|
|
position_ids_header, start_layer=None, **kwargs):
|
|
self.calls += 1
|
|
return " x" if self.calls < 3 else ""
|
|
|
|
head_backend = _CachedHeadBackend()
|
|
tail_backend = _LegacyTailBackend()
|
|
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
content = _chat_once(head_port, tail_port, max_tokens=5)
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert content == " x x"
|
|
# No token_id from the tail → every step is a full prefill (legacy cost),
|
|
# never a decode against a cache the tail doesn't keep.
|
|
assert head_backend.decode_calls == []
|
|
assert len(head_backend.prefills) == 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Golden test on a real two-shard split (env-gated: loads Qwen2.5-0.5B twice)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_GOLDEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
|
|
|
|
requires_real_model = pytest.mark.skipif(
|
|
os.environ.get("MESHNET_REAL_MODEL_TESTS") != "1",
|
|
reason="set MESHNET_REAL_MODEL_TESTS=1 to run the real-model golden test",
|
|
)
|
|
|
|
|
|
@requires_real_model
|
|
def test_cached_distributed_generation_matches_stateless_golden():
|
|
pytest.importorskip("torch")
|
|
from meshnet_node.model_backend import TorchModelShard
|
|
|
|
head = TorchModelShard(_GOLDEN_MODEL, 0, 11)
|
|
tail = TorchModelShard(_GOLDEN_MODEL, 12, 23)
|
|
steps = 12
|
|
prompt = head.tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": "Count from 1 to 5."}],
|
|
add_generation_prompt=True,
|
|
tokenize=False,
|
|
)
|
|
|
|
# Reference: today's stateless path — re-encode the full sequence each step.
|
|
stateless_ids: list[int] = []
|
|
text = prompt
|
|
for _ in range(steps):
|
|
payload = head.encode_prompt(text)
|
|
result = tail.forward_bytes(
|
|
payload.body, payload.shape,
|
|
payload.attention_mask_header, payload.position_ids_header,
|
|
start_layer=12,
|
|
)
|
|
stateless_ids.append(result.token_id)
|
|
text += result.text
|
|
|
|
# Cached path: one prefill, then single-token decode steps.
|
|
session = "golden-session"
|
|
cached_ids: list[int] = []
|
|
payload = head.encode_prompt(prompt, session_id=session)
|
|
result = tail.forward_bytes(
|
|
payload.body, payload.shape,
|
|
payload.attention_mask_header, payload.position_ids_header,
|
|
start_layer=12, session_id=session, cache_mode="prefill",
|
|
)
|
|
cached_ids.append(result.token_id)
|
|
for _ in range(steps - 1):
|
|
payload = head.encode_next_token(cached_ids[-1], session)
|
|
assert payload.shape[1] == 1, "decode payload must be a single token"
|
|
result = tail.forward_bytes(
|
|
payload.body, payload.shape,
|
|
None, payload.position_ids_header,
|
|
start_layer=12, session_id=session, cache_mode="decode",
|
|
past_len=payload.past_len,
|
|
)
|
|
cached_ids.append(result.token_id)
|
|
|
|
assert cached_ids == stateless_ids
|