diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index 38a595cc7..eff3d429e 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -41,6 +41,74 @@ def _wait_for_instance_list( return None +def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None: + """Drain queued SUB messages for a short window.""" + deadline = time.time() + max_seconds + while time.time() < deadline: + _recv_instance_list(sub, timeout_ms=50) + + +def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): + """Rapid updates should be coalesced into fewer PUB messages.""" + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + sub_ctx = zmq.Context.instance() + sub = sub_ctx.socket(zmq.SUB) + sub.connect(pub_addr) + sub.setsockopt(zmq.SUBSCRIBE, b"") + + time.sleep(0.3) # PUB/SUB slow-joiner + + client = OmniCoordClientForStage( + router_addr, + "tcp://stage:coalesce", + "tcp://stage:coalesce-out", + 0, + ) + + # Wait for initial registration broadcast and clear any queued messages. + msg = _wait_for_instance_list(sub, expected_count=1) + assert msg is not None + _drain_sub_messages(sub) + + # Burst many queue updates in a short period. + update_count = 80 + for i in range(update_count): + client.update_info(queue_length=i) + + # With publish_min_interval=0.1s, received messages over ~1s should be + # much smaller than update_count (coalescing effect). + window_s = 1.1 + deadline = time.time() + window_s + recv_count = 0 + while time.time() < deadline: + if _recv_instance_list(sub, timeout_ms=100) is not None: + recv_count += 1 + + assert recv_count < update_count // 2, ( + f"expected coalesced PUB traffic, got {recv_count} for {update_count} updates" + ) + + client.close() + coordinator.close() + sub.close(0) + sub_ctx.term() + + def test_omni_coordinator_registration_broadcast(): """Verify that after multiple OmniCoordClientForStage instances register, OmniCoordinator publishes an InstanceList containing all registered instances. diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index 7ff608f2f..2c7c8fbb9 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -61,6 +61,7 @@ class OmniCoordinator: self._publish_min_interval: float = 0.1 # seconds self._pending_broadcast: bool = False + self._pending_lock = threading.Lock() self._running = True self._closed = False @@ -84,13 +85,13 @@ class OmniCoordinator: """Add a new instance based on an incoming event.""" with self._lock: self._add_new_instance_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() def update_instance_info(self, event: InstanceEvent) -> None: """Update an existing instance based on an incoming event.""" with self._lock: self._update_instance_info_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() def remove_instance(self, event: InstanceEvent) -> None: """Mark an instance as removed / down based on an incoming event. @@ -101,10 +102,15 @@ class OmniCoordinator: """ with self._lock: self._remove_instance_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() - def publish_instance_list_update(self) -> None: - """Publish the current active instance list to all subscribers.""" + def publish_instance_list_update(self) -> bool: + """Publish the current active instance list to all subscribers. + + Returns: + True if the PUB send succeeded, False if it was dropped (e.g. + socket not ready when using ``zmq.NOBLOCK``). + """ active_list = self.get_active_instances() payload = asdict(active_list) data = json.dumps(payload).encode("utf-8") @@ -113,20 +119,18 @@ class OmniCoordinator: try: # PUB socket is best-effort; drop update if not ready. self._pub.send(data, flags=zmq.NOBLOCK) + return True except (zmq.Again, zmq.ZMQError): # Silently ignore send failures; next update will catch up. - return + return False - def _schedule_broadcast(self, force: bool) -> None: - """Schedule a broadcast, optionally bypassing throttling. + def _schedule_broadcast(self) -> None: + """Request a broadcast to be flushed by the periodic loop. - When ``force`` is True, publish immediately. Otherwise, mark a pending - broadcast that will be flushed by the periodic loop at most once per - ``_publish_min_interval``. + All broadcast requests are coalesced via ``_pending_broadcast`` and + flushed at most once per ``_publish_min_interval``. """ - if force: - self.publish_instance_list_update() - else: + with self._pending_lock: self._pending_broadcast = True def _mark_instance_error_locked(self, info: InstanceInfo) -> None: @@ -152,8 +156,8 @@ class OmniCoordinator: for input_addr in to_delete: del self._instances[input_addr] if timed_out: - # Instance liveness changed; force immediate broadcast. - self._schedule_broadcast(force=True) + # Instance liveness changed; request broadcast. + self._schedule_broadcast() def close(self) -> None: """Shut down background threads and close all ZMQ sockets.""" @@ -228,9 +232,9 @@ class OmniCoordinator: def _periodic_loop(self) -> None: """Periodic loop to check heartbeat timeouts and flush broadcasts. - Heartbeat timeouts are checked on their original cadence, while - queue_length / non-liveness updates are coalesced and flushed at - most once per ``_publish_min_interval``. + Heartbeat timeouts are checked on their original cadence, while all + broadcast requests are coalesced and flushed at most once per + ``_publish_min_interval``. """ heartbeat_interval = max(1.0, min(self._heartbeat_timeout / 2.0, 5.0)) loop_interval = self._publish_min_interval @@ -243,9 +247,18 @@ class OmniCoordinator: self._check_heartbeat_timeouts() last_heartbeat_check = now - if self._pending_broadcast: - self.publish_instance_list_update() - self._pending_broadcast = False + with self._pending_lock: + has_pending_broadcast = self._pending_broadcast + + if not has_pending_broadcast: + if self._stop_event.wait(timeout=loop_interval): + break + continue + + # Publish outside lock. Clear pending only on success. + if self.publish_instance_list_update(): + with self._pending_lock: + self._pending_broadcast = False if self._stop_event.wait(timeout=loop_interval): break @@ -267,27 +280,23 @@ class OmniCoordinator: info.status = StageStatus.UP promote = True if promote: - self._schedule_broadcast(force=True) + self._schedule_broadcast() return # Check-and-act under single lock to avoid TOCTOU race (duplicate # registration when concurrent events arrive for the same instance). with self._lock: - force_broadcast = False if input_addr not in self._instances: self._add_new_instance_locked(event) - force_broadcast = True else: if event.status == StageStatus.DOWN: self._remove_instance_locked(event) - force_broadcast = True else: self._update_instance_info_locked(event) - # New instances / DOWN events are broadcast immediately; other - # updates (e.g. queue_length changes) are throttled via the - # periodic loop. - self._schedule_broadcast(force=force_broadcast) + # Any non-heartbeat state change that affects the active list + # is coalesced and flushed via the periodic loop. + self._schedule_broadcast() except (KeyError, ValueError, TypeError) as e: logger.warning("Dropping malformed event: %s", e)