mirrors/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/llm/_internal/serve/routing_policies/__init__.py
View File
python/ray/llm/_internal/serve/routing_policies/kv_aware/__init__.py
View File
python/ray/llm/_internal/serve/routing_policies/kv_aware/constants.py
+18
View File
@@ -0,0 +1,18 @@
# experimental_configs key overriding the per-node base port.
KV_EVENTS_PORT_BASE_KEY = "KV_EVENTS_PORT_BASE"
DEFAULT_KV_EVENTS_PORT_BASE = 5557
# experimental_configs key overriding the selection service's KV-indexer thread count.
KV_INDEXER_THREADS_KEY = "KV_INDEXER_THREADS"
DEFAULT_KV_INDEXER_THREADS = 4
# The engine's KV-event replay (ROUTER) socket sits this many ports above its PUB
# port, a separate range so it never collides with the PUB ports of colocated
# replicas (PORT_BASE + replica rank). Dynamo's selection service dials it to recover
# events missed before its SUB connected.
DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET = 1000
# TTL for a request's lifecycle tracking on the KV router actor. A live
# replica whose completion event was lost (e.g. a batch dropped on a
# transient actor outage) would otherwise leave its entry tracked forever.
REQUEST_TRACKING_TTL_S = 3600
python/ray/llm/_internal/serve/routing_policies/kv_aware/kv_aware_actor.py
+483
View File
@@ -0,0 +1,483 @@
import asyncio
import hashlib
import logging
import math
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, TypedDict
import ray
from ray import serve
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
DEFAULT_KV_INDEXER_THREADS,
REQUEST_TRACKING_TTL_S,
)
from ray.serve._private.common import DeploymentTargetInfo
from ray.serve._private.constants import (
SERVE_CONTROLLER_NAME,
SERVE_LOGGER_NAME,
SERVE_NAMESPACE,
)
from ray.serve._private.long_poll import LongPollClient, LongPollNamespace
logger = logging.getLogger(SERVE_LOGGER_NAME)
KV_ROUTER_ACTOR_NAME = "serve_llm_kv_router"
# Dynamo's selection service keys all worker, indexer, and load state by
# (model_name, tenant_id). KVRouterActor is a deployment-scoped actor that
# instantiates a selection service and serves exactly one model, so a single
# fixed key scopes all of its workers together.
_MODEL_NAME = "default"
_TENANT_ID = "default"
# Hooks a replica may invoke through ``KVRouterActor.on_lifecycle_events``.
LIFECYCLE_HOOKS = frozenset(
{
"on_request_added",
"on_prefill_complete",
"on_decode_progress",
"on_request_completed",
}
)
def get_worker_id(replica_unique_id: str) -> int:
"""Deterministically derive a Dynamo worker id from a replica's unique id."""
return int.from_bytes(
hashlib.blake2b(replica_unique_id.encode(), digest_size=8).digest(), "big"
)
@dataclass
class RequestLifecycle:
"""In-flight request load state while the request is served by a replica."""
worker_id: int
prompt_tokens: int = 0
# Client-provided output-length estimate (``sampling_params.max_tokens``);
# weights each decode block's load by how much generation remains.
expected_output_tokens: Optional[int] = None
prefill_completed: bool = False
output_tokens: int = 0
# Running count of KV blocks (prompt + output) the request occupies; the
# cursor for booking each newly crossed decode block.
total_blocks: int = 0
# Monotonic admission time, for the TTL eviction sweep.
created_at: float = field(default_factory=time.monotonic)
class WorkerSelection(TypedDict):
"""The worker chosen by ``KVRouterActor.select_worker`` for a request."""
# The chosen worker.
worker_id: int
# Data-parallel rank within the worker.
dp_rank: int
# Matched prompt tokens available on the selected worker.
overlap_tokens: int
# Prompt tokens that still need prefill on the selected worker.
effective_prefill_tokens: int
class KVRouterActor:
"""Deployment-scoped Ray actor backing KV-aware routing.
Attached to the LLMServer deployment via Serve's ``DeploymentActorConfig``,
independent of any replica's lifetime.
1. Created once per deployment, attached to the LLMServer deployment via
Serve's ``DeploymentActorConfig`` (independent of any replica's lifetime).
2. Owns an in-process Dynamo ``SelectionService``.
3. Tracks live replicas via a ``LongPollClient`` on ``DEPLOYMENT_TARGETS``,
mapping each running replica to a Dynamo worker id.
4. The ``SelectionService`` maintains a global KV index radix tree, fed by
every replica's KV events; each node records which workers hold that KV block.
5. Scoring (``select_worker``) ranks candidate workers by KV-cache overlap
and prefill/decode load.
6. Books each request's lifecycle into the service's active-load tracker, so
in-flight load feeds back into scoring for subsequent requests.
"""
def __init__(self, indexer_threads: int = DEFAULT_KV_INDEXER_THREADS):
# KV-cache block size, learned once from the first replica's reported
# engine config and passed to the selection service, which uses it to
# track the worker's active load and index its KV blocks for overlap.
self._block_size: Optional[int] = None
self._indexer_threads = indexer_threads
# _replica_id_by_worker maps a Dynamo worker id to the running replica's full
# id string, kept in sync with the deployment's live replicas over LongPoll.
# NOTE (jeffreywang): _replica_id_by_worker is later used by select_worker
# to get candidate workers to route among.
self._replica_id_by_worker: Dict[int, str] = {}
# Per-request state that the lifecycle hooks need, keyed by request id, serves
# the following purposes:
# 1. Block cursor: Turn cumulative decode tokens into add_output_block deltas.
# 2. expected_output_tokens for decode-block decay weighting.
# 3. In-flight request set: Free reservation exactly once.
# Ordered oldest-first so the TTL sweep pops stale entries off the front.
self._requests: "OrderedDict[str, RequestLifecycle]" = OrderedDict()
# Reverse index of in-flight request ids per worker, kept in lockstep with
# _requests, so remove_worker is O(k) in the worker's requests, not O(N).
self._request_ids_by_worker: Dict[int, Set[str]] = {}
# Carries the effective prefill tokens select() computed at routing time to
# on_request_added, which books them via the explicit create_reservation.
# TODO(jeffreywang): this map is only needed because create_reservation
# requires the effective prefill tokens to be passed in explicitly. Once the
# selection service caches each select() result and create_reservation can
# look it up by request id, Ray no longer needs to forward it.
self._effective_prefill_tokens_by_request: Dict[str, int] = {}
self._pending_tasks: Set[asyncio.Task] = set()
self._long_poll_client: Optional[LongPollClient] = None
self._create_selection_service()
self._start_replica_tracking()
async def ready(self) -> None:
"""Readiness probe for KVAwareRouter to confirm KVRouterActor is initialized
before it starts routing requests to it.
"""
def get_block_size(self) -> int:
"""Return the KV-cache block size used for decode-block accounting."""
return self._block_size
def _create_selection_service(self) -> None:
"""Create the in-process Dynamo selection service for this deployment."""
# Imported here, not at module scope: Ray pickles this actor class by
# value, and Dynamo's pyo3 classes cannot be pickled as its globals.
try:
from dynamo.llm import SelectionService
except ImportError:
self._svc = None
logger.warning(
"ai-dynamo is not installed; KV-aware routing requires ai-dynamo."
)
return
self._svc = SelectionService(indexer_threads=self._indexer_threads)
logger.info(
"Dynamo SelectionService created (indexer threads %d).",
self._indexer_threads,
)
def _start_replica_tracking(self) -> None:
"""Subscribe to this deployment's running replicas via LongPollClient."""
deployment_id = serve.get_deployment_actor_context().deployment_id
controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
self._long_poll_client = LongPollClient(
controller,
{
(
LongPollNamespace.DEPLOYMENT_TARGETS,
deployment_id,
): self._on_deployment_targets,
},
# Relies on KVRouterActor being an async actor (it defines async
# methods), so Ray runs __init__ inside the actor's event loop.
call_in_event_loop=asyncio.get_running_loop(),
client_id=f"{type(self).__name__}:{deployment_id}",
)
def _schedule(self, coro) -> None:
"""Run a coroutine on the actor's event loop, holding a reference until
it completes.
"""
task = asyncio.ensure_future(coro)
self._pending_tasks.add(task)
task.add_done_callback(self._pending_tasks.discard)
def _register_block_size(self, block_size: int, replica_id: str) -> None:
"""Pin the deployment's KV-cache block size from the first replica's
reported engine config.
"""
if self._block_size is None:
self._block_size = block_size
logger.info("KV router block size set to %d.", block_size)
elif block_size != self._block_size:
# Replicas of a deployment are expected to resolve the same block
# size, so a mismatch is unexpected. We still register the worker so
# the selection service spawns its KV-event listener, but the indexer
# only ingests blocks whose size matches the pinned block size, so a
# genuinely mismatched replica's KV events would be dropped (its KV
# cache never indexed).
logger.error(
"Replica %s reports KV block size %d but the KV router is "
"pinned at %d; registering it at the pinned size (replicas of a "
"deployment are expected to agree).",
replica_id,
block_size,
self._block_size,
)
def _on_deployment_targets(self, target_info: DeploymentTargetInfo) -> None:
"""LongPoll listener: reconcile tracked workers against the running-replica
snapshot.
Each replica advertises its KV-events endpoint via ``record_routing_stats``
(carried in ``RunningReplicaInfo.routing_stats``); newly advertised replicas
are registered with the selection service and departed ones evicted.
"""
members: Dict[int, tuple] = {}
for replica in target_info.running_replicas:
worker_id = get_worker_id(replica.replica_id.unique_id)
kv_event_metadata = replica.routing_stats.get("kv_event_metadata")
if kv_event_metadata is not None:
members[worker_id] = (
replica.replica_id.to_full_id_str(),
kv_event_metadata,
)
registered = set(self._replica_id_by_worker)
added = members.keys() - registered
removed = registered - members.keys()
for worker_id in removed:
self.remove_worker(worker_id)
self._replica_id_by_worker.pop(worker_id, None)
for worker_id in added:
replica_id, kv_event_metadata = members[worker_id]
self._register_block_size(kv_event_metadata["block_size"], replica_id)
self._replica_id_by_worker[worker_id] = replica_id
self._schedule(
self._upsert_worker(worker_id, replica_id, kv_event_metadata)
)
if added or removed:
logger.info(
"KV router replica membership updated: +%d -%d, tracking %d worker(s).",
len(added),
len(removed),
len(self._replica_id_by_worker),
)
def remove_worker(self, worker_id: int) -> None:
"""Evict a departed replica's worker and its KV blocks from the
selection service.
"""
# Drop the departed replica's in-flight requests; their completions can
# never arrive, so they would otherwise leak. delete_worker below frees
# their load in the service, so no per-request free_reservation is needed.
for request_id in self._request_ids_by_worker.pop(worker_id, set()):
self._requests.pop(request_id, None)
if self._svc is None:
return
self._schedule(self._svc.delete_worker(worker_id))
async def _upsert_worker(
self, worker_id: int, replica_id: str, kv_event_metadata: Dict[str, Any]
) -> None:
"""Register a replica's KV-event endpoint with the selection service.
The selection service spawns a connect-out ZMQ listener to the
replica's ``endpoint`` and indexes its live KV events.
"""
if self._svc is None:
return
dp_rank = kv_event_metadata["dp_rank"]
await self._svc.upsert_worker(
{
"worker_id": worker_id,
"model_name": _MODEL_NAME,
"tenant_id": _TENANT_ID,
# NOTE: SelectionService requires endpoint to be non-empty although it's left
# unused under an external runtime like Ray Serve LLM.
# TODO (jeffreywang): Allow empty endpoints upstream.
"endpoint": f"ray://{replica_id}",
"block_size": self._block_size,
# NOTE: max_num_batched_tokens is a proxy of load capacity for load-based
# scoring in the selection service.
"max_num_batched_tokens": kv_event_metadata["max_num_batched_tokens"],
"data_parallel_start_rank": dp_rank,
# TODO (jeffreywang): Support KV-aware routing for data parallel deployments.
"data_parallel_size": 1,
"kv_events_endpoints": {dp_rank: kv_event_metadata["endpoint"]},
# The listener dials this on a sequence gap (slow-joiner) to replay
# the events it missed before its SUB connected; without it those
# events are dropped and never indexed.
"replay_endpoint": kv_event_metadata.get("replay_endpoint"),
}
)
logger.info(
"Registered KV event worker %d for replica %s at %s.",
worker_id,
replica_id,
kv_event_metadata["endpoint"],
)
async def select_worker(
self,
request_id: str,
token_ids: List[int],
allowed_worker_ids: List[int],
) -> WorkerSelection:
"""Score the allowed workers for a request based on KV-cache overlap and
load and pick the best one.
Args:
request_id: Unique identifier for the request being routed.
token_ids: Prompt token ids used to compute KV-cache overlap.
allowed_worker_ids: Candidate worker ids the router may select from.
Returns:
The selected worker (see ``WorkerSelection``).
"""
if token_ids is None or len(token_ids) == 0:
raise ValueError("KV aware routing requires non-empty token_ids.")
if self._svc is None:
# ai-dynamo is not installed, so this deployment cannot score requests.
# Fail fast and surface RuntimeError to the client as a 503 via LLMRouter.
raise RuntimeError(
"KV-aware routing is unavailable because ai-dynamo is not "
"installed in the deployment's environment."
)
selection = await self._svc.select(
{
"model_name": _MODEL_NAME,
"tenant_id": _TENANT_ID,
"selection_id": request_id,
"token_ids": token_ids,
"allowed_worker_ids": allowed_worker_ids,
}
)
self._effective_prefill_tokens_by_request[request_id] = selection[
"effective_prefill_tokens"
]
return {
"worker_id": selection["worker_id"],
"dp_rank": selection["dp_rank"],
"overlap_tokens": selection["overlap"]["longest_matched"],
"effective_prefill_tokens": selection["effective_prefill_tokens"],
}
async def on_lifecycle_events(self, events: List[tuple]) -> None:
"""Apply a replica's ``(hook_name, args)`` lifecycle events in order.
The hooks are order-sensitive (e.g. a completion arriving before its
admission would resurrect an evicted request) so a replica sends its
events in submission order, batched into one call.
"""
if self._svc is None or self._block_size is None:
return
for hook_name, args in events:
if hook_name not in LIFECYCLE_HOOKS:
logger.warning("Ignoring unknown lifecycle hook %s", hook_name)
continue
try:
await getattr(self, hook_name)(*args)
except Exception:
# One hook raising must not abort the batch and drop other events.
logger.exception(
"KV lifecycle hook %s failed; skipping it and continuing.",
hook_name,
)
async def on_request_added(
self,
request_id: str,
worker_id: int,
token_ids: List[int],
expected_output_tokens: Optional[int] = None,
) -> None:
"""Admit a routed request into ``worker_id``'s active load, booking it
into the selection service which computes the worker's KV overlap from
``token_ids``, so the recorded prefill excludes the cached prefix."""
await self._evict_stale_requests()
prompt_tokens = len(token_ids)
self._requests[request_id] = RequestLifecycle(
worker_id=worker_id,
prompt_tokens=prompt_tokens,
expected_output_tokens=expected_output_tokens,
total_blocks=math.ceil(prompt_tokens / self._block_size),
)
self._request_ids_by_worker.setdefault(worker_id, set()).add(request_id)
effective_prefill_tokens = self._effective_prefill_tokens_by_request.pop(
request_id, None
)
await self._svc.create_reservation(
{
"model_name": _MODEL_NAME,
"tenant_id": _TENANT_ID,
"reservation_id": request_id,
"worker_id": worker_id,
"token_ids": token_ids,
"expected_output_tokens": expected_output_tokens,
"effective_prefill_tokens": effective_prefill_tokens,
}
)
if request_id not in self._requests:
await self._svc.free_reservation(request_id)
async def on_prefill_complete(self, request_id: str) -> None:
"""Record a request's prefill -> decode transition, dropping its prefill
load in the selection service."""
state = self._requests.get(request_id)
if state is None:
return
state.prefill_completed = True
await self._svc.prefill_complete(request_id)
async def on_decode_progress(
self, request_id: str, cumulative_output_tokens: int
) -> None:
"""Advance ``request_id`` to an exact cumulative output-token count,
booking one decode block in the selection service per crossed boundary.
"""
state = self._requests.get(request_id)
if state is None:
return
state.output_tokens = cumulative_output_tokens
new_total_blocks = math.ceil(
(state.prompt_tokens + cumulative_output_tokens) / self._block_size
)
decay_fraction = self._get_decay_fraction(state)
while new_total_blocks > state.total_blocks:
state.total_blocks += 1
self._svc.add_output_block(request_id, decay_fraction=decay_fraction)
async def on_request_completed(self, request_id: str) -> None:
"""Free ``request_id`` from the selection service's active load and the
local view."""
self._effective_prefill_tokens_by_request.pop(request_id, None)
state = self._requests.pop(request_id, None)
if state is not None:
self._untrack_worker_request(request_id, state.worker_id)
await self._svc.free_reservation(request_id)
def _untrack_worker_request(self, request_id: str, worker_id: int) -> None:
"""Drop a request from the per-worker reverse index, keeping it in
lockstep with ``_requests``."""
request_ids = self._request_ids_by_worker.get(worker_id)
if request_ids is not None:
request_ids.discard(request_id)
if not request_ids:
del self._request_ids_by_worker[worker_id]
async def _evict_stale_requests(self) -> None:
"""Backstop for a lost completion on a live replica: evict requests tracked
past ``REQUEST_TRACKING_TTL_S``, freeing their reservations.
"""
cutoff = time.monotonic() - REQUEST_TRACKING_TTL_S
while self._requests:
request_id, state = next(iter(self._requests.items()))
if state.created_at > cutoff:
break
self._requests.popitem(last=False)
self._untrack_worker_request(request_id, state.worker_id)
self._effective_prefill_tokens_by_request.pop(request_id, None)
logger.warning(
"Evicting stale KV request %s (tracked > %ds without completion); "
"freeing its reservation.",
request_id,
REQUEST_TRACKING_TTL_S,
)
await self._svc.free_reservation(request_id)
def _get_decay_fraction(self, state: RequestLifecycle) -> Optional[float]:
"""Fraction of output still expected, or ``None`` without an estimate;
weights each decode block by how much generation remains."""
if not state.expected_output_tokens:
return None
return max(0.0, 1.0 - state.output_tokens / state.expected_output_tokens)
python/ray/llm/_internal/serve/routing_policies/kv_aware/kv_aware_router.py
+122
View File
@@ -0,0 +1,122 @@
import logging
import random
from typing import List, Optional
import ray
from ray.actor import ActorHandle
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
from ray.llm._internal.serve.core.ingress.tokenizer import REQUEST_TOKEN_IDS_KWARG
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
KV_ROUTER_ACTOR_NAME,
get_worker_id,
)
from ray.serve._private.constants import (
SERVE_DEPLOYMENT_ACTOR_PREFIX,
SERVE_LOGGER_NAME,
SERVE_NAMESPACE,
)
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
from ray.serve.config import RequestRouterConfig
logger = logging.getLogger(SERVE_LOGGER_NAME)
class KVAwareRouter(RequestRouter):
"""Routes each request to the candidate that best balances expected KV-cache
overlap against the worker's current prefill/decode load.
Scoring is delegated to the deployment-scoped ``KVRouterActor`` (which owns the
Dynamo selection service and the global KV index); this per-handle router stays
thin and simply maps candidate replicas to/from Dynamo worker ids.
"""
def initialize_state(self):
"""Resolve the deployment's ``KVRouterActor``.
The actor is attached to this deployment via ``DeploymentActorConfig``
whenever the request router is a ``KVAwareRouter``, so it exists by the time
requests route. We resolve its Serve-generated name and block on a cheap
call to confirm it finished initializing, so the first routed request finds
a ready scorer.
"""
self._kv_router_actor = self._discover_kv_router_actor()
# Synchronization barrier: Ray defers actor methods until __init__ completes,
# so awaiting any method blocks until KVRouterActor is constructed.
ray.get(self._kv_router_actor.ready.remote())
def _discover_kv_router_actor(self) -> ActorHandle:
"""Handle to this deployment's ``KVRouterActor`` by its Serve-scoped name."""
prefix = (
f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}"
f"{self._deployment_id.app_name}::{self._deployment_id.name}::"
)
suffix = f"::{KV_ROUTER_ACTOR_NAME}"
for entry in ray.util.list_named_actors(all_namespaces=True):
name = entry.get("name") or ""
if (
entry.get("namespace") == SERVE_NAMESPACE
and name.startswith(prefix)
and name.endswith(suffix)
):
return ray.get_actor(name, namespace=SERVE_NAMESPACE)
raise RuntimeError(
f"KVRouterActor for deployment {self._deployment_id} not found; it must "
"be attached via DeploymentActorConfig when using KVAwareRouter."
)
async def choose_replicas(
self,
candidate_replicas: List[RunningReplica],
pending_request: Optional[PendingRequest] = None,
) -> List[List[RunningReplica]]:
"""Choose the candidate replica(s) to route ``pending_request`` to.
Maps the candidate replicas to their Dynamo worker ids, asks the
``KVRouterActor`` to rank them via ``select_worker``, and routes to
the chosen worker's replica. With direct streaming enabled, HAProxy
then forwards the original request to that replica.
Requests with no prompt token ids have nothing to score on, so they route
to a random candidate. This covers the pre-routing ``/tokenize`` RPC (routed
before token ids exist) and token-less fallbacks (batch prompts,
truncated/unparseable bodies).
TODO (jeffreywang): Move pre-routing tokenization to KVRouterActor while
ensuring tokenization correctness.
Args:
candidate_replicas: The replicas eligible to serve the request.
pending_request: The request being routed.
Returns:
Ranked groups of replicas.
"""
token_ids = (
pending_request.kwargs.get(REQUEST_TOKEN_IDS_KWARG)
if pending_request is not None
else None
)
if not token_ids:
return [[random.choice(candidate_replicas)]] if candidate_replicas else []
worker_id_to_replica = {
get_worker_id(replica.replica_id.unique_id): replica
for replica in candidate_replicas
}
selection = await self._kv_router_actor.select_worker.remote(
pending_request.metadata.request_id,
token_ids,
list(worker_id_to_replica),
)
return [[worker_id_to_replica[selection["worker_id"]]]]
def is_kv_aware(llm_config: LLMConfig) -> bool:
"""Whether ``llm_config`` selects a ``KVAwareRouter`` for replica selection."""
request_router_config = llm_config.deployment_config.get("request_router_config")
if isinstance(request_router_config, dict):
request_router_config = RequestRouterConfig(**request_router_config)
return isinstance(request_router_config, RequestRouterConfig) and issubclass(
request_router_config.get_request_router_class(), KVAwareRouter
)
python/ray/llm/_internal/serve/routing_policies/kv_aware/utils.py
+66
View File
@@ -0,0 +1,66 @@
"""Helpers for wiring KV-aware routing into an LLM deployment."""
import logging
import ray
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
DEFAULT_KV_INDEXER_THREADS,
KV_INDEXER_THREADS_KEY,
)
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
KV_ROUTER_ACTOR_NAME,
KVRouterActor,
)
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
is_kv_aware,
)
from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
configure_kv_events_for_kv_routing,
)
from ray.serve._private.constants import SERVE_LOGGER_NAME
from ray.serve.config import DeploymentActorConfig
logger = logging.getLogger(SERVE_LOGGER_NAME)
def _maybe_setup_kv_aware_routing(
deployment_options: dict, llm_config: LLMConfig
) -> None:
"""Set up KV-aware routing when the deployment's request router is a
KVAwareRouter.
Attaches the KVRouterActor, which owns the deployment's global KV radix
tree, and enables the engine KV events that feed it.
"""
if not is_kv_aware(llm_config):
if llm_config.engine_kwargs.get("kv_events_config") is not None:
logger.warning(
"engine_kwargs['kv_events_config'] is set but the deployment's "
"request router is not a KVAwareRouter, so the engine's KV events "
"will not be consumed. To use them, configure KVAwareRouter via "
"deployment_config.request_router_config."
)
return
# Keep the engine's token-tracking gate which reads llm_config consistent
# with the router resolved here from the merged deployment options.
llm_config.deployment_config["request_router_config"] = deployment_options[
"request_router_config"
]
deployment_options["deployment_actors"] = [
*deployment_options.get("deployment_actors", []),
DeploymentActorConfig(
name=KV_ROUTER_ACTOR_NAME,
actor_class=ray.remote(KVRouterActor),
actor_options={"num_cpus": 0},
init_kwargs={
"indexer_threads": llm_config.experimental_configs.get(
KV_INDEXER_THREADS_KEY, DEFAULT_KV_INDEXER_THREADS
),
},
),
]
configure_kv_events_for_kv_routing(llm_config)
python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/__init__.py
+1
View File
@@ -0,0 +1 @@
"""vLLM-specific KV-aware routing helpers."""
python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/kv_events.py
+163
View File
@@ -0,0 +1,163 @@
import logging
from typing import Any, Dict, Optional
import ray
from ray import serve
from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
DEFAULT_KV_EVENTS_PORT_BASE,
DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET,
KV_EVENTS_PORT_BASE_KEY,
)
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
is_kv_aware,
)
from ray.serve._private.constants import SERVE_LOGGER_NAME
logger = logging.getLogger(SERVE_LOGGER_NAME)
def configure_kv_events_for_kv_routing(llm_config: LLMConfig) -> None:
"""
Enable engine KV-cache events for a KV-aware-routed deployment.
"""
engine_kwargs = llm_config.engine_kwargs
if engine_kwargs.get("enable_prefix_caching") is False:
logger.warning(
"KV-aware routing is configured but enable_prefix_caching is False; "
"the engine will not emit KV-cache events."
)
llm_config.update_engine_kwargs(
kv_events_config={
"enable_kv_cache_events": True,
"publisher": "zmq",
"endpoint": _default_kv_events_endpoint(llm_config),
# Bind a replay (ROUTER) socket so the selection service can recover
# the events it missed before its SUB connected (slow-joiner gap).
"replay_endpoint": _default_kv_events_replay_endpoint(llm_config),
}
)
_pin_block_hash_seed(llm_config)
def _pin_block_hash_seed(llm_config: LLMConfig) -> None:
"""Make engine block hashes content-deterministic across replicas.
The KV router's global indexer chains and dedups blocks by the engines'
block hashes, so identical content must hash identically on every
replica. vLLM salts its block-hash chain root per process unless
``PYTHONHASHSEED`` is set, so pin it deployment-wide.
"""
runtime_env = dict(llm_config.runtime_env or {})
env_vars = dict(runtime_env.get("env_vars") or {})
env_vars.setdefault("PYTHONHASHSEED", "0")
runtime_env["env_vars"] = env_vars
llm_config.runtime_env = runtime_env
def assign_replica_kv_events_endpoint(llm_config: LLMConfig) -> None:
"""Pin the engine's KV-events endpoint to a per-replica port.
Replicas of a deployment share one ``engine_kwargs``, so colocated
replicas would otherwise bind the same ZMQ port. Offsets the configured
base port by the replica's rank.
"""
if not is_kv_aware(llm_config):
return
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
if kv_events_config is None:
return
updated = dict(kv_events_config)
endpoint = updated.pop("endpoint")
replay_endpoint = updated.pop("replay_endpoint")
# With data parallelism the engine offsets ports by dp_rank under ZmqEventPublisher;
# otherwise offset by the replica's node-local rank so colocated replicas don't collide.
if llm_config.engine_kwargs.get("data_parallel_rank") is not None:
offset = 0
else:
offset = _get_replica_rank()
updated["endpoint"] = _get_offset_endpoint_port(endpoint, offset)
updated["replay_endpoint"] = _get_offset_endpoint_port(replay_endpoint, offset)
llm_config.update_engine_kwargs(kv_events_config=updated)
def resolve_kv_event_source_endpoint(llm_config: LLMConfig) -> Optional[str]:
"""This replica's node-routable KV-events endpoint, for the selection
service to connect out to.
The engine's KV-events endpoint at the replica's node IP; ``None`` when
KV-cache events are not enabled.
"""
if not is_kv_aware(llm_config):
return None
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
if kv_events_config is None:
return None
return _get_node_routable_endpoint(llm_config, kv_events_config["endpoint"])
def get_kv_event_routing_stats(
llm_config: LLMConfig, block_size: int, max_num_batched_tokens: int
) -> Dict[str, Any]:
"""Returns this replica's routing-stats payload advertising its KV-events
endpoint and the engine's resolved KV-cache block size."""
if not is_kv_aware(llm_config):
return {}
kv_events_config = llm_config.engine_kwargs.get("kv_events_config")
if kv_events_config is None:
return {}
kv_event_metadata = {
"endpoint": _get_node_routable_endpoint(
llm_config, kv_events_config["endpoint"]
),
"block_size": block_size,
"max_num_batched_tokens": max_num_batched_tokens,
"dp_rank": llm_config.engine_kwargs.get("data_parallel_rank") or 0,
"replay_endpoint": _get_node_routable_endpoint(
llm_config, kv_events_config["replay_endpoint"]
),
}
return {"kv_event_metadata": kv_event_metadata}
def _get_node_routable_endpoint(llm_config: LLMConfig, endpoint: str) -> str:
"""Rewrite a wildcard-bound engine endpoint to this replica's node IP.
Offsets by ``data_parallel_rank`` to match the engine's own per-rank offset,
so the selection service dials the right socket.
"""
dp_rank = llm_config.engine_kwargs.get("data_parallel_rank")
if dp_rank is not None:
endpoint = _get_offset_endpoint_port(endpoint, dp_rank)
port = endpoint.rsplit(":", 1)[1]
return f"tcp://{ray.util.get_node_ip_address()}:{port}"
def _get_replica_rank() -> int:
return serve.get_replica_context().rank.local_rank
def _default_kv_events_endpoint(llm_config: LLMConfig) -> str:
port_base = int(
llm_config.experimental_configs.get(
KV_EVENTS_PORT_BASE_KEY, DEFAULT_KV_EVENTS_PORT_BASE
)
)
return f"tcp://*:{port_base}"
def _default_kv_events_replay_endpoint(llm_config: LLMConfig) -> str:
port_base = int(
llm_config.experimental_configs.get(
KV_EVENTS_PORT_BASE_KEY, DEFAULT_KV_EVENTS_PORT_BASE
)
)
return f"tcp://*:{port_base + DEFAULT_KV_EVENTS_REPLAY_PORT_OFFSET}"
def _get_offset_endpoint_port(endpoint: str, offset: int) -> str:
"""Offset a TCP endpoint's port with vLLM's ZmqEventPublisher convention."""
base, port = endpoint.rsplit(":", 1)
return f"{base}:{int(port) + offset}"
python/ray/llm/_internal/serve/routing_policies/kv_aware/vllm/token_tracking.py
+192
View File
@@ -0,0 +1,192 @@
import asyncio
from typing import Any, List, Optional, Type
from vllm.outputs import RequestOutput
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine.async_llm import AsyncLLM
from ray import serve
from ray.actor import ActorHandle
from ray.exceptions import RayActorError, RayTaskError
from ray.llm._internal.serve.observability.logging import get_logger
from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
KV_ROUTER_ACTOR_NAME,
get_worker_id,
)
from ray.llm._internal.serve.utils.server_utils import get_serve_request_id
logger = get_logger(__name__)
def _get_prompt_token_ids(prompt: Any) -> List[int]:
"""The prompt's pre-tokenized token ids."""
try:
return list(prompt["prompt_token_ids"])
except (KeyError, TypeError) as e:
raise ValueError(
"KV-aware token tracking requires a pre-tokenized prompt "
f"(dict with 'prompt_token_ids'); got {type(prompt).__name__}"
) from e
class LifecycleEventForwarder:
"""Ordered, non-blocking bridge from the engine to the KV router actor
which maintains per-replica token load statistics.
``report`` only enqueues locally, so generation never blocks on the actor.
A single delivery task per replica drains the queue to the actor, awaiting
one ``on_lifecycle_events`` call at a time so events arrive in the order they
were reported. Events that pile up during a call are sent together in the
next one.
"""
def __init__(self, actor: ActorHandle, worker_id: int):
self.actor = actor
self.worker_id = worker_id
self._events: asyncio.Queue = asyncio.Queue()
self._delivery_task: Optional[asyncio.Task] = None
def report(self, method_name: str, *args) -> None:
if self._delivery_task is None or self._delivery_task.done():
self._delivery_task = asyncio.get_running_loop().create_task(
self._deliver()
)
self._events.put_nowait((method_name, args))
async def _deliver(self) -> None:
while True:
# Wait for the next event, then drain whatever queued up behind it
# into the same batch.
batch = [await self._events.get()]
while not self._events.empty():
batch.append(self._events.get_nowait())
try:
await self.actor.on_lifecycle_events.remote(batch)
except (RayActorError, RayTaskError) as e:
logger.warning("Dropping KV lifecycle events: %s", e)
finally:
for _ in batch:
self._events.task_done()
async def flush(self) -> None:
"""Wait until every reported event has been delivered."""
await self._events.join()
def close(self) -> None:
"""Cancel the delivery task on engine shutdown."""
if self._delivery_task is not None:
self._delivery_task.cancel()
self._delivery_task = None
class RequestTokenTracker:
"""Drives the request lifecycle hooks for one ``generate()`` stream."""
def __init__(
self,
forwarder: LifecycleEventForwarder,
request_id: str,
prompt_token_ids: List[int],
expected_output_tokens: Optional[int],
):
self._forwarder = forwarder
self._request_id = request_id
self._cumulative = 0
self._prefill_marked = False
self._finished = False
forwarder.report(
"on_request_added",
request_id,
forwarder.worker_id,
prompt_token_ids,
expected_output_tokens,
)
def on_output(self, output: RequestOutput) -> None:
"""Observe one engine ``RequestOutput`` (forwarded to the caller as-is).
vLLM streams either DELTA chunks or a single FINAL_ONLY chunk; both
carry only new tokens, so output progress simply accumulates.
"""
step_tokens = sum(len(o.token_ids or []) for o in output.outputs)
if step_tokens == 0:
# No new tokens this step (e.g. a finish-only chunk).
return
self._cumulative += step_tokens
if not self._prefill_marked:
# The first output token signals prefill completion.
self._prefill_marked = True
self._forwarder.report("on_prefill_complete", self._request_id)
self._forwarder.report("on_decode_progress", self._request_id, self._cumulative)
def finish(self) -> None:
"""Report completion exactly once."""
if not self._finished:
self._finished = True
self._forwarder.report("on_request_completed", self._request_id)
def enable_token_tracking(engine_cls: Type[AsyncLLM]) -> Type[AsyncLLM]:
"""Decorator adding KV-router request lifecycle tracking."""
class TokenTrackingEngine(engine_cls):
_lifecycle_forwarder: Optional[LifecycleEventForwarder] = None
_resolve_warned: bool = False
def _resolve_lifecycle_forwarder(self) -> Optional[LifecycleEventForwarder]:
if self._lifecycle_forwarder is None:
try:
actor = serve.get_deployment_actor(KV_ROUTER_ACTOR_NAME)
worker_id = get_worker_id(
serve.get_replica_context().replica_id.unique_id
)
self._lifecycle_forwarder = LifecycleEventForwarder(
actor, worker_id
)
except Exception as e:
# Warn once: resolution is retried per request until it succeeds.
if not self._resolve_warned:
self._resolve_warned = True
logger.warning("KV token tracking disabled: %s", e)
return self._lifecycle_forwarder
def shutdown(self, *args, **kwargs):
if self._lifecycle_forwarder is not None:
self._lifecycle_forwarder.close()
return super().shutdown(*args, **kwargs)
async def generate(self, prompt, sampling_params, request_id, *args, **kwargs):
stream = super().generate(
prompt, sampling_params, request_id, *args, **kwargs
)
forwarder = self._resolve_lifecycle_forwarder()
# CUMULATIVE repeats output-so-far per chunk; our accounting sums
# deltas, so skip it rather than over-count. vLLM's OpenAI layer only
# uses DELTA/FINAL_ONLY (*Request.to_sampling_params):
# https://github.com/vllm-project/vllm/tree/main/vllm/entrypoints/openai
if forwarder is None or (
sampling_params.output_kind == RequestOutputKind.CUMULATIVE
):
async for output in stream:
yield output
return
lifecycle_request_id = get_serve_request_id() or request_id
tracker = RequestTokenTracker(
forwarder,
lifecycle_request_id,
_get_prompt_token_ids(prompt),
# The request's own output cap is its expected length; weights
# the selection service's decode-block decay.
# TODO(jeffreywang): Use an agent-provided expected-OSL hint for
# more accurate decode-load estimation.
sampling_params.max_tokens,
)
try:
async for output in stream:
tracker.on_output(output)
yield output
finally:
tracker.finish()
return TokenTrackingEngine
python/ray/llm/_internal/serve/routing_policies/prefix_aware/__init__.py
View File
python/ray/llm/_internal/serve/routing_policies/prefix_aware/prefix_aware_router.py
+420
View File
@@ -0,0 +1,420 @@
# These imports are used for metrics tracking, will remove for PR
import logging
import time
from typing import (
Any,
List,
Optional,
)
import ray
from ray.actor import ActorHandle
from ray.llm._internal.serve.routing_policies.prefix_aware.prefix_tree import (
PrefixTreeActor,
)
from ray.serve._private.common import ReplicaID
from ray.serve._private.constants import (
SERVE_CONTROLLER_NAME,
SERVE_LOGGER_NAME,
SERVE_NAMESPACE,
)
from ray.serve._private.replica_result import ReplicaResult
from ray.serve._private.request_router import (
PowerOfTwoChoicesRequestRouter,
)
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 (
LocalityMixin,
MultiplexMixin,
RequestRouter,
)
logger = logging.getLogger(SERVE_LOGGER_NAME)
class PrefixCacheAffinityRouter(LocalityMixin, MultiplexMixin, RequestRouter):
"""Extends the PowerOfTwoChoicesRequestRouter with prefix-matching capabilities.
This request router optimizes replica selection by considering input text prefixes:
1. Mixes between three strategies to balance prefix cache hit rate and load balancing:
- When load is balanced (queue length difference < threshold), it selects replicas
with the highest prefix match rate for the input text
- When load is balanced but match rate is below 10%, it falls back to the smallest tenants
- When load is imbalanced, it uses the default Power of Two selection
2. Maintains a prefix tree to track which replicas have processed similar inputs:
- Inserts prompt text into the prefix tree after routing
- Uses this history to inform future routing decisions
This approach improves performance by routing related requests to the same replicas,
increasing cache locality and reducing overhead for language model inference.
"""
def initialize_state(
self,
imbalanced_threshold: Optional[float] = float("inf"),
match_rate_threshold: Optional[float] = 0.1,
do_eviction: Optional[bool] = False,
eviction_threshold_chars: Optional[int] = 400_000,
eviction_target_chars: Optional[int] = 360_000,
eviction_interval_secs: Optional[int] = 10,
tree_actor: Optional[ActorHandle] = None,
):
"""Initialize the prefix-aware routing state and configuration.
Args:
imbalanced_threshold: Threshold for queue length difference to consider
load balanced. When the difference between replica queue lengths is
less than this value, prefix-aware routing is used.
match_rate_threshold: Minimum prefix match rate (0.0-1.0) required to
use prefix-aware routing. If match rate is below this threshold,
falls back to smallest tenant selection.
do_eviction: Whether to enable automatic eviction of old prefix tree
entries to manage memory usage.
eviction_threshold_chars: Maximum number of characters in the prefix
tree before eviction is triggered.
eviction_target_chars: Target number of characters to reduce the
prefix tree to during eviction.
eviction_interval_secs: Interval in seconds between eviction checks
when eviction is enabled.
tree_actor: The actor to use for the prefix tree in a test environment.
If None, a detached actor will be created/retrieved.
"""
# === Prefix-aware routing logic hyperparameters ===
self._imbalanced_threshold = imbalanced_threshold
self._match_rate_threshold = match_rate_threshold
# === Eviction policy ===
self._do_eviction = do_eviction
self._eviction_loop_running = False
self._eviction_threshold_chars = eviction_threshold_chars
# Default eviction_target_chars to eviction_threshold_chars if not specified
self._eviction_target_chars = (
eviction_target_chars
if eviction_target_chars is not None
else eviction_threshold_chars
)
self._eviction_interval_secs = eviction_interval_secs
if tree_actor is None:
# Create deployment-specific detached actor to avoid replica ID conflicts
# in multi-deployment scenarios (e.g., PD disaggregation with DP)
deployment_name = self._deployment_id.name if self._deployment_id else None
app_name = self._deployment_id.app_name if self._deployment_id else None
actor_name = "LlmPrefixTreeActor"
actor_namespace_components = [SERVE_NAMESPACE]
if app_name:
actor_namespace_components.append(app_name)
if deployment_name:
actor_namespace_components.append(deployment_name)
actor_namespace = "::".join(actor_namespace_components)
self._tree_actor = PrefixTreeActor.options(
name=actor_name,
namespace=actor_namespace,
get_if_exists=True,
lifetime="detached",
).remote()
# Register the actor with the controller for cleanup on serve.shutdown()
controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
ray.get(
controller._register_shutdown_cleanup_actor.remote(self._tree_actor)
)
else:
self._tree_actor = tree_actor
def _extract_text_from_request(self, pending_request: PendingRequest) -> str:
"""Extracts the text content from a pending request for prefix matching.
Searches through request arguments for either 'messages' or 'prompt' attributes,
then normalizes the content to a single string representation that can be used
for prefix tree operations.
Args:
pending_request: The request to extract text from
Returns:
A string containing the prompt text or concatenated message contents
Raises:
ValueError: If no prompt or messages attribute is found in the request
"""
prompt = None
for arg in pending_request.args:
valid_input_types = ["messages", "prompt"]
for valid_input_type in valid_input_types:
if hasattr(arg, valid_input_type):
prompt = (
arg.prompt if valid_input_type == "prompt" else arg.messages
)
break
if prompt is not None:
break
if prompt is None:
raise ValueError(
"No request with message or prompt attribute found in pending_request.args"
)
return self._normalize_prompt_to_string(prompt)
def _coerce_to_text(self, value: Any) -> str:
if value is None:
return ""
if isinstance(value, str):
return value
if isinstance(value, list):
return "".join(self._coerce_to_text(item) for item in value)
if isinstance(value, dict):
text_value = value.get("text")
if isinstance(text_value, str):
return text_value
if "content" in value:
return self._coerce_to_text(value["content"])
return ""
def _normalize_prompt_to_string(self, prompt: Any) -> str:
"""Normalize prompt/messages a single string of characters.
This is not exhaustive (e.g. thinking parts, multimodal are not supported).
TODO(seiji): find a more maintainable way to normalize the prompt/messages.
Supported:
- string → return as-is
- list of strings → concat
- list of message dicts with 'content' as string → concat
- list of message dicts with 'content' as list of dicts → concat the 'text' fields from those parts
"""
if isinstance(prompt, str):
return prompt
if isinstance(prompt, list):
return "".join(
self._coerce_to_text(
message.get("content") if isinstance(message, dict) else message
)
for message in prompt
)
return ""
async def _prefix_match_best_replicas(
self,
pending_request: Optional[PendingRequest],
candidate_replicas: List[RunningReplica],
) -> List[RunningReplica]:
"""
Returns a set of candidate replicas, of which the one with the smallest replica queue will be chosen.
0. Default: same as pow 2 request router, return 2 replicas at random.
1. If load is balanced, choose replica(s) with highest prefix match rate. If highest hit rate is below 10% or no match found, use replicas with smallest KV cache usage.
2. If load is imbalanced, use default.
"""
chosen_replica_id_strings = []
if (
pending_request is not None
and pending_request.args is not None
and len(pending_request.args) > 0
):
input_text = self._extract_text_from_request(pending_request)
if input_text is not None:
# Start Sphinx tag: __begin_load_balance_component__
# Check for imbalanced load.
highest_queue_len = 0
lowest_queue_len = float("inf")
not_in_cache: List[ReplicaID] = []
if self._use_replica_queue_len_cache:
# Populate available queue lens from the cache.
for r in candidate_replicas:
queue_len = self._replica_queue_len_cache.get(r.replica_id)
if queue_len is None or queue_len >= r.max_ongoing_requests:
not_in_cache.append(r)
else:
highest_queue_len = max(highest_queue_len, queue_len)
lowest_queue_len = min(lowest_queue_len, queue_len)
else:
not_in_cache = candidate_replicas
if len(not_in_cache) > 0:
for r, queue_len in await self._probe_queue_lens(
not_in_cache,
0,
):
if queue_len is None:
continue
highest_queue_len = max(highest_queue_len, queue_len)
lowest_queue_len = min(lowest_queue_len, queue_len)
is_imbalanced = (
highest_queue_len - lowest_queue_len > self._imbalanced_threshold
)
# End Sphinx tag: __end_load_balance_component__
# Start Sphinx tag: __begin_prefix_match_component__
if not is_imbalanced:
# Convert candidate replica IDs to strings for prefix matching.
candidate_replica_ids_strings = [
r.replica_id.to_full_id_str() for r in candidate_replicas
]
(matched_text, matched_tenant_id_strings,) = ray.get(
self._tree_actor.prefix_match.remote(
input_text, candidate_replica_ids_strings
)
)
match_rate = len(matched_text) / len(input_text)
if match_rate < self._match_rate_threshold:
smallest_tenants_id_strings = ray.get(
self._tree_actor.get_smallest_tenants.remote()
)
if (
smallest_tenants_id_strings is not None
and len(smallest_tenants_id_strings) > 0
):
chosen_replica_id_strings = smallest_tenants_id_strings
else:
if (
matched_tenant_id_strings is not None
and len(matched_tenant_id_strings) > 0
):
chosen_replica_id_strings = matched_tenant_id_strings
# End Sphinx tag: __end_prefix_match_component__
return [
[
self._replicas[ReplicaID.from_full_id_str(chosen_id_string)]
for chosen_id_string in chosen_replica_id_strings
]
]
# Start Sphinx tag: __begin_on_replica_actor_died__
def on_replica_actor_died(self, replica_id: ReplicaID):
"""Drop replica from replica set so it's not considered for future requests."""
super().on_replica_actor_died(replica_id)
ray.get(self._tree_actor.remove_tenants.remote([replica_id.to_full_id_str()]))
# End Sphinx tag: __end_on_replica_actor_died__
def update_replicas(self, replicas: List[RunningReplica]):
"""Update the set of available replicas to be considered for routing.
When the set of replicas changes, we may spawn additional routing tasks
if there are pending requests.
"""
# 1) Record the old replica IDs
old_ids = set(self._replica_id_set)
# 2) Run the default update_replicas logic
super().update_replicas(replicas)
# 3) Figure out which replicas were added / removed
new_ids = set(self._replica_id_set)
added = new_ids - old_ids
removed = old_ids - new_ids
# 4) Update the prefix tree with the changes
if added:
added_strings = [rid.to_full_id_str() for rid in added]
ray.get(self._tree_actor.add_tenants.remote(added_strings, time.time()))
if removed:
removed_strings = [rid.to_full_id_str() for rid in removed]
ray.get(self._tree_actor.remove_tenants.remote(removed_strings))
# === Start tasks (if enabled and not already running) ===
if self._do_eviction and not self._eviction_loop_running:
ray.get(
self._tree_actor.start_eviction_loop.remote(
self._eviction_threshold_chars,
self._eviction_target_chars,
self._eviction_interval_secs,
)
)
self._eviction_loop_running = True
async def choose_replicas(
self,
candidate_replicas: List[RunningReplica],
pending_request: Optional[PendingRequest] = None,
) -> List[RunningReplica]:
"""One iteration of the power of two choices procedure that chooses
(at most) two random available replicas.
For multiplexing, this will first attempt to choose replicas that have the
requested model ID for a configured timeout. If no replicas with the matching
model ID are available after that timeout, it will fall back to the regular
procedure.
"""
# Start Sphinx tag: __begin_pow2_router_base__
# Get fallback replicas from PowerOfTwoChoicesRequestRouter
fallback_replicas = await PowerOfTwoChoicesRequestRouter.choose_replicas(
self,
candidate_replicas=candidate_replicas,
pending_request=pending_request,
)
if pending_request is None or not fallback_replicas:
return fallback_replicas
# End Sphinx tag: __end_pow2_router_base__
if (
pending_request is not None
and pending_request.metadata.multiplexed_model_id
):
# Get candidates for multiplexed model ID.
candidate_replica_ids = self.apply_multiplex_routing(
pending_request=pending_request,
)
else:
# Get candidates for locality preference.
candidate_replica_ids = self.apply_locality_routing(
pending_request=pending_request,
)
if not candidate_replica_ids:
return fallback_replicas
# Convert candidate replica IDs to RunningReplica objects.
replica_id_to_replica_map = {
replica.replica_id: replica for replica in candidate_replicas
}
candidate_replicas = [
replica_id_to_replica_map[candidate_replica_id]
for candidate_replica_id in candidate_replica_ids
]
chosen_replicas = await self._prefix_match_best_replicas(
pending_request, candidate_replicas
)
if chosen_replicas[0]:
return chosen_replicas
return fallback_replicas
# Start Sphinx tag: __begin_on_request_routed__
def on_request_routed(
self,
pending_request: PendingRequest,
replica_id: ReplicaID,
result: ReplicaResult,
):
"""Called when a request is routed to a replica.
This is used as a callback to update the state of the request router
after a response is generated.
"""
# Right now this only inserts the prompt into the prefix tree, not the response (streaming response makes things complicated)
if (
pending_request is not None
and pending_request.args is not None
and len(pending_request.args) > 0
):
input_text = self._extract_text_from_request(pending_request)
if input_text is not None:
# Insert into prefix tree
ray.get(
self._tree_actor.insert.remote(
input_text, replica_id.to_full_id_str(), time.time()
)
)
# End Sphinx tag: __end_on_request_routed__
python/ray/llm/_internal/serve/routing_policies/prefix_aware/prefix_tree.py
+620
View File
@@ -0,0 +1,620 @@
from __future__ import annotations
import logging
import os
import threading
from threading import RLock
from typing import Any, Dict, List, Optional, Tuple
import ray
from ray.serve._private.constants import (
SERVE_LOGGER_NAME,
)
logger = logging.getLogger(SERVE_LOGGER_NAME)
class Node:
"""
Node in a prefix tree that represents a segment of text and can belong to multiple tenants.
Each node also tracks the last access time for each tenant.
Simple example of root node connected to two children Nodes:
root = Node(text="", parent=None, edge_label_to_child={"f": fooNode, "b": barNode}, tenant_to_last_access_time={"tenant_1": 2})
fooNode = Node(text="foo", parent=root, edge_label_to_child={}, tenant_to_last_access_time={"tenant_1": 1})
barNode = Node(text="bar", parent=root, edge_label_to_child={}, tenant_to_last_access_time={"tenant_1": 2})
In the above example, "foo" was inserted at time 1, and "bar" was inserted at time 2.
It follows that root was last accessed at time 2.
"""
def __init__(self, text: str = "", parent: Optional[Node] = None) -> None:
"""
Initialize a node in the prefix tree.
Args:
text: The text segment this node represents
parent: The parent node of this node
"""
self.text: str = text
self.parent: Optional[Node] = parent
# Maps first character to child node
self.edge_label_to_child: Dict[str, Node] = {}
# For each tenant that has inserted text matching this node, track its last access timestamp (in seconds)
self.tenant_to_last_access_time: Dict[str, float] = {}
# Doubly linked list pointers for LRU tracking per tenant
# Points to the less recently used node (toward tail)
self.tenant_to_older_node: Dict[str, Optional[Node]] = {}
# Points to the more recently used node (toward head)
self.tenant_to_newer_node: Dict[str, Optional[Node]] = {}
class PrefixTree:
"""
Thread-safe multi-tenant prefix tree (approximate radix tree).
Features:
1. Stores data for multiple tenants in the same tree structure
2. Thread-safe with node-level locking for concurrent access
3. LRU eviction based on tenant access time
4. Efficient prefix matching across multiple tenants
Example tree structure:
Representing the strings inserted in order:
- "helloworld" at time 1 by tenant_1
- "hellothere" at time 2 by tenant_2
- "hellothomas" at time 3 by tenant_2
root: [] {tenant_1: 1, tenant_2: 3}
(h) → [hello] {tenant_1: 1, tenant_2: 3}
(w) → [world] {tenant_1: 1}
(t) → [th] {tenant_2: 3}
(e) → [ere] {tenant_2: 2}
(o) → [omas] {tenant_2: 3}
Legend for each node:
- [text] = Node.text
- {tenant, timestamp} = Node.tenant_to_last_access_time
- (x) = edge label (first character used as key for parent's children)
PrefixTree instance variables:
self.tenant_to_char_count = {"tenant_1": 10, "tenant_2": 14}
self.tenant_to_lru_tail = {"tenant_1": Node("world"), "tenant_2": Node("ere")}
"""
def __init__(self) -> None:
"""Initialize an empty prefix tree."""
self.lock: RLock = RLock()
# Root is always the head of the LRU list for each tenant.
self.root: Node = Node()
# Tracks total character count per tenant. Can be used by the replica request router to determine which tenant to evict, and by how much.
# Also uses the keys to track the active tenants in the tree.
self.tenant_to_char_count: Dict[str, int] = {}
# LRU tracking - root is always the head, tail is the least recently used.
self.tenant_to_lru_tail: Dict[str, Optional[Node]] = {}
self._eviction_thread: Optional[threading.Thread] = None
self._eviction_stop_event: threading.Event = threading.Event()
@staticmethod
def _shared_prefix_count(a: str, b: str) -> int:
"""
Count the number of shared characters at the beginning of two strings.
Args:
a: First string
b: Second string
Returns:
Number of matching characters at the beginning
"""
return len(os.path.commonprefix([a, b]))
def _get_lru_chain(self, tenant: str) -> List[Node]:
"""
Get the LRU chain for a given tenant by traversing from the root to the oldest node.
Note: This method is intended to be used only in tests.
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
return []
nodes = []
current_node = self.root
while current_node:
nodes.append(current_node)
current_node = current_node.tenant_to_older_node.get(tenant)
return nodes
def _insert_node_into_linked_list(
self,
node: Node,
newer_neighbor: Optional[Node],
older_neighbor: Optional[Node],
tenant: str,
) -> None:
"""
Insert a node into the LRU list between two neighbors. Updates the neighbors' pointers and the tail pointer, if that changes.
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
return
# Skip if node is the root
if node == self.root:
return
node.tenant_to_newer_node[tenant] = newer_neighbor
node.tenant_to_older_node[tenant] = older_neighbor
if newer_neighbor:
newer_neighbor.tenant_to_older_node[tenant] = node
if older_neighbor:
older_neighbor.tenant_to_newer_node[tenant] = node
if self.tenant_to_lru_tail[tenant] == newer_neighbor:
self.tenant_to_lru_tail[tenant] = node
def _remove_node_from_linked_list(self, node: Node, tenant: str) -> None:
"""
Remove a node from the LRU list. Updates the neighbors' pointers and the tail pointer, if that changes.
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
return
# Skip if node is the root
if node == self.root:
return
# Connect older and newer neighbors
older_neighbor = node.tenant_to_older_node.get(tenant)
newer_neighbor = node.tenant_to_newer_node.get(tenant)
if older_neighbor:
older_neighbor.tenant_to_newer_node[tenant] = newer_neighbor
if newer_neighbor:
newer_neighbor.tenant_to_older_node[tenant] = older_neighbor
# Update tail pointer if necessary
if self.tenant_to_lru_tail[tenant] == node:
self.tenant_to_lru_tail[tenant] = newer_neighbor
# Remove node from list
node.tenant_to_newer_node.pop(tenant, None)
node.tenant_to_older_node.pop(tenant, None)
def _remove_tenant_single_node(self, tenant: str, node: Node) -> int:
"""
Remove a tenant from a single node.
Args:
tenant: Tenant to remove
node: Node to remove tenant from
Returns:
Number of characters removed (0 if preconditions not met)
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
logger.debug(f"Tenant '{tenant}' does not exist. No action taken.")
return 0
if tenant not in node.tenant_to_last_access_time:
logger.debug(
f"Tenant '{tenant}' does not have node '{node.text}'. No action taken."
)
return 0
removed_chars_len: int = len(node.text)
self.tenant_to_char_count[tenant] -= removed_chars_len
node.tenant_to_last_access_time.pop(tenant, None)
self._remove_node_from_linked_list(node, tenant)
# Clean up empty nodes
if not node.tenant_to_last_access_time and node.parent:
if (
node.text and node.text[0] in node.parent.edge_label_to_child
): # Defensive check
node.parent.edge_label_to_child.pop(node.text[0], None)
return removed_chars_len
def add_tenants(self, tenants: List[str], time_s: float) -> None:
"""
Add multiple new tenants to the tree. Also inserts an empty string for each tenant into the tree.
For each tenant that already exists, a warning is logged and that tenant is skipped.
Args:
tenants: List of tenants to add
time_s: Current timestamp in seconds
"""
with self.lock:
for tenant in tenants:
if tenant in self.tenant_to_char_count:
logger.debug(f"Tenant '{tenant}' already exists. Skipping.")
continue
self.tenant_to_char_count[tenant] = 0
self.tenant_to_lru_tail[tenant] = self.root
# Initialize the root node as the head of the LRU list for this tenant
self.root.tenant_to_newer_node[tenant] = None
self.root.tenant_to_older_node[tenant] = None
self.insert("", tenant, time_s)
def insert(self, text: str, tenant: str, time_s: float) -> None:
"""
Insert text into tree for a specific tenant, but only if the tenant already exists.
If the tenant doesn't exist in the tree, this will log a warning and return without
inserting anything. Use add_tenants() first to add a new tenant.
Args:
text: Text to insert
tenant: Tenant
time_s: Current timestamp in seconds
Loop structure:
1. We update the current node at the start of each iteration of the while loop.
This includes updating tenant_to_char_count and tenant_to_last_access_time, and moving the node to the front of the LRU list.
2. Each iteration then either:
a. Breaks (if we've processed the entire string).
b. Processes the next segment of text by:
1. If no child exists for the first character, create a new leaf node that matches the current text.
2. Then, match the current text with the child's text:
a. If they share a prefix (partial match), split the node and traverse into the new parent.
b. If they fully match, traverse into the child node.
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
logger.debug(
f"Tenant '{tenant}' does not exist. Use add_tenants() first."
)
return
curr_node: Node = self.root
i: int = 0
while i <= len(text):
# Invariant at beginning of each iteration: assume curr_node has not been visited by tenant yet.
# Update tenant info for current node.
if tenant not in curr_node.tenant_to_last_access_time:
self.tenant_to_char_count[tenant] += len(curr_node.text)
curr_node.tenant_to_last_access_time[tenant] = time_s
if curr_node != self.root:
self._remove_node_from_linked_list(curr_node, tenant)
self._insert_node_into_linked_list(
curr_node,
self.root,
self.root.tenant_to_older_node.get(tenant),
tenant,
)
if i == len(text):
break
first_char: str = text[i]
curr_text: str = text[i:]
if first_char not in curr_node.edge_label_to_child:
# No match, create new node. Don't update new node as "visited" by tenant yet; it will be done at the beginning of the next iteration.
# e.g. curr_node.edge_label_to_child = {}, curr_text = "hello" -> curr_node.edge_label_to_child = {"h": Node("hello")}
new_node: Node = Node(text=curr_text, parent=curr_node)
curr_node.edge_label_to_child[first_char] = new_node
# Add the node to the back of the LRU list; it will be moved to the front in the next iteration.
self._insert_node_into_linked_list(
new_node, self.tenant_to_lru_tail[tenant], None, tenant
)
# Match found, check if we need to split
matched_node: Node = curr_node.edge_label_to_child[first_char]
shared_count: int = self._shared_prefix_count(
matched_node.text, curr_text
)
if shared_count < len(matched_node.text):
# Partial match, split node at matched point
# Example:
## Before update:
### curr_node.edge_label_to_child = {"h": Node("helloworld")}, curr_text = "hellothere" -> shared_count = 5
### matched_node = Node("helloworld")
## After update:
### curr_node.edge_label_to_child = {"h": Node("hello", edge_label_to_child = {"w": Node("world")})}
### parent_node = Node("hello"), matched_node = Node("world")
### Copy matched_node.tenant_to_last_access_time to parent_node.tenant_to_last_access_time
### Insert parent_node into the back of the LRU list; it will be moved to the front in the next iteration. (for the current tenant)
### Insert parent_node between matched_node and matched_node's newer neighbor (for all other tenants)
### (new) curr_text = "there", (new) curr_node = parent_node
### Continue adding "there" to tree in next iteration
matched_text: str = matched_node.text[:shared_count]
remaining_text: str = matched_node.text[shared_count:]
# Create new intermediate node
# Note that we don't update new_parent.tenant_to_last_access_time yet; it will be done at the beginning of the next iteration.
new_parent: Node = Node(text=matched_text, parent=curr_node)
new_parent.tenant_to_last_access_time = (
matched_node.tenant_to_last_access_time.copy()
)
# Insert new_parent into the back of the LRU list; it will be moved to the front in the next iteration. (for the current tenant)
self._insert_node_into_linked_list(
new_parent, self.tenant_to_lru_tail[tenant], None, tenant
)
# Insert new_parent between matched_node and matched_node's newer neighbor (for all other tenants)
for existing_tenant in new_parent.tenant_to_last_access_time:
if existing_tenant != tenant:
self._insert_node_into_linked_list(
new_parent,
matched_node.tenant_to_newer_node.get(existing_tenant),
matched_node,
existing_tenant,
)
# Update existing matched node
matched_node.text = remaining_text
matched_node.parent = new_parent
# Connect nodes
new_parent.edge_label_to_child[remaining_text[0]] = matched_node
curr_node.edge_label_to_child[first_char] = new_parent
# Continue traversal
curr_node = new_parent
i += shared_count
else:
# Full match, continue down the tree
curr_node = matched_node
i += shared_count
def prefix_match(
self, text: str, available_tenants: Optional[List[str]] = None
) -> Tuple[str, Optional[List[str]]]:
"""
Match text against tree and return matched text and matching tenants.
Args:
text: Text to match
available_tenants: List of tenants to match against (or None for all)
Returns:
Tuple of (matched_text, matched_tenants):
If the list of available tenants doesn't match any tenants in the tree: returns ("", None)
When no prefix match is found (does not traverse further than the root node): returns ("", list of available tenants)
When a prefix match is found: returns (matched_prefix, list of tenants that own the matched node)
"""
with self.lock:
if available_tenants:
# Filter available_tenants to only include those in the tree
available_tenants = [
tenant
for tenant in available_tenants
if tenant in self.tenant_to_char_count
]
if not available_tenants:
return "", None
else:
available_tenants = list(self.tenant_to_char_count.keys())
curr_node: Node = self.root
i: int = 0
text_len: int = len(text)
while i < text_len:
first_char: str = text[i]
curr_text: str = text[i:]
if first_char in curr_node.edge_label_to_child:
matched_node: Node = curr_node.edge_label_to_child[first_char]
# Check if any available tenants match this node
if not any(
tenant in matched_node.tenant_to_last_access_time
for tenant in available_tenants
):
break
shared_count: int = self._shared_prefix_count(
matched_node.text, curr_text
)
i += shared_count
curr_node = matched_node
if shared_count < len(matched_node.text):
# Partial match, stop here
break
else:
# No match found, stop here
break
# Find tenants in current node that match available tenants
matched_tenants = [
tenant
for tenant in available_tenants
if tenant in curr_node.tenant_to_last_access_time
] or None
matched_text: str = text[:i]
return matched_text, matched_tenants
def remove_tenants(self, tenants: List[str]) -> Dict[str, int]:
"""
Remove multiple tenants and all their nodes from the tree.
Time complexity: O(n) where n is the total number of nodes owned by all tenants.
Args:
tenants: List of tenants to remove
Returns:
Dictionary mapping each tenant to the number of characters removed
(0 if tenant doesn't exist)
"""
chars_removed: Dict[str, int] = {}
with self.lock:
for tenant in tenants:
if tenant not in self.tenant_to_char_count:
logger.debug(f"Tenant '{tenant}' does not exist. Skipping.")
chars_removed[tenant] = 0
continue
tenant_chars_removed: int = 0
# Start from the tail and remove all nodes
current_tail = self.tenant_to_lru_tail.get(tenant)
while current_tail:
newer_neighbor = current_tail.tenant_to_newer_node.get(tenant)
tenant_chars_removed += self._remove_tenant_single_node(
tenant, current_tail
)
current_tail = newer_neighbor
# Clean up tenant references
self.tenant_to_char_count.pop(tenant, None)
self.tenant_to_lru_tail.pop(tenant, None)
chars_removed[tenant] = tenant_chars_removed
return chars_removed
def evict_tenant_by_lru(self, tenant: str, min_remove_size: int) -> int:
"""
Evict least recently used nodes for a tenant until minimum size is freed.
Time complexity: O(m) where m is the number of nodes removed.
Args:
tenant: The tenant to evict nodes from
min_remove_size: Minimum number of characters to remove
Returns:
Actual number of characters removed (0 if tenant doesn't exist)
Note:
- All nodes with the same oldest access time are removed together to maintain tree integrity, even if only removing a subset of them satisfies the min_remove_size.
- For more predictable eviction, use unique timestamps for each insertion.
- The root node is never evicted as it serves as the permanent head of the LRU list.
"""
with self.lock:
if tenant not in self.tenant_to_char_count:
logger.debug(
f"Cannot evict tenant '{tenant}': tenant does not exist. No action taken."
)
return 0
if self.tenant_to_char_count[tenant] < min_remove_size:
logger.debug(
f"Cannot evict {min_remove_size} characters from tenant '{tenant}', which has only "
f"{self.tenant_to_char_count[tenant]} characters. Will remove all available characters."
)
min_remove_size = self.tenant_to_char_count[tenant]
total_chars_removed: int = 0
# Start removing from the tail (least recently used)
current_tail = self.tenant_to_lru_tail.get(tenant)
# Continue until we've freed enough space or run out of nodes
while total_chars_removed < min_remove_size and current_tail:
# Stop if we've reached the root - the root is never evicted
if current_tail == self.root:
break
# Get the current timestamp to remove all nodes with this timestamp
current_timestamp = current_tail.tenant_to_last_access_time[tenant]
# Collect all nodes with the same timestamp (guaranteed to be contiguous in our LRU list)
while (
current_tail != self.root # Never include the root
and current_tail.tenant_to_last_access_time[tenant]
== current_timestamp
):
newer_neighbor = current_tail.tenant_to_newer_node.get(tenant)
total_chars_removed += self._remove_tenant_single_node(
tenant, current_tail
)
current_tail = newer_neighbor
return total_chars_removed
def get_smallest_tenants(self) -> Optional[List[str]]:
"""
Get the tenants with the smallest total character count.
Returns:
Tenants with smallest character count, or None if no tenants
"""
with self.lock:
if not self.tenant_to_char_count:
return None
min_count = min(self.tenant_to_char_count.values())
return [
tenant
for tenant, count in self.tenant_to_char_count.items()
if count == min_count
]
def start_eviction_loop(
self, eviction_threshold: int, eviction_target: int, interval_secs: float
) -> bool:
"""Start a single eviction loop within the actor itself.
Args:
eviction_threshold: Minimum number of characters a tenant must have to be evicted
eviction_target: The maximum number of characters a tenant should have after eviction
interval_secs: Number of seconds between eviction checks
Returns:
True if the loop was started, False if it was already running
"""
self._eviction_stop_event.clear()
with self.lock:
if self._eviction_thread is None:
self._eviction_thread = threading.Thread(
target=self._run_eviction_loop,
args=(eviction_threshold, eviction_target, interval_secs),
daemon=True,
)
self._eviction_thread.start()
return True
else:
logger.debug("Eviction loop already running")
return False
def _run_eviction_loop(self, eviction_threshold, eviction_target, interval_secs):
while not self._eviction_stop_event.is_set():
if self._eviction_stop_event.wait(interval_secs):
# Stop event was set, exit loop asap
break
with self.lock:
for tenant, char_count in self.tenant_to_char_count.items():
if char_count > eviction_threshold:
excess = char_count - eviction_target
self.evict_tenant_by_lru(tenant, excess)
def stop_eviction_loop(self):
self._eviction_stop_event.set()
if self._eviction_thread:
self._eviction_thread.join()
self._eviction_thread = None
@ray.remote
class PrefixTreeActor(PrefixTree):
def getattr(self, attribute: str) -> Any:
"""
Get an attribute of the PrefixTree.
Note: This method is intended to be used only in tests.
"""
return getattr(self, attribute)
def setattr(self, attribute: str, value: Any) -> None:
setattr(self, attribute, value)