fix: attempt to fix hanging test

Signed-off-by: alec-flowers <aflowers@nvidia.com>

Author: alec-flowers

tests/v1/engine/test_engine_core_client.py

@@ -8,6 +8,7 @@
 
 import psutil
 import pytest
+import zmq
 from transformers import AutoTokenizer
 
 from vllm import SamplingParams
@@ -250,7 +251,6 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
         assert str(e_info.value) == "Call to echo method failed: help!"
 
 
-@create_new_process_for_each_test()
 @pytest.mark.parametrize(
     "multiprocessing_mode,publisher_config",
     [(True, "tcp"), (False, "inproc")],
@@ -332,8 +332,11 @@ def test_kv_cache_events(
             assert event.token_ids == custom_tokens, (
                 "Token ids should be the same as the custom tokens")
         finally:
-            subscriber.close()
             client.shutdown()
+            subscriber.close()
+            # TODO hack to try and fix CI hang
+            ctx = zmq.Context.instance()
+            ctx.term()
         return
 
 

vllm/distributed/kv_events.py

@@ -119,7 +119,7 @@ def __init__(
         topic: str = "",
     ) -> None:
         # Storage
-        self._event_queue = Queue[EventBatch](maxsize=max_queue_size)
+        self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)
         self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
 
         # ZMQ sockets
@@ -151,9 +151,9 @@ def publish(self, events: EventBatch) -> None:
     def shutdown(self) -> None:
         """Stop the publisher thread and clean up resources."""
         self._running = False
+        self._event_queue.put_nowait(None)
 
         start = time.time()
-
         pending_items = True
         while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT):
             pending_items = not self._event_queue.empty()
@@ -177,8 +177,7 @@ def shutdown(self) -> None:
             if self._replay is not None:
                 self._replay.close(linger=0)
         finally:
-            # Do not terminate context; other sockets may use it
-            pass
+            pass  # Do not terminate context; other sockets may use it
 
     def _socket_setup(self) -> None:
         """Initialize sockets
@@ -211,7 +210,7 @@ def _publisher_thread(self) -> None:
 
         assert self._pub is not None  # narrows type for mypy
 
-        while self._running or not self._event_queue.empty():
+        while self._running or self._event_queue.qsize() > 0:
             # --- replay (non-critical) ---------------------------------
             if self._replay is not None and self._replay.poll(0):
                 try:
@@ -222,6 +221,8 @@ def _publisher_thread(self) -> None:
             # --- main queue (critical) ---------------------------------
             try:
                 event = self._event_queue.get(timeout=0.1)
+                if event is None:
+                    break  # Sentinel received, exit thread
             except queue.Empty:
                 continue