Skip to content

Commit f04d3f6

Browse files
Etelisorozery
andauthored
[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 <oro@il.ibm.com>
1 parent 0a9396a commit f04d3f6

1 file changed

Lines changed: 14 additions & 19 deletions

File tree

tests/v1/kv_connector/unit/test_offloading_connector.py

Lines changed: 14 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -251,23 +251,19 @@ def test_cpu_offloading(
251251
kv_connector_extra_config=extra_config,
252252
)
253253

254-
# KV events are incompatible with HMA (setting kv_events_config
255-
# would force HMA off), so only enable them for non-HMA models.
256254
subscriber: MockSubscriber | None = None
257-
kv_events_config: KVEventsConfig | None = None
258-
if not uses_hma:
259-
port: int
260-
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
261-
s.bind(("0.0.0.0", 0))
262-
port = s.getsockname()[1]
263-
264-
events_endpoint = f"tcp://*:{port}"
265-
kv_events_config = KVEventsConfig(
266-
enable_kv_cache_events=True,
267-
publisher="zmq",
268-
endpoint=events_endpoint,
269-
topic="test",
270-
)
255+
port: int
256+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
257+
s.bind(("0.0.0.0", 0))
258+
port = s.getsockname()[1]
259+
260+
events_endpoint = f"tcp://*:{port}"
261+
kv_events_config = KVEventsConfig(
262+
enable_kv_cache_events=True,
263+
publisher="zmq",
264+
endpoint=events_endpoint,
265+
topic="test",
266+
)
271267

272268
# Attention-free / hybrid models disable prefix caching by default
273269
# (ModelConfig.is_prefix_caching_supported returns False). Without it,
@@ -289,9 +285,8 @@ def test_cpu_offloading(
289285
**({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
290286
)
291287

292-
if kv_events_config is not None:
293-
events_endpoint = events_endpoint.replace("*", "127.0.0.1")
294-
subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
288+
events_endpoint = events_endpoint.replace("*", "127.0.0.1")
289+
subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
295290

296291
try:
297292
_latency_test(llm, subscriber)

0 commit comments

Comments
 (0)