chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,328 @@
|
||||
# flake8: noqa
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Deque, Dict, List, Optional, Set, Tuple
|
||||
|
||||
import ray
|
||||
from ray import serve
|
||||
from ray._common.utils import get_or_create_event_loop
|
||||
from ray.serve._private.common import DeploymentID, DeploymentTargetInfo
|
||||
from ray.serve._private.constants import SERVE_CONTROLLER_NAME, SERVE_NAMESPACE
|
||||
from ray.serve._private.long_poll import LongPollClient, LongPollNamespace
|
||||
|
||||
logger = logging.getLogger("ray.serve")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapacityQueueStats:
|
||||
"""Statistics for a CapacityQueue."""
|
||||
|
||||
queue_size: int
|
||||
num_waiters: int
|
||||
num_replicas: int
|
||||
total_capacity: int
|
||||
total_in_flight: int
|
||||
total_acquires: int
|
||||
total_releases: int
|
||||
total_timeouts: int
|
||||
max_waiters_seen: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReplicaCapacityInfo:
|
||||
"""Tracks capacity information for a single replica."""
|
||||
|
||||
replica_id: str
|
||||
max_capacity: int
|
||||
in_flight: int = 0
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class CapacityQueue:
|
||||
"""Centralized capacity manager that routes requests to the least-loaded replica.
|
||||
|
||||
A Ray actor that tracks per-replica in-flight counts and hands out capacity
|
||||
tokens. Routers call ``acquire()`` to get a replica_id, route to it, then
|
||||
call ``release()`` when done. If no capacity is available, acquire blocks
|
||||
until ``acquire_timeout_s`` expires (returns None).
|
||||
|
||||
Replica membership is kept in sync via long-poll subscription to the Serve
|
||||
controller. After a crash/restart, CapacityQueue may temporarily over-provision;
|
||||
this self-heals because replicas enforce their own limits and reject excess.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
acquire_timeout_s: float,
|
||||
token_ttl_s: Optional[float],
|
||||
_enable_long_poll: bool = True,
|
||||
):
|
||||
self._acquire_timeout_s = acquire_timeout_s
|
||||
self._token_ttl_s = token_ttl_s
|
||||
|
||||
# Waiters: (asyncio.Future, timestamp) for requests waiting for capacity
|
||||
self._waiters: Deque[Tuple[asyncio.Future, float]] = deque()
|
||||
|
||||
# Registered replicas and their capacity info
|
||||
self._replicas: Dict[str, ReplicaCapacityInfo] = {}
|
||||
|
||||
# Track when each in-flight token was acquired: {replica_id -> [timestamp, ...]}
|
||||
self._in_flight_timestamps: Dict[str, Deque[float]] = {}
|
||||
|
||||
# Statistics
|
||||
self._total_acquires: int = 0
|
||||
self._total_releases: int = 0
|
||||
self._total_timeouts: int = 0
|
||||
self._total_ttl_reclaims: int = 0
|
||||
self._max_waiters_seen: int = 0
|
||||
|
||||
# Subscribe to replica updates from the Serve controller so the queue
|
||||
# automatically registers new replicas and unregisters dead ones.
|
||||
self._long_poll_client: Optional[LongPollClient] = None
|
||||
if _enable_long_poll:
|
||||
deployment_context = serve.get_deployment_actor_context()
|
||||
deployment_id = deployment_context.deployment_id
|
||||
controller_handle = ray.get_actor(
|
||||
SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE
|
||||
)
|
||||
self._long_poll_client = LongPollClient(
|
||||
controller_handle,
|
||||
{
|
||||
(
|
||||
LongPollNamespace.DEPLOYMENT_TARGETS,
|
||||
deployment_id,
|
||||
): self._update_deployment_targets,
|
||||
},
|
||||
call_in_event_loop=get_or_create_event_loop(),
|
||||
client_id=f"{type(self).__name__}:{ray.get_runtime_context().get_actor_id()}",
|
||||
)
|
||||
|
||||
# Start background TTL reaper if token_ttl_s is set.
|
||||
if self._token_ttl_s is not None and self._token_ttl_s > 0:
|
||||
self._ttl_task = get_or_create_event_loop().create_task(
|
||||
self._reap_expired_tokens()
|
||||
)
|
||||
|
||||
async def _reap_expired_tokens(self) -> None:
|
||||
"""Periodically reclaim tokens that have exceeded their TTL."""
|
||||
interval = self._token_ttl_s / 2
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
now = time.time()
|
||||
reclaimed = 0
|
||||
for replica_id, timestamps in list(self._in_flight_timestamps.items()):
|
||||
while timestamps and (now - timestamps[0]) > self._token_ttl_s:
|
||||
timestamps.popleft()
|
||||
if replica_id in self._replicas:
|
||||
self._decrement_in_flight(replica_id)
|
||||
reclaimed += 1
|
||||
if reclaimed > 0:
|
||||
self._total_ttl_reclaims += reclaimed
|
||||
logger.info(f"TTL reaper reclaimed {reclaimed} expired token(s).")
|
||||
self._fulfill_waiters()
|
||||
|
||||
def _decrement_in_flight(self, replica_id: str) -> None:
|
||||
"""Decrement in-flight count for a replica, clamping at zero."""
|
||||
info = self._replicas[replica_id]
|
||||
info.in_flight = max(0, info.in_flight - 1)
|
||||
|
||||
def register_replica(self, replica_id: str, capacity: int) -> None:
|
||||
"""Register a replica with its capacity.
|
||||
|
||||
If the replica is already registered, this is a no-op.
|
||||
"""
|
||||
if replica_id in self._replicas:
|
||||
return
|
||||
self._replicas[replica_id] = ReplicaCapacityInfo(
|
||||
replica_id=replica_id,
|
||||
max_capacity=capacity,
|
||||
in_flight=0,
|
||||
)
|
||||
logger.debug(
|
||||
f"Registered replica {replica_id} with capacity {capacity}. "
|
||||
f"Total replicas: {len(self._replicas)}."
|
||||
)
|
||||
self._fulfill_waiters()
|
||||
|
||||
def unregister_replica(self, replica_id: str) -> None:
|
||||
"""Unregister a replica and remove its capacity."""
|
||||
if replica_id not in self._replicas:
|
||||
return
|
||||
info = self._replicas.pop(replica_id)
|
||||
self._in_flight_timestamps.pop(replica_id, None)
|
||||
logger.debug(
|
||||
f"Unregistered replica {replica_id} "
|
||||
f"(had {info.in_flight}/{info.max_capacity} in-flight). "
|
||||
f"Total replicas: {len(self._replicas)}."
|
||||
)
|
||||
|
||||
def _update_deployment_targets(
|
||||
self, deployment_target_info: DeploymentTargetInfo
|
||||
) -> None:
|
||||
"""Handle deployment target updates from the controller (via long poll).
|
||||
|
||||
Automatically registers new replicas and unregisters removed ones so the
|
||||
queue always reflects the set of live replicas.
|
||||
"""
|
||||
running_replicas = deployment_target_info.running_replicas
|
||||
current_ids: Set[str] = {r.replica_id.unique_id for r in running_replicas}
|
||||
registered_ids: Set[str] = set(self._replicas.keys())
|
||||
|
||||
for rid in registered_ids - current_ids:
|
||||
self.unregister_replica(rid)
|
||||
|
||||
replica_by_uid = {r.replica_id.unique_id: r for r in running_replicas}
|
||||
for uid in current_ids - registered_ids:
|
||||
self.register_replica(uid, replica_by_uid[uid].max_ongoing_requests)
|
||||
|
||||
def _get_least_loaded_replica(self) -> Optional[str]:
|
||||
"""Find the replica with fewest in-flight requests that has capacity.
|
||||
|
||||
Ties are broken randomly.
|
||||
"""
|
||||
best_replicas: List[str] = []
|
||||
best_in_flight: float = float("inf")
|
||||
|
||||
for replica_id, info in self._replicas.items():
|
||||
available = info.max_capacity - info.in_flight
|
||||
if available > 0:
|
||||
if info.in_flight < best_in_flight:
|
||||
best_replicas = [replica_id]
|
||||
best_in_flight = info.in_flight
|
||||
elif info.in_flight == best_in_flight:
|
||||
best_replicas.append(replica_id)
|
||||
|
||||
if not best_replicas:
|
||||
return None
|
||||
return random.choice(best_replicas)
|
||||
|
||||
def _get_total_available_capacity(self) -> int:
|
||||
"""Get total available capacity across all replicas."""
|
||||
return sum(
|
||||
max(0, info.max_capacity - info.in_flight)
|
||||
for info in self._replicas.values()
|
||||
)
|
||||
|
||||
def _has_available_capacity(self) -> bool:
|
||||
"""Check if any replica has available capacity."""
|
||||
return any(
|
||||
info.max_capacity > info.in_flight for info in self._replicas.values()
|
||||
)
|
||||
|
||||
async def acquire(self, timeout_s: Optional[float] = None) -> Optional[str]:
|
||||
"""Acquire a capacity token from the least loaded replica.
|
||||
|
||||
Returns the replica_id to route to, or None on timeout.
|
||||
Caller MUST call release() when the request completes.
|
||||
"""
|
||||
self._total_acquires += 1
|
||||
timeout = timeout_s if timeout_s is not None else self._acquire_timeout_s
|
||||
|
||||
# Fast path: find least loaded replica
|
||||
replica_id = self._get_least_loaded_replica()
|
||||
if replica_id is not None:
|
||||
self._replicas[replica_id].in_flight += 1
|
||||
self._record_acquire_timestamp(replica_id)
|
||||
return replica_id
|
||||
|
||||
# Slow path: wait for capacity
|
||||
loop = asyncio.get_running_loop()
|
||||
future: asyncio.Future[str] = loop.create_future()
|
||||
waiter_entry = (future, time.time())
|
||||
self._waiters.append(waiter_entry)
|
||||
self._max_waiters_seen = max(self._max_waiters_seen, len(self._waiters))
|
||||
|
||||
try:
|
||||
replica_id = await asyncio.wait_for(future, timeout=timeout)
|
||||
self._record_acquire_timestamp(replica_id)
|
||||
return replica_id
|
||||
except asyncio.TimeoutError:
|
||||
self._total_timeouts += 1
|
||||
self._waiters = deque((f, t) for f, t in self._waiters if f is not future)
|
||||
return None
|
||||
except asyncio.CancelledError:
|
||||
self._waiters = deque((f, t) for f, t in self._waiters if f is not future)
|
||||
raise
|
||||
|
||||
def release(self, replica_id: str) -> None:
|
||||
"""Release a capacity token when a request completes.
|
||||
|
||||
MUST be called after acquire() when the request finishes.
|
||||
"""
|
||||
self._total_releases += 1
|
||||
|
||||
if replica_id not in self._replicas:
|
||||
logger.warning(f"Release called for unknown replica {replica_id}.")
|
||||
return
|
||||
|
||||
self._decrement_in_flight(replica_id)
|
||||
# Remove the oldest timestamp for this replica.
|
||||
if replica_id in self._in_flight_timestamps:
|
||||
ts = self._in_flight_timestamps[replica_id]
|
||||
if ts:
|
||||
ts.popleft()
|
||||
self._fulfill_waiters()
|
||||
|
||||
def _record_acquire_timestamp(self, replica_id: str) -> None:
|
||||
"""Record the timestamp of an acquired token for TTL tracking."""
|
||||
if self._token_ttl_s is not None:
|
||||
if replica_id not in self._in_flight_timestamps:
|
||||
self._in_flight_timestamps[replica_id] = deque()
|
||||
self._in_flight_timestamps[replica_id].append(time.time())
|
||||
|
||||
def _fulfill_waiters(self) -> None:
|
||||
"""Try to fulfill waiting requests with available capacity."""
|
||||
while self._waiters and self._has_available_capacity():
|
||||
future, _ = self._waiters.popleft()
|
||||
if future.done():
|
||||
continue
|
||||
|
||||
replica_id = self._get_least_loaded_replica()
|
||||
if replica_id is None:
|
||||
self._waiters.appendleft((future, time.time()))
|
||||
break
|
||||
|
||||
self._replicas[replica_id].in_flight += 1
|
||||
self._record_acquire_timestamp(replica_id)
|
||||
future.set_result(replica_id)
|
||||
|
||||
def get_stats(self) -> CapacityQueueStats:
|
||||
"""Get queue statistics."""
|
||||
total_capacity = sum(r.max_capacity for r in self._replicas.values())
|
||||
total_in_flight = sum(r.in_flight for r in self._replicas.values())
|
||||
|
||||
return CapacityQueueStats(
|
||||
queue_size=self._get_total_available_capacity(),
|
||||
num_waiters=len(self._waiters),
|
||||
num_replicas=len(self._replicas),
|
||||
total_capacity=total_capacity,
|
||||
total_in_flight=total_in_flight,
|
||||
total_acquires=self._total_acquires,
|
||||
total_releases=self._total_releases,
|
||||
total_timeouts=self._total_timeouts,
|
||||
max_waiters_seen=self._max_waiters_seen,
|
||||
)
|
||||
|
||||
def get_queue_length(self) -> int:
|
||||
"""Get total available capacity."""
|
||||
return self._get_total_available_capacity()
|
||||
|
||||
def get_num_waiters(self) -> int:
|
||||
"""Get number of requests waiting for capacity."""
|
||||
return len(self._waiters)
|
||||
|
||||
def get_registered_replicas(self) -> List[str]:
|
||||
"""Get list of registered replica IDs."""
|
||||
return list(self._replicas.keys())
|
||||
|
||||
def get_replica_in_flight(self) -> Dict[str, Tuple[int, int]]:
|
||||
"""Get per-replica (in_flight, max_capacity) for convergence checks."""
|
||||
return {
|
||||
rid: (info.in_flight, info.max_capacity)
|
||||
for rid, info in self._replicas.items()
|
||||
}
|
||||
@@ -0,0 +1,303 @@
|
||||
# flake8: noqa
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.exceptions import RayActorError
|
||||
from ray.serve._private.common import ReplicaID
|
||||
from ray.serve._private.constants import (
|
||||
SERVE_DEPLOYMENT_ACTOR_PREFIX,
|
||||
SERVE_NAMESPACE,
|
||||
)
|
||||
from ray.serve._private.request_router.pow_2_router import (
|
||||
PowerOfTwoChoicesRequestRouter,
|
||||
)
|
||||
from ray.serve.request_router import (
|
||||
LocalityMixin,
|
||||
MultiplexMixin,
|
||||
PendingRequest,
|
||||
ReplicaResult,
|
||||
RequestRouter,
|
||||
RunningReplica,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("ray.serve")
|
||||
|
||||
|
||||
DEFAULT_MAX_FAULT_RETRIES = 3
|
||||
|
||||
|
||||
class CapacityQueueRouter(LocalityMixin, MultiplexMixin, RequestRouter):
|
||||
"""Custom request router that uses a CapacityQueue deployment actor.
|
||||
|
||||
Acquires capacity tokens from a centralized CapacityQueue before routing
|
||||
requests. Each token guarantees the target replica has capacity.
|
||||
|
||||
If the capacity queue is unavailable (dead, not yet discovered, or times out),
|
||||
the router retries with exponential backoff up to MAX_FAULT_RETRIES times,
|
||||
then falls back to the standard power-of-two-choices algorithm for that
|
||||
request.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._capacity_queue = None
|
||||
self._capacity_queue_actor_name: Optional[str] = None
|
||||
self._capacity_queue_full_name: Optional[str] = None
|
||||
|
||||
# Tokens acquired but not yet accepted by the replica (pending the
|
||||
# rejection protocol handshake). On rejection the token is
|
||||
# intentionally leaked so the CQ's in_flight stays elevated, teaching
|
||||
# it that the replica is busy.
|
||||
self._pending_tokens: Dict[str, str] = {}
|
||||
|
||||
# Tokens whose replica accepted the request. Released when the
|
||||
# request completes or is cancelled.
|
||||
self._acquired_tokens: Dict[str, str] = {}
|
||||
|
||||
# Reverse index: unique_id string -> ReplicaID for O(1) lookup.
|
||||
# Rebuilt by update_replicas whenever the replica set changes.
|
||||
self._uid_to_replica_id: Dict[str, ReplicaID] = {}
|
||||
|
||||
def initialize_state(self, **kwargs):
|
||||
if "capacity_queue_actor_name" not in kwargs:
|
||||
raise ValueError(
|
||||
"CapacityQueueRouter requires 'capacity_queue_actor_name' in "
|
||||
"request_router_kwargs."
|
||||
)
|
||||
self._capacity_queue_actor_name = kwargs["capacity_queue_actor_name"]
|
||||
self._max_fault_retries = kwargs.get(
|
||||
"max_fault_retries", DEFAULT_MAX_FAULT_RETRIES
|
||||
)
|
||||
|
||||
def _try_discover_capacity_queue(self) -> bool:
|
||||
"""Try to discover the CapacityQueue deployment actor.
|
||||
|
||||
TODO (jeffreywang): Find a better way to discover the CapacityQueue actor.
|
||||
|
||||
Returns True if discovered, False otherwise (does not raise).
|
||||
"""
|
||||
if self._capacity_queue_full_name is not None:
|
||||
try:
|
||||
self._capacity_queue = ray.get_actor(
|
||||
self._capacity_queue_full_name, namespace=SERVE_NAMESPACE
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
prefix = (
|
||||
f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}"
|
||||
f"{self._deployment_id.app_name}::"
|
||||
f"{self._deployment_id.name}::"
|
||||
)
|
||||
suffix = f"::{self._capacity_queue_actor_name}"
|
||||
|
||||
try:
|
||||
actors = ray.util.list_named_actors(all_namespaces=True)
|
||||
for actor_info in actors:
|
||||
if actor_info["namespace"] == SERVE_NAMESPACE:
|
||||
name = actor_info["name"]
|
||||
if name.startswith(prefix) and name.endswith(suffix):
|
||||
self._capacity_queue = ray.get_actor(
|
||||
name, namespace=SERVE_NAMESPACE
|
||||
)
|
||||
self._capacity_queue_full_name = name
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
def update_replicas(self, replicas: List[RunningReplica]):
|
||||
super().update_replicas(replicas)
|
||||
self._uid_to_replica_id = {rid.unique_id: rid for rid in self._replicas}
|
||||
|
||||
def _safe_release(self, replica_id: str) -> None:
|
||||
"""Release a token, handling queue actor death."""
|
||||
if self._capacity_queue is None:
|
||||
return
|
||||
try:
|
||||
self._capacity_queue.release.remote(replica_id)
|
||||
except RayActorError:
|
||||
self._capacity_queue = None
|
||||
|
||||
async def _choose_replica_for_request(
|
||||
self, pending_request: PendingRequest, *, is_retry: bool = False
|
||||
) -> RunningReplica:
|
||||
"""Choose a replica by acquiring a token from the capacity queue.
|
||||
|
||||
This overrides the base-class method to bypass the routing-task
|
||||
machinery (choose_replicas -> probe -> retry loop). The CQ acquire
|
||||
runs directly in the caller's coroutine so the event loop stays
|
||||
responsive.
|
||||
|
||||
On transient CQ faults (actor death, discovery failure), retries with
|
||||
exponential backoff up to MAX_FAULT_RETRIES times, then falls back to
|
||||
power-of-two-choices for this request.
|
||||
"""
|
||||
internal_request_id = pending_request.metadata.internal_request_id
|
||||
|
||||
# On retry (replica rejected), drop the old pending token WITHOUT
|
||||
# releasing it. The leaked token keeps in_flight elevated on the CQ,
|
||||
# teaching it that the replica is busy. The TTL reaper eventually
|
||||
# reclaims it.
|
||||
if is_retry:
|
||||
self._pending_tokens.pop(internal_request_id, None)
|
||||
|
||||
# Wait for at least one replica.
|
||||
while len(self._replicas) == 0:
|
||||
self._replicas_updated_event.clear()
|
||||
await self._replicas_updated_event.wait()
|
||||
|
||||
fault_attempt = 0
|
||||
capacity_depleted_attempt = 0
|
||||
acquired_replica_id = None
|
||||
try:
|
||||
while True:
|
||||
# Too many consecutive faults — fall back to Pow2.
|
||||
if fault_attempt >= self._max_fault_retries:
|
||||
return await super()._choose_replica_for_request(
|
||||
pending_request, is_retry=is_retry
|
||||
)
|
||||
|
||||
# Discover the CQ actor if we don't have a handle.
|
||||
if self._capacity_queue is None:
|
||||
if not self._try_discover_capacity_queue():
|
||||
await self._backoff(fault_attempt)
|
||||
fault_attempt += 1
|
||||
continue
|
||||
|
||||
# Acquire a token from the CQ.
|
||||
acquire_ref = None
|
||||
try:
|
||||
acquire_ref = self._capacity_queue.acquire.remote()
|
||||
acquired_replica_id = await acquire_ref
|
||||
except asyncio.CancelledError:
|
||||
if acquire_ref is not None:
|
||||
ray.cancel(acquire_ref)
|
||||
raise
|
||||
except RayActorError:
|
||||
logger.warning(
|
||||
"Lost connection to CapacityQueue actor: "
|
||||
f"{self._capacity_queue_full_name}."
|
||||
)
|
||||
self._capacity_queue = None
|
||||
await self._backoff(fault_attempt)
|
||||
fault_attempt += 1
|
||||
continue
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Unexpected error acquiring from CapacityQueue: "
|
||||
f"{self._capacity_queue_full_name}.",
|
||||
exc_info=True,
|
||||
)
|
||||
await self._backoff(fault_attempt)
|
||||
fault_attempt += 1
|
||||
continue
|
||||
|
||||
# CQ returned None — capacity exhausted. This is normal
|
||||
# backpressure, not a fault, so it doesn't increment
|
||||
# fault_attempt or trigger Pow2 fallback. Apply backoff
|
||||
# to avoid a hot loop when acquire_timeout_s is low.
|
||||
if acquired_replica_id is None:
|
||||
await self._backoff(capacity_depleted_attempt)
|
||||
capacity_depleted_attempt += 1
|
||||
continue
|
||||
|
||||
# O(1) lookup via reverse index.
|
||||
rid = self._uid_to_replica_id.get(acquired_replica_id)
|
||||
replica = self._replicas.get(rid) if rid is not None else None
|
||||
if replica is not None:
|
||||
# Store as pending until on_request_routed confirms
|
||||
# acceptance. If the replica rejects, the token is
|
||||
# intentionally NOT released.
|
||||
self._pending_tokens[internal_request_id] = acquired_replica_id
|
||||
acquired_replica_id = None
|
||||
fault_attempt = 0
|
||||
capacity_depleted_attempt = 0
|
||||
return replica
|
||||
|
||||
# Replica not found — release token and wait for replica update.
|
||||
# This happens during replica transitions when the CQ knows about
|
||||
# a new replica before the local router does. Without waiting,
|
||||
# the CQ fast path creates a hot loop returning the same unknown
|
||||
# replica.
|
||||
logger.debug(
|
||||
f"CapacityQueue returned unknown replica {acquired_replica_id}; "
|
||||
"releasing token and waiting for replica update."
|
||||
)
|
||||
self._safe_release(acquired_replica_id)
|
||||
acquired_replica_id = None
|
||||
|
||||
# Wait for the router's local replica state to be updated via long polling.
|
||||
# This ensures we don't create a hot loop when the queue is ahead of us.
|
||||
self._replicas_updated_event.clear()
|
||||
await self._replicas_updated_event.wait()
|
||||
except asyncio.CancelledError:
|
||||
# If cancelled after acquiring a token but before storing it in
|
||||
# _pending_tokens, release it to avoid leaking capacity.
|
||||
if acquired_replica_id is not None:
|
||||
self._safe_release(acquired_replica_id)
|
||||
raise
|
||||
|
||||
async def choose_replicas(
|
||||
self,
|
||||
candidate_replicas: List[RunningReplica],
|
||||
pending_request: Optional[PendingRequest] = None,
|
||||
) -> List[List[RunningReplica]]:
|
||||
"""Pow2 candidate selection — no CQ calls.
|
||||
|
||||
Only reached when _choose_replica_for_request falls back to the
|
||||
base class after MAX_FAULT_RETRIES consecutive CQ faults.
|
||||
"""
|
||||
return await PowerOfTwoChoicesRequestRouter.choose_replicas(
|
||||
self,
|
||||
candidate_replicas=candidate_replicas,
|
||||
pending_request=pending_request,
|
||||
)
|
||||
|
||||
def on_request_routed(
|
||||
self,
|
||||
pending_request: PendingRequest,
|
||||
replica_id: ReplicaID,
|
||||
result: ReplicaResult,
|
||||
):
|
||||
"""Promote a pending token to acquired once the replica accepts.
|
||||
|
||||
Only called when the rejection protocol is active and the replica
|
||||
accepts. On rejection this is NOT called, so the token stays in
|
||||
_pending_tokens and is intentionally leaked on retry to teach the
|
||||
CQ that the replica is busy.
|
||||
"""
|
||||
rid = pending_request.metadata.internal_request_id
|
||||
token = self._pending_tokens.pop(rid, None)
|
||||
if token is not None:
|
||||
self._acquired_tokens[rid] = token
|
||||
|
||||
def on_request_completed(
|
||||
self,
|
||||
replica_id: ReplicaID,
|
||||
internal_request_id: str,
|
||||
):
|
||||
"""
|
||||
Release the capacity token when a request completes.
|
||||
"""
|
||||
token = self._acquired_tokens.pop(internal_request_id, None)
|
||||
if token is None:
|
||||
token = self._pending_tokens.pop(internal_request_id, None)
|
||||
if token is not None:
|
||||
self._safe_release(token)
|
||||
|
||||
def on_request_cancelled(
|
||||
self,
|
||||
internal_request_id: str,
|
||||
):
|
||||
"""Release the capacity token when a request is cancelled."""
|
||||
token = self._acquired_tokens.pop(internal_request_id, None)
|
||||
if token is None:
|
||||
token = self._pending_tokens.pop(internal_request_id, None)
|
||||
if token is not None:
|
||||
self._safe_release(token)
|
||||
@@ -0,0 +1,356 @@
|
||||
"""Consistent-hash request router to honor session stickiness.
|
||||
|
||||
ConsistentHashRouter implements a consistent-hash ring with virtual nodes
|
||||
to map session IDs to replicas. When the assigned replica rejects the request
|
||||
due to backpressure, the router falls back to the next replica in the ring,
|
||||
with at most DEFAULT_NUM_FALLBACK_REPLICAS replicas before backing off the request.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import bisect
|
||||
import logging
|
||||
import time
|
||||
from typing import FrozenSet, List, Optional
|
||||
|
||||
import mmh3
|
||||
|
||||
from ray.serve._private.common import ReplicaID
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME, SERVE_SESSION_ID
|
||||
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 RequestRouter
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
# Fixed seed for consistent hashing. A seed change invalidates every running
|
||||
# affinity map in the fleet and would silently remap every session the instant
|
||||
# a new router picks up the controller broadcast.
|
||||
CONSISTENT_HASH_SEED = 0xF9B4CA77
|
||||
|
||||
DEFAULT_NUM_VIRTUAL_NODES = 100
|
||||
DEFAULT_NUM_FALLBACK_REPLICAS = 2
|
||||
|
||||
|
||||
def _hash_bytes(key: bytes) -> int:
|
||||
"""Hash ``key`` with MurmurHash3 and return the low 64 bits."""
|
||||
return mmh3.hash64(key, seed=CONSISTENT_HASH_SEED, signed=False)[0]
|
||||
|
||||
|
||||
class ConsistentHashRouter(RequestRouter):
|
||||
"""Routes each request to a replica chosen by consistent hashing over
|
||||
``session_id`` (falling back to ``internal_request_id``).
|
||||
|
||||
Every ``ConsistentHashRouter`` instance owns a private ring built from
|
||||
the controller's long-poll replica broadcast. Routers on different
|
||||
processes converge because they consume the same broadcast and use the
|
||||
same hash + vnode count; router restarts are free because a fresh
|
||||
router rebuilds the identical ring.
|
||||
|
||||
Consistent hashing is a deterministic, affinity-first strategy: the chosen
|
||||
replica is a pure function of session ID and ring state. Mixing in queue-depth,
|
||||
locality, or multiplexed model signals breaks determinism and therefore breaks
|
||||
affinity.
|
||||
|
||||
Kwargs (passed via ``RequestRouterConfig.request_router_kwargs``):
|
||||
|
||||
- ``num_virtual_nodes`` (int, default ``100``): vnodes per replica on
|
||||
the hash ring. Higher values spread sessions more evenly across
|
||||
replicas at the cost of a larger ring.
|
||||
- ``num_fallback_replicas`` (int, default ``2``): number of clockwise
|
||||
successor replicas tried after the primary if it rejects (e.g.
|
||||
backpressure). Set to ``0`` for strict affinity (no fallback).
|
||||
|
||||
Session affinity:
|
||||
The routing key is ``RequestMetadata.session_id`` — populated upstream
|
||||
from the HTTP header named by ``SERVE_SESSION_ID`` (set via env var
|
||||
``RAY_SERVE_SESSION_ID_HEADER_KEY``; default ``x-session-id``).
|
||||
Requests without that header fall back to ``internal_request_id``
|
||||
(a fresh UUID per request) and therefore distribute uniformly across
|
||||
replicas — same code path, no affinity. See
|
||||
``ray.serve._private.constants.SERVE_SESSION_ID``.
|
||||
|
||||
Usage (Ray Serve LLM):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray.serve.config import RequestRouterConfig
|
||||
from ray.serve.llm import LLMConfig, ModelLoadingConfig
|
||||
|
||||
LLMConfig(
|
||||
model_loading_config=ModelLoadingConfig(
|
||||
model_id="Qwen/Qwen3-0.6B-FP8",
|
||||
model_source="Qwen/Qwen3-0.6B-FP8",
|
||||
),
|
||||
deployment_config=dict(
|
||||
request_router_config=RequestRouterConfig(
|
||||
request_router_class=(
|
||||
"ray.serve.experimental.consistent_hash_router."
|
||||
"ConsistentHashRouter"
|
||||
),
|
||||
request_router_kwargs={
|
||||
"num_virtual_nodes": 100,
|
||||
"num_fallback_replicas": 2,
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
The session header itself is set globally on the cluster:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Clients sending sessions on x-correlation-id (e.g. aiperf):
|
||||
export RAY_SERVE_SESSION_ID_HEADER_KEY=x-correlation-id
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# Ring state
|
||||
self._ring_hashes: List[int] = []
|
||||
self._ring_replicas: List[ReplicaID] = []
|
||||
|
||||
# Change-detection cache: Compared against the incoming set on
|
||||
# each update_replicas tick.
|
||||
self._replica_set_snapshot: FrozenSet[ReplicaID] = frozenset()
|
||||
|
||||
# Tunables populated by ``initialize_state``.
|
||||
self._num_virtual_nodes: int = DEFAULT_NUM_VIRTUAL_NODES
|
||||
self._num_fallback_replicas: int = DEFAULT_NUM_FALLBACK_REPLICAS
|
||||
|
||||
def initialize_state(self, **kwargs) -> None:
|
||||
num_virtual_nodes = kwargs.get("num_virtual_nodes", DEFAULT_NUM_VIRTUAL_NODES)
|
||||
num_fallback_replicas = kwargs.get(
|
||||
"num_fallback_replicas", DEFAULT_NUM_FALLBACK_REPLICAS
|
||||
)
|
||||
|
||||
if not isinstance(num_virtual_nodes, int) or num_virtual_nodes <= 0:
|
||||
raise ValueError(
|
||||
f"num_virtual_nodes must be a positive int, got {num_virtual_nodes!r}."
|
||||
)
|
||||
if not isinstance(num_fallback_replicas, int) or num_fallback_replicas < 0:
|
||||
raise ValueError(
|
||||
"num_fallback_replicas must be a non-negative int, got "
|
||||
f"{num_fallback_replicas!r}."
|
||||
)
|
||||
|
||||
self._num_virtual_nodes = num_virtual_nodes
|
||||
self._num_fallback_replicas = num_fallback_replicas
|
||||
|
||||
logger.info(
|
||||
f"ConsistentHashRouter initialized for {self._deployment_id}: "
|
||||
f"num_virtual_nodes={self._num_virtual_nodes}, "
|
||||
f"num_fallback_replicas={self._num_fallback_replicas}, "
|
||||
f"session_id_header={SERVE_SESSION_ID!r} "
|
||||
"(set via RAY_SERVE_SESSION_ID_HEADER_KEY)."
|
||||
)
|
||||
|
||||
def update_replicas(self, replicas: List[RunningReplica]) -> None:
|
||||
"""
|
||||
Update available replicas and rebuild the ring when the set changes.
|
||||
"""
|
||||
super().update_replicas(replicas)
|
||||
|
||||
new_snapshot = frozenset(self._replica_id_set)
|
||||
if new_snapshot == self._replica_set_snapshot:
|
||||
return
|
||||
|
||||
self._rebuild_ring(new_snapshot)
|
||||
|
||||
def _rebuild_ring(self, new_snapshot: FrozenSet[ReplicaID]) -> None:
|
||||
"""
|
||||
Rebuild ring state from the current replica set.
|
||||
|
||||
Every vnode is hashed from f"{unique_id}:{vnode_index}". Using
|
||||
the replica's unique_id rather than the full ReplicaID keeps the
|
||||
key stable across controller restarts.
|
||||
"""
|
||||
new_hashes: List[int] = []
|
||||
new_replicas: List[ReplicaID] = []
|
||||
|
||||
for replica_id in new_snapshot:
|
||||
for i in range(self._num_virtual_nodes):
|
||||
key = f"{replica_id.unique_id}:{i}".encode()
|
||||
new_hashes.append(_hash_bytes(key))
|
||||
new_replicas.append(replica_id)
|
||||
|
||||
if new_hashes:
|
||||
paired = sorted(zip(new_hashes, new_replicas), key=lambda p: p[0])
|
||||
new_hashes = [p[0] for p in paired]
|
||||
new_replicas = [p[1] for p in paired]
|
||||
|
||||
self._ring_hashes = new_hashes
|
||||
self._ring_replicas = new_replicas
|
||||
self._replica_set_snapshot = new_snapshot
|
||||
|
||||
logger.info(
|
||||
f"Rebuilt consistent-hash ring for {self._deployment_id}: "
|
||||
f"{len(new_snapshot)} replicas, "
|
||||
f"{len(new_hashes)} vnodes."
|
||||
)
|
||||
|
||||
def _routing_key(self, pending_request: Optional[PendingRequest]) -> Optional[str]:
|
||||
"""Prefer session_id; fall back to internal_request_id.
|
||||
|
||||
Hashing on internal_request_id (a fresh UUID per request)
|
||||
spreads session-less traffic uniformly through the same code
|
||||
path, avoiding mixing with another routing strategy (e.g. pow-2).
|
||||
"""
|
||||
if pending_request is None:
|
||||
return None
|
||||
metadata = pending_request.metadata
|
||||
if metadata.session_id:
|
||||
return metadata.session_id
|
||||
# internal_request_id is always populated by the proxy or by
|
||||
# get_request_metadata's fallback to generate_request_id().
|
||||
return metadata.internal_request_id
|
||||
|
||||
def _lookup_ranked_replicas(self, routing_key: str) -> List[ReplicaID]:
|
||||
"""
|
||||
Walk the ring and return the primary plus up to K distinct
|
||||
clockwise successor replicas.
|
||||
"""
|
||||
if not self._ring_hashes:
|
||||
return []
|
||||
|
||||
key_hash = _hash_bytes(routing_key.encode())
|
||||
start_idx = bisect.bisect_right(self._ring_hashes, key_hash)
|
||||
if start_idx == len(self._ring_hashes):
|
||||
# Wrapped past the last vnode
|
||||
start_idx = 0
|
||||
|
||||
# Walk clockwise from the owner, collecting the primary and up to K
|
||||
# fallback replicas, ensuring that each replica only appears once
|
||||
# even if it owns multiple virtual nodes.
|
||||
ranked: List[ReplicaID] = []
|
||||
seen: set = set()
|
||||
max_ranked = 1 + self._num_fallback_replicas
|
||||
n = len(self._ring_hashes)
|
||||
|
||||
num_replicas = len(self._replica_set_snapshot)
|
||||
for offset in range(n):
|
||||
replica_id = self._ring_replicas[(start_idx + offset) % n]
|
||||
if replica_id in seen:
|
||||
continue
|
||||
seen.add(replica_id)
|
||||
ranked.append(replica_id)
|
||||
if len(ranked) == max_ranked or len(seen) == num_replicas:
|
||||
break
|
||||
|
||||
return ranked
|
||||
|
||||
async def choose_replicas(
|
||||
self,
|
||||
candidate_replicas: List[RunningReplica],
|
||||
pending_request: Optional[PendingRequest] = None,
|
||||
) -> List[List[RunningReplica]]:
|
||||
"""Return ranked candidates for the outer retry loop.
|
||||
|
||||
Each rank is a single-element list containing one replica, so the
|
||||
retry loop walks primary -> fallback_1 -> fallback_2 in strict
|
||||
hash order with backoff.
|
||||
"""
|
||||
if not candidate_replicas:
|
||||
return []
|
||||
|
||||
if pending_request is not None:
|
||||
# Enable exponential-backoff sleep between outer retry iterations.
|
||||
pending_request.routing_context.should_backoff = True
|
||||
|
||||
routing_key = self._routing_key(pending_request)
|
||||
if routing_key is None:
|
||||
# Returning [] here would livelock the framework's retry loop.
|
||||
# Yield candidates so the task can loop back for real requests.
|
||||
return [candidate_replicas]
|
||||
|
||||
ranked_ids = self._lookup_ranked_replicas(routing_key)
|
||||
|
||||
ranks: List[List[RunningReplica]] = []
|
||||
for replica_id in ranked_ids:
|
||||
replica = self._replicas.get(replica_id)
|
||||
if replica is not None:
|
||||
ranks.append([replica])
|
||||
|
||||
return ranks
|
||||
|
||||
async def _fulfill_pending_requests(self):
|
||||
"""Overrides the base loop with two consistent-hash-specific
|
||||
invariants:
|
||||
|
||||
1. Exit immediately when there is nothing to route (rather than
|
||||
routing for a None pending request, which the base class does
|
||||
assuming a FIFO fallback that ConsistentHashRouter does not
|
||||
provide).
|
||||
2. The routing task that owns a popped pending request must keep
|
||||
retrying the inner loop until the request is fulfilled, even if
|
||||
there are now more routing tasks than the target. Without this,
|
||||
concurrent same-session bursts would orphan their pending
|
||||
requests under load.
|
||||
"""
|
||||
try:
|
||||
while len(self._routing_tasks) <= self.target_num_routing_tasks:
|
||||
start_time = time.time()
|
||||
backoff_index = 0
|
||||
pending_request = self._get_next_pending_request_to_route()
|
||||
if pending_request is None:
|
||||
# No work for this task; let it exit. Future calls to
|
||||
# _maybe_start_routing_tasks (on new requests or replica
|
||||
# updates) will spawn fresh tasks. Drain done entries
|
||||
# from the front of the fulfill deque first so that
|
||||
# ``num_pending_requests`` accurately reflects active
|
||||
# work for the next round of task accounting.
|
||||
while (
|
||||
len(self._pending_requests_to_fulfill) > 0
|
||||
and self._pending_requests_to_fulfill[0].future.done()
|
||||
):
|
||||
self._pending_requests_to_fulfill.popleft()
|
||||
return
|
||||
request_metadata = pending_request.metadata
|
||||
gen_choose_replicas_with_backoff = self._choose_replicas_with_backoff(
|
||||
pending_request
|
||||
)
|
||||
try:
|
||||
async for candidates in gen_choose_replicas_with_backoff:
|
||||
while (
|
||||
len(self._pending_requests_to_fulfill) > 0
|
||||
and self._pending_requests_to_fulfill[0].future.done()
|
||||
):
|
||||
self._pending_requests_to_fulfill.popleft()
|
||||
|
||||
# Note: unlike the base class, we deliberately do NOT
|
||||
# break on `len(routing_tasks) > target` here when our
|
||||
# popped pending_request hasn't been fulfilled yet, so
|
||||
# that the popping task always finishes its work.
|
||||
if (
|
||||
pending_request.future.done()
|
||||
and len(self._routing_tasks) > self.target_num_routing_tasks
|
||||
):
|
||||
break
|
||||
|
||||
replica = await self._select_from_candidate_replicas(
|
||||
candidates, backoff_index
|
||||
)
|
||||
if replica is not None:
|
||||
self._fulfill_next_pending_request(
|
||||
replica, request_metadata
|
||||
)
|
||||
break
|
||||
|
||||
backoff_index += 1
|
||||
if backoff_index >= 50 and backoff_index % 50 == 0:
|
||||
routing_time_elapsed = time.time() - start_time
|
||||
warning_log = (
|
||||
"Failed to route request after "
|
||||
f"{backoff_index} attempts over "
|
||||
f"{routing_time_elapsed:.2f}s. Retrying. "
|
||||
f"Request ID: {request_metadata.request_id}."
|
||||
)
|
||||
logger.warning(warning_log)
|
||||
finally:
|
||||
await gen_choose_replicas_with_backoff.aclose()
|
||||
|
||||
except Exception:
|
||||
logger.exception("Unexpected error in _fulfill_pending_requests.")
|
||||
finally:
|
||||
self._routing_tasks.remove(asyncio.current_task(loop=self._event_loop))
|
||||
self.num_routing_tasks_gauge.set(self.curr_num_routing_tasks)
|
||||
@@ -0,0 +1,90 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user