Skip to content

Commit c528b90

Browse files
authored
Fix EventPublisherFactory logic for disabled KV cache events (#27419)
Signed-off-by: Bradley <[email protected]>
1 parent 85fee74 commit c528b90

File tree

2 files changed

+54
-1
lines changed

2 files changed

+54
-1
lines changed

tests/distributed/test_events.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,3 +263,52 @@ def test_data_parallel_rank_tagging(publisher_config):
263263
pub_1.shutdown()
264264
sub_0.close()
265265
sub_1.close()
266+
267+
268+
def test_event_publisher_factory():
269+
"""Test event publisher factory creation behavior under different configurations"""
270+
from vllm.config.kv_events import KVEventsConfig
271+
from vllm.distributed.kv_events import ZmqEventPublisher
272+
273+
# test config is None
274+
publisher = EventPublisherFactory.create(None, DP_RANK)
275+
assert isinstance(publisher, NullEventPublisher)
276+
publisher.shutdown()
277+
278+
# test disable kv cache events
279+
config = KVEventsConfig(
280+
enable_kv_cache_events=False,
281+
publisher="zmq", # Even if zmq is specified, should return NullEventPublisher
282+
endpoint="tcp://localhost:5557",
283+
)
284+
publisher = EventPublisherFactory.create(config, DP_RANK)
285+
assert isinstance(publisher, NullEventPublisher)
286+
publisher.shutdown()
287+
288+
# test zmq publisher
289+
config = KVEventsConfig(
290+
enable_kv_cache_events=True,
291+
publisher="zmq",
292+
endpoint="inproc://test-factory-true",
293+
)
294+
publisher = EventPublisherFactory.create(config, DP_RANK)
295+
assert isinstance(publisher, ZmqEventPublisher)
296+
publisher.shutdown()
297+
298+
# test unknown publisher
299+
with pytest.raises(ValueError, match="Input should be"):
300+
KVEventsConfig(
301+
enable_kv_cache_events=True,
302+
publisher="unknown_publisher",
303+
endpoint="tcp://localhost:5557",
304+
)
305+
306+
# test publisher not specified
307+
config = KVEventsConfig(
308+
enable_kv_cache_events=True,
309+
# publisher not specified, should default to "zmq"
310+
endpoint="tcp://localhost:5557",
311+
)
312+
publisher = EventPublisherFactory.create(config, DP_RANK)
313+
assert isinstance(publisher, ZmqEventPublisher)
314+
publisher.shutdown()

vllm/distributed/kv_events.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,11 @@ def create(
353353
cls, config: KVEventsConfig | None, data_parallel_rank: int = 0
354354
) -> EventPublisher:
355355
"""Create publisher from a config mapping."""
356-
if config is None or config.publisher == "null":
356+
if (
357+
config is None
358+
or not config.enable_kv_cache_events
359+
or config.publisher == "null"
360+
):
357361
return NullEventPublisher()
358362

359363
config_dict = asdict(config)

0 commit comments

Comments
 (0)