[large-scale-serving] Integrate OmniCoordinator into stage engine pipeline (#3569)

Signed-off-by: chickeyton <ngton2014@gmail.com>
Signed-off-by: herotai214 <herotai214@gmail.com>
Co-authored-by: herotai214 <herotai214@gmail.com>

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -6,18 +6,18 @@ from time import time
 import pytest
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceInfo,
     LeastQueueLengthBalancer,
     RandomBalancer,
+    ReplicaInfo,
+    ReplicaStatus,
     RoundRobinBalancer,
-    StageStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_load_balancer_select_returns_valid_index():
-    """Verify RandomBalancer.select() returns a valid index for instances."""
+    """Verify RandomBalancer.select() returns a valid index for replicas."""
     # Task structure mirrors async_omni; RandomBalancer ignores task contents.
     task: dict = {
         "request_id": "test",
@@ -26,30 +26,30 @@ def test_load_balancer_select_returns_valid_index():
     }
 
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
@@ -58,38 +58,38 @@ def test_load_balancer_select_returns_valid_index():
 
     balancer = RandomBalancer()
 
-    index = balancer.select(task, instances)
+    index = balancer.select(task, replicas)
 
     assert isinstance(index, int)
-    assert 0 <= index < len(instances)
+    assert 0 <= index < len(replicas)
 
 
-def test_round_robin_balancer_cycles_instances():
+def test_round_robin_balancer_cycles_replicas():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -97,35 +97,35 @@ def test_round_robin_balancer_cycles_instances():
     ]
 
     balancer = RoundRobinBalancer()
-    results = [balancer.select({}, instances) for _ in range(5)]
+    results = [balancer.select({}, replicas) for _ in range(5)]
 
     # Default start_index=0 => 0,1,2,0,1
     assert results == [0, 1, 2, 0, 1]
 
 
-def test_round_robin_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_round_robin_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         RoundRobinBalancer().select({}, [])
 
 
 def test_round_robin_balancer_after_large_index_and_shorter_list():
-    """Large start_index % len(instances) then counter wraps with shorter list."""
+    """Large start_index % len(replicas) then counter wraps with shorter list."""
     now = time()
     two = [
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -138,30 +138,30 @@ def test_round_robin_balancer_after_large_index_and_shorter_list():
 
 def test_least_queue_length_balancer_picks_min_queue():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=5,
             last_heartbeat=now,
             registered_at=now,
@@ -169,41 +169,41 @@ def test_least_queue_length_balancer_picks_min_queue():
     ]
 
     balancer = LeastQueueLengthBalancer()
-    index = balancer.select({}, instances)
+    index = balancer.select({}, replicas)
     assert index == 1
 
 
-def test_least_queue_length_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_least_queue_length_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         LeastQueueLengthBalancer().select({}, [])
 
 
 def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
@@ -214,21 +214,21 @@ def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
         "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
         return_value=2,
     )
-    assert balancer.select({}, instances) == 2
+    assert balancer.select({}, replicas) == 2
 
 
 def test_least_queue_length_balancer_negative_queue_raises():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=-1,
             last_heartbeat=now,
             registered_at=now,
         ),
     ]
     with pytest.raises(ValueError, match="queue_length must be non-negative"):
-        LeastQueueLengthBalancer().select({}, instances)
+        LeastQueueLengthBalancer().select({}, replicas)

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -8,8 +8,8 @@ import pytest
 import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceList,
     OmniCoordClientForHub,
+    ReplicaList,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,8 +32,8 @@ def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> b
     return False
 
 
-def test_hub_client_caches_instance_list_from_pub():
-    """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list()."""
+def test_hub_client_caches_replica_list_from_pub():
+    """Verify OmniCoordClientForHub receives replica list updates from OmniCoordinator and caches for get_replica_list()."""
     ctx, pub, endpoint = _bind_pub()
 
     client = OmniCoordClientForHub(endpoint)
@@ -41,7 +41,7 @@ def test_hub_client_caches_instance_list_from_pub():
     time.sleep(0.2)
 
     now = time.time()
