[Feature] Load Balancer - Add LeastQueueLengthBalancer RoundRobinBalancer (#2448)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -7,7 +7,9 @@ import pytest
 
 from vllm_omni.distributed.omni_coordinator import (
     InstanceInfo,
+    LeastQueueLengthBalancer,
     RandomBalancer,
+    RoundRobinBalancer,
     StageStatus,
 )
 
@@ -60,3 +62,173 @@ def test_load_balancer_select_returns_valid_index():
 
     assert isinstance(index, int)
     assert 0 <= index < len(instances)
+
+
+def test_round_robin_balancer_cycles_instances():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=2,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=1,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+
+    balancer = RoundRobinBalancer()
+    results = [balancer.select({}, instances) 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"):
+        RoundRobinBalancer().select({}, [])
+
+
+def test_round_robin_balancer_after_large_index_and_shorter_list():
+    """Large start_index % len(instances) then counter wraps with shorter list."""
+    now = time()
+    two = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    balancer = RoundRobinBalancer(start_index=7)
+    assert balancer.select({}, two) == 1  # 7 % 2
+    assert balancer.select({}, two) == 0  # next index wrapped to 0
+
+
+def test_least_queue_length_balancer_picks_min_queue():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=2,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=0,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=5,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+
+    balancer = LeastQueueLengthBalancer()
+    index = balancer.select({}, instances)
+    assert index == 1
+
+
+def test_least_queue_length_balancer_empty_instances_raises():
+    with pytest.raises(ValueError, match="instances must not be empty"):
+        LeastQueueLengthBalancer().select({}, [])
+
+
+def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10002",
+            output_addr="tcp://host:10002-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+        InstanceInfo(
+            input_addr="tcp://host:10003",
+            output_addr="tcp://host:10003-out",
+            stage_id=1,
+            status=StageStatus.UP,
+            queue_length=3,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    balancer = LeastQueueLengthBalancer()
+    mocker.patch(
+        "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
+        return_value=2,
+    )
+    assert balancer.select({}, instances) == 2
+
+
+def test_least_queue_length_balancer_negative_queue_raises():
+    now = time()
+    instances = [
+        InstanceInfo(
+            input_addr="tcp://host:10001",
+            output_addr="tcp://host:10001-out",
+            stage_id=0,
+            status=StageStatus.UP,
+            queue_length=-1,
+            last_heartbeat=now,
+            registered_at=now,
+        ),
+    ]
+    with pytest.raises(ValueError, match="queue_length must be non-negative"):
+        LeastQueueLengthBalancer().select({}, instances)

vllm_omni/distributed/omni_coordinator/__init__.py

@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from .load_balancer import (
+    LeastQueueLengthBalancer,
     LoadBalancer,
     LoadBalancingPolicy,
     RandomBalancer,
+    RoundRobinBalancer,
     Task,
 )
 from .messages import InstanceEvent, InstanceInfo, InstanceList, StageStatus
@@ -24,4 +26,6 @@ __all__ = [
     "LoadBalancer",
     "LoadBalancingPolicy",
     "RandomBalancer",
+    "RoundRobinBalancer",
+    "LeastQueueLengthBalancer",
 ]

vllm_omni/distributed/omni_coordinator/load_balancer.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import random
+import threading
 from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Any, TypedDict
@@ -26,11 +27,16 @@ class Task(TypedDict, total=False):
 class LoadBalancingPolicy(str, Enum):
     """Enumeration for load balancing policies.
 
-    Only ``RANDOM`` is implemented. Additional policies (e.g. round-robin,
-    least-connections) can be added in the future.
+    These policies are used by :class:`LoadBalancer` implementations to route
+    tasks to a subset of available instances.
+
+    TODO(NumberWan): Map enum values to balancer classes when OmniCoordinator
+    integration lands. Tracked in https://github.com/vllm-project/vllm-omni/pull/2448
     """
 
     RANDOM = "random"
+    ROUND_ROBIN = "round-robin"
+    LEAST_QUEUE_LENGTH = "least-queue-length"
 
 
 class LoadBalancer(ABC):
@@ -61,10 +67,10 @@ class LoadBalancer(ABC):
 class RandomBalancer(LoadBalancer):
     """Load balancer that selects an instance uniformly at random.
 
-    This is the initial and only policy supported. It intentionally ignores
-    the task payload and chooses a random index from the provided instance
-    list. More sophisticated policies (e.g. round-robin, least-connections)
-    can be implemented as additional subclasses of :class:`LoadBalancer`.
+    It intentionally ignores the task payload and chooses a random index from
+    the provided instance list. More sophisticated policies (e.g. round-robin,
+    least-queue-length) can be implemented as additional subclasses of
+    :class:`LoadBalancer`.
     """
 
     def select(self, task: Task, instances: list[InstanceInfo]) -> int:  # noqa: ARG002
@@ -74,9 +80,69 @@ class RandomBalancer(LoadBalancer):
         return random.randrange(len(instances))
 
 
+class RoundRobinBalancer(LoadBalancer):
+    """Load balancer that selects instances in a round-robin fashion.
+
+    This implementation keeps a running index modulo ``len(instances)``. It
+    therefore depends on the **order and stable meaning** of the ``instances``
+    list between calls. If the list length or ordering changes, the sequence
+    of picks may skip or repeat entries relative to a fixed set of backends.
+
+    When instance membership changes dynamically, callers should reset routing
+    state—for example by constructing a new ``RoundRobinBalancer`` or resetting
+    ``_next_index``—similar to rebuilding ``itertools.cycle`` after mutating
+    the instance list (as in vLLM's disaggregated proxy examples).
+
+    Concurrency: ``select`` is synchronous and is expected to run on the
+    coordinator asyncio event loop thread without ``await`` inside this
+    method, so a single invocation is not interleaved with another on that
+    thread. A :class:`threading.Lock` still serializes updates to
+    ``_next_index`` for callers that might invoke ``select`` from multiple
+    threads or alongside threaded infrastructure (e.g. ZMQ receive threads).
+    """
+
+    def __init__(self, start_index: int = 0) -> None:
+        self._next_index = start_index
+        self._lock = threading.Lock()
+
+    def select(self, task: Task, instances: list[InstanceInfo]) -> int:  # noqa: ARG002
+        if not instances:
+            raise ValueError("instances must not be empty")
+
+        n = len(instances)
+        with self._lock:
+            idx = self._next_index % n
+            self._next_index = (self._next_index + 1) % n
+        return idx
+
+
+class LeastQueueLengthBalancer(LoadBalancer):
+    """Select the instance with the smallest ``queue_length``.
+
+    If multiple instances share the same minimum queue length, one of them is
+    chosen uniformly at random.
+
+    Raises:
+        ValueError: If any instance has a negative ``queue_length``.
+    """
+
+    def select(self, task: Task, instances: list[InstanceInfo]) -> int:  # noqa: ARG002
+        if not instances:
+            raise ValueError("instances must not be empty")
+
+        queue_lengths = [inst.queue_length for inst in instances]
+        if any(q < 0 for q in queue_lengths):
+            raise ValueError("queue_length must be non-negative for all instances")
+        min_q = min(queue_lengths)
+        candidates = [i for i, q in enumerate(queue_lengths) if q == min_q]
+        return random.choice(candidates)
+
+
 __all__ = [
     "Task",
     "LoadBalancingPolicy",
     "LoadBalancer",
     "RandomBalancer",
+    "RoundRobinBalancer",
+    "LeastQueueLengthBalancer",
 ]