neuron-tai/tests/test_kv_cache_distributed.py

"""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.error
import urllib.request

import pytest

from meshnet_node.model_backend import (
    KVCacheMiss,
    SessionCacheStore,
    TailTokenResult,
    TensorPayload,
    TorchModelShard,
)
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")


def test_prefill_cache_triton_cpu_failure_disables_cache_and_retries_stateless():
    """CPU shards must recover when hybrid model cache path dispatches Triton."""
    shard = object.__new__(TorchModelShard)
    shard.model_id = "fake-hybrid"
    shard.supports_kv_cache = True
    shard._effective_start = lambda start_layer=None: 22
    shard._new_session_cache = lambda: object()

    calls = []

    def fake_run_layers(hidden_states, attention_mask, position_ids, *, start_layer=None, cache=None, past_len=0):
        calls.append({"cache": cache, "past_len": past_len})
        if cache is not None:
            raise RuntimeError("Pointer argument cannot be accessed from Triton (cpu tensor?)")
        return "stateless-ok"

    shard._run_layers = fake_run_layers

    result = TorchModelShard._run_layers_session(
        shard,
        hidden_states=object(),
        attention_mask=None,
        position_ids=None,
        session_id="session-1",
        cache_mode="prefill",
    )

    assert result == "stateless-ok"
    assert shard.supports_kv_cache is False
    assert len(calls) == 2
    assert calls[0]["cache"] is not None
    assert calls[1]["cache"] is None


# ---------------------------------------------------------------------------
# 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]


class _BrokenTailBackend(_CachedTailBackend):
    """Tail whose forward always fails (e.g. missing compiler, OOM)."""

    def forward_bytes(self, *args, **kwargs):
        raise RuntimeError("Failed to find C compiler (test)")


def test_pipeline_failure_before_first_token_returns_502():
    """A dead hop must surface as an error, not an empty 200 completion."""
    head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BrokenTailBackend([]))
    head_port = head.start()
    tail_port = tail.start()
    try:
        try:
            _chat_once(head_port, tail_port, max_tokens=3)
        except urllib.error.HTTPError as exc:
            assert exc.code == 502
            body = json.loads(exc.read())
            assert "pipeline error" in body["error"]["message"]
            assert "C compiler" in body["error"]["message"]
        else:
            raise AssertionError("expected HTTP 502 from a failed pipeline")
    finally:
        head.stop()
        tail.stop()


def test_pipeline_failure_in_stream_emits_error_frame():
    """Streaming requests get an OpenAI-style error frame before [DONE]."""
    head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BrokenTailBackend([]))
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 3,
            "stream": True,
        }).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:
            assert resp.status == 200
            events = resp.read().decode().strip().split("\n\n")
    finally:
        head.stop()
        tail.stop()

    assert events[-1] == "data: [DONE]"
    error_frame = json.loads(events[-2][len("data: "):])
    assert error_frame["error"]["type"] == "upstream_error"
    assert "C compiler" in error_frame["error"]["message"]


def test_large_prefill_activation_survives_zstd_compressed_hop():
    """A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed.

    The head compresses and sets X-Meshnet-Encoding; the tail's /forward must
    decompress before shape validation, so a passing generation proves the
    compressed round trip (a mishandled encoding fails validation with 400).
    """

    class _BigHeadBackend(_CachedHeadBackend):
        def encode_prompt(self, prompt, session_id=None):
            self.prefills.append(session_id)
            if session_id:
                self._seq[session_id] = 2048
            return TensorPayload(
                body=b"\x00" * (1 * 2048 * 32 * 2),  # 128 KiB, above the zstd threshold
                shape=[1, 2048, 32],
                attention_mask_header=None,
                position_ids_header=None,
            )

    head_backend = _BigHeadBackend()
    tail_backend = _CachedTailBackend([(" a", 1), (" b", 2)])
    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=2)
    finally:
        head.stop()
        tail.stop()

    assert content == " a b"
    assert tail_backend.calls[0]["mode"] == "prefill"
    assert tail_backend.calls[0]["shape"] == [1, 2048, 32]


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