-    instances_payload = [
+    replicas_payload = [
         {
             "input_addr": "tcp://stage:10001",
             "output_addr": "tcp://stage:10001-out",
@@ -71,37 +71,37 @@ def test_hub_client_caches_instance_list_from_pub():
         },
     ]
 
-    payload = {"instances": instances_payload, "timestamp": now}
+    payload = {"replicas": replicas_payload, "timestamp": now}
     pub.send(json.dumps(payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3)
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 3)
 
-    inst_list = client.get_instance_list()
-    assert isinstance(inst_list, InstanceList)
-    assert len(inst_list.instances) == 3
+    rep_list = client.get_replica_list()
+    assert isinstance(rep_list, ReplicaList)
+    assert len(rep_list.replicas) == 3
 
-    for src, inst in zip(instances_payload, inst_list.instances, strict=True):
-        assert inst.input_addr == src["input_addr"]
-        assert inst.output_addr == src["output_addr"]
-        assert inst.stage_id == src["stage_id"]
-        assert inst.status.value == src["status"]
+    for src, rep in zip(replicas_payload, rep_list.replicas, strict=True):
+        assert rep.input_addr == src["input_addr"]
+        assert rep.output_addr == src["output_addr"]
+        assert rep.stage_id == src["stage_id"]
+        assert rep.status.value == src["status"]
 
-    stage0 = client.get_instances_for_stage(0)
-    stage1 = client.get_instances_for_stage(1)
+    stage0 = client.get_replicas_for_stage(0)
+    stage1 = client.get_replicas_for_stage(1)
 
-    assert all(inst.stage_id == 0 for inst in stage0.instances)
-    assert all(inst.stage_id == 1 for inst in stage1.instances)
+    assert all(rep.stage_id == 0 for rep in stage0.replicas)
+    assert all(rep.stage_id == 1 for rep in stage1.replicas)
 
-    # Send an updated list with fewer instances and verify cache refresh.
+    # Send an updated list with fewer replicas and verify cache refresh.
     updated_payload = {
-        "instances": instances_payload[:2],
+        "replicas": replicas_payload[:2],
         "timestamp": now + 1.0,
     }
     pub.send(json.dumps(updated_payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2)
-    updated_list = client.get_instance_list()
-    assert len(updated_list.instances) == 2
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 2)
+    updated_list = client.get_replica_list()
+    assert len(updated_list.replicas) == 2
 
     client.close()
     pub.close(0)

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -9,7 +9,7 @@ import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
-    StageStatus,
+    ReplicaStatus,
 )
 from vllm_omni.distributed.omni_coordinator import (
     omni_coord_client_for_stage as stage_client_module,
@@ -47,7 +47,7 @@ def test_stage_client_auto_register_on_init():
     event = _recv_event(router)
 
     assert event["event_type"] == "update"
-    assert event["status"] == StageStatus.UP.value
+    assert event["status"] == ReplicaStatus.UP.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -70,13 +70,13 @@ def test_stage_client_update_info_sends_correct_event():
     # Discard initial registration event.
     _recv_event(router)
 
-    client.update_info(status=StageStatus.ERROR)
+    client.update_info(status=ReplicaStatus.ERROR)
     client.update_info(queue_length=10)
 
     first = _recv_event(router)
     second = _recv_event(router)
 
-    assert first["status"] == StageStatus.ERROR.value
+    assert first["status"] == ReplicaStatus.ERROR.value
     assert first["stage_id"] == stage_id
     assert first["input_addr"] == input_addr
     assert first["output_addr"] == output_addr
@@ -107,7 +107,7 @@ def test_stage_client_close_sends_down_status():
     client.close()
 
     event = _recv_event(router)
-    assert event["status"] == StageStatus.DOWN.value
+    assert event["status"] == ReplicaStatus.DOWN.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -274,6 +274,7 @@ def test_heartbeat_loop_retries_after_transient_send_failure():
     client._closed = False
     client._heartbeat_interval = 0.0
     client._stop_event = _FakeStopEvent()
+    client._on_heartbeat = None
 
     calls = {"count": 0}
 

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -11,14 +11,14 @@ from vllm.v1.utils import get_engine_client_zmq_addr
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
     OmniCoordinator,
-    StageStatus,
+    ReplicaStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
-    """Receive InstanceList JSON from SUB socket. Returns None on timeout."""
+def _recv_replica_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
+    """Receive ReplicaList JSON from SUB socket. Returns None on timeout."""
     sub.setsockopt(zmq.RCVTIMEO, timeout_ms)
     try:
         data = sub.recv()
@@ -27,16 +27,16 @@ def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
         return None
 
 
-def _wait_for_instance_list(
+def _wait_for_replica_list(
     sub: zmq.Socket,
     expected_count: int,
     timeout: float = 3.0,
 ) -> dict | None:
-    """Wait until received InstanceList with expected_count active instances."""
+    """Wait until received ReplicaList with expected_count active replicas."""
     start = time.time()
     while time.time() - start < timeout:
-        msg = _recv_instance_list(sub, timeout_ms=500)
-        if msg is not None and len(msg.get("instances", [])) == expected_count:
+        msg = _recv_replica_list(sub, timeout_ms=500)
+        if msg is not None and len(msg.get("replicas", [])) == expected_count:
             return msg
     return None
 
@@ -45,7 +45,7 @@ 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)
+        _recv_replica_list(sub, timeout_ms=50)
 
 
 def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
@@ -81,7 +81,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     )
 
     # Wait for initial registration broadcast and clear any queued messages.
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
     _drain_sub_messages(sub)
 
