wehub-resource-sync
357 lines
15 KiB
Python
357 lines
15 KiB
Python
"""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)
|