diff --git a/vllm/config/kv_events.py b/vllm/config/kv_events.py index dc829113a8aa8..ce46cc03c39fe 100644 --- a/vllm/config/kv_events.py +++ b/vllm/config/kv_events.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Literal + +from pydantic import Field from pydantic.dataclasses import dataclass from vllm.config.utils import config @@ -17,7 +20,7 @@ class KVEventsConfig: Events can be published externally by zmq using the event publisher config. """ - publisher: str = "null" + publisher: Literal["null", "zmq"] = Field(default=None) """The publisher to use for publishing kv events. Can be "null", "zmq". """ @@ -47,3 +50,7 @@ class KVEventsConfig: """The topic to use for the event publisher. Consumers can subscribe to this topic to receive events. """ + + def __post_init__(self): + if self.publisher is None: + self.publisher = "zmq" if self.enable_kv_cache_events else "null" diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 4711467dafbdc..2fd46b23249ff 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -353,12 +353,12 @@ class EventPublisherFactory: cls, config: KVEventsConfig | None, data_parallel_rank: int = 0 ) -> EventPublisher: """Create publisher from a config mapping.""" - if not config: + if config is None or config.publisher == "null": return NullEventPublisher() config_dict = asdict(config) - kind = config_dict.pop("publisher", "null") + kind = config_dict.pop("publisher") config_dict.pop("enable_kv_cache_events") try: constructor = cls._registry[kind]