popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test descriptions

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-11 22:25:30 +03:00
parent 7d259d7c9b
commit 7cf8d9bcf3
43 changed files with 876 additions and 265 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
tests/test_kv_cache_distributed.py
View File

@@ -37,6 +37,7 @@ class _Clock:
def test_store_lookup_roundtrip_advances_lru():
"Store lookup roundtrip advances lru\n\nTags: cache, model, node"
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)
@@ -46,12 +47,14 @@ def test_store_lookup_roundtrip_advances_lru():
def test_lookup_unknown_session_raises_cache_miss():
"Lookup unknown session raises cache miss\n\nTags: cache, model, node"
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():
"Seq len mismatch drops entry and raises\n\nTags: cache, model, node"
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=6, effective_start=0)
with pytest.raises(KVCacheMiss):
@@ -62,6 +65,7 @@ def test_seq_len_mismatch_drops_entry_and_raises():
def test_effective_start_mismatch_raises():
"Effective start mismatch raises\n\nTags: cache, model, node"
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=6, effective_start=12)
with pytest.raises(KVCacheMiss):
@@ -69,6 +73,7 @@ def test_effective_start_mismatch_raises():
def test_ttl_expiry_evicts_stale_sessions():
"Ttl expiry evicts stale sessions\n\nTags: cache, model, node"
clock = _Clock()
store = SessionCacheStore(max_sessions=4, ttl_seconds=60.0, clock=clock)
store.store("s1", cache=object(), seq_len=6, effective_start=0)
@@ -79,6 +84,7 @@ def test_ttl_expiry_evicts_stale_sessions():
def test_lru_eviction_bounds_session_count():
"Lru eviction bounds session count\n\nTags: cache, model, node"
clock = _Clock()
store = SessionCacheStore(max_sessions=2, ttl_seconds=1000.0, clock=clock)
store.store("s1", cache=object(), seq_len=1, effective_start=0)
@@ -93,6 +99,7 @@ def test_lru_eviction_bounds_session_count():
def test_drop_removes_session():
"Drop removes session\n\nTags: cache, model, node"
store = SessionCacheStore(max_sessions=4, ttl_seconds=100.0)
store.store("s1", cache=object(), seq_len=1, effective_start=0)
store.drop("s1")
@@ -101,7 +108,8 @@ def test_drop_removes_session():
def test_prefill_cache_triton_cpu_failure_disables_cache_and_retries_stateless():
"""CPU shards must recover when hybrid model cache path dispatches Triton."""
"CPU shards must recover when hybrid model cache path dispatches Triton.\n\nTags: cache, model, node"
shard = object.__new__(TorchModelShard)
shard.model_id = "fake-hybrid"
shard.supports_kv_cache = True
@@ -248,6 +256,7 @@ def _chat_once(head_port: int, tail_port: int, max_tokens: int) -> str:
def test_session_is_stable_and_decode_payloads_are_single_token():
"Session is stable and decode payloads are single token\n\nTags: cache, model, node"
head_backend = _CachedHeadBackend()
tail_backend = _CachedTailBackend([(" a", 1), (" b", 2), (" c", 3)])
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
@@ -287,7 +296,8 @@ class _BrokenTailBackend(_CachedTailBackend):
def test_pipeline_failure_before_first_token_returns_502():
"""A dead hop must surface as an error, not an empty 200 completion."""
"A dead hop must surface as an error, not an empty 200 completion.\n\nTags: cache, model, node"
head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
tail = TorchNodeServer(backend=_BrokenTailBackend([]))
head_port = head.start()
@@ -308,7 +318,8 @@ def test_pipeline_failure_before_first_token_returns_502():
def test_pipeline_failure_in_stream_emits_error_frame():
"""Streaming requests get an OpenAI-style error frame before [DONE]."""
"Streaming requests get an OpenAI-style error frame before [DONE].\n\nTags: cache, model, node, streaming"
head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
tail = TorchNodeServer(backend=_BrokenTailBackend([]))
head_port = head.start()
@@ -345,12 +356,8 @@ def test_pipeline_failure_in_stream_emits_error_frame():
def test_large_prefill_activation_survives_zstd_compressed_hop():
"""A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed.
"A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed.\n\nTags: cache, model, node"
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):
@@ -382,6 +389,7 @@ def test_large_prefill_activation_survives_zstd_compressed_hop():
def test_eos_token_id_stops_generation():
"Eos token id stops generation\n\nTags: cache, model, node"
head_backend = _CachedHeadBackend()
tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])
head = TorchNodeServer(backend=head_backend, tracker_mode=True)
@@ -399,8 +407,8 @@ def test_eos_token_id_stops_generation():
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."""
"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.\n\nTags: cache, model, node"
class _StatelessHead(_CachedHeadBackend):
supports_kv_cache = False
@@ -424,7 +432,8 @@ def test_stateless_fallback_stops_at_eos_token_id():
def test_decode_forward_logging_is_rate_limited():
"""Shard nodes log a per-session decode summary, not one line per token."""
"Shard nodes log a per-session decode summary, not one line per token.\n\nTags: cache, model, node"
tail_backend = _CachedTailBackend([])
tail = TorchNodeServer(backend=tail_backend)
tail.start()
@@ -442,6 +451,7 @@ def test_decode_forward_logging_is_rate_limited():
def test_downstream_cache_miss_falls_back_to_full_reprefill():
"Downstream cache miss falls back to full reprefill\n\nTags: cache, model, node, streaming"
head_backend = _CachedHeadBackend()
# Call 1 (the first decode) raises KVCacheMiss → node answers 409 →
# head re-prefills the full sequence and keeps generating.
@@ -467,7 +477,8 @@ def test_downstream_cache_miss_falls_back_to_full_reprefill():
def test_kv_head_with_legacy_tail_reprefills_every_step():
"""Mixed fleet: tail predates the protocol and returns no token_id."""
"Mixed fleet: tail predates the protocol and returns no token_id.\n\nTags: cache, model, node"
class _LegacyTailBackend:
model_id = "fake-model"
@@ -516,6 +527,7 @@ requires_real_model = pytest.mark.skipif(
@requires_real_model
def test_cached_distributed_generation_matches_stateless_golden():
"Cached distributed generation matches stateless golden\n\nTags: cache, model, node"
pytest.importorskip("torch")
from meshnet_node.model_backend import TorchModelShard