ray-project--ray/python/ray/serve/_private/router.py

import asyncio
import concurrent.futures
import logging
import sys
import threading
import time
import weakref
from abc import ABC, abstractmethod
from asyncio import AbstractEventLoop, ensure_future, futures
from collections import defaultdict
from collections.abc import MutableMapping
from contextlib import asynccontextmanager, contextmanager
from dataclasses import replace
from functools import lru_cache, partial
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Coroutine,
    DefaultDict,
    Dict,
    List,
    Optional,
    Tuple,
    Union,
)

import ray
from ray.actor import ActorHandle
from ray.exceptions import (
    ActorDiedError,
    ActorUnavailableError,
    RayError,
    RayTaskError,
    TaskCancelledError,
)
from ray.serve._private.common import (
    RUNNING_REQUESTS_KEY,
    DeploymentHandleSource,
    DeploymentID,
    DeploymentTargetInfo,
    HandleMetricReport,
    ReplicaID,
    RequestMetadata,
    RunningReplicaInfo,
)
from ray.serve._private.config import DeploymentConfig
from ray.serve._private.constants import (
    DEFAULT_LATENCY_BUCKET_MS,
    RAY_SERVE_AUTOSCALING_METRIC_RECORD_INTERVAL_FACTOR,
    RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE,
    RAY_SERVE_METRICS_EXPORT_INTERVAL_MS,
    RAY_SERVE_PROXY_PREFER_LOCAL_AZ_ROUTING,
    SERVE_LOGGER_NAME,
)
from ray.serve._private.constants_utils import warn_if_deprecated_env_var_set
from ray.serve._private.event_loop_monitoring import EventLoopMonitor
from ray.serve._private.long_poll import LongPollClient, LongPollNamespace
from ray.serve._private.metrics_utils import (
    QUEUED_REQUESTS_KEY,
    InMemoryMetricsStore,
    MetricsPusher,
    TimeStampedValue,
)
from ray.serve._private.replica_result import ReplicaResult
from ray.serve._private.request_router import PendingRequest, RequestRouter
from ray.serve._private.request_router.pow_2_router import (
    PowerOfTwoChoicesRequestRouter,
)
from ray.serve._private.request_router.replica_wrapper import (
    ReplicaSelection,
    RunningReplica,
)
from ray.serve._private.tracing_utils import (
    create_propagated_context,
    is_span_recording,
    set_http_span_attributes,
    set_rpc_span_attributes,
    set_span_attributes,
    set_span_exception,
    set_span_name,
    tracing_decorator_factory,
)
from ray.serve._private.usage import ServeUsageTag
from ray.serve._private.utils import (
    check_obj_ref_ready_nowait,
    compress_metric_report,
    generate_request_id,
    resolve_deployment_response,
)
from ray.serve.config import AutoscalingConfig
from ray.serve.exceptions import (
    BackPressureError,
    DeploymentUnavailableError,
    RayServeException,
    ReplicaUnavailableError,
)
from ray.types import ObjectRef
from ray.util import metrics

logger = logging.getLogger(SERVE_LOGGER_NAME)


