From b60a661b3f60890f0ddf24e72f9741e575d6ce90 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Wed, 22 Apr 2026 02:17:28 +0800 Subject: [PATCH] [Feature] Load Balancer - Add LeastQueueLengthBalancer RoundRobinBalancer (#2448) Signed-off-by: NumberWan --- .../omni_coordinator/test_load_balancer.py | 172 ++++++++++++++++++ .../distributed/omni_coordinator/__init__.py | 4 + .../omni_coordinator/load_balancer.py | 78 +++++++- 3 files changed, 248 insertions(+), 6 deletions(-) diff --git a/tests/distributed/omni_coordinator/test_load_balancer.py b/tests/distributed/omni_coordinator/test_load_balancer.py index b2d1f3ee8..8350b33d3 100644 --- a/tests/distributed/omni_coordinator/test_load_balancer.py +++ b/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) diff --git a/vllm_omni/distributed/omni_coordinator/__init__.py b/vllm_omni/distributed/omni_coordinator/__init__.py index cbef920d4..6894e3113 100644 --- a/vllm_omni/distributed/omni_coordinator/__init__.py +++ b/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", ] diff --git a/vllm_omni/distributed/omni_coordinator/load_balancer.py b/vllm_omni/distributed/omni_coordinator/load_balancer.py index 15a079b0a..41b03be16 100644 --- a/vllm_omni/distributed/omni_coordinator/load_balancer.py +++ b/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", ]