@@ -96,7 +96,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     deadline = time.time() + window_s
     recv_count = 0
     while time.time() < deadline:
-        if _recv_instance_list(sub, timeout_ms=100) is not None:
+        if _recv_replica_list(sub, timeout_ms=100) is not None:
             recv_count += 1
 
     assert recv_count < update_count // 2, (
@@ -110,8 +110,8 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
 
 
 def test_omni_coordinator_registration_broadcast():
-    """Verify that after multiple OmniCoordClientForStage instances register,
-    OmniCoordinator publishes an InstanceList containing all registered instances.
+    """Verify that after multiple OmniCoordClientForStage replicas register,
+    OmniCoordinator publishes a ReplicaList containing all registered replicas.
     """
     router_addr = get_engine_client_zmq_addr(
         local_only=False,
@@ -144,12 +144,12 @@ def test_omni_coordinator_registration_broadcast():
         OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
     ]
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected InstanceList with 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected ReplicaList with 3 replicas"
+    assert len(msg["replicas"]) == 3
     assert isinstance(msg["timestamp"], (int, float))
 
-    input_addrs = {inst["input_addr"] for inst in msg["instances"]}
+    input_addrs = {rep["input_addr"] for rep in msg["replicas"]}
     assert "tcp://stage:10001" in input_addrs
     assert "tcp://stage:10002" in input_addrs
     assert "tcp://stage:10003" in input_addrs
@@ -162,7 +162,7 @@ def test_omni_coordinator_registration_broadcast():
 
 
 def test_omni_coordinator_heartbeat_timeout_handling():
-    """Verify that when a stage instance stops sending heartbeats,
+    """Verify that when a stage replica stops sending heartbeats,
     OmniCoordinator marks it as unhealthy and excludes it from the active list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -201,23 +201,23 @@ def test_omni_coordinator_heartbeat_timeout_handling():
         "output_addr": "tcp://stage:c-out",
         "stage_id": 0,
         "event_type": "update",
-        "status": StageStatus.UP.value,
+        "status": ReplicaStatus.UP.value,
         "queue_length": 0,
     }
     dealer_c.send(json.dumps(reg_event).encode("utf-8"))
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected initial 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected initial 3 replicas"
+    assert len(msg["replicas"]) == 3
 
     # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s).
     time.sleep(8.0)
 
     # Receive the update (C should be ERROR and excluded from active list).
-    msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0)
-    assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout"
-    instances = msg_after_timeout.get("instances", [])
-    input_addrs = {inst["input_addr"] for inst in instances}
+    msg_after_timeout = _wait_for_replica_list(sub, expected_count=2, timeout=5.0)
+    assert msg_after_timeout is not None, "Expected ReplicaList with 2 replicas after timeout"
+    replicas = msg_after_timeout.get("replicas", [])
+    input_addrs = {rep["input_addr"] for rep in replicas}
 
     assert "tcp://stage:a" in input_addrs
     assert "tcp://stage:b" in input_addrs
@@ -232,8 +232,8 @@ def test_omni_coordinator_heartbeat_timeout_handling():
     sub_ctx.term()
 
 
-def test_omni_coordinator_instance_shutdown_handling():
-    """Verify that when a stage instance sends status='down',
+def test_omni_coordinator_replica_shutdown_handling():
+    """Verify that when a stage replica sends status='down',
     OmniCoordinator removes it from the active list and broadcasts an updated list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -261,18 +261,18 @@ def test_omni_coordinator_instance_shutdown_handling():
 
     client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
 
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
-    assert len(msg["instances"]) == 1
-    assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown"
+    assert len(msg["replicas"]) == 1
+    assert msg["replicas"][0]["input_addr"] == "tcp://stage:shutdown"
 
     # Send down status (simulating graceful shutdown).
-    client.update_info(status=StageStatus.DOWN)
+    client.update_info(status=ReplicaStatus.DOWN)
 
-    # Receive updated list (should have 0 active instances).
-    msg = _wait_for_instance_list(sub, expected_count=0)
+    # Receive updated list (should have 0 active replicas).
+    msg = _wait_for_replica_list(sub, expected_count=0)
     assert msg is not None
-    assert len(msg["instances"]) == 0
+    assert len(msg["replicas"]) == 0
 
     client.close()
     coordinator.close()