class RouterMetricsManager:
    """Manages metrics for the router."""

    PUSH_METRICS_TO_CONTROLLER_TASK_NAME = "push_metrics_to_controller"
    RECORD_METRICS_TASK_NAME = "record_metrics"

    def __init__(
        self,
        deployment_id: DeploymentID,
        handle_id: str,
        self_actor_id: str,
        handle_source: DeploymentHandleSource,
        controller_handle: ActorHandle,
        router_requests_counter: metrics.Counter,
        queued_requests_gauge: metrics.Gauge,
        running_requests_gauge: metrics.Gauge,
        event_loop: asyncio.BaseEventLoop,
    ):
        self._handle_id = handle_id
        self._deployment_id = deployment_id
        self._self_actor_id = self_actor_id
        self._handle_source = handle_source
        self._controller_handle = controller_handle

        # Exported metrics
        self.num_router_requests = router_requests_counter
        self.num_router_requests.set_default_tags(
            {
                "deployment": deployment_id.name,
                "application": deployment_id.app_name,
                "handle": self._handle_id,
                "actor_id": self._self_actor_id,
            }
        )

        self.num_queued_requests = 0
        self.num_queued_requests_gauge = queued_requests_gauge
        self.num_queued_requests_gauge.set_default_tags(
            {
                "deployment": deployment_id.name,
                "application": deployment_id.app_name,
                "handle": self._handle_id,
                "actor_id": self._self_actor_id,
            }
        )
        self.num_queued_requests_gauge.set(0)

        # Track queries sent to replicas for the autoscaling algorithm.
        self.num_requests_sent_to_replicas: DefaultDict[ReplicaID, int] = defaultdict(
            int
        )
        self.num_running_requests_gauge = running_requests_gauge
        self.num_running_requests_gauge.set_default_tags(
            {
                "deployment": deployment_id.name,
                "application": deployment_id.app_name,
                "handle": self._handle_id,
                "actor_id": self._self_actor_id,
            }
        )
        # We use Ray object ref callbacks to update state when tracking
        # number of requests running on replicas. The callbacks will be
        # called from a C++ thread into the router's async event loop,
        # so non-atomic read and write operations need to be guarded by
        # this thread-safe lock.
        self._queries_lock = threading.Lock()

        # Track reserved slots for choose_replica operations
        self._num_reserved_slots = 0
        self._reserved_slots_gauge = metrics.Gauge(
            "serve_reserved_slots_active",
            description=(
                "The current number of reserved slots for choose_replica operations."
            ),
            tag_keys=("deployment", "application", "handle", "actor_id"),
        )
        self._reserved_slots_gauge.set_default_tags(
            {
                "deployment": deployment_id.name,
                "application": deployment_id.app_name,
                "handle": self._handle_id,
                "actor_id": self._self_actor_id,
            }
        )
        self._reserved_slots_gauge.set(0)
        # Regularly aggregate and push autoscaling metrics to controller
        self.metrics_pusher = MetricsPusher()
        self.metrics_store = InMemoryMetricsStore()
        # The config for the deployment this router sends requests to will be broadcast
        # by the controller. That means it is not available until we get the first
        # update. This includes an optional autoscaling config.
        self._deployment_config: Optional[DeploymentConfig] = None
        # Track whether the metrics manager has been shutdown
        self._shutdown: bool = False

        # Tracks in-flight metrics push to controller. Skip if new one is sent.
        self._pending_metrics_push_ref: Optional[ObjectRef] = None
        self._metrics_push_lock = threading.Lock()

        # If the interval is set to 0, eagerly sets all metrics.
        self._cached_metrics_enabled = RAY_SERVE_METRICS_EXPORT_INTERVAL_MS != 0
        self._cached_metrics_interval_s = RAY_SERVE_METRICS_EXPORT_INTERVAL_MS / 1000
        self._cached_metrics_task: Optional[asyncio.Task] = None

        if self._cached_metrics_enabled:
            self._cached_num_router_requests = defaultdict(int)

            def create_metrics_task():
                self._cached_metrics_task = event_loop.create_task(
                    self._report_cached_metrics_forever()
                )

            # the constructor is called in the user thread, but its trying to create a task on the event loop
            # which is running in the router thread. This is not thread safe, so we need to use call_soon_threadsafe
            # to create the task on the event loop thread safely.
            event_loop.call_soon_threadsafe(create_metrics_task)

    @contextmanager
    def wrap_request_assignment(self, request_meta: RequestMetadata):
        max_queued_requests = (
            self._deployment_config.max_queued_requests
            if self._deployment_config is not None
            else -1
        )
        if (
            max_queued_requests != -1
            and self.num_queued_requests >= max_queued_requests
        ):
            # Due to the async nature of request handling, we may reject more requests
            # than strictly necessary. This is more likely to happen during
            # high concurrency. Here's why:
            #
            # When multiple requests arrive simultaneously with max_queued_requests=1:
            # 1. First request increments num_queued_requests to 1
            # 2. Before that request gets assigned to a replica and decrements the counter,
            #    we yield to the event loop
            # 3. Other requests see num_queued_requests=1 and get rejected, even though
            #    the first request will soon free up the queue slot
            #
            # For example, with max_queued_requests=1 and 4 simultaneous requests:
            # - Request 1 gets queued (num_queued_requests=1)
            # - Requests 2,3,4 get rejected since queue appears full
            # - Request 1 gets assigned and frees queue slot (num_queued_requests=0)
            # - But we already rejected Request 2 which could have been queued
            e = BackPressureError(
                num_queued_requests=self.num_queued_requests,
                max_queued_requests=max_queued_requests,
            )
            logger.warning(e.message)
            raise e

        self.inc_num_total_requests(request_meta.route)
        yield

    @contextmanager
    def wrap_queued_request(self, is_retry: bool, num_curr_replicas: int):
        """Increment queued requests gauge and maybe push autoscaling metrics to controller."""
        try:
            self.inc_num_queued_requests()
            # Optimization: if there are currently zero replicas for a deployment,
            # push handle metric to controller to allow for fast cold start time.
            # Only do this on the first attempt to route the request.
            if not is_retry and self.should_send_scaled_to_zero_optimized_push(
                curr_num_replicas=num_curr_replicas
            ):
                self.push_autoscaling_metrics_to_controller()

            yield
        finally:
            # If the request is disconnected before assignment, this coroutine
            # gets cancelled by the caller and an asyncio.CancelledError is
            # raised. The finally block ensures that num_queued_requests
            # is correctly decremented in this case.
            self.dec_num_queued_requests()

    def _update_running_replicas(self, running_replicas: List[RunningReplicaInfo]):
        """Prune list of replica ids in self.num_queries_sent_to_replicas.

        We want to avoid self.num_queries_sent_to_replicas from growing
        in memory as the deployment upscales and downscales over time.
        """

        running_replica_set = {replica.replica_id for replica in running_replicas}
        with self._queries_lock:
            self.num_requests_sent_to_replicas = defaultdict(
                int,
                {
                    id: self.num_requests_sent_to_replicas[id]
                    for id, num_queries in self.num_requests_sent_to_replicas.items()
                    if num_queries or id in running_replica_set
                },
            )

    @property
    def autoscaling_config(self) -> Optional[AutoscalingConfig]:
        if self._deployment_config is None:
            return None

        return self._deployment_config.autoscaling_config

    def update_deployment_config(
        self, deployment_config: DeploymentConfig, curr_num_replicas: int
    ):
        """Update the config for the deployment this router sends requests to."""

        if self._shutdown:
            return

        self._deployment_config = deployment_config

        # Start the metrics pusher if autoscaling is enabled.
        autoscaling_config = self.autoscaling_config
        if autoscaling_config:
            self.metrics_pusher.start()
            # Optimization for autoscaling cold start time. If there are
            # currently 0 replicas for the deployment, and there is at
            # least one queued request on this router, then immediately
            # push handle metric to the controller.
            if self.should_send_scaled_to_zero_optimized_push(curr_num_replicas):
                self.push_autoscaling_metrics_to_controller()

            # Record number of queued + ongoing requests at regular
            # intervals into the in-memory metrics store
            record_interval_s = (
                autoscaling_config.look_back_period_s
                * RAY_SERVE_AUTOSCALING_METRIC_RECORD_INTERVAL_FACTOR
            )
            self.metrics_pusher.register_or_update_task(
                self.RECORD_METRICS_TASK_NAME,
                self._add_autoscaling_metrics_point,
                min(record_interval_s, autoscaling_config.metrics_interval_s),
            )
            # Push metrics to the controller periodically.
            self.metrics_pusher.register_or_update_task(
                self.PUSH_METRICS_TO_CONTROLLER_TASK_NAME,
                self.push_autoscaling_metrics_to_controller,
                autoscaling_config.metrics_interval_s,
            )

        else:
            if self.metrics_pusher:
                self.metrics_pusher.stop_tasks()

    def _report_cached_metrics(self):
        for route, count in self._cached_num_router_requests.items():
            self.num_router_requests.inc(count, tags={"route": route})
        self._cached_num_router_requests.clear()

        self.num_queued_requests_gauge.set(self.num_queued_requests)

        self.num_running_requests_gauge.set(
            sum(self.num_requests_sent_to_replicas.values())
        )

    async def _report_cached_metrics_forever(self):
        assert self._cached_metrics_interval_s > 0

        consecutive_errors = 0
        while True:
            try:
                await asyncio.sleep(self._cached_metrics_interval_s)
                self._report_cached_metrics()
                consecutive_errors = 0
            except Exception:
                logger.exception("Unexpected error reporting metrics.")

                # Exponential backoff starting at 1s and capping at 10s.
                backoff_time_s = min(10, 2**consecutive_errors)
                consecutive_errors += 1
                await asyncio.sleep(backoff_time_s)

    def inc_num_total_requests(self, route: str):
        if self._cached_metrics_enabled:
            self._cached_num_router_requests[route] += 1
        else:
            self.num_router_requests.inc(tags={"route": route})

    def inc_num_queued_requests(self):
        self.num_queued_requests += 1
        if not self._cached_metrics_enabled:
            self.num_queued_requests_gauge.set(self.num_queued_requests)

    def dec_num_queued_requests(self):
        self.num_queued_requests -= 1
        if not self._cached_metrics_enabled:
            self.num_queued_requests_gauge.set(self.num_queued_requests)

    def inc_reserved_slots(self):
        self._num_reserved_slots += 1
        self._reserved_slots_gauge.set(self._num_reserved_slots)

    def dec_reserved_slots(self):
        self._num_reserved_slots -= 1
        self._reserved_slots_gauge.set(self._num_reserved_slots)

    def inc_num_running_requests_for_replica(self, replica_id: ReplicaID):
        with self._queries_lock:
            self.num_requests_sent_to_replicas[replica_id] += 1
            if not self._cached_metrics_enabled:
                self.num_running_requests_gauge.set(
                    sum(self.num_requests_sent_to_replicas.values())
                )

    def dec_num_running_requests_for_replica(self, replica_id: ReplicaID):
        with self._queries_lock:
            self.num_requests_sent_to_replicas[replica_id] -= 1
            if not self._cached_metrics_enabled:
                self.num_running_requests_gauge.set(
                    sum(self.num_requests_sent_to_replicas.values())
                )

    def should_send_scaled_to_zero_optimized_push(self, curr_num_replicas: int) -> bool:
        return (
            self.autoscaling_config is not None
            and curr_num_replicas == 0
            and self.num_queued_requests > 0
        )

    def push_autoscaling_metrics_to_controller(self):
        """Pushes queued and running request metrics to the controller.

        These metrics are used by the controller for autoscaling.
        If a previous push is already in flight, skips this push (will try again next interval).
        """
        with self._metrics_push_lock:
            if self._pending_metrics_push_ref is not None:
                if not check_obj_ref_ready_nowait(self._pending_metrics_push_ref):
                    return  # Previous push still in flight, skip and try again later
            self._pending_metrics_push_ref = (
                self._controller_handle.record_autoscaling_metrics_from_handle.remote(
                    compress_metric_report(self._get_metrics_report())
                )
            )

    def _add_autoscaling_metrics_point(self):
        """Adds metrics point for queued and running requests at replicas.

        Also prunes keys in the in memory metrics store with outdated datapoints.

        ┌─────────────────────────────────────────────────────────────────┐
        │  Handle-based metrics collection                                │
        ├─────────────────────────────────────────────────────────────────┤
        │                                                                 │
        │  Client                Handle              Replicas             │
        │  ┌──────┐            ┌────────┐          ┌─────────┐           │
        │  │  App │───────────>│ Handle │─────────>│ Replica │           │
        │  │      │  Requests  │        │ Forwards │    1    │           │
        │  └──────┘            │ Tracks │          └─────────┘           │
        │                      │ Queued │                                │
        │                      │   +    │          ┌─────────┐           │
        │                      │Running │─────────>│ Replica │           │
        │                      │Requests│ Forwards │    2    │           │
        │                      └────────┘          └─────────┘           │
        │                          │                                      │
        │                          │ Push metrics                         │
        │                          └─────────────────> Controller         │
        │                                                                 │
        └─────────────────────────────────────────────────────────────────┘

        :::{note}
        The long-term plan is to deprecate handle-based metrics collection in favor of
        replica-based collection. Replica-based collection will become the default in a
        future release. Queued requests will be continues to be tracked at the handle.
        :::
        """

        timestamp = time.time()
        self.metrics_store.add_metrics_point(
            {QUEUED_REQUESTS_KEY: self.num_queued_requests}, timestamp
        )
        if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE:
            self.metrics_store.add_metrics_point(
                self.num_requests_sent_to_replicas, timestamp
            )

        # Prevent in memory metrics store memory from growing
        start_timestamp = time.time() - self.autoscaling_config.look_back_period_s
        self.metrics_store.prune_keys_and_compact_data(start_timestamp)

    def _get_metrics_report(self) -> HandleMetricReport:
        timestamp = time.time()
        running_requests = dict()
        avg_running_requests = dict()
        look_back_period = self.autoscaling_config.look_back_period_s
        self.metrics_store.prune_keys_and_compact_data(time.time() - look_back_period)
        avg_queued_requests = self.metrics_store.aggregate_avg([QUEUED_REQUESTS_KEY])[0]
        if avg_queued_requests is None:
            # If the queued requests timeseries is empty, we set the
            # average to the current number of queued requests.
            avg_queued_requests = self.num_queued_requests
        # If the queued requests timeseries is empty, we set the number of data points to 1.
        # This is to avoid division by zero.
        num_data_points = self.metrics_store.timeseries_count(QUEUED_REQUESTS_KEY) or 1
        queued_requests = self.metrics_store.data.get(
            QUEUED_REQUESTS_KEY, [TimeStampedValue(timestamp, self.num_queued_requests)]
        )
        if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE and self.autoscaling_config:
            for replica_id, num_requests in self.num_requests_sent_to_replicas.items():
                # Calculate avg running requests.
                # NOTE (abrar): The number of data points from queued requests is often higher than
                # those from running requests. This is because replica metrics are only collected
                # once a replica is up, whereas queued request metrics are collected continuously
                # as long as the handle is alive. To approximate the true average of ongoing requests,
                # we should normalize by using the same number of data points for both queued and
                # running request time series.
                running_requests_sum = self.metrics_store.aggregate_sum([replica_id])[0]
                if running_requests_sum is None:
                    # If the running requests timeseries is empty, we set the sum
                    # to the current number of requests.
                    running_requests_sum = num_requests
                replica_str = replica_id.to_full_id_str()
                avg_running_requests[replica_str] = (
                    running_requests_sum / num_data_points
                )
                # Get running requests data
                running_requests[replica_str] = self.metrics_store.data.get(
                    replica_id, [TimeStampedValue(timestamp, num_requests)]
                )
        handle_metric_report = HandleMetricReport(
            deployment_id=self._deployment_id,
            handle_id=self._handle_id,
            actor_id=self._self_actor_id,
            handle_source=self._handle_source,
            aggregated_queued_requests=avg_queued_requests,
            queued_requests=queued_requests,
            aggregated_metrics={
                RUNNING_REQUESTS_KEY: avg_running_requests,
            },
            metrics={
                RUNNING_REQUESTS_KEY: running_requests,
            },
            timestamp=timestamp,
        )

        return handle_metric_report

    async def shutdown(self):
        """Shutdown metrics manager gracefully."""

        if self.metrics_pusher:
            await self.metrics_pusher.graceful_shutdown()

        self._shutdown = True

        if self._cached_metrics_task is not None:
            self._cached_metrics_task.cancel()
            try:
                await self._cached_metrics_task
            except asyncio.CancelledError:
                pass


class Router(ABC):
    @property
    def event_loop(self) -> Optional[AbstractEventLoop]:
        """The event loop the router runs on, or None (e.g. local testing)."""
        return getattr(self, "_asyncio_loop", None)

    @abstractmethod
    def running_replicas_populated(self) -> bool:
        pass

    @abstractmethod
    def assign_request(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> concurrent.futures.Future[ReplicaResult]:
        pass

    @abstractmethod
    @asynccontextmanager
    async def choose_replica(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> AsyncIterator[ReplicaSelection]:
        yield

    @abstractmethod
    def dispatch(
        self,
        selection: ReplicaSelection,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> concurrent.futures.Future[ReplicaResult]:
        pass

    @abstractmethod
    async def broadcast(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> List[ReplicaResult]:
        pass

    @abstractmethod
    def shutdown(self) -> concurrent.futures.Future:
        pass


async def create_event() -> asyncio.Event:
    """Helper to create an asyncio event in the current event loop."""
    return asyncio.Event()


class AsyncioRouter:
    def __init__(
        self,
        controller_handle: ActorHandle,
        deployment_id: DeploymentID,
        handle_id: str,
        self_actor_id: str,
        handle_source: DeploymentHandleSource,
        event_loop: asyncio.BaseEventLoop,
        enable_strict_max_ongoing_requests: bool,
        node_id: str,
        availability_zone: Optional[str],
        prefer_local_node_routing: bool,
        resolve_request_arg_func: Coroutine = resolve_deployment_response,
        request_router_class: Optional[Callable] = None,
        request_router_kwargs: Optional[Dict[str, Any]] = None,
        request_router: Optional[RequestRouter] = None,
        _request_router_initialized_event: Optional[asyncio.Event] = None,
    ):
        """Used to assign requests to downstream replicas for a deployment.

        The routing behavior is delegated to a RequestRouter; this is a thin
        wrapper that adds metrics and logging.
        """
        self._controller_handle = controller_handle
        self.deployment_id = deployment_id
        self._self_actor_id = self_actor_id
        self._handle_source = handle_source
        self._event_loop = event_loop
        self._request_router_class = request_router_class
        self._request_router_kwargs = (
            request_router_kwargs if request_router_kwargs else {}
        )
        self._enable_strict_max_ongoing_requests = enable_strict_max_ongoing_requests
        self._node_id = node_id
        self._availability_zone = availability_zone
        self._prefer_local_node_routing = prefer_local_node_routing
        # By default, deployment is available unless we receive news
        # otherwise through a long poll broadcast from the controller.
        self._deployment_available = True

        # The request router will be lazy loaded to decouple form the initialization.
        self._request_router: Optional[RequestRouter] = request_router

        if _request_router_initialized_event:
            self._request_router_initialized = _request_router_initialized_event
        else:
            future = asyncio.run_coroutine_threadsafe(create_event(), self._event_loop)
            self._request_router_initialized = future.result()

        if self._request_router:
            self._request_router_initialized.set()
        self._resolve_request_arg_func = resolve_request_arg_func
        self._running_replicas: Optional[List[RunningReplicaInfo]] = None

        # Flipped to `True` once the router has received a non-empty
        # replica set at least once.
        self._running_replicas_populated: bool = False

        self._initial_backoff_s: Optional[float] = None
        self._backoff_multiplier: Optional[float] = None
        self._max_backoff_s: Optional[float] = None

        # Initializing `self._metrics_manager` before `self.long_poll_client` is
        # necessary to avoid race condition where `self.update_deployment_config()`
        # might be called before `self._metrics_manager` instance is created.
        self._metrics_manager = RouterMetricsManager(
            deployment_id,
            handle_id,
            self_actor_id,
            handle_source,
            controller_handle,
            metrics.Counter(
                "serve_num_router_requests",
                description="The number of requests processed by the router.",
                tag_keys=("deployment", "route", "application", "handle", "actor_id"),
            ),
            metrics.Gauge(
                "serve_deployment_queued_queries",
                description=(
                    "The current number of queries to this deployment waiting"
                    " to be assigned to a replica."
                ),
                tag_keys=("deployment", "application", "handle", "actor_id"),
            ),
            metrics.Gauge(
                "serve_num_ongoing_requests_at_replicas",
                description=(
                    "The current number of requests to this deployment that "
                    "have been submitted to a replica."
                ),
                tag_keys=("deployment", "application", "handle", "actor_id"),
            ),
            event_loop,
        )

        self._objref_resolution_latency_ms = metrics.Histogram(
            "serve_router_args_resolution_latency_ms",
            description=(
                "Time in milliseconds spent resolving upstream ObjectRef or "
                "DeploymentResponse arguments before a request enters the "
                "routing queue."
            ),
            boundaries=DEFAULT_LATENCY_BUCKET_MS,
            tag_keys=("deployment", "application", "handle", "actor_id"),
        ).set_default_tags(
            {
                "deployment": deployment_id.name,
                "application": deployment_id.app_name,
                "handle": handle_id,
                "actor_id": self_actor_id,
            }
        )

        # The Router needs to stay informed about changes to the target deployment's
        # running replicas and deployment config. We do this via the long poll system.
        # However, for efficiency, we don't want to create a LongPollClient for every
        # DeploymentHandle, so we use a shared LongPollClient that all Routers
        # register themselves with. But first, the router needs to get a fast initial
        # update so that it can start serving requests, which we do with a dedicated
        # LongPollClient that stops running once the shared client takes over.

        self.long_poll_client = LongPollClient(
            controller_handle,
            {
                (
                    LongPollNamespace.DEPLOYMENT_TARGETS,
                    deployment_id,
                ): self.update_deployment_targets,
                (
                    LongPollNamespace.DEPLOYMENT_CONFIG,
                    deployment_id,
                ): self.update_deployment_config,
            },
            call_in_event_loop=self._event_loop,
            # Multiple AsyncioRouters can share an actor (one per downstream
            # handle), so include the deployment id to disambiguate.
            client_id=f"{type(self).__name__}:{self_actor_id}:{deployment_id}",
        )

        shared = SharedRouterLongPollClient.get_or_create(
            controller_handle, self._event_loop
        )
        shared.register(self)

    @property
    def request_router(self) -> Optional[RequestRouter]:
        """Get and lazy loading request router.

        If the request_router_class not provided, and the request router is not
        yet initialized, then it will return None. Otherwise, if request router
        is not yet initialized, it will be initialized and returned. Also,
        setting `self._request_router_initialized` to signal that the request
        router is initialized.
        """
        if not self._request_router and self._request_router_class:
            backoff_kwargs = {}
            if self._initial_backoff_s is not None:
                backoff_kwargs["initial_backoff_s"] = self._initial_backoff_s
            if self._backoff_multiplier is not None:
                backoff_kwargs["backoff_multiplier"] = self._backoff_multiplier
            if self._max_backoff_s is not None:
                backoff_kwargs["max_backoff_s"] = self._max_backoff_s

            request_router = self._request_router_class(
                deployment_id=self.deployment_id,
                handle_source=self._handle_source,
                self_node_id=self._node_id,
                self_actor_id=self._self_actor_id,
                self_actor_handle=ray.get_runtime_context().current_actor
                if ray.get_runtime_context().get_actor_id()
                else None,
                # Streaming ObjectRefGenerators are not supported in Ray Client
                use_replica_queue_len_cache=self._enable_strict_max_ongoing_requests,
                create_replica_wrapper_func=lambda r: RunningReplica(r),
                prefer_local_node_routing=self._prefer_local_node_routing,
                prefer_local_az_routing=RAY_SERVE_PROXY_PREFER_LOCAL_AZ_ROUTING,
                self_availability_zone=self._availability_zone,
                **backoff_kwargs,
            )
            request_router.initialize_state(**(self._request_router_kwargs))

            # Populate the running replicas if they are already available.
            if self._running_replicas is not None:
                request_router._update_running_replicas(self._running_replicas)

            self._request_router = request_router
            self._request_router_initialized.set()

            # Log usage telemetry to indicate that custom request router
            # feature is being used in this cluster.
            if (
                self._request_router_class.__name__
                != PowerOfTwoChoicesRequestRouter.__name__
            ):
                ServeUsageTag.CUSTOM_REQUEST_ROUTER_USED.record("1")
        return self._request_router

    def running_replicas_populated(self) -> bool:
        return self._running_replicas_populated

    def update_deployment_targets(self, deployment_target_info: DeploymentTargetInfo):
        self._deployment_available = deployment_target_info.is_available

        running_replicas = deployment_target_info.running_replicas
        if self.request_router:
            self.request_router._update_running_replicas(running_replicas)
        else:
            # In this case, the request router hasn't been initialized yet.
            # Store the running replicas so that we can update the request
            # router once it is initialized.
            self._running_replicas = running_replicas
        self._metrics_manager._update_running_replicas(running_replicas)

        if running_replicas:
            self._running_replicas_populated = True

    def update_deployment_config(self, deployment_config: DeploymentConfig):
        self._request_router_class = (
            deployment_config.request_router_config.get_request_router_class()
        )
        self._request_router_kwargs = (
            deployment_config.request_router_config.request_router_kwargs
        )

        # Warn if deprecated env vars are set
        warn_if_deprecated_env_var_set("RAY_SERVE_ROUTER_RETRY_INITIAL_BACKOFF_S")
        warn_if_deprecated_env_var_set("RAY_SERVE_ROUTER_RETRY_BACKOFF_MULTIPLIER")
        warn_if_deprecated_env_var_set("RAY_SERVE_ROUTER_RETRY_MAX_BACKOFF_S")

        self._initial_backoff_s = (
            deployment_config.request_router_config.initial_backoff_s
        )
        self._backoff_multiplier = (
            deployment_config.request_router_config.backoff_multiplier
        )
        self._max_backoff_s = deployment_config.request_router_config.max_backoff_s

        if self._request_router:
            self._request_router.update_backoff_params(
                initial_backoff_s=self._initial_backoff_s,
                backoff_multiplier=self._backoff_multiplier,
                max_backoff_s=self._max_backoff_s,
            )

        # Guard against the case where request_router is None (e.g., when
        # request_router_class is None and lazy initialization has not yet
        # occurred). In that scenario there are no active replicas yet, so
        # passing 0 is the correct and safe value.
        _router = self.request_router
        curr_num_replicas = len(_router.curr_replicas) if _router is not None else 0
        self._metrics_manager.update_deployment_config(
            deployment_config,
            curr_num_replicas=curr_num_replicas,
        )

    async def _resolve_request_arguments(
        self,
        pr: PendingRequest,
    ) -> None:
        """Asynchronously resolve and replace top-level request args and kwargs."""
        if pr.resolved:
            return

        new_args = list(pr.args)
        new_kwargs = pr.kwargs.copy()

        # Map from index -> task for resolving positional arg
        resolve_arg_tasks = {}
        for i, obj in enumerate(pr.args):
            task = await self._resolve_request_arg_func(obj, pr.metadata)
            if task is not None:
                resolve_arg_tasks[i] = task

        # Map from key -> task for resolving key-word arg
        resolve_kwarg_tasks = {}
        for k, obj in pr.kwargs.items():
            task = await self._resolve_request_arg_func(obj, pr.metadata)
            if task is not None:
                resolve_kwarg_tasks[k] = task

        # Gather all argument resolution tasks concurrently.
        if resolve_arg_tasks or resolve_kwarg_tasks:
            all_tasks = list(resolve_arg_tasks.values()) + list(
                resolve_kwarg_tasks.values()
            )
            await asyncio.wait(all_tasks)

        # Update new args and new kwargs with resolved arguments
        for index, task in resolve_arg_tasks.items():
            new_args[index] = task.result()
        for key, task in resolve_kwarg_tasks.items():
            new_kwargs[key] = task.result()

        pr.args = new_args
        pr.kwargs = new_kwargs
        pr.resolved = True

    def _process_finished_request(
        self,
        replica_id: ReplicaID,
        internal_request_id: str,
        replica_actor_id: Optional[ray.ActorID],
        result: Union[Any, RayError],
    ) -> None:
        if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE:
            self._metrics_manager.dec_num_running_requests_for_replica(replica_id)

        # Notify request router that request completed (for cleanup, e.g., token release)
        if self.request_router:
            self.request_router.on_request_completed(replica_id, internal_request_id)

        actor_died_error = self._get_actor_died_error(result)
        if actor_died_error is not None:
            self._handle_actor_died_error(
                replica_id, replica_actor_id, actor_died_error
            )
        elif isinstance(result, ActorUnavailableError):
            # There are network issues, or replica has died but GCS is down so
            # ActorUnavailableError will be raised until GCS recovers. For the
            # time being, invalidate the cache entry so that we don't try to
            # send requests to this replica without actively probing, and retry
            # routing request.
            if self.request_router:
                self.request_router.on_replica_actor_unavailable(replica_id)
            logger.warning(
                f"Request failed because {replica_id} is temporarily unavailable."
            )

    def _get_actor_died_error(
        self, result: Union[Any, RayError]
    ) -> Optional[ActorDiedError]:
        if isinstance(result, ActorDiedError):
            return result

        if isinstance(result, RayTaskError) and isinstance(
            getattr(result, "cause", None), ActorDiedError
        ):
            # RayTaskError wrapping ActorDiedError (e.g., from failed object ref
            # resolution in chained deployment calls).
            return result.cause

        return None

    def _handle_actor_died_error(
        self,
        replica_id: ReplicaID,
        replica_actor_id: Optional[ray.ActorID],
        actor_died_error: ActorDiedError,
    ) -> bool:
        """Handle an ActorDiedError from a replica request.

        Returns True if the error is from this replica (i.e., this replica
        died and should be retried on another replica). Returns False if the
        error is from an upstream dependency (i.e., this replica is healthy
        but the request failed due to a bad input).
        """
        # Only mark the replica as dead if the ActorDiedError refers to this
        # replica. With chained DeploymentResponses, the error may come from
        # an upstream deployment that was passed as an object ref to this
        # replica. In that case, this replica is still healthy.
        error_actor_id = getattr(actor_died_error, "actor_id", None)
        replica_actor_id_hex = (
            replica_actor_id.hex() if replica_actor_id is not None else None
        )
        # When error_actor_id or replica_actor_id_hex is None, we cannot
        # definitively compare. Treat as match to preserve conservative
        # behavior: mark replica dead rather than leaving it in rotation.
        if (
            error_actor_id is None
            or replica_actor_id_hex is None
            or error_actor_id == replica_actor_id_hex
        ):
            # Replica has died but controller hasn't notified the router yet.
            if self.request_router:
                self.request_router.on_replica_actor_died(replica_id)
            logger.warning(
                f"{replica_id} will not be considered for future "
                "requests because it has died."
            )
            return True
        else:
            # Error from upstream dependency, not from this replica.
            logger.debug(
                f"ActorDiedError from upstream (actor_id={error_actor_id}), "
                f"not from {replica_id} (actor_id={replica_actor_id_hex}). "
                "Replica remains in rotation."
            )
            return False

    def _make_upstream_crash_error(self, e: ActorDiedError) -> RayServeException:
        """Surface a clear Serve error while preserving Ray's actor death details."""
        msg = (
            f"Request to deployment '{self.deployment_id.name}' failed because "
            f"an upstream actor died before finishing a dependent task. "
            f"Ray reported:\n{e}"
        )
        wrapped = RayServeException(msg)
        wrapped.__cause__ = e
        return wrapped

    async def _route_and_send_request_once(
        self,
        pr: PendingRequest,
        response_id: str,
        is_retry: bool,
    ) -> Optional[ReplicaResult]:
        result: Optional[ReplicaResult] = None
        replica: Optional[RunningReplica] = None
        callback_registered = False
        try:
            # Resolve request arguments BEFORE incrementing queued requests.
            # This ensures that queue metrics reflect actual pending work,
            # not time spent waiting for upstream DeploymentResponse arguments.
            # See: https://github.com/ray-project/ray/issues/60624
            await self._resolve_args_with_metrics(pr)

            num_curr_replicas = len(self.request_router.curr_replicas)
            with self._metrics_manager.wrap_queued_request(is_retry, num_curr_replicas):
                replica = await self.request_router._choose_replica_for_request(
                    pr, is_retry=is_retry
                )

                # If the queue len cache is disabled or we're sending a request to Java,
                # then directly send the query and hand the response back. The replica will
                # never reject requests in this code path.
                with_rejection = (
                    self._enable_strict_max_ongoing_requests
                    and not replica.is_cross_language
                )
                result = replica.try_send_request(pr, with_rejection=with_rejection)
                # Proactively update the queue length cache.
                self.request_router.on_send_request(replica.replica_id)

            self._register_completion_callback(result, replica, pr)
            callback_registered = True

            if not with_rejection:
                self._register_decrement_queue_len_cache_callback(
                    result, replica.replica_id
                )
                return result

            queue_info = await result.get_rejection_response()
            self.request_router.on_new_queue_len_info(
                replica.replica_id, queue_info.num_ongoing_requests
            )
            if queue_info.accepted:
                self.request_router.on_request_routed(pr, replica.replica_id, result)
                self._register_decrement_queue_len_cache_callback(
                    result, replica.replica_id
                )
                return result

            # Request was rejected: cancel so done callbacks fire.
            # Without this, same-loop (means the router is running on the main event loop,
            # where the DeploymentHandle lives) gRPC streaming results are
            # never consumed (no background drain task exists in that mode).
            # The call stays open, the running request counter is never decremented,
            # and the autoscaler sees load that blocks downscaling.
            result.cancel()

        except asyncio.CancelledError:
            # NOTE(edoakes): this is not strictly necessary because there are
            # currently no `await` statements between getting the ref and returning,
            # but I'm adding it defensively.
            if result is not None:
                result.cancel()

            raise
        except ActorDiedError as e:
            if replica is not None:
                is_from_this_replica = self._handle_actor_died_error(
                    replica.replica_id, replica.actor_id, e
                )
                if not is_from_this_replica and callback_registered:
                    # Error from an upstream dependency during request
                    # execution (e.g., a chained DeploymentResponse whose
                    # source actor died). The request was already accepted
                    # by this (healthy) replica, but the input is
                    # permanently failed — retrying with another replica
                    # won't help. Propagate immediately so the caller gets
                    # a fast error.
                    raise self._make_upstream_crash_error(e)
            elif not pr.resolved:
                # ActorDiedError during argument resolution — same upstream
                # cause as above, caught before a replica was even chosen.
                raise self._make_upstream_crash_error(e)
        except ActorUnavailableError:
            # There are network issues, or replica has died but GCS is down so
            # ActorUnavailableError will be raised until GCS recovers. For the
            # time being, invalidate the cache entry so that we don't try to
            # send requests to this replica without actively probing, and retry
            # routing request.
            if replica is not None:
                self.request_router.on_replica_actor_unavailable(replica.replica_id)
                logger.warning(f"{replica.replica_id} is temporarily unavailable.")
        finally:
            # Only release if callback wasn't registered (callback handles release).
            if replica is not None and not callback_registered:
                self.request_router.on_request_completed(
                    replica.replica_id, pr.metadata.internal_request_id
                )

        return None

    async def route_and_send_request(
        self,
        pr: PendingRequest,
        response_id: str,
    ) -> ReplicaResult:
        """Choose a replica for the request and send it.

        This will block indefinitely if no replicas are available to handle the
        request, so it's up to the caller to time out or cancel the request.
        """
        # Wait for the router to be initialized before sending the request.
        await self._request_router_initialized.wait()

        is_retry = False
        while True:
            result = await self._route_and_send_request_once(
                pr,
                response_id,
                is_retry,
            )
            if result is not None:
                return result

            # If the replica rejects the request, retry the routing process. The
            # request will be placed on the front of the queue to avoid tail latencies.
            # TODO(edoakes): this retry procedure is not perfect because it'll reset the
            # process of choosing candidates replicas (i.e., for locality-awareness).
            is_retry = True

    @tracing_decorator_factory(
        trace_name="route_to_replica",
    )
    async def assign_request(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> ReplicaResult:
        """Assign a request to a replica and return the resulting object_ref."""
        if is_span_recording():
            trace_attributes = {
                "request_id": request_meta.request_id,
                "deployment": self.deployment_id.name,
                "app": self.deployment_id.app_name,
                "call_method": request_meta.call_method,
                "route": request_meta.route,
                "multiplexed_model_id": request_meta.multiplexed_model_id,
                "is_streaming": request_meta.is_streaming,
                "is_http_request": request_meta.is_http_request,
                "is_grpc_request": request_meta.is_grpc_request,
            }
            set_span_attributes(trace_attributes)
            set_span_name(
                f"route_to_replica {self.deployment_id.name} {request_meta.call_method}"
            )
            # Add context to request meta to link
            # traces collected in the replica.
            propagate_context = create_propagated_context()
            request_meta.tracing_context = propagate_context
        else:
            request_meta.tracing_context = None

        if not self._deployment_available:
            raise DeploymentUnavailableError(self.deployment_id)

        response_id = generate_request_id()
        assign_request_task = asyncio.current_task()
        ray.serve.context._add_request_pending_assignment(
            request_meta.internal_request_id, response_id, assign_request_task
        )
        assign_request_task.add_done_callback(
            lambda _: ray.serve.context._remove_request_pending_assignment(
                request_meta.internal_request_id, response_id
            )
        )

        # Wait for the router to be initialized before sending the request.
        await self._request_router_initialized.wait()

        with self._metrics_manager.wrap_request_assignment(request_meta):
            replica_result = None
            exc = None
            try:
                replica_result = await self.route_and_send_request(
                    PendingRequest(
                        args=list(request_args),
                        kwargs=request_kwargs,
                        metadata=request_meta,
                    ),
                    response_id,
                )
                return replica_result
            except asyncio.CancelledError as e:
                exc = e
                # NOTE(edoakes): this is not strictly necessary because
                # there are currently no `await` statements between
                # getting the ref and returning, but I'm adding it defensively.
                if replica_result is not None:
                    replica_result.cancel()

                raise
            finally:
                if is_span_recording():
                    if request_meta.is_http_request:
                        set_http_span_attributes(
                            method=request_meta.call_method,
                            route=request_meta.route,
                        )
                    else:
                        set_rpc_span_attributes(
                            system=request_meta._request_protocol,
                            method=request_meta.call_method,
                            service=self.deployment_id.name,
                        )
                    if exc:
                        set_span_exception(exc, escaped=True)

    @asynccontextmanager
    async def choose_replica(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> AsyncIterator[ReplicaSelection]:
        """Pick a replica, hold capacity on it, and yield a selection.

        The yielded selection can be dispatched (consuming the held
        capacity) or allowed to drop, in which case the capacity is
        released automatically when the context exits.

        Behavior:
          1. Refuses immediately if the deployment is unavailable.
          2. Waits for the router to finish initial setup before routing.
          3. Refuses with backpressure if the queue limit is already
             reached; otherwise the request counts against router-level
             request and queue metrics for the duration of routing.
          4. Resolves the request arguments (which may block on
             upstream responses).
          5. Picks a replica and asks it to hold capacity. Retries on
             actor failure, transient unavailability, or replica-side
             capacity rejection until a replica accepts.
          6. Yields the selection while the reservation is held; on exit,
             releases the capacity if dispatch never consumed it.
             Reservation metrics stay accurate even if the release call
             itself fails.
        """
        # TODO(jeffreywang): Add tracing support
        if not self._deployment_available:
            raise DeploymentUnavailableError(self.deployment_id)

        # Internal opt-out for pick-only callers (e.g. ingress router that
        # forwards traffic out-of-band via HAProxy). Skips the replica-side
        # reserve_slot RPC and the rejection-retry loop; the configured
        # RequestRouter still drives ordering.
        reserve = request_kwargs.pop("_reserve", True)

        await self._request_router_initialized.wait()

        with self._metrics_manager.wrap_request_assignment(request_meta):
            pr = PendingRequest(
                args=list(request_args),
                kwargs=request_kwargs,
                metadata=request_meta,
            )
            try:
                await self._resolve_args_with_metrics(pr)
            except ActorDiedError as e:
                raise self._make_upstream_crash_error(e)
            if reserve:
                replica, slot_token = await self._pick_and_reserve_replica(pr)
            else:
                # Fast path: synchronously ask the configured RequestRouter to
                # pick a replica from the current snapshot, bypassing the
                # _pending_requests_to_fulfill queue and the routing-task
                # workers. Safe because there's no reservation -> no rejection
                # -> no retry, so the queue's ordering/backoff guarantees are
                # unused.
                ranks = await self.request_router.choose_replicas(
                    candidate_replicas=self.request_router._replicas_list,
                    pending_request=pr,
                )
                replica = next((r for rank in ranks for r in rank), None)
                if replica is None:
                    raise RuntimeError(
                        f"no replicas available for {self.deployment_id}"
                    )
                slot_token = None

            selection = ReplicaSelection(
                replica_id=replica.replica_id.unique_id,
                node_ip=replica._replica_info.node_ip,
                port=replica._replica_info.port,
                node_id=replica.node_id,
                availability_zone=replica.availability_zone,
                replica_metadata=replica.replica_metadata,
                _replica=replica,
                _deployment_id=None,  # Injected by DeploymentHandle for dispatch-time validation.
                _request_metadata=request_meta,
                _method_name=request_meta.call_method,
                _slot_token=slot_token,
            )

        if not reserve:
            yield selection
            return

        try:
            self._metrics_manager.inc_reserved_slots()
            yield selection
        finally:
            try:
                await self._release_slot_and_refresh_cache(selection, replica)
                # Every dispatched request must fire on_request_completed exactly
                # once. When dispatch wires up the result's done-callback it fires
                # there; otherwise (dispatch raised before registering the
                # callback) fire it here. Skipped entirely when the caller exits
                # without dispatch, since no request was sent to the replica.
                if (
                    selection._dispatched
                    and not selection._completion_callback_registered
                    and self.request_router
                ):
                    self.request_router.on_request_completed(
                        replica.replica_id, request_meta.internal_request_id
                    )
            finally:
                # Decrement reserved slots metric even if release failed,
                # otherwise the gauge leaks for the lifetime of the process.
                self._metrics_manager.dec_reserved_slots()

    async def dispatch(
        self,
        selection: ReplicaSelection,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> ReplicaResult:
        """Dispatch to a specific replica, consuming the reserved slot.

        Args:
            selection: The replica selection from choose_replica().
            request_meta: Request metadata.
            *request_args: Request positional arguments.
            **request_kwargs: Request keyword arguments.

        Returns:
            ReplicaResult for the dispatched request.

        Raises:
            RuntimeError: If the selection has already been dispatched.
            ReplicaUnavailableError: If the replica is no longer available.
        """
        selection._mark_dispatched()
        return await self._dispatch_to_marked_selection(
            selection, request_meta, *request_args, **request_kwargs
        )

    async def _dispatch_to_marked_selection(
        self,
        selection: ReplicaSelection,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> ReplicaResult:
        """Send a request to the replica chosen by `choose_replica`.

        Mirrors the send half of `_route_and_send_request_once`: build a
        PendingRequest, resolve args, send with `with_rejection=False`
        (the reservation guarantees acceptance), and wire up the same
        completion callbacks. On failure, releases the held slot before
        re-raising.
        """
        replica = selection._replica
        if selection._slot_token is None:
            raise RuntimeError(
                "Cannot dispatch a ReplicaSelection that was created without a "
                "reservation (_reserve=False)."
            )
        # Inject the slot token so the replica skips re-acquiring its semaphore.
        # Args are re-resolved here because dispatch may carry augmented args.
        pr = PendingRequest(
            args=list(request_args),
            kwargs=request_kwargs,
            metadata=replace(request_meta, _reserved_slot_token=selection._slot_token),
        )

        try:
            if replica.replica_id not in self.request_router.curr_replicas:
                raise ReplicaUnavailableError(
                    f"Replica {selection.replica_id} is no longer available"
                )
            try:
                await self._resolve_args_with_metrics(pr)
            except ActorDiedError as e:
                raise self._make_upstream_crash_error(e)
            result = replica.try_send_request(pr, with_rejection=False)
        except BaseException:
            # Dispatch failed; release the reservation before re-raising.
            await self._release_slot_and_refresh_cache(selection, replica, force=True)
            raise

        self._register_completion_callback(result, replica, pr)
        # The done-callback registered above will fire on_request_completed
        # when the result completes; flag the selection so choose_replica's
        # finally skips its manual call (exactly one per reservation).
        selection._completion_callback_registered = True
        self._register_decrement_queue_len_cache_callback(result, replica.replica_id)

        # Only the cancellation path needs a follow-up release_slot RPC: on
        # success or replica-side failure, `_start_request` already consumed
        # the slot. Skipping it on the happy path saves one actor RPC per
        # dispatched request.
        def _maybe_release_slot(fut, sel=selection):
            cancelled = isinstance(fut, TaskCancelledError) or (
                callable(getattr(fut, "cancelled", None)) and fut.cancelled()
            )
            if cancelled:
                self._event_loop.call_soon_threadsafe(
                    lambda: self._event_loop.create_task(
                        self._release_slot_if_still_reserved(sel)
                    )
                )

        result.add_done_callback(_maybe_release_slot)

        return result

    async def _release_slot_if_still_reserved(
        self, selection: ReplicaSelection
    ) -> None:
        """Best-effort release_slot when a dispatched request was cancelled
        before the replica could consume the reservation. A no-op on the
        replica side if the slot was already consumed."""
        try:
            await selection._release_slot(force=True)
        except Exception:
            logger.debug("Failed to release reserved replica slot.", exc_info=True)

    async def _resolve_args_with_metrics(self, pr: PendingRequest) -> None:
        """Resolve a PendingRequest's arguments, observing resolution latency."""
        if pr.resolved:
            return
        resolution_start = time.monotonic()
        await self._resolve_request_arguments(pr)
        resolution_ms = (time.monotonic() - resolution_start) * 1000
        self._objref_resolution_latency_ms.observe(resolution_ms)

    async def _pick_and_reserve_replica(
        self, pr: PendingRequest
    ) -> Tuple[RunningReplica, str]:
        """Pick a replica and hold capacity on it.

        Loops until a replica accepts the reservation, retrying on
        capacity rejection, replica actor death, or transient
        unavailability. Updates the queue-length cache from the replica's
        reported count.
        """
        is_retry = False
        while True:
            num_curr_replicas = len(self.request_router.curr_replicas)
            with self._metrics_manager.wrap_queued_request(
                is_retry=is_retry, num_curr_replicas=num_curr_replicas
            ):
                replica = await self.request_router._choose_replica_for_request(
                    pr, is_retry=is_retry
                )
                try:
                    slot_token, queue_info = await replica.reserve_slot(pr.metadata)
                except ActorDiedError as e:
                    self._handle_actor_died_error(
                        replica.replica_id, replica.actor_id, e
                    )
                    is_retry = True
                    continue
                except ActorUnavailableError:
                    self.request_router.on_replica_actor_unavailable(replica.replica_id)
                    logger.warning(f"{replica.replica_id} is temporarily unavailable.")
                    is_retry = True
                    continue

                self.request_router.on_new_queue_len_info(
                    replica.replica_id, queue_info.num_ongoing_requests
                )
                if queue_info.accepted:
                    return replica, slot_token

            is_retry = True

    def _register_completion_callback(
        self, result: ReplicaResult, replica: RunningReplica, pr: PendingRequest
    ) -> None:
        """Count this request as in-flight and register the cleanup
        callback fired when the result completes or is cancelled.
        """
        if RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE:
            self._metrics_manager.inc_num_running_requests_for_replica(
                replica.replica_id
            )
        # Always register callback to notify router when request completes
        # (needed for token release in queue-based routing, metrics tracking, etc.)
        # NOTE: add_done_callback fires from a C++ worker thread (for actor
        # ObjectRefs) or a gRPC callback thread. _process_finished_request
        # and decrement_queue_len_cache both access shared router state
        # (e.g., _replica_queue_len_cache) that is not thread-safe, so we
        # schedule them on the router's event loop.
        callback = partial(
            self._process_finished_request,
            replica.replica_id,
            pr.metadata.internal_request_id,
            replica.actor_id,
        )
        result.add_done_callback(
            lambda _, cb=callback: self._event_loop.call_soon_threadsafe(cb, _)
        )

    def _register_decrement_queue_len_cache_callback(
        self, result: ReplicaResult, replica_id: ReplicaID
    ) -> None:
        """Schedule a queue-length-cache decrement for when this result completes."""
        result.add_done_callback(
            lambda _: self._event_loop.call_soon_threadsafe(
                self.request_router.decrement_queue_len_cache,
                replica_id,
            )
        )

    async def _release_slot_and_refresh_cache(
        self,
        selection: ReplicaSelection,
        replica: RunningReplica,
        *,
        force: bool = False,
    ) -> None:
        """Release a reserved slot and feed the replica's reported queue
        length back into the cache. Swallows release failures so callers
        can put this in a finally without masking other errors."""
        try:
            num_ongoing_requests = await selection._release_slot(force=force)
            if num_ongoing_requests is not None:
                self.request_router.on_new_queue_len_info(
                    replica.replica_id, num_ongoing_requests
                )
        except Exception:
            logger.exception("Failed to release reserved replica slot.")

    async def broadcast(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> List[ReplicaResult]:
        """Send a request to all running replicas in parallel.

        Bypasses the normal load-balancing path and sends the request
        directly to every replica. Waits for the request router to be
        initialized so the replica set is populated.
        """
        # Propagate tracing context, matching assign_request behavior.
        if is_span_recording():
            propagate_context = create_propagated_context()
            request_meta.tracing_context = propagate_context
        else:
            request_meta.tracing_context = None

        if not self._deployment_available:
            raise DeploymentUnavailableError(self.deployment_id)

        await self._request_router_initialized.wait()

        if not self._deployment_available:
            raise DeploymentUnavailableError(self.deployment_id)

        replicas: List[RunningReplica] = list(
            self.request_router.curr_replicas.values()
        )
        if not replicas:
            raise DeploymentUnavailableError(self.deployment_id)

        # Resolve arguments (e.g. DeploymentResponse objects) before sending.
        pr = PendingRequest(
            args=list(request_args),
            kwargs=dict(request_kwargs),
            metadata=request_meta,
        )
        await self._resolve_request_arguments(pr)

        results: List[ReplicaResult] = []
        for replica in replicas:
            replica_pr = PendingRequest(
                args=list(pr.args),
                kwargs=dict(pr.kwargs),
                metadata=replace(
                    request_meta,
                    internal_request_id=generate_request_id(),
                ),
            )
            replica_pr.resolved = True
            try:
                result = replica.try_send_request(replica_pr, with_rejection=False)
            except ActorDiedError:
                # Replica has died but controller hasn't notified the router yet.
                # Skip this replica and continue broadcasting to healthy replicas.
                self.request_router.on_replica_actor_died(replica.replica_id)
                logger.warning(
                    f"{replica.replica_id} will not be considered for future "
                    "requests because it has died."
                )
                continue
            except ActorUnavailableError:
                # Replica is temporarily unavailable. Invalidate the cache entry
                # and continue broadcasting to other replicas.
                self.request_router.on_replica_actor_unavailable(replica.replica_id)
                logger.warning(f"{replica.replica_id} is temporarily unavailable.")
                continue

            # Proactively update the queue length cache.
            self.request_router.on_send_request(replica.replica_id)

            self._register_completion_callback(result, replica, replica_pr)
            self._register_decrement_queue_len_cache_callback(
                result, replica.replica_id
            )

            results.append(result)

        if not results:
            raise DeploymentUnavailableError(self.deployment_id)

        return results

    async def shutdown(self):
        await self._metrics_manager.shutdown()


class SingletonThreadRouter(Router):
    """Wrapper class that runs an AsyncioRouter on a separate thread.

    The motivation for this is to avoid user code blocking the event loop and
    preventing the router from making progress.

    Maintains a singleton event loop running in a daemon thread that is shared by
    all AsyncioRouters.
    """

    _asyncio_loop: Optional[asyncio.AbstractEventLoop] = None
    _asyncio_loop_creation_lock = threading.Lock()
    _event_loop_monitor: Optional[EventLoopMonitor] = None

    def __init__(self, **passthrough_kwargs):
        assert (
            "event_loop" not in passthrough_kwargs
        ), "SingletonThreadRouter manages the router event loop."

        if passthrough_kwargs.get("handle_source") == DeploymentHandleSource.REPLICA:
            component = EventLoopMonitor.COMPONENT_REPLICA
        elif passthrough_kwargs.get("handle_source") == DeploymentHandleSource.PROXY:
            component = EventLoopMonitor.COMPONENT_PROXY
        else:
            component = EventLoopMonitor.COMPONENT_UNKNOWN

        self._asyncio_router = AsyncioRouter(
            event_loop=self._get_singleton_asyncio_loop(component), **passthrough_kwargs
        )

    @classmethod
    def _get_singleton_asyncio_loop(cls, component: str) -> asyncio.AbstractEventLoop:
        """Get singleton asyncio loop running in a daemon thread.

        This method is thread safe.
        """
        with cls._asyncio_loop_creation_lock:
            if cls._asyncio_loop is None:
                cls._asyncio_loop = asyncio.new_event_loop()

                # Create event loop monitor for the router loop.
                # This is shared across all replicas in this process.
                actor_id = ray.get_runtime_context().get_actor_id()
                cls._event_loop_monitor = EventLoopMonitor(
                    component=component,
                    loop_type=EventLoopMonitor.LOOP_TYPE_ROUTER,
                    # actor_id is None when using DeploymentHandle.remote()
                    # from the driver.
                    actor_id=actor_id or "",
                )

                def _run_router_event_loop():
                    asyncio.set_event_loop(cls._asyncio_loop)
                    # Start monitoring before run_forever so the task is scheduled.
                    cls._event_loop_monitor.start(cls._asyncio_loop)
                    cls._asyncio_loop.run_forever()

                thread = threading.Thread(
                    daemon=True,
                    target=_run_router_event_loop,
                )
                thread.start()

        return cls._asyncio_loop

    def running_replicas_populated(self) -> bool:
        return self._asyncio_router.running_replicas_populated()

    def assign_request(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> concurrent.futures.Future[ReplicaResult]:
        """Routes assign_request call on the internal asyncio loop.

        This method uses `run_coroutine_threadsafe` to execute the actual request
        assignment logic (`_asyncio_router.assign_request`) on the dedicated
        asyncio event loop thread. It returns a `concurrent.futures.Future` that
        can be awaited or queried from the calling thread.

        Args:
            request_meta: Metadata describing the inbound request.
            *request_args: Positional arguments forwarded to the replica handler.
            **request_kwargs: Keyword arguments forwarded to the replica handler.

        Returns:
            A concurrent.futures.Future resolving to the ReplicaResult representing
            the assigned request.
        """
        return self._wrap_asyncio_call_in_future(
            self._asyncio_router.assign_request(
                request_meta, *request_args, **request_kwargs
            )
        )

    @asynccontextmanager
    async def choose_replica(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> AsyncIterator[ReplicaSelection]:
        """Bridge async context manager to router event loop.

        This ensures choose_replica runs on the singleton router loop,
        maintaining thread safety for all state modifications.
        """
        # Enter context on router loop
        async def enter_context():
            cm = self._asyncio_router.choose_replica(
                request_meta, *request_args, **request_kwargs
            )
            selection = await cm.__aenter__()
            return selection, cm

        async def exit_context(cm, exc_type, exc_val, exc_tb):
            return await cm.__aexit__(exc_type, exc_val, exc_tb)

        async def release_entered_context(entry_future):
            """Release a slot whose caller was cancelled before owning the selection.

            Awaits the entry on the router loop and drives ``__aexit__`` on the
            entered context manager. The inner CM is parked at ``yield`` holding
            the slot, and only an explicit ``__aexit__`` releases it reliably.
            """
            try:
                _, cm = await asyncio.wrap_future(entry_future)
            except BaseException:
                # __aenter__ was cancelled/failed: nothing entered to release.
                return
            try:
                await cm.__aexit__(None, None, None)
            except Exception:
                logger.exception("Failed to release reserved replica slot.")

        future = asyncio.run_coroutine_threadsafe(enter_context(), self._asyncio_loop)
        # Shield so a caller cancellation does not propagate through wrap_future
        # and cancel ``enter_context``: __aenter__ finishes and returns, so the
        # entered CM stays reachable for release. Otherwise the CM is discarded
        # mid-entry and the slot leaks until GC.
        try:
            selection, context_manager = await asyncio.shield(
                asyncio.wrap_future(future)
            )
        except BaseException:
            # Honor the cancellation now; release the orphaned slot on the
            # router loop instead of making the cancelled caller wait on it.
            asyncio.run_coroutine_threadsafe(
                release_entered_context(future), self._asyncio_loop
            )
            raise

        try:
            yield selection
        finally:
            exc_info = sys.exc_info()
            future = asyncio.run_coroutine_threadsafe(
                exit_context(context_manager, *exc_info), self._asyncio_loop
            )
            # Shielded so a cancel landing during cleanup doesn't propagate
            # through wrap_future and abort __aexit__ on the router loop.
            await asyncio.shield(asyncio.wrap_future(future))

    def dispatch(
        self,
        selection: ReplicaSelection,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> concurrent.futures.Future[ReplicaResult]:
        """Dispatch request to a previously selected replica."""
        try:
            selection._mark_dispatched()
        except Exception as exc:
            future = concurrent.futures.Future()
            future.set_exception(exc)
            return future

        return self._wrap_asyncio_call_in_future(
            self._asyncio_router._dispatch_to_marked_selection(
                selection, request_meta, *request_args, **request_kwargs
            )
        )

    def _wrap_asyncio_call_in_future(
        self,
        coro: Coroutine,
    ) -> concurrent.futures.Future[ReplicaResult]:
        """Wrap an async call in a concurrent.futures.Future for cross-thread execution.

        This is a helper method to execute AsyncioRouter's async methods on the dedicated asyncio event loop thread.

        Args:
            coro: The coroutine to execute (e.g., _asyncio_router.assign_request(...))

        Returns:
            A concurrent.futures.Future that resolves to the ReplicaResult.
        """
        # Extract operation name from coroutine for logging
        operation_name = coro.__name__

        def asyncio_future_callback(
            asyncio_future: asyncio.Future, concurrent_future: concurrent.futures.Future
        ):
            """Callback attached to the asyncio Task running assign_request.

            This runs when the asyncio Task finishes (completes, fails, or is cancelled).
            Its primary goal is to propagate cancellation initiated via the
            `concurrent_future` back to the `ReplicaResult` in situations where
            asyncio_future didn't see the cancellation event in time. Think of it
            like a second line of defense for cancellation of replica results.
            """
            # Check if the cancellation originated from the concurrent.futures.Future
            if (
                concurrent_future.cancelled()
                and not asyncio_future.cancelled()
                and asyncio_future.exception() is None
            ):
                result: ReplicaResult = asyncio_future.result()
                logger.info(
                    f"Asyncio task completed despite cancellation attempt during {operation_name}. "
                    "Attempting to cancel the request."
                )
                result.cancel()

        concurrent_future = concurrent.futures.Future()

        def create_task_and_setup():
            task = self._asyncio_loop.create_task(coro)

            # Set up your cancellation callback
            task.add_done_callback(
                lambda _: asyncio_future_callback(_, concurrent_future)
            )

            try:
                # chain the two futures to handle direction channel of cancellation
                futures._chain_future(
                    ensure_future(task, loop=self._asyncio_loop), concurrent_future
                )
            except (SystemExit, KeyboardInterrupt):
                raise
            except BaseException as exc:
                if concurrent_future.set_running_or_notify_cancel():
                    concurrent_future.set_exception(exc)
                raise

        # Schedule on the event loop thread
        self._asyncio_loop.call_soon_threadsafe(create_task_and_setup)
        return concurrent_future

    async def broadcast(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> List[ReplicaResult]:
        return await self._asyncio_router.broadcast(
            request_meta, *request_args, **request_kwargs
        )

    def shutdown(self) -> concurrent.futures.Future:
        return asyncio.run_coroutine_threadsafe(
            self._asyncio_router.shutdown(), loop=self._asyncio_loop
        )


class SharedRouterLongPollClient:
    def __init__(self, controller_handle: ActorHandle, event_loop: AbstractEventLoop):
        self.controller_handler = controller_handle
        self.event_loop = event_loop

        # We use a WeakSet to store the Routers so that we don't prevent them
        # from being garbage-collected.
        self.routers: MutableMapping[
            DeploymentID, weakref.WeakSet[AsyncioRouter]
        ] = defaultdict(weakref.WeakSet)

        # Creating the LongPollClient implicitly starts it
        self.long_poll_client = LongPollClient(
            controller_handle,
            key_listeners={},
            call_in_event_loop=self.event_loop,
            client_id=f"{type(self).__name__}:{ray.get_runtime_context().get_worker_id()}",
        )

    @classmethod
    @lru_cache(maxsize=None)
    def get_or_create(
        cls, controller_handle: ActorHandle, event_loop: AbstractEventLoop
    ) -> "SharedRouterLongPollClient":
        shared = cls(controller_handle=controller_handle, event_loop=event_loop)
        logger.info(f"Started {shared}.")
        return shared

    def update_deployment_targets(
        self,
        deployment_target_info: DeploymentTargetInfo,
        deployment_id: DeploymentID,
    ) -> None:
        for router in self.routers[deployment_id]:
            router.update_deployment_targets(deployment_target_info)
            router.long_poll_client.stop()

    def update_deployment_config(
        self, deployment_config: DeploymentConfig, deployment_id: DeploymentID
    ) -> None:
        for router in self.routers[deployment_id]:
            router.update_deployment_config(deployment_config)
            router.long_poll_client.stop()

    def register(self, router: AsyncioRouter) -> None:
        # We need to run the underlying method in the same event loop that runs
        # the long poll loop, because we need to mutate the mapping of routers,
        # which are also being iterated over by the key listener callbacks.
        # If those happened concurrently in different threads,
        # we could get a `RuntimeError: Set changed size during iteration`.
        # See https://github.com/ray-project/ray/pull/53613 for more details.
        self.event_loop.call_soon_threadsafe(self._register, router)

    def _register(self, router: AsyncioRouter) -> None:
        self.routers[router.deployment_id].add(router)

        # Remove the entries for any deployment ids that no longer have any routers.
        # The WeakSets will automatically lose track of Routers that get GC'd,
        # but the outer dict will keep the key around, so we need to clean up manually.
        # Note the list(...) to avoid mutating self.routers while iterating over it.
        for deployment_id, routers in list(self.routers.items()):
            if not routers:
                self.routers.pop(deployment_id)

        # Register the new listeners on the long poll client.
        # Some of these listeners may already exist, but it's safe to add them again.
        key_listeners = {
            (LongPollNamespace.DEPLOYMENT_TARGETS, deployment_id): partial(
                self.update_deployment_targets, deployment_id=deployment_id
            )
            for deployment_id in self.routers.keys()
        } | {
            (LongPollNamespace.DEPLOYMENT_CONFIG, deployment_id): partial(
                self.update_deployment_config, deployment_id=deployment_id
            )
            for deployment_id in self.routers.keys()
        }
        self.long_poll_client.add_key_listeners(key_listeners)


class CurrentLoopRouter(Router):
    """Wrapper class that runs an AsyncioRouter on the current asyncio loop.
    Note that this class is NOT THREAD-SAFE, and all methods are expected to be
    invoked from a single asyncio event loop.
    """

    def __init__(self, **passthrough_kwargs):
        assert (
            "event_loop" not in passthrough_kwargs
        ), "CurrentLoopRouter uses the current event loop."

        self._asyncio_loop = asyncio.get_running_loop()
        self._asyncio_router = AsyncioRouter(
            event_loop=self._asyncio_loop,
            _request_router_initialized_event=asyncio.Event(),
            **passthrough_kwargs,
        )

    def running_replicas_populated(self) -> bool:
        return self._asyncio_router.running_replicas_populated()

    def assign_request(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> asyncio.Future[ReplicaResult]:
        return self._asyncio_loop.create_task(
            self._asyncio_router.assign_request(
                request_meta, *request_args, **request_kwargs
            ),
        )

    @asynccontextmanager
    async def choose_replica(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> AsyncIterator[ReplicaSelection]:
        """Delegate to AsyncioRouter's choose_replica."""
        async with self._asyncio_router.choose_replica(
            request_meta, *request_args, **request_kwargs
        ) as selection:
            yield selection

    def dispatch(
        self,
        selection: ReplicaSelection,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> asyncio.Future[ReplicaResult]:
        """Dispatch request to a previously selected replica.

        Returns an asyncio.Future wrapping the async dispatch call.
        """
        try:
            selection._mark_dispatched()
        except Exception as exc:
            future = self._asyncio_loop.create_future()
            future.set_exception(exc)
            return future

        return self._asyncio_loop.create_task(
            self._asyncio_router._dispatch_to_marked_selection(
                selection, request_meta, *request_args, **request_kwargs
            )
        )

    async def broadcast(
        self,
        request_meta: RequestMetadata,
        *request_args,
        **request_kwargs,
    ) -> List[ReplicaResult]:
        return await self._asyncio_router.broadcast(
            request_meta, *request_args, **request_kwargs
        )

    def shutdown(self) -> asyncio.Future:
        return self._asyncio_loop.create_task(self._asyncio_router.shutdown())