From 8ed8f1728c745d4e91ba9eaa09ed81159aa57e41 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Tue, 11 Aug 2026 11:25:33 -0500 Subject: [PATCH] fix(router): decode ATOM KV events with map schema Signed-off-by: Cam Quilici --- infera/router/kv_event/events.py | 6 ++--- tests/unit/router/test_kv_event_client.py | 28 +++++++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/infera/router/kv_event/events.py b/infera/router/kv_event/events.py index ed797fb2..ef6cb47a 100644 --- a/infera/router/kv_event/events.py +++ b/infera/router/kv_event/events.py @@ -132,11 +132,11 @@ class SglangKVEventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc def batch_type_for_engine(engine) -> type: """KVEventBatch schema matching a worker's engine wire format. - vLLM emits tagged-MAP events; SGLang (and anything else, the historical - default) emits tagged-ARRAY events. + vLLM and Infera's ATOM event hook emit tagged-MAP events; SGLang (the + historical default) emits tagged-ARRAY events. """ from infera.common.worker_pool import EngineType - if engine == EngineType.VLLM: + if engine in (EngineType.VLLM, EngineType.ATOM): return KVEventBatch return SglangKVEventBatch diff --git a/tests/unit/router/test_kv_event_client.py b/tests/unit/router/test_kv_event_client.py index 17161c43..7ad106f2 100644 --- a/tests/unit/router/test_kv_event_client.py +++ b/tests/unit/router/test_kv_event_client.py @@ -42,8 +42,10 @@ AllBlocksCleared, BlockRemoved, BlockStored, + KVEventBatch, SglangBlockStored, SglangKVEventBatch, + batch_type_for_engine, ) from infera.router.kv_event.hasher import ROUTER_SEED, hash_chunk @@ -335,6 +337,32 @@ def test_on_worker_removed_unknown_id_is_noop(): # ---------------------------------------------------------------------- +def test_atom_decoder_matches_infera_atom_hook_wire_format(): + """ATOM's Infera hook publishes the same tagged-map schema as vLLM. + + Selecting SGLang's tagged-array decoder for an ATOM worker rejects every + event and leaves KV-aware routing with an empty cache view. + """ + payload = msgspec.msgpack.encode( + KVEventBatch( + ts=1.0, + events=[ + _make_stored( + block_hashes=[111], + parent_block_hash=None, + token_ids=[1, 2, 3, 4], + ) + ], + ) + ) + + decoded = msgspec.msgpack.decode(payload, type=batch_type_for_engine(EngineType.ATOM)) + + assert isinstance(decoded, KVEventBatch) + assert isinstance(decoded.events[0], BlockStored) + assert decoded.events[0].block_hashes == [111] + + @pytest.mark.asyncio async def test_real_zmq_publisher_drives_cache_view(): """End-to-end: a real ZMQ PUB on `tcp://127.0.0.1:` emits a