From f04d3f640e41ff217d4a793c315438aa2878f193 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:22:27 +0300 Subject: [PATCH] [Test] Enable KV cache events for HMA models in CPU offloading test (#47754) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../unit/test_offloading_connector.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 7d139f6ccd..16420b164b 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -251,23 +251,19 @@ def test_cpu_offloading( kv_connector_extra_config=extra_config, ) - # KV events are incompatible with HMA (setting kv_events_config - # would force HMA off), so only enable them for non-HMA models. subscriber: MockSubscriber | None = None - kv_events_config: KVEventsConfig | None = None - if not uses_hma: - port: int - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("0.0.0.0", 0)) - port = s.getsockname()[1] + port: int + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("0.0.0.0", 0)) + port = s.getsockname()[1] - events_endpoint = f"tcp://*:{port}" - kv_events_config = KVEventsConfig( - enable_kv_cache_events=True, - publisher="zmq", - endpoint=events_endpoint, - topic="test", - ) + events_endpoint = f"tcp://*:{port}" + kv_events_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=events_endpoint, + topic="test", + ) # Attention-free / hybrid models disable prefix caching by default # (ModelConfig.is_prefix_caching_supported returns False). Without it, @@ -289,9 +285,8 @@ def test_cpu_offloading( **({"max_num_seqs": 1} if current_platform.is_rocm() else {}), ) - if kv_events_config is not None: - events_endpoint = events_endpoint.replace("*", "127.0.0.1") - subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic) + events_endpoint = events_endpoint.replace("*", "127.0.0.1") + subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic) try: _latency_test(llm, subscriber)