| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import queue |
| 5 | import threading |
| 6 | import time |
| 7 | from abc import ABC, abstractmethod |
| 8 | from collections import Counter, deque |
| 9 | from collections.abc import Callable |
| 10 | from dataclasses import asdict |
| 11 | from itertools import count |
| 12 | from queue import Queue |
| 13 | from typing import Any |
| 14 | |
| 15 | import msgspec |
| 16 | import zmq |
| 17 | |
| 18 | from vllm.config.kv_events import KVEventsConfig |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.utils.network_utils import ( |
| 21 | get_ip, |
| 22 | get_tcp_uri, |
| 23 | is_valid_ipv6_address, |
| 24 | split_zmq_path, |
| 25 | ) |
| 26 | from vllm.v1.core.kv_cache_utils import ExternalBlockHash |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | class EventBatch( |
| 32 | msgspec.Struct, |
| 33 | array_like=True, # type: ignore[call-arg] |
| 34 | omit_defaults=True, # type: ignore[call-arg] |
| 35 | gc=False, # type: ignore[call-arg] |
| 36 | ): |
| 37 | ts: float |
| 38 | events: list[Any] |
| 39 | data_parallel_rank: int | None = None |
| 40 | |
| 41 | |
| 42 | class KVCacheEvent( |
| 43 | msgspec.Struct, |
| 44 | omit_defaults=True, # type: ignore[call-arg] |
| 45 | gc=False, # type: ignore[call-arg] |
| 46 | tag=True, |
| 47 | ): |
| 48 | """Base class for all KV cache-related events.""" |
| 49 | |
| 50 | |
| 51 | MEDIUM_GPU = "GPU" |
| 52 | MEDIUM_CPU = "CPU" |
| 53 | MEDIUM_STORAGE = "STORAGE" |
| 54 | |
| 55 | |
| 56 | class BlockStored(KVCacheEvent): |
| 57 | block_hashes: list[ExternalBlockHash] |
| 58 | parent_block_hash: ExternalBlockHash | None |
| 59 | token_ids: list[int] |
| 60 | block_size: int |
| 61 | |
| 62 | """Deprecated: use `lora_name` for KV block key hash. |
| 63 | Retained for backward compatibility. |
| 64 | """ |
| 65 | lora_id: int | None |
| 66 | |
| 67 | medium: str | None |
| 68 | lora_name: str | None |
| 69 | |
| 70 | """Extra keys used in block hash computation, one entry per block in |
| 71 | block_hashes. Each entry contains, in order, the LoRA name (if |
| 72 | `lora_name` is set), `(mm_identifier, offset)` pairs, the cache_salt |
| 73 | string and prompt embedding digests for that specific block. These are |
| 74 | published untagged; the block hash itself uses each value tagged with |
| 75 | its source, e.g. `("lora", name)` or `("mm", mm_identifier, offset)`. |
| 76 | """ |
| 77 | extra_keys: list[tuple[Any, ...] | None] | None = None |
| 78 | |
| 79 | """Store events carry cache-spec metadata so consumers can classify and |
| 80 | filter groups as they are learned. Remove events only need group_idx+hash. |
| 81 | """ |
| 82 | group_idx: int | None = None |
| 83 | kv_cache_spec_kind: str | None = None |
| 84 | kv_cache_spec_sliding_window: int | None = None |
| 85 | """LOCAL or REMOTE relative to the publisher; None means unspecified.""" |
| 86 | locality: str | None = None |
| 87 | """Secondary offloading tier identifier, if generated by one.""" |
| 88 | ownership: str | None = None |
| 89 | """Session that triggered this store or reuse report. |
| 90 | |
| 91 | This identifies the request context that emitted the event, not exclusive |
| 92 | ownership of the underlying block, which may be shared across sessions. |
| 93 | """ |
| 94 | session_id: str | None = None |
| 95 | |
| 96 | def __hash__(self) -> int: |
| 97 | return hash( |
| 98 | ( |
| 99 | tuple(self.block_hashes), |
| 100 | self.parent_block_hash, |
| 101 | tuple(self.token_ids), |
| 102 | self.block_size, |
| 103 | self.lora_id, |
| 104 | self.medium, |
| 105 | tuple(self.extra_keys) if self.extra_keys else None, |
| 106 | self.group_idx, |
| 107 | self.kv_cache_spec_kind, |
| 108 | self.kv_cache_spec_sliding_window, |
| 109 | self.locality, |
| 110 | self.ownership, |
| 111 | self.session_id, |
| 112 | ) |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | class BlockRemoved(KVCacheEvent): |
| 117 | block_hashes: list[ExternalBlockHash] |
| 118 | medium: str | None |
| 119 | group_idx: int | None = None |
| 120 | """LOCAL or REMOTE relative to the publisher; None means unspecified.""" |
| 121 | locality: str | None = None |
| 122 | """Secondary offloading tier identifier, if generated by one.""" |
| 123 | ownership: str | None = None |
| 124 | |
| 125 | def __hash__(self) -> int: |
| 126 | return hash( |
| 127 | ( |
| 128 | tuple(self.block_hashes), |
| 129 | self.medium, |
| 130 | self.group_idx, |
| 131 | self.locality, |
| 132 | self.ownership, |
| 133 | ) |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | class AllBlocksCleared(KVCacheEvent): |
| 138 | pass |
| 139 | |
| 140 | |
| 141 | class KVEventBatch(EventBatch): |
| 142 | events: list[BlockStored | BlockRemoved | AllBlocksCleared] |
| 143 | |
| 144 | |
| 145 | class KVEventAggregator: |
| 146 | """Aggregates KV events across multiple workers. |
| 147 | Tracks how many times each event appears and returns only those |
| 148 | that were emitted by all workers. |
| 149 | """ |
| 150 | |
| 151 | __slots__ = ("_event_counter", "_num_workers") |
| 152 | |
| 153 | def __init__(self, num_workers: int) -> None: |
| 154 | if num_workers <= 0: |
| 155 | raise ValueError("num_workers must be greater than zero.") |
| 156 | self._event_counter: Counter[KVCacheEvent] = Counter() |
| 157 | self._num_workers: int = num_workers |
| 158 | |
| 159 | def add_events(self, events: list[KVCacheEvent]) -> None: |
| 160 | """Add events from a worker batch. |
| 161 | |
| 162 | Args: |
| 163 | events: List of KVCacheEvent objects. |
| 164 | |
| 165 | """ |
| 166 | if not isinstance(events, list): |
| 167 | raise TypeError("events must be a list of KVCacheEvent.") |
| 168 | self._event_counter.update(events) |
| 169 | |
| 170 | def get_common_events(self) -> list[KVCacheEvent]: |
| 171 | """Return events that appeared in all workers. |
| 172 | |
| 173 | Returns: |
| 174 | List of events present in all workers. |
| 175 | |
| 176 | """ |
| 177 | return [ |
| 178 | event |
| 179 | for event, count in self._event_counter.items() |
| 180 | if count == self._num_workers |
| 181 | ] |
| 182 | |
| 183 | def get_all_events(self) -> list[KVCacheEvent]: |
| 184 | """Return all events for all workers. |
| 185 | |
| 186 | Returns: |
| 187 | List of events for all workers. |
| 188 | |
| 189 | """ |
| 190 | return list(self._event_counter.elements()) |
| 191 | |
| 192 | def clear_events(self) -> None: |
| 193 | """Clear all tracked events.""" |
| 194 | self._event_counter.clear() |
| 195 | |
| 196 | def increment_workers(self, count: int = 1) -> None: |
| 197 | """Increment the number of workers contributing events. |
| 198 | |
| 199 | Args: |
| 200 | count: Number to increment the workers by. |
| 201 | |
| 202 | """ |
| 203 | if count <= 0: |
| 204 | raise ValueError("count must be positive.") |
| 205 | self._num_workers += count |
| 206 | |
| 207 | def reset_workers(self) -> None: |
| 208 | """Reset the number of workers to 1.""" |
| 209 | self._num_workers = 1 |
| 210 | |
| 211 | def get_number_of_workers(self) -> int: |
| 212 | """Return the number of workers. |
| 213 | |
| 214 | Returns: |
| 215 | int number of workers. |
| 216 | |
| 217 | """ |
| 218 | return self._num_workers |
| 219 | |
| 220 | def __repr__(self) -> str: |
| 221 | return ( |
| 222 | f"<KVEventAggregator workers={self._num_workers}, " |
| 223 | f"events={len(self._event_counter)}>" |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | class KVConnectorKVEvents(ABC): |
| 228 | """Abstract base class for KV events. |
| 229 | Acts as a container for KV events from the connector. |
| 230 | """ |
| 231 | |
| 232 | @abstractmethod |
| 233 | def add_events(self, events: list[KVCacheEvent]) -> None: |
| 234 | raise NotImplementedError |
| 235 | |
| 236 | @abstractmethod |
| 237 | def aggregate(self) -> "KVConnectorKVEvents": |
| 238 | raise NotImplementedError |
| 239 | |
| 240 | @abstractmethod |
| 241 | def increment_workers(self, count: int = 1) -> None: |
| 242 | raise NotImplementedError |
| 243 | |
| 244 | @abstractmethod |
| 245 | def get_all_events(self) -> list[KVCacheEvent]: |
| 246 | raise NotImplementedError |
| 247 | |
| 248 | @abstractmethod |
| 249 | def get_number_of_workers(self) -> int: |
| 250 | raise NotImplementedError |
| 251 | |
| 252 | @abstractmethod |
| 253 | def clear_events(self) -> None: |
| 254 | raise NotImplementedError |
| 255 | |
| 256 | def merge(self, other: "KVConnectorKVEvents") -> "KVConnectorKVEvents": |
| 257 | self.add_events(other.get_all_events()) |
| 258 | return self |
| 259 | |
| 260 | |
| 261 | class EventPublisher(ABC): |
| 262 | """Lightweight publisher for EventBatch batches with data parallelism |
| 263 | support. |
| 264 | |
| 265 | In data parallel setups, each DP rank runs its own EventPublisher instance |
| 266 | to avoid duplicate events and ensure proper event attribution: |
| 267 | |
| 268 | - Each DP rank creates a separate publisher |
| 269 | - Publishers automatically annotate events with their data_parallel_rank |
| 270 | - This allows consumers to distinguish events from different DP ranks |
| 271 | |
| 272 | The publisher is responsible for adding DP metadata since the scheduler |
| 273 | operates independently of DP topology and shouldn't need DP awareness. |
| 274 | """ |
| 275 | |
| 276 | def __init__(self, data_parallel_rank: int = 0) -> None: |
| 277 | self._data_parallel_rank = data_parallel_rank |
| 278 | |
| 279 | @abstractmethod |
| 280 | def publish(self, events: EventBatch) -> None: |
| 281 | """Emit events in order. |
| 282 | |
| 283 | Implementations should guarantee at-least-once delivery and |
| 284 | monotonic ordering (e.g., via sequence numbers). |
| 285 | """ |
| 286 | |
| 287 | @abstractmethod |
| 288 | def shutdown(self) -> None: |
| 289 | """Shutdown the publisher.""" |
| 290 | |
| 291 | def get_publisher_config(self) -> KVEventsConfig | None: |
| 292 | """Return the publisher's resolved runtime configuration.""" |
| 293 | return None |
| 294 | |
| 295 | |
| 296 | class NullEventPublisher(EventPublisher): |
| 297 | """No-op implementation (default when disabled).""" |
| 298 | |
| 299 | def publish(self, events) -> None: |
| 300 | return |
| 301 | |
| 302 | def shutdown(self) -> None: |
| 303 | return |
| 304 | |
| 305 | |
| 306 | class ZmqEventPublisher(EventPublisher): |
| 307 | """Reliable PUB/ROUTER publisher with an in-memory replay buffer. |
| 308 | |
| 309 | Spawns a separate thread to handle publishing from a queue. |
| 310 | |
| 311 | Parameters |
| 312 | ---------- |
| 313 | endpoint: |
| 314 | PUB address. Use `tcp://*:5557` to bind or `tcp://host:5557` to |
| 315 | connect. |
| 316 | replay_endpoint: |
| 317 | Optional ROUTER address for replay requests. When given, subscribers can |
| 318 | request missed batches by sending the starting sequence number as an |
| 319 | 8-byte big-endian integer. |
| 320 | buffer_steps: |
| 321 | Number of past batches to keep for replay. |
| 322 | hwm: |
| 323 | ZeroMQ high-water-mark for PUB socket. |
| 324 | max_queue_size: |
| 325 | Maximum number of events to buffer in memory. |
| 326 | topic: |
| 327 | Topic to publish events to. |
| 328 | |
| 329 | """ |
| 330 | |
| 331 | SHUTDOWN_TIMEOUT: float = 1.0 |
| 332 | END_SEQ = (-1).to_bytes(8, "big", signed=True) |
| 333 | |
| 334 | def __init__( |
| 335 | self, |
| 336 | data_parallel_rank: int, |
| 337 | endpoint: str = "tcp://*:5557", |
| 338 | replay_endpoint: str | None = None, |
| 339 | buffer_steps: int = 10_000, |
| 340 | hwm: int = 100_000, |
| 341 | max_queue_size: int = 100_000, |
| 342 | topic: str = "", |
| 343 | ) -> None: |
| 344 | # Storage |
| 345 | super().__init__(data_parallel_rank) |
| 346 | self._event_queue = Queue[EventBatch | None](maxsize=max_queue_size) |
| 347 | self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps) |
| 348 | |
| 349 | # ZMQ sockets |
| 350 | self._ctx = zmq.Context.instance() |
| 351 | self._pub: zmq.Socket | None = None |
| 352 | self._replay: zmq.Socket | None = None |
| 353 | self._dp_rank = data_parallel_rank |
| 354 | |
| 355 | self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank) |
| 356 | self._replay_endpoint = self.offset_endpoint_port( |
| 357 | replay_endpoint, self._dp_rank |
| 358 | ) |
| 359 | assert self._endpoint is not None |
| 360 | self._hwm = hwm |
| 361 | self._socket_setup() |
| 362 | self._publisher_config = KVEventsConfig( |
| 363 | enable_kv_cache_events=True, |
| 364 | publisher="zmq", |
| 365 | endpoint=self._endpoint, |
| 366 | replay_endpoint=self._replay_endpoint, |
| 367 | buffer_steps=buffer_steps, |
| 368 | hwm=hwm, |
| 369 | max_queue_size=max_queue_size, |
| 370 | topic=topic, |
| 371 | ) |
| 372 | |
| 373 | # Payload |
| 374 | self._seq_gen = count() |
| 375 | self._topic_bytes = topic.encode("utf-8") |
| 376 | |
| 377 | # Thread |
| 378 | self._running = True |
| 379 | logger.info("Starting ZMQ publisher thread") |
| 380 | |
| 381 | self._thread = threading.Thread( |
| 382 | target=self._publisher_thread, daemon=True, name="zmq-publisher" |
| 383 | ) |
| 384 | self._thread.start() |
| 385 | |
| 386 | def get_publisher_config(self) -> KVEventsConfig: |
| 387 | return self._publisher_config |
| 388 | |
| 389 | def publish(self, events: EventBatch) -> None: |
| 390 | if not self._running: |
| 391 | raise RuntimeError("Publisher is closed") |
| 392 | if events.data_parallel_rank is None: |
| 393 | events.data_parallel_rank = self._data_parallel_rank |
| 394 | self._event_queue.put(events) |
| 395 | |
| 396 | def shutdown(self) -> None: |
| 397 | """Stop the publisher thread and clean up resources.""" |
| 398 | self._running = False |
| 399 | self._event_queue.put_nowait(None) |
| 400 | |
| 401 | start = time.time() |
| 402 | pending_items = True |
| 403 | while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT): |
| 404 | pending_items = not self._event_queue.empty() |
| 405 | if pending_items: |
| 406 | time.sleep(0.1) |
| 407 | |
| 408 | if pending_items: |
| 409 | logger.warning( |
| 410 | "Warning: Queue still has %s items after %s seconds timeout", |
| 411 | self._event_queue.qsize(), |
| 412 | self.SHUTDOWN_TIMEOUT, |
| 413 | ) |
| 414 | |
| 415 | if self._thread.is_alive(): |
| 416 | self._thread.join(timeout=self.SHUTDOWN_TIMEOUT) |
| 417 | |
| 418 | # Clean up ZMQ resources |
| 419 | try: |
| 420 | if self._pub is not None: |
| 421 | self._pub.close(linger=0) |
| 422 | if self._replay is not None: |
| 423 | self._replay.close(linger=0) |
| 424 | finally: |
| 425 | pass # Do not terminate context; other sockets may use it |
| 426 | |
| 427 | def _socket_setup(self) -> None: |
| 428 | """Initialize sockets |
| 429 | https://pyzmq.readthedocs.io/en/v19.0.0/morethanbindings.html#thread-safety |
| 430 | """ |
| 431 | if self._pub is None: |
| 432 | self._pub = self._ctx.socket(zmq.PUB) |
| 433 | self._pub.set_hwm(self._hwm) |
| 434 | assert self._endpoint is not None |
| 435 | if self._is_bind_endpoint(self._endpoint): |
| 436 | self._endpoint = self._bind(self._pub, self._endpoint) |
| 437 | else: |
| 438 | self._pub.connect(self._endpoint) |
| 439 | |
| 440 | # Set up replay socket: use ROUTER |
| 441 | # 1) handles multiple REQ clients (identities) |
| 442 | # 2) lets us send back one request → many replies (streamed events) |
| 443 | # 3) works in our non‑blocking poll loop alongside PUB |
| 444 | if self._replay_endpoint is not None: |
| 445 | self._replay = self._ctx.socket(zmq.ROUTER) |
| 446 | self._replay_endpoint = self._bind(self._replay, self._replay_endpoint) |
| 447 | |
| 448 | @staticmethod |
| 449 | def _is_bind_endpoint(endpoint: str) -> bool: |
| 450 | """Bind wildcard and ipc/inproc endpoints; connect to the rest.""" |
| 451 | return ( |
| 452 | "*" in endpoint |
| 453 | or "::" in endpoint |
| 454 | or endpoint.startswith(("ipc://", "inproc://")) |
| 455 | ) |
| 456 | |
| 457 | @staticmethod |
| 458 | def _bind(sock: zmq.Socket, endpoint: str) -> str: |
| 459 | """Bind; tcp://*:0 binds this node's IP and returns the bound endpoint.""" |
| 460 | if endpoint != "tcp://*:0": |
| 461 | sock.bind(endpoint) |
| 462 | return endpoint |
| 463 | ip = get_ip() |
| 464 | if ip in ("0.0.0.0", "::"): |
| 465 | raise ValueError("tcp://*:0 needs a node IP; set VLLM_HOST_IP") |
| 466 | if is_valid_ipv6_address(ip): |
| 467 | sock.setsockopt(zmq.IPV6, 1) |
| 468 | sock.bind(get_tcp_uri(ip, 0)) |
| 469 | _, _, port = split_zmq_path(sock.getsockopt_string(zmq.LAST_ENDPOINT)) |
| 470 | return get_tcp_uri(ip, int(port)) |
| 471 | |
| 472 | def _publisher_thread(self) -> None: |
| 473 | """Background thread that processes the event queue.""" |
| 474 | self._pack = msgspec.msgpack.Encoder() |
| 475 | |
| 476 | assert self._pub is not None # narrows type for mypy |
| 477 | |
| 478 | while self._running or self._event_queue.qsize() > 0: |
| 479 | # --- replay (non-critical) --------------------------------- |
| 480 | if self._replay is not None and self._replay.poll(0): |
| 481 | try: |
| 482 | self._service_replay() |
| 483 | except Exception as e: |
| 484 | logger.exception("Error in replay: %s", e) |
| 485 | |
| 486 | # --- main queue (critical) --------------------------------- |
| 487 | try: |
| 488 | event = self._event_queue.get(timeout=0.1) |
| 489 | if event is None: |
| 490 | break # Sentinel received, exit thread |
| 491 | except queue.Empty: |
| 492 | continue |
| 493 | |
| 494 | try: |
| 495 | seq = next(self._seq_gen) |
| 496 | |
| 497 | payload = self._pack.encode(event) |
| 498 | seq_bytes = seq.to_bytes(8, "big") |
| 499 | self._pub.send_multipart((self._topic_bytes, seq_bytes, payload)) |
| 500 | |
| 501 | self._buffer.append((seq, payload)) |
| 502 | self._event_queue.task_done() |
| 503 | |
| 504 | except Exception as e: |
| 505 | # Publishing failed; back-off a bit to avoid a tight error loop |
| 506 | logger.exception("Error in publisher thread: %s", e) |
| 507 | time.sleep(0.1) |
| 508 | |
| 509 | def _service_replay(self) -> None: |
| 510 | """If a replay request is waiting, send buffered batches.""" |
| 511 | assert self._replay is not None # narrows type for mypy |
| 512 | |
| 513 | frame = self._replay.recv_multipart() |
| 514 | if len(frame) != 3: |
| 515 | logger.warning("Invalid replay request: %s", frame) |
| 516 | return |
| 517 | client_id, _, start_seq_bytes = frame |
| 518 | start_seq = int.from_bytes(start_seq_bytes, "big") |
| 519 | |
| 520 | for seq, buf in self._buffer: |
| 521 | if seq >= start_seq: |
| 522 | # Subscriber receives (topic, seq_bytes, payload) |
| 523 | self._replay.send_multipart( |
| 524 | (client_id, b"", self._topic_bytes, seq.to_bytes(8, "big"), buf) |
| 525 | ) |
| 526 | # Send end of sequence marker |
| 527 | self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b"")) |
| 528 | |
| 529 | @staticmethod |
| 530 | def offset_endpoint_port( |
| 531 | endpoint: str | None, data_parallel_rank: int |
| 532 | ) -> str | None: |
| 533 | """Helper function to offset the port in an endpoint by |
| 534 | the data parallel rank. |
| 535 | |
| 536 | Args: |
| 537 | endpoint: The endpoint string |
| 538 | (e.g., "tcp://*:5557" or "inproc://cache") |
| 539 | data_parallel_rank: The data parallel rank to offset by |
| 540 | |
| 541 | Returns: |
| 542 | The endpoint with the port offset by data_parallel_rank |
| 543 | or suffix appended; a tcp port of 0 is left for the OS. |
| 544 | |
| 545 | """ |
| 546 | # Do nothing if input is None or data_parallel_rank is 0 |
| 547 | if not endpoint or data_parallel_rank == 0: |
| 548 | return endpoint |
| 549 | |
| 550 | if "inproc" in endpoint: |
| 551 | return f"{endpoint}_dp{data_parallel_rank}" |
| 552 | if "tcp" in endpoint: |
| 553 | if endpoint and ":" in endpoint: |
| 554 | # Get everything after the last colon (the port) |
| 555 | last_colon_idx = endpoint.rfind(":") |
| 556 | base_addr = endpoint[:last_colon_idx] |
| 557 | base_port = int(endpoint[last_colon_idx + 1 :]) |
| 558 | if base_port == 0: |
| 559 | return endpoint |
| 560 | new_port = base_port + data_parallel_rank |
| 561 | return f"{base_addr}:{new_port}" |
| 562 | return endpoint |
| 563 | raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'") |
| 564 | |
| 565 | |
| 566 | class EventPublisherFactory: |
| 567 | _registry: dict[str, Callable[..., EventPublisher]] = { |
| 568 | "null": NullEventPublisher, |
| 569 | "zmq": ZmqEventPublisher, |
| 570 | } |
| 571 | |
| 572 | @classmethod |
| 573 | def register_publisher(cls, name: str, ctor: Callable[..., EventPublisher]) -> None: |
| 574 | if name in cls._registry: |
| 575 | raise KeyError(f"publisher '{name}' already registered") |
| 576 | cls._registry[name] = ctor |
| 577 | |
| 578 | @classmethod |
| 579 | def create( |
| 580 | cls, config: KVEventsConfig | None, data_parallel_rank: int = 0 |
| 581 | ) -> EventPublisher: |
| 582 | """Create publisher from a config mapping.""" |
| 583 | if ( |
| 584 | config is None |
| 585 | or not config.enable_kv_cache_events |
| 586 | or config.publisher == "null" |
| 587 | ): |
| 588 | return NullEventPublisher() |
| 589 | |
| 590 | config_dict = asdict(config) |
| 591 | |
| 592 | kind = config_dict.pop("publisher") |
| 593 | config_dict.pop("enable_kv_cache_events") |
| 594 | try: |
| 595 | constructor = cls._registry[kind] |
| 596 | except KeyError as exc: |
| 597 | raise ValueError(f"Unknown event publisher '{kind}'") from exc |
| 598 | return constructor(data_parallel_rank=data_parallel_rank, **config_dict) |