wehub-resource-sync
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
"""Round-robin request router.
|
|
|
|
RoundRobinRouter cycles through the candidate replicas passed in by Serve and
|
|
returns strictly ordered singleton ranks. Each request starts at the current
|
|
round-robin cursor; if that replica cannot fulfill the request, Serve tries the
|
|
next replica in order, wrapping around the candidate list.
|
|
"""
|
|
|
|
import random
|
|
from collections.abc import Sequence
|
|
from typing import List, Optional
|
|
|
|
from ray.serve._private.request_router.common import (
|
|
PendingRequest,
|
|
)
|
|
from ray.serve._private.request_router.replica_wrapper import (
|
|
RunningReplica,
|
|
)
|
|
from ray.serve._private.request_router.request_router import (
|
|
FIFOMixin,
|
|
RequestRouter,
|
|
)
|
|
|
|
|
|
class _RoundRobinReplicaRanks(Sequence[List[RunningReplica]]):
|
|
"""Lazy ordered singleton ranks for a strict round-robin attempt."""
|
|
|
|
def __init__(
|
|
self,
|
|
replicas: List[RunningReplica],
|
|
start_index: int,
|
|
):
|
|
self._replicas = replicas
|
|
self._start_index = start_index
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._replicas)
|
|
|
|
def __getitem__(self, index):
|
|
if isinstance(index, slice):
|
|
return list(self)[index]
|
|
|
|
num_replicas = len(self._replicas)
|
|
if index < 0:
|
|
index += num_replicas
|
|
if index < 0 or index >= num_replicas:
|
|
raise IndexError(index)
|
|
|
|
return [self._replicas[(self._start_index + index) % num_replicas]]
|
|
|
|
def __iter__(self):
|
|
num_replicas = len(self._replicas)
|
|
for offset in range(num_replicas):
|
|
yield [self._replicas[(self._start_index + offset) % num_replicas]]
|
|
|
|
|
|
class RoundRobinRouter(FIFOMixin, RequestRouter):
|
|
"""Routes requests by cycling through candidate replicas.
|
|
|
|
Each call to ``choose_replicas`` advances a shared cursor by one position
|
|
and returns ordered singleton ranks starting from that cursor. This upholds
|
|
strict round-robin ordering while still allowing Serve's existing selector
|
|
to continue to the next replica if the current one is full.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._round_robin_counter = random.randrange(2**31)
|
|
|
|
def initialize_state(self, **kwargs) -> None:
|
|
pass
|
|
|
|
async def choose_replicas(
|
|
self,
|
|
candidate_replicas: List[RunningReplica],
|
|
pending_request: Optional[PendingRequest] = None,
|
|
) -> Sequence[List[RunningReplica]]:
|
|
if not candidate_replicas:
|
|
return []
|
|
|
|
if pending_request is not None:
|
|
# Enable exponential-backoff sleep between outer retry iterations.
|
|
# Without this, the base class tight-loops calling choose_replicas
|
|
# when every replica is at capacity.
|
|
pending_request.routing_context.should_backoff = True
|
|
|
|
index = self._round_robin_counter % len(candidate_replicas)
|
|
self._round_robin_counter += 1
|
|
|
|
return _RoundRobinReplicaRanks(candidate_replicas, index)
|