import asyncio import concurrent.futures import errno import functools import inspect import logging import math import os import pickle import threading import time import traceback import warnings from collections import defaultdict, deque from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass from functools import wraps from importlib import import_module from typing import ( Any, AsyncGenerator, Callable, Dict, Generator, List, Optional, Set, Tuple, Union, ) import grpc import starlette.responses from anyio import to_thread from fastapi import Request from starlette.types import ASGIApp, Receive, Scope, Send import ray from ray import cloudpickle from ray._common.filters import CoreContextFilter from ray._common.utils import get_or_create_event_loop from ray.actor import ActorClass, ActorHandle from ray.dag.py_obj_scanner import _PyObjScanner from ray.remote_function import RemoteFunction from ray.serve import metrics from ray.serve._private.common import ( RUNNING_REQUESTS_KEY, DeploymentID, ReplicaID, ReplicaMetricReport, ReplicaQueueLengthInfo, RequestMetadata, RequestProtocol, ServeComponentType, StreamingHTTPRequest, gRPCRequest, gRPCStreamingRequest, ) from ray.serve._private.config import DeploymentConfig from ray.serve._private.constants import ( GRPC_CONTEXT_ARG_NAME, HEALTH_CHECK_METHOD, HEALTHY_MESSAGE, RAY_SERVE_AUTOSCALING_METRIC_RECORD_INTERVAL_FACTOR, RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE, RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S, RAY_SERVE_DIRECT_INGRESS_PORT_RETRY_COUNT, RAY_SERVE_ENABLE_DIRECT_INGRESS, RAY_SERVE_ENABLE_HA_PROXY, RAY_SERVE_HAPROXY_METRICS_ENABLED, RAY_SERVE_METRICS_EXPORT_INTERVAL_MS, RAY_SERVE_RECORD_AUTOSCALING_STATS_TIMEOUT_S, RAY_SERVE_REPLICA_GRPC_MAX_MESSAGE_LENGTH, RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_NUM_BUCKETS, RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_REPORT_INTERVAL_S, RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_WINDOW_S, RAY_SERVE_REPLICA_UTILIZATION_NUM_BUCKETS, RAY_SERVE_REPLICA_UTILIZATION_REPORT_INTERVAL_S, RAY_SERVE_REPLICA_UTILIZATION_WINDOW_S, RAY_SERVE_REQUEST_PATH_LOG_BUFFER_SIZE, RAY_SERVE_RUN_SYNC_IN_THREADPOOL, RAY_SERVE_RUN_SYNC_IN_THREADPOOL_WARNING, RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD, RECONFIGURE_METHOD, RECORD_REPLICA_METADATA_METHOD, REQUEST_LATENCY_BUCKETS_MS, REQUEST_ROUTING_STATS_METHOD, SERVE_CONTROLLER_NAME, SERVE_HTTP_REQUEST_ID_HEADER, SERVE_LOG_APPLICATION, SERVE_LOG_COMPONENT, SERVE_LOG_DEPLOYMENT, SERVE_LOG_REPLICA, SERVE_LOG_REQUEST_ID, SERVE_LOG_ROUTE, SERVE_LOGGER_NAME, SERVE_NAMESPACE, USER_HEALTH_CHECK_PROBE_INTERVAL_S, USER_HEALTH_CHECK_PROBE_MAX_FAIL, USER_HEALTH_CHECK_PROBE_TIMEOUT_S, ) from ray.serve._private.default_impl import ( create_replica_impl, create_replica_metrics_manager, ) from ray.serve._private.direct_ingress_grpc_util import gRPCDIReceiveStream from ray.serve._private.direct_ingress_http_util import ASGIDIReceiveProxy from ray.serve._private.event_loop_monitoring import EventLoopMonitor from ray.serve._private.grpc_util import ( get_grpc_response_status, set_grpc_code_and_details, start_grpc_server, ) from ray.serve._private.http_util import ( ASGIAppReplicaWrapper, ASGIArgs, ASGIReceiveProxy, MessageQueue, Response, configure_http_middlewares, configure_http_options_with_defaults, convert_object_to_asgi_messages, parse_disconnect_disabled_header, parse_request_timeout_header, parse_session_id_header, start_asgi_http_server, ) from ray.serve._private.logging_utils import ( access_log_msg, configure_component_logger, configure_component_memory_profiler, format_client_address, format_grpc_peer_address, get_component_logger_file_path, ) from ray.serve._private.metrics_utils import InMemoryMetricsStore, MetricsPusher from ray.serve._private.proxy_request_response import ResponseStatus, gRPCStreamingType from ray.serve._private.replica_response_generator import ReplicaResponseGenerator from ray.serve._private.request_ingress_metrics import RequestIngressMetrics from ray.serve._private.rolling_window import ( RollingWindowAccumulator, RollingWindowMax, ) from ray.serve._private.serialization import RPCSerializer from ray.serve._private.task_consumer import TaskConsumerWrapper from ray.serve._private.thirdparty.get_asgi_route_name import ( extract_route_patterns, get_asgi_route_name, ) from ray.serve._private.tracing_utils import ( TraceContextManager, extract_propagated_context, is_span_recording, is_tracing_enabled, set_http_span_attributes, set_rpc_span_attributes, set_span_attributes, set_span_exception, setup_tracing, ) from ray.serve._private.usage import ServeUsageTag from ray.serve._private.utils import ( Semaphore, _callable_uses_multiplexing, asyncio_grpc_exception_handler, check_obj_ref_ready_nowait, compress_metric_report, generate_request_id, get_component_file_name, # noqa: F401 is_grpc_enabled, parse_import_path, ) from ray.serve._private.version import DeploymentVersion from ray.serve.config import ( AutoscalingConfig, HTTPOptions, ProxyLocation, gRPCOptions, ) from ray.serve.context import _get_in_flight_requests from ray.serve.deployment import Deployment from ray.serve.exceptions import ( BackPressureError, DeploymentUnavailableError, RayServeException, gRPCStatusError, ) from ray.serve.gang import GangContext from ray.serve.generated.serve_pb2 import ( ASGIRequest, ASGIResponse, HealthzResponse, ListApplicationsResponse, ) from ray.serve.generated.serve_pb2_grpc import add_ASGIServiceServicer_to_server from ray.serve.grpc_util import ( RayServegRPCContext, gRPCInputStream, ) from ray.serve.handle import DeploymentHandle from ray.serve.schema import EncodingType, LoggingConfig, ReplicaRank from ray.types import ObjectRef logger = logging.getLogger(SERVE_LOGGER_NAME) SERVE_BUILD_ASGI_APP_METHOD = "__serve_build_asgi_app__" def _validate_replica_metadata(metadata: Any) -> Dict[str, Any]: """Validate the return value of a user ``record_replica_metadata`` hook. Returns an empty dict for ``None``; raises ``TypeError`` if the hook returned something other than a dict (it must be a JSON-serializable mapping that the controller can propagate to routers). """ if metadata is None: return {} if not isinstance(metadata, dict): raise TypeError( f"{RECORD_REPLICA_METADATA_METHOD} must return a dict, got " f"{type(metadata).__name__}." ) return metadata def _wrap_grpc_call(f): """Decorator that processes grpc methods.""" def serialize(result, metadata): if metadata.is_streaming and metadata.is_http_request: return result else: # Use cached serializer to avoid per-request instantiation overhead serializer = RPCSerializer.get_cached_serializer( metadata.request_serialization, metadata.response_serialization, ) return serializer.dumps_response(result) @wraps(f) async def wrapper( self, request: ASGIRequest, context: grpc.aio.ServicerContext, ): request_metadata = pickle.loads(request.pickled_request_metadata) # Get cached serializer with options from metadata serializer = RPCSerializer.get_cached_serializer( request_metadata.request_serialization, request_metadata.response_serialization, ) request_args = serializer.loads_request(request.request_args) request_kwargs = serializer.loads_request(request.request_kwargs) if request_metadata.is_http_request or request_metadata.is_grpc_request: request_args = (pickle.loads(request_args[0]),) try: result = await f( self, context, request_metadata, *request_args, **request_kwargs ) return ASGIResponse(serialized_message=serialize(result, request_metadata)) except (Exception, asyncio.CancelledError) as e: return ASGIResponse( serialized_message=serializer.dumps_response(e), is_error=True, ) @wraps(f) async def gen_wrapper( self, request: ASGIRequest, context: grpc.aio.ServicerContext, ): request_metadata = pickle.loads(request.pickled_request_metadata) # Get cached serializer with options from metadata serializer = RPCSerializer.get_cached_serializer( request_metadata.request_serialization, request_metadata.response_serialization, ) request_args = serializer.loads_request(request.request_args) request_kwargs = serializer.loads_request(request.request_kwargs) if request_metadata.is_http_request or request_metadata.is_grpc_request: request_args = (pickle.loads(request_args[0]),) try: async for result in f( self, context, request_metadata, *request_args, **request_kwargs ): yield ASGIResponse( serialized_message=serialize(result, request_metadata) ) except (Exception, asyncio.CancelledError) as e: yield ASGIResponse( serialized_message=serializer.dumps_response(e), is_error=True, ) if inspect.isasyncgenfunction(f): return gen_wrapper else: return wrapper ReplicaMetadata = Tuple[ DeploymentConfig, DeploymentVersion, Optional[float], Optional[int], Optional[str], int, int, ReplicaRank, # rank Optional[List[str]], # route_patterns Optional[List[DeploymentID]], # outbound_deployments bool, # has_user_routing_stats_method Optional[GangContext], # gang_context Dict[str, Any], # replica_metadata ] def _load_deployment_def_from_import_path(import_path: str) -> Callable: module_name, attr_name = parse_import_path(import_path) deployment_def = getattr(import_module(module_name), attr_name) # For ray or serve decorated class or function, strip to return # original body. if isinstance(deployment_def, RemoteFunction): deployment_def = deployment_def._function elif isinstance(deployment_def, ActorClass): deployment_def = deployment_def.__ray_metadata__.modified_class elif isinstance(deployment_def, Deployment): logger.warning( f'The import path "{import_path}" contains a ' "decorated Serve deployment. The decorator's settings " "are ignored when deploying via import path." ) deployment_def = deployment_def.func_or_class return deployment_def class ReplicaMetricsManager: """Manages metrics for the replica. A variety of metrics are managed: - Fine-grained metrics are set for every request. - Autoscaling statistics are periodically pushed to the controller. - Queue length metrics are periodically recorded as user-facing gauges. """ PUSH_METRICS_TO_CONTROLLER_TASK_NAME = "push_metrics_to_controller" RECORD_METRICS_TASK_NAME = "record_metrics" SET_REPLICA_REQUEST_METRIC_GAUGE_TASK_NAME = "set_replica_request_metric_gauge" def __init__( self, replica_id: ReplicaID, event_loop: asyncio.BaseEventLoop, autoscaling_config: Optional[AutoscalingConfig], ingress: bool, max_ongoing_requests: int, ): self._replica_id = replica_id self._deployment_id = replica_id.deployment_id self._metrics_pusher = MetricsPusher() self._metrics_store = InMemoryMetricsStore() self._ingress = ingress self._controller_handle = ray.get_actor( SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE ) self._num_ongoing_requests = 0 self._max_ongoing_requests = max_ongoing_requests # Store event loop for scheduling async tasks from sync context self._event_loop = event_loop or asyncio.get_event_loop() # Cache user_callable_wrapper initialization state to avoid repeated runtime checks self._custom_metrics_enabled = False # On first call to _fetch_custom_autoscaling_metrics. Failing validation disables _custom_metrics_enabled self._checked_custom_metrics = False self._record_autoscaling_stats_fn = None # 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 # Request counter (only set on replica startup). self._restart_counter = metrics.Counter( "serve_deployment_replica_starts", description=( "The number of times this replica has been restarted due to failure." ), ) self._restart_counter.inc() # Per-request metrics. self._request_counter = metrics.Counter( "serve_deployment_request_counter", description=( "The number of queries that have been processed in this replica." ), tag_keys=("route",), ) if self._cached_metrics_enabled: self._cached_request_counter = defaultdict(int) self._error_counter = metrics.Counter( "serve_deployment_error_counter", description=( "The number of exceptions that have occurred in this replica." ), tag_keys=("route", "exception_type"), ) if self._cached_metrics_enabled: self._cached_error_counter = defaultdict(int) # log REQUEST_LATENCY_BUCKET_MS logger.debug(f"REQUEST_LATENCY_BUCKETS_MS: {REQUEST_LATENCY_BUCKETS_MS}") self._processing_latency_tracker = metrics.Histogram( "serve_deployment_processing_latency_ms", description="The latency for queries to be processed.", boundaries=REQUEST_LATENCY_BUCKETS_MS, tag_keys=("route",), ) if self._cached_metrics_enabled: self._cached_latencies = defaultdict(deque) self._event_loop.create_task(self._report_cached_metrics_forever()) # Track maximum processing latency over a rolling window. self._max_processing_latency_trackers = defaultdict( lambda: RollingWindowMax( window_duration_s=RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_WINDOW_S, num_buckets=RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_NUM_BUCKETS, ) ) self._max_processing_latency_gauge = metrics.Gauge( "serve_deployment_max_processing_latency_ms", description="The maximum observed time spent processing a query.", tag_keys=("route",), ) self._max_processing_latency_report_interval_s = ( RAY_SERVE_REPLICA_MAX_PROCESSING_LATENCY_REPORT_INTERVAL_S ) self._event_loop.create_task(self._report_max_processing_latency_forever()) self._num_ongoing_requests_gauge = metrics.Gauge( "serve_replica_processing_queries", description="The current number of queries being processed.", ) self.record_autoscaling_stats_failed_counter = metrics.Counter( "serve_record_autoscaling_stats_failed", tag_keys=("exception_name",), description="The number of errored record_autoscaling_stats invocations.", ) self.user_autoscaling_stats_latency_tracker = metrics.Histogram( "serve_user_autoscaling_stats_latency_ms", description=( "Time taken to execute the user-defined autoscaling stats function " "in milliseconds." ), boundaries=REQUEST_LATENCY_BUCKETS_MS, ) # Replica utilization tracking with rolling window. # Tracks total user code execution time over a rolling window to calculate # utilization as: user_code_time / (window_duration * max_ongoing_requests). self._user_code_time_accumulator = RollingWindowAccumulator( window_duration_s=RAY_SERVE_REPLICA_UTILIZATION_WINDOW_S, num_buckets=RAY_SERVE_REPLICA_UTILIZATION_NUM_BUCKETS, ) self._replica_utilization_gauge = metrics.Gauge( "serve_replica_utilization_percent", description=( "Percentage of replica capacity utilized by user code execution " "over a rolling window. Calculated as: " "user_code_time / (window_duration * max_ongoing_requests)." ), ) self._utilization_report_interval_s = ( RAY_SERVE_REPLICA_UTILIZATION_REPORT_INTERVAL_S ) self._event_loop.create_task(self._report_utilization_forever()) self.set_autoscaling_config(autoscaling_config) # Only populated if direct ingress is enabled. self._ingress_metrics: Dict[RequestProtocol, RequestIngressMetrics] = {} self._ingress_ongoing_requests: Dict[RequestProtocol, int] = {} if self._is_direct_ingress: # These ingress metrics share the same names, tag keys, and emission # logic as those collected by the proxy (see RequestIngressMetrics). When # direct ingress is enabled traffic bypasses the proxy, so a given # request is recorded by exactly one of the two. self._ingress_node_id = ray.get_runtime_context().get_node_id() self._ingress_node_ip_address = ray.util.get_node_ip_address() # gRPC ingress metrics are allocated lazily via # `enable_grpc_ingress_metrics()` once the gRPC config has been fetched from # the controller, which is not available at construction time. if self._should_emit_request_ingress_metrics(RequestProtocol.HTTP): self._add_ingress_metrics(RequestProtocol.HTTP) if self._cached_metrics_enabled: # Mapping from protocol -> {request_tags -> value}. self._cached_ingress_request_counter = defaultdict( lambda: defaultdict(int) ) self._cached_ingress_request_error_counter = defaultdict( lambda: defaultdict(int) ) self._cached_deployment_request_error_counter = defaultdict( lambda: defaultdict(int) ) self._cached_ingress_processing_latencies = defaultdict( lambda: defaultdict(deque) ) @property def _is_direct_ingress(self) -> bool: return self._ingress and RAY_SERVE_ENABLE_DIRECT_INGRESS def _should_emit_request_ingress_metrics(self, protocol: RequestProtocol) -> bool: # When HAProxy is enabled, http ingress request metrics are emitted by # the HAProxyManager. return self._is_direct_ingress and not ( RAY_SERVE_ENABLE_HA_PROXY and RAY_SERVE_HAPROXY_METRICS_ENABLED and protocol == RequestProtocol.HTTP ) def _add_ingress_metrics(self, protocol: RequestProtocol): """Allocate metric objects and ongoing-request counter for a protocol.""" self._ingress_metrics[protocol] = RequestIngressMetrics( protocol, source="ingress", node_id=self._ingress_node_id, node_ip_address=self._ingress_node_ip_address, ) self._ingress_ongoing_requests[protocol] = 0 def enable_grpc_ingress_metrics(self): """Allocate gRPC ingress metrics for a direct-ingress replica. gRPC config is fetched from the controller after this manager is constructed, so gRPC ingress metrics are allocated here (from the replica's server-start path) rather than in `__init__`. """ if not self._is_direct_ingress: return if RequestProtocol.GRPC not in self._ingress_metrics: self._add_ingress_metrics(RequestProtocol.GRPC) def _report_cached_metrics(self): for route, count in self._cached_request_counter.items(): self._request_counter.inc(count, tags={"route": route}) self._cached_request_counter.clear() for (route, exception_type), count in self._cached_error_counter.items(): self._error_counter.inc( count, tags={"route": route, "exception_type": exception_type} ) self._cached_error_counter.clear() for route, latencies in self._cached_latencies.items(): for latency_ms in latencies: self._processing_latency_tracker.observe( latency_ms, tags={"route": route} ) self._cached_latencies.clear() self._num_ongoing_requests_gauge.set(self._num_ongoing_requests) if not self._is_direct_ingress: return for protocol in self._ingress_metrics: protocol_metrics = self._ingress_metrics[protocol] protocol_metrics.set_num_ongoing_requests( self._ingress_ongoing_requests[protocol] ) for request_tags, count in self._cached_ingress_request_counter[ protocol ].items(): protocol_metrics.request_counter.inc(count, tags=dict(request_tags)) for request_tags, count in self._cached_ingress_request_error_counter[ protocol ].items(): protocol_metrics.request_error_counter.inc( count, tags=dict(request_tags) ) for request_tags, count in self._cached_deployment_request_error_counter[ protocol ].items(): protocol_metrics.deployment_request_error_counter.inc( count, tags=dict(request_tags) ) for latency_tags, latencies in self._cached_ingress_processing_latencies[ protocol ].items(): for latency_ms in latencies: protocol_metrics.processing_latency_tracker.observe( latency_ms, tags=dict(latency_tags) ) self._cached_ingress_request_counter.clear() self._cached_ingress_request_error_counter.clear() self._cached_deployment_request_error_counter.clear() self._cached_ingress_processing_latencies.clear() 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) async def shutdown(self): """Stop periodic background tasks.""" await self._metrics_pusher.graceful_shutdown() def start_metrics_pusher(self): self._metrics_pusher.start() # Push autoscaling metrics to the controller periodically. self._metrics_pusher.register_or_update_task( self.PUSH_METRICS_TO_CONTROLLER_TASK_NAME, self._push_autoscaling_metrics, self._autoscaling_config.metrics_interval_s, ) # Collect autoscaling metrics locally periodically. record_interval_s = ( self._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_async, min(record_interval_s, self._autoscaling_config.metrics_interval_s), ) def should_collect_ongoing_requests(self) -> bool: """Determine if replicas should collect ongoing request metrics. ┌────────────────────────────────────────────────────────────────┐ │ Replica-based metrics collection │ ├────────────────────────────────────────────────────────────────┤ │ │ │ Client Handle Replicas │ │ ┌──────┐ ┌────────┐ │ │ │ App │─────>│ Handle │────┬───>┌─────────┐ │ │ │ │ │ Tracks │ │ │ Replica │ │ │ └──────┘ │ Queued │ │ │ 1 │ │ │ │Requests│ │ │ Tracks │ │ │ └────────┘ │ │ Running │ │ │ │ │ └─────────┘ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ┌─────────┐ │ │ │ └───>│ Replica │ │ │ │ │ 2 │ │ │ │ │ Tracks │ │ │ │ │ Running │ │ │ │ └─────────┘ │ │ │ │ │ │ │ │ │ │ ▼ ▼ │ │ ┌──────────────────────────────┐ │ │ │ Controller │ │ │ │ • Queued metrics (handle) │ │ │ │ • Running metrics (replica1)│ │ │ │ • Running metrics (replica2)│ │ │ └──────────────────────────────┘ │ │ │ └────────────────────────────────────────────────────────────────┘ For direct ingress deployments, metrics must be collected from replicas regardless of whether autoscaling metrics are being collected via handles. This is necessary because direct ingress traffic bypasses deployment handles and goes directly to the replicas. """ if self._is_direct_ingress and self._autoscaling_config: return True return not RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE def set_autoscaling_config(self, autoscaling_config: Optional[AutoscalingConfig]): """Dynamically update autoscaling config.""" self._autoscaling_config = autoscaling_config if self._autoscaling_config and self.should_collect_ongoing_requests(): self.start_metrics_pusher() def enable_custom_autoscaling_metrics( self, custom_metrics_enabled: bool, record_autoscaling_stats_fn: Callable[[], Optional[concurrent.futures.Future]], ): """Runs after the user callable wrapper is initialized to enable autoscaling metrics collection.""" if custom_metrics_enabled: self._custom_metrics_enabled = custom_metrics_enabled self._record_autoscaling_stats_fn = record_autoscaling_stats_fn if self._autoscaling_config: self.start_metrics_pusher() def _change_num_ongoing_requests( self, request_metadata: RequestMetadata, delta: int ) -> None: self._num_ongoing_requests += delta protocol = request_metadata.protocol if ( self._is_direct_ingress and request_metadata.is_direct_ingress and protocol in self._ingress_metrics ): self._ingress_ongoing_requests[protocol] += delta if not self._cached_metrics_enabled: self._num_ongoing_requests_gauge.set(self._num_ongoing_requests) if ( self._is_direct_ingress and request_metadata.is_direct_ingress and protocol in self._ingress_metrics ): self._ingress_metrics[protocol].set_num_ongoing_requests( self._ingress_ongoing_requests[protocol] ) def inc_num_ongoing_requests(self, request_metadata: RequestMetadata) -> None: self._change_num_ongoing_requests(request_metadata, 1) def dec_num_ongoing_requests(self, request_metadata: RequestMetadata) -> None: self._change_num_ongoing_requests(request_metadata, -1) def get_num_ongoing_requests(self) -> int: """Get current total queue length of requests for this replica.""" return self._num_ongoing_requests def set_max_ongoing_requests(self, max_ongoing_requests: int) -> None: """Update max_ongoing_requests when deployment config changes.""" self._max_ongoing_requests = max_ongoing_requests async def _report_max_processing_latency_forever(self) -> None: """Background task to emit max processing latency gauge continuously.""" consecutive_errors = 0 while True: try: await asyncio.sleep(self._max_processing_latency_report_interval_s) for route, tracker in list( self._max_processing_latency_trackers.items() ): max_latency = tracker.get_max() self._max_processing_latency_gauge.set( max_latency, tags={"route": route} ) consecutive_errors = 0 except Exception: logger.exception( "Unexpected error reporting max processing latency 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) async def _report_utilization_forever(self) -> None: """Background task to emit utilization gauge continuously.""" consecutive_errors = 0 while True: try: await asyncio.sleep(self._utilization_report_interval_s) utilization = self._calculate_utilization() self._replica_utilization_gauge.set(utilization) consecutive_errors = 0 except Exception: logger.exception("Unexpected error reporting utilization 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 _calculate_utilization(self) -> float: """Calculate current utilization percentage based on rolling window. Utilization is calculated as: user_code_time / (window_duration * max_ongoing_requests) This represents the percentage of the replica's theoretical maximum capacity that was used for executing user code. """ total_user_code_time_ms = self._user_code_time_accumulator.get_total() # Max capacity = window_duration_ms * max_ongoing_requests window_duration_ms = RAY_SERVE_REPLICA_UTILIZATION_WINDOW_S * 1000 max_capacity_ms = window_duration_ms * self._max_ongoing_requests if max_capacity_ms > 0: utilization_percent = (total_user_code_time_ms / max_capacity_ms) * 100 # Cap at 100% (can theoretically exceed if requests overlap heavily) utilization_percent = min(utilization_percent, 100.0) else: utilization_percent = 0.0 return utilization_percent def record_request_metrics( self, *, route: str, latency_ms: float, is_error: bool, exception_type: Optional[str] = None, ): """Records per-request metrics.""" # Track latency for utilization calculation (rolling window). self._user_code_time_accumulator.add(latency_ms) self._max_processing_latency_trackers[route].add(latency_ms) if self._cached_metrics_enabled: self._cached_latencies[route].append(latency_ms) if is_error: exc_type = exception_type or "Unknown" self._cached_error_counter[(route, exc_type)] += 1 else: self._cached_request_counter[route] += 1 else: self._processing_latency_tracker.observe(latency_ms, tags={"route": route}) if is_error: exc_type = exception_type or "Unknown" self._error_counter.inc( tags={"route": route, "exception_type": exc_type} ) else: self._request_counter.inc(tags={"route": route}) def record_ingress_request_metrics( self, *, protocol: RequestProtocol, method: str, route: str, app_name: str, deployment_name: str, latency_ms: float, is_error: bool, status_code: str, ): """Record per-request metrics.""" if not self._should_emit_request_ingress_metrics(protocol): return if self._cached_metrics_enabled: # Cached path: accumulate per-tag-set counts/latencies keyed by the same # canonical tag schemas used by the direct emit path, and flush them in # `_report_cached_metrics`. request_tags = RequestIngressMetrics.request_tags( route=route, method=method, application=app_name, status_code=status_code, ) self._cached_ingress_request_counter[protocol][ frozenset(request_tags.items()) ] += 1 self._cached_ingress_processing_latencies[protocol][ frozenset(request_tags.items()) ].append(latency_ms) if is_error: request_error_tags = RequestIngressMetrics.request_error_tags( route=route, method=method, application=app_name, status_code=status_code, ) deployment_error_tags = RequestIngressMetrics.deployment_error_tags( route=route, method=method, application=app_name, status_code=status_code, deployment=deployment_name, ) self._cached_ingress_request_error_counter[protocol][ frozenset(request_error_tags.items()) ] += 1 self._cached_deployment_request_error_counter[protocol][ frozenset(deployment_error_tags.items()) ] += 1 else: self._ingress_metrics[protocol].record_request( route=route, method=method, application=app_name, status_code=status_code, latency_ms=latency_ms, is_error=is_error, deployment_name=deployment_name, ) def _push_autoscaling_metrics(self) -> Dict[str, Any]: look_back_period = self._autoscaling_config.look_back_period_s self._metrics_store.prune_keys_and_compact_data(time.time() - look_back_period) new_aggregated_metrics = {} new_metrics = {**self._metrics_store.data} if self.should_collect_ongoing_requests(): # Keep the legacy window_avg ongoing requests in the merged metrics dict window_avg = ( self._metrics_store.aggregate_avg([RUNNING_REQUESTS_KEY])[0] or 0.0 ) new_aggregated_metrics.update({RUNNING_REQUESTS_KEY: window_avg}) replica_metric_report = ReplicaMetricReport( replica_id=self._replica_id, timestamp=time.time(), aggregated_metrics=new_aggregated_metrics, metrics=new_metrics, ) 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_replica.remote( compress_metric_report(replica_metric_report) ) ) async def _fetch_custom_autoscaling_metrics( self, ) -> Optional[Dict[str, Union[int, float]]]: try: start_time = time.time() res = await asyncio.wait_for( self._record_autoscaling_stats_fn(), timeout=RAY_SERVE_RECORD_AUTOSCALING_STATS_TIMEOUT_S, ) latency_ms = (time.time() - start_time) * 1000 self.user_autoscaling_stats_latency_tracker.observe(latency_ms) # Perform validation only first call if not self._checked_custom_metrics: # Enforce return type to be Dict[str, Union[int, float]] if not isinstance(res, dict): logger.error( f"User autoscaling stats method returned {type(res).__name__}, " f"expected Dict[str, Union[int, float]]. Disabling autoscaling stats." ) self._custom_metrics_enabled = False return None for key, value in res.items(): if not isinstance(value, (int, float)): logger.error( f"User autoscaling stats method returned invalid value type " f"{type(value).__name__} for key '{key}', expected int or float. " f"Disabling autoscaling stats." ) self._custom_metrics_enabled = False return None self._checked_custom_metrics = True return res except asyncio.TimeoutError as e: logger.error( f"Replica autoscaling stats timed out after {RAY_SERVE_RECORD_AUTOSCALING_STATS_TIMEOUT_S}s." ) self.record_autoscaling_stats_failed_counter.inc( tags={"exception_name": e.__class__.__name__} ) except Exception as e: logger.error(f"Replica autoscaling stats failed. {e}") self.record_autoscaling_stats_failed_counter.inc( tags={"exception_name": e.__class__.__name__} ) return None async def _add_autoscaling_metrics_point_async(self) -> None: metrics_dict = {} if self.should_collect_ongoing_requests(): metrics_dict = {RUNNING_REQUESTS_KEY: self._num_ongoing_requests} # Use cached availability flag to avoid repeated runtime checks if self._custom_metrics_enabled: custom_metrics = await self._fetch_custom_autoscaling_metrics() if custom_metrics: metrics_dict.update(custom_metrics) self._metrics_store.add_metrics_point( metrics_dict, time.time(), ) StatusCodeCallback = Callable[[str], None] class Replica: def __init__( self, replica_id: ReplicaID, deployment_def: Callable, init_args: Tuple, init_kwargs: Dict, deployment_config: DeploymentConfig, version: DeploymentVersion, ingress: bool, route_prefix: str, is_ingress_request_router: bool = False, ): self._version = version self._replica_id = replica_id self._deployment_id = replica_id.deployment_id self._deployment_config = deployment_config self._ingress = ingress self._is_ingress_request_router = is_ingress_request_router self._route_prefix = route_prefix self._component_name = f"{self._deployment_id.name}" if self._deployment_id.app_name: self._component_name = ( f"{self._deployment_id.app_name}_" + self._component_name ) self._component_id = self._replica_id.unique_id self._configure_logger_and_profilers(self._deployment_config.logging_config) self._event_loop = get_or_create_event_loop() actor_id = ray.get_runtime_context().get_actor_id() self._user_callable_wrapper = UserCallableWrapper( deployment_def, init_args, init_kwargs, deployment_id=self._deployment_id, run_sync_methods_in_threadpool=RAY_SERVE_RUN_SYNC_IN_THREADPOOL, run_user_code_in_separate_thread=RAY_SERVE_RUN_USER_CODE_IN_SEPARATE_THREAD, local_testing_mode=False, deployment_config=deployment_config, actor_id=actor_id, ray_actor_options=self._version.ray_actor_options, ) self._semaphore = Semaphore(lambda: self.max_ongoing_requests) # Guards against calling the user's callable constructor multiple times. self._user_callable_initialized = False self._user_callable_initialized_lock = asyncio.Lock() self._initialization_latency: Optional[float] = None # Track deployment handles created dynamically via get_deployment_handle() self._dynamically_created_handles: Set[DeploymentID] = set() # Flipped to `True` when health checks pass and `False` when they fail. May be # used by replica subclass implementations. self._healthy = False # Flipped to `True` once graceful shutdown is initiated. May be used by replica # subclass implementations. self._shutting_down = False # Gang context for this replica. self._gang_context: Optional[GangContext] = None # Static, immutable per-replica metadata captured once at init time # via the user's `record_replica_metadata` hook (if defined). self._replica_metadata: Dict[str, Any] = {} # Will be populated with the wrapped ASGI app if the user callable is an # `ASGIAppReplicaWrapper` (i.e., they are using the FastAPI integration). self._user_callable_asgi_app: Optional[ASGIApp] = None # Set metadata for logs and metrics. # servable_object will be populated in `initialize_and_get_metadata`. self._set_internal_replica_context(servable_object=None, rank=None) self._metrics_manager = create_replica_metrics_manager( replica_id=replica_id, event_loop=self._event_loop, autoscaling_config=self._deployment_config.autoscaling_config, ingress=ingress, max_ongoing_requests=self._deployment_config.max_ongoing_requests, ) # Start event loop monitoring for the replica's main event loop. self._main_loop_monitor = EventLoopMonitor( component=EventLoopMonitor.COMPONENT_REPLICA, loop_type=EventLoopMonitor.LOOP_TYPE_MAIN, actor_id=actor_id, extra_tags={ "deployment": self._deployment_id.name, "application": self._deployment_id.app_name, }, ) self._main_loop_monitor.start(self._event_loop) self._internal_grpc_port: Optional[int] = None self._docs_path: Optional[str] = None self._http_port: Optional[int] = None self._grpc_port: Optional[int] = None self._rank: Optional[ReplicaRank] = None # gRPC server for inter-deployment communication self._server = grpc.aio.server( options=[ ( "grpc.max_receive_message_length", RAY_SERVE_REPLICA_GRPC_MAX_MESSAGE_LENGTH, ) ] ) # Silence spammy false positive errors from gRPC Python self._event_loop.set_exception_handler(asyncio_grpc_exception_handler) try: is_tracing_setup_successful = setup_tracing( component_type=ServeComponentType.REPLICA, component_name=self._component_name, component_id=self._component_id, ) if is_tracing_setup_successful: logger.info("Successfully set up tracing for replica") except Exception as e: logger.warning( f"Failed to set up tracing: {e}. " "The replica will continue running, but traces will not be exported." ) self._controller_handle = ray.get_actor( SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE ) # get node ID self._node_id = ray.get_runtime_context().get_node_id() self._http_options: Optional[HTTPOptions] = None self._grpc_options: Optional[gRPCOptions] = None self._direct_ingress_http_server_task: Optional[asyncio.Task] = None self._direct_ingress_grpc_server_task: Optional[asyncio.Task] = None # Server objects backing the tasks above, used for graceful shutdown # (uvicorn.Server and gRPCGenericServer, respectively). self._direct_ingress_http_server = None self._direct_ingress_grpc_server = None # Set after the graceful shutdown drain completes; new handle-path # requests are then rejected (the router retries them elsewhere). self._quiescing = False self._num_queued_requests = 0 self._reserved_slots: Set[str] = set() @property def max_ongoing_requests(self) -> int: return self._deployment_config.max_ongoing_requests def get_num_ongoing_requests(self) -> int: return self._metrics_manager.get_num_ongoing_requests() + len( self._reserved_slots ) async def reserve_slot( self, request_metadata: RequestMetadata, slot_token: str ) -> Tuple[bool, int]: """Reserve replica capacity for a future dispatch call.""" if request_metadata.is_direct_ingress: raise RuntimeError( "Slot reservation is not supported for direct-ingress requests." ) if not self._can_accept_request(request_metadata): return False, self.get_num_ongoing_requests() await self._semaphore.acquire() self._reserved_slots.add(slot_token) return True, self.get_num_ongoing_requests() def release_slot(self, slot_token: str) -> Tuple[bool, int]: """Release replica capacity reserved by choose_replica().""" if slot_token not in self._reserved_slots: return False, self.get_num_ongoing_requests() self._reserved_slots.remove(slot_token) self._semaphore.release() return True, self.get_num_ongoing_requests() def get_metadata(self) -> ReplicaMetadata: current_rank = ray.serve.context._get_internal_replica_context().rank # Extract route patterns from ASGI app if available route_patterns = None if self._user_callable_asgi_app is not None: # _user_callable_asgi_app is the actual ASGI app (FastAPI/Starlette) # It's set when initialize_callable() returns an ASGI app if hasattr(self._user_callable_asgi_app, "routes"): route_patterns = extract_route_patterns(self._user_callable_asgi_app) has_user_routing_stats_method = ( self._user_callable_wrapper is not None and self._user_callable_wrapper.has_user_routing_stats_method ) return ( self._version.deployment_config, self._version, self._initialization_latency, self._internal_grpc_port, self._docs_path, self._http_port, self._grpc_port, current_rank, route_patterns, self.list_outbound_deployments(), has_user_routing_stats_method, self._gang_context, self._replica_metadata, ) def get_dynamically_created_handles(self) -> Set[DeploymentID]: return self._dynamically_created_handles def list_outbound_deployments(self) -> List[DeploymentID]: """List all outbound deployment IDs this replica calls into. This includes: - Handles created via get_deployment_handle() - Handles passed as init args/kwargs to the deployment constructor This is used to determine which deployments are reachable from this replica. The list of DeploymentIDs can change over time as new handles can be created at runtime. Also its not guaranteed that the list of DeploymentIDs are identical across replicas because it depends on user code. Returns: A list of DeploymentIDs that this replica calls into. """ seen_deployment_ids: Set[DeploymentID] = set() # First, collect dynamically created handles for deployment_id in self.get_dynamically_created_handles(): seen_deployment_ids.add(deployment_id) # Get the init args/kwargs init_args = self._user_callable_wrapper._init_args init_kwargs = self._user_callable_wrapper._init_kwargs # Use _PyObjScanner to find all DeploymentHandle objects in: # The init_args and init_kwargs (handles might be passed as init args) scanner = _PyObjScanner(source_type=DeploymentHandle) try: handles = scanner.find_nodes((init_args, init_kwargs)) for handle in handles: deployment_id = handle.deployment_id seen_deployment_ids.add(deployment_id) finally: scanner.clear() return list(seen_deployment_ids) def _set_internal_replica_context( self, *, servable_object: Callable = None, rank: ReplicaRank = None ): # Calculate world_size from deployment config instead of storing it world_size = self._deployment_config.num_replicas # Create callback for registering dynamically created handles def register_handle_callback(deployment_id: DeploymentID) -> None: self._dynamically_created_handles.add(deployment_id) code_version = self._version.code_version ray.serve.context._set_internal_replica_context( replica_id=self._replica_id, servable_object=servable_object, _deployment_config=self._deployment_config, rank=rank, world_size=world_size, handle_registration_callback=register_handle_callback, gang_context=self._gang_context, code_version=code_version, ) def _configure_logger_and_profilers( self, logging_config: Union[None, Dict, LoggingConfig] ): if logging_config is None: logging_config = {} if isinstance(logging_config, dict): logging_config = LoggingConfig(**logging_config) configure_component_logger( component_type=ServeComponentType.REPLICA, component_name=self._component_name, component_id=self._component_id, logging_config=logging_config, buffer_size=RAY_SERVE_REQUEST_PATH_LOG_BUFFER_SIZE, ) configure_component_memory_profiler( component_type=ServeComponentType.REPLICA, component_name=self._component_name, component_id=self._component_id, ) if logging_config.encoding == EncodingType.JSON: # Create logging context for access logs as a performance optimization. # While logging_utils can automatically add Ray core and Serve access log context, # we pre-compute it here since context evaluation is expensive and this context # will be reused for multiple access log entries. ray_core_logging_context = CoreContextFilter.get_ray_core_logging_context() # remove task level log keys from ray core logging context, it would be nice # to have task level log keys here but we are letting those go in favor of # performance optimization. Also we cannot include task level log keys here because # they would referance the current task (__init__) and not the task that is logging. for key in CoreContextFilter.TASK_LEVEL_LOG_KEYS: ray_core_logging_context.pop(key, None) self._access_log_context = { **ray_core_logging_context, SERVE_LOG_DEPLOYMENT: self._component_name, SERVE_LOG_REPLICA: self._component_id, SERVE_LOG_COMPONENT: ServeComponentType.REPLICA, SERVE_LOG_APPLICATION: self._deployment_id.app_name, "skip_context_filter": True, "serve_access_log": True, } else: self._access_log_context = { "skip_context_filter": True, "serve_access_log": True, } def _is_replica_quiescing(self, request_metadata: RequestMetadata) -> bool: # During graceful shutdown, reject new handle-path requests so the # router retries them on another replica. Direct ingress requests are # not rejected: their servers are shutting down gracefully and anything # that still arrives is served to completion. return self._quiescing and not request_metadata.is_direct_ingress def _can_accept_request(self, request_metadata: RequestMetadata) -> bool: if self._is_replica_quiescing(request_metadata): return False if request_metadata.is_direct_ingress: limit = self.max_queued_requests if limit != -1 and self._num_queued_requests >= limit: return False return True # This replica gates concurrent request handling with an asyncio.Semaphore. # Each in-flight request acquires the semaphore. When the number of ongoing # requests reaches max_ongoing_requests, the semaphore becomes locked. # A new request can be accepted if the semaphore is currently unlocked. return not self._semaphore.locked() @contextmanager def _handle_errors_and_metrics( self, request_metadata: RequestMetadata ) -> Generator[StatusCodeCallback, None, None]: start_time = time.time() user_exception = None status_code = None def _status_code_callback(s: str): nonlocal status_code status_code = s try: yield _status_code_callback except asyncio.CancelledError as e: user_exception = e self._on_request_cancelled(request_metadata, e) except Exception as e: user_exception = e logger.exception("Request failed.") self._on_request_failed(request_metadata, e) latency_ms = (time.time() - start_time) * 1000 self._record_errors_and_metrics( user_exception, status_code, latency_ms, request_metadata ) if user_exception is not None and not request_metadata.is_direct_ingress: raise user_exception from None def _record_errors_and_metrics( self, user_exception: Optional[BaseException], status_code: Optional[str], latency_ms: float, request_metadata: RequestMetadata, ): http_method = request_metadata._http_method route = request_metadata.route call_method = request_metadata.call_method if user_exception is None: status_str = "OK" elif isinstance(user_exception, asyncio.CancelledError): status_str = "CANCELLED" else: status_str = "ERROR" # Mutating self._access_log_context is not thread safe, but since this # is only called from the same thread, it is safe. Mutating the same object # because creating a new dict is expensive. self._access_log_context[SERVE_LOG_ROUTE] = route self._access_log_context[SERVE_LOG_REQUEST_ID] = request_metadata.request_id logger.info( access_log_msg( method=http_method or "CALL", route=route if self._ingress and route else call_method, # Prefer the HTTP status code if it was populated. status=status_code or status_str, latency_ms=latency_ms, client=request_metadata._client, ), extra=self._access_log_context, ) exception_type = ( type(user_exception).__name__ if user_exception is not None else None ) self._metrics_manager.record_request_metrics( route=route, latency_ms=latency_ms, is_error=user_exception is not None, exception_type=exception_type, ) if is_span_recording(): if request_metadata.is_http_request: set_http_span_attributes( method=request_metadata._http_method, status_code=status_code, route=request_metadata.route, ) else: set_rpc_span_attributes( system=request_metadata._request_protocol, method=request_metadata.call_method, status_code=status_code, service=self._deployment_id.name, ) if user_exception is not None: set_span_exception(user_exception, escaped=False) # Record ingress metrics for direct ingress requests if request_metadata.is_direct_ingress and status_code is not None: if request_metadata.is_grpc_request: protocol = RequestProtocol.GRPC # Match the proxy's `method` tag (the full gRPC service method). # gRPC status codes are names (e.g. "OK", "UNAVAILABLE"); anything # other than OK is an error, mirroring ResponseStatus.is_error. method = request_metadata._grpc_service_method is_error = status_code != grpc.StatusCode.OK.name else: protocol = RequestProtocol.HTTP method = request_metadata._http_method is_error = status_code.startswith(("4", "5")) self._metrics_manager.record_ingress_request_metrics( protocol=protocol, method=method, route=route, app_name=self._deployment_id.app_name, deployment_name=self._deployment_id.name, latency_ms=latency_ms, is_error=is_error, status_code=status_code, ) def _unpack_proxy_args( self, request_metadata: RequestMetadata, request_args: Tuple[Any], request_kwargs: Dict[str, Any], ) -> Tuple[Tuple[Any], Dict[str, Any], Any]: # Extract _ray_trace_ctx from kwargs at the entry point. # # Context: When tracing is enabled, Ray's tracing decorators inject # _ray_trace_ctx into ServeReplica actor method calls. The ServeReplica # actor methods properly handle this, but we # need to extract it before calling user-defined deployment methods. # # Design: We return it so it can be passed to _wrap_request() which # stores it in _RequestContext. Users can then access it via serve.context # if needed (advanced use case), while keeping it out of their method signatures. # # For gRPC requests (when using RAY_SERVE_USE_GRPC_BY_DEFAULT), # the _ray_trace_ctx is not injected into kwargs since there's no Ray actor # call. Instead, the tracing context is passed via request_metadata.tracing_context. # We fall back to using that if _ray_trace_ctx is not in kwargs. ray_trace_ctx = request_kwargs.pop("_ray_trace_ctx", None) if ray_trace_ctx is None: ray_trace_ctx = request_metadata.tracing_context if request_metadata.is_http_request: assert len(request_args) == 1 and isinstance( request_args[0], StreamingHTTPRequest ) request: StreamingHTTPRequest = request_args[0] scope = request.asgi_scope receive = ASGIReceiveProxy( scope, request_metadata, request.receive_asgi_messages ) request_metadata._http_method = scope.get("method", "WS").upper() request_args = (scope, receive) elif request_metadata.is_grpc_request: if len(request_args) == 1 and isinstance( request_args[0], gRPCStreamingRequest ): # Handle client streaming or bidirectional streaming request streaming_request: gRPCStreamingRequest = request_args[0] request_args, request_kwargs = self._setup_grpc_streaming_args( request_metadata, streaming_request ) else: # Handle unary or server streaming request assert len(request_args) == 1 and isinstance( request_args[0], gRPCRequest ) request: gRPCRequest = request_args[0] method_info = self._user_callable_wrapper.get_user_method_info( request_metadata.call_method ) request_args = (request.user_request_proto,) request_kwargs = ( {GRPC_CONTEXT_ARG_NAME: request_metadata.grpc_context} if method_info.takes_grpc_context_kwarg else {} ) return request_args, request_kwargs, ray_trace_ctx def _setup_grpc_streaming_args( self, request_metadata: RequestMetadata, streaming_request: gRPCStreamingRequest, ) -> Tuple[Tuple[Any], Dict[str, Any]]: """Set up request args for gRPC client/bidirectional streaming. Creates a gRPCInputStream that wraps the callback to the proxy, allowing the user method to iterate over incoming request messages. """ # Look up the proxy actor fresh per-request to avoid stale handles # if the proxy actor restarts (same name, new actor ID). This mirrors # the per-request lookup in StreamingHTTPRequest.receive_asgi_messages. proxy_actor = ray.get_actor( streaming_request.proxy_actor_name, namespace=SERVE_NAMESPACE ) # Create a cancel event that will be set when the client cancels cancel_event = asyncio.Event() # Create an async iterator that fetches messages from the proxy async def request_message_iterator(): while True: # Ray handles serialization - no manual pickle needed ( has_more, message, is_cancelled, ) = await proxy_actor.receive_grpc_messages.remote( streaming_request.session_id ) if is_cancelled: # Set the cancel event so is_cancelled() returns True cancel_event.set() if not has_more: break yield message # Create the gRPCInputStream wrapper with the cancel event input_stream = gRPCInputStream( request_message_iterator(), cancel_event=cancel_event ) method_info = self._user_callable_wrapper.get_user_method_info( request_metadata.call_method ) request_args = (input_stream,) request_kwargs = ( {GRPC_CONTEXT_ARG_NAME: request_metadata.grpc_context} if method_info.takes_grpc_context_kwarg else {} ) return request_args, request_kwargs async def handle_request( self, request_metadata: RequestMetadata, *request_args, **request_kwargs ) -> Tuple[bytes, Any]: request_args, request_kwargs, ray_trace_ctx = self._unpack_proxy_args( request_metadata, request_args, request_kwargs ) with self._wrap_request(request_metadata, ray_trace_ctx): async with self._start_request(request_metadata): try: return await self._user_callable_wrapper.call_user_method( request_metadata, request_args, request_kwargs ) except Exception as e: # For gRPC requests, wrap exception with user-set status code. # Non-gRPC requests should preserve the original exception # without creating a self-referential __cause__ chain. self._raise_user_exception(e, request_metadata) async def handle_request_streaming( self, request_metadata: RequestMetadata, *request_args, **request_kwargs ) -> AsyncGenerator[Any, None]: """Generator that is the entrypoint for all `stream=True` handle calls.""" request_args, request_kwargs, ray_trace_ctx = self._unpack_proxy_args( request_metadata, request_args, request_kwargs ) with self._wrap_request( request_metadata, ray_trace_ctx ) as status_code_callback: async with self._start_request(request_metadata): try: if request_metadata.is_http_request: scope, receive = request_args async for msgs in self._user_callable_wrapper.call_http_entrypoint( request_metadata, status_code_callback, scope, receive, ): yield pickle.dumps(msgs) else: async for result in self._user_callable_wrapper.call_user_generator( request_metadata, request_args, request_kwargs, ): yield result except Exception as e: self._raise_user_exception(e, request_metadata) def _maybe_wrap_grpc_exception( self, e: BaseException, request_metadata: RequestMetadata ) -> BaseException: """Wrap exception with gRPCStatusError if user set a status code. For gRPC requests, if the user set a status code on the grpc_context before raising an exception, we wrap the exception with gRPCStatusError to preserve the user's intended status code through the error handling path. """ if request_metadata.is_grpc_request: grpc_context = request_metadata.grpc_context if grpc_context and grpc_context.code(): return gRPCStatusError( original_exception=e, code=grpc_context.code(), details=grpc_context.details(), ) return e def _raise_user_exception( self, e: BaseException, request_metadata: RequestMetadata ) -> None: wrapped_exception = self._maybe_wrap_grpc_exception(e, request_metadata) if wrapped_exception is e: raise e raise wrapped_exception from e async def handle_request_with_rejection( self, request_metadata: RequestMetadata, *request_args, **request_kwargs ): # Reject new requests once quiescing so the router retries them on # another replica. if self._is_replica_quiescing(request_metadata): logger.info( "Replica is shutting down, rejecting request " f"{request_metadata.request_id}.", extra={"log_to_stderr": False}, ) yield ReplicaQueueLengthInfo(False, self.get_num_ongoing_requests()) return # Check if the replica has capacity for the request. if not self._can_accept_request(request_metadata): limit = self.max_ongoing_requests logger.warning( f"Replica at capacity of max_ongoing_requests={limit}, " f"rejecting request {request_metadata.request_id}.", extra={"log_to_stderr": False}, ) yield ReplicaQueueLengthInfo(False, self.get_num_ongoing_requests()) return request_args, request_kwargs, ray_trace_ctx = self._unpack_proxy_args( request_metadata, request_args, request_kwargs ) with self._wrap_request( request_metadata, ray_trace_ctx ) as status_code_callback: async with self._start_request(request_metadata): yield ReplicaQueueLengthInfo( accepted=True, # NOTE(edoakes): `_wrap_request` will increment the number # of ongoing requests to include this one, so re-fetch the value. num_ongoing_requests=self.get_num_ongoing_requests(), ) try: if request_metadata.is_http_request: scope, receive = request_args async for msgs in self._user_callable_wrapper.call_http_entrypoint( request_metadata, status_code_callback, scope, receive, ): yield pickle.dumps(msgs) elif request_metadata.is_streaming: async for result in self._user_callable_wrapper.call_user_generator( request_metadata, request_args, request_kwargs, ): yield result else: yield await self._user_callable_wrapper.call_user_method( request_metadata, request_args, request_kwargs ) except Exception as e: self._raise_user_exception(e, request_metadata) async def _on_initialized(self): await self._maybe_start_direct_ingress_servers() current_rank = ray.serve.context._get_internal_replica_context().rank self._set_internal_replica_context( servable_object=self._user_callable_wrapper.user_callable, rank=current_rank, ) # Start the gRPC server for inter-deployment communication add_ASGIServiceServicer_to_server(self, self._server) self._internal_grpc_port = self._server.add_insecure_port("[::]:0") await self._server.start() logger.debug( f"Started inter-deployment gRPC server on port {self._internal_grpc_port}" ) # Save the initialization latency if the replica is initializing # for the first time. if self._initialization_latency is None: self._initialization_latency = time.time() - self._initialization_start_time def _raise_if_multiplexing_with_direct_ingress(self): """Reject model multiplexing on the ingress deployment under direct ingress. Model multiplexing relies on the multiplexed model ID being propagated through the proxy, which direct ingress bypasses (the model ID is never populated). This runs after the user callable is initialized so it also catches multiplexing that is wired up dynamically in the constructor (e.g. `self._load_model = serve.multiplexed(...)(fn)`), which is invisible to the static check performed at deploy time. """ if self._ingress and RAY_SERVE_ENABLE_DIRECT_INGRESS: if _callable_uses_multiplexing(self._user_callable_wrapper._callable): raise RuntimeError( "Model multiplexing (`@serve.multiplexed`) is not supported on the " "ingress deployment when direct ingress or HAProxy is enabled " "(RAY_SERVE_ENABLE_DIRECT_INGRESS)." ) async def initialize( self, deployment_config: Optional[DeploymentConfig], rank: Optional[ReplicaRank], gang_context: Optional[GangContext] = None, ): if gang_context is not None: self._gang_context = gang_context if rank is not None: self._rank = rank self._set_internal_replica_context( servable_object=self._user_callable_wrapper.user_callable, rank=rank ) try: # Ensure that initialization is only performed once. # When controller restarts, it will call this method again. async with self._user_callable_initialized_lock: self._initialization_start_time = time.time() if not self._user_callable_initialized: self._user_callable_asgi_app = ( await self._user_callable_wrapper.initialize_callable() ) self._raise_if_multiplexing_with_direct_ingress() self._user_callable_wrapper.start_user_loop_watchdog( self._event_loop ) # Capture static per-replica metadata exactly once, now that # the user callable is initialized. Unlike routing stats, # this is never polled and is treated as immutable. if self._user_callable_wrapper.has_user_replica_metadata_method: self._replica_metadata = _validate_replica_metadata( await self._user_callable_wrapper.call_user_record_replica_metadata() ) if self._user_callable_asgi_app: self._docs_path = ( self._user_callable_wrapper._callable.docs_path ) await self._on_initialized() self._user_callable_initialized = True if self._user_callable_wrapper is not None: initialized = ( hasattr( self._user_callable_wrapper, "_user_autoscaling_stats" ) and self._user_callable_wrapper._user_autoscaling_stats is not None ) self._metrics_manager.enable_custom_autoscaling_metrics( custom_metrics_enabled=initialized, record_autoscaling_stats_fn=self._user_callable_wrapper.call_record_autoscaling_stats, ) if deployment_config is not None: await self._user_callable_wrapper.set_sync_method_threadpool_limit( deployment_config.max_ongoing_requests ) rank = ray.serve.context._get_internal_replica_context().rank await self._user_callable_wrapper.call_reconfigure( deployment_config.user_config, rank=rank, ) # A new replica should not be considered healthy until it passes # an initial health check. If an initial health check fails, # consider it an initialization failure. await self.check_health() except Exception: raise RuntimeError(traceback.format_exc()) from None async def reconfigure( self, deployment_config: DeploymentConfig, rank: ReplicaRank, route_prefix: Optional[str] = None, ): try: user_config_changed = ( deployment_config.user_config != self._deployment_config.user_config ) rank_changed = rank != self._rank self._rank = rank logging_config_changed = ( deployment_config.logging_config != self._deployment_config.logging_config ) self._deployment_config = deployment_config self._version = DeploymentVersion.from_deployment_version( self._version, deployment_config, route_prefix ) self._metrics_manager.set_autoscaling_config( deployment_config.autoscaling_config ) self._metrics_manager.set_max_ongoing_requests( deployment_config.max_ongoing_requests ) if logging_config_changed: self._configure_logger_and_profilers(deployment_config.logging_config) await self._user_callable_wrapper.set_sync_method_threadpool_limit( deployment_config.max_ongoing_requests ) if user_config_changed or rank_changed: await self._user_callable_wrapper.call_reconfigure( deployment_config.user_config, rank=rank, ) # We need to update internal replica context to reflect the new # deployment_config and rank. self._set_internal_replica_context( servable_object=self._user_callable_wrapper.user_callable, rank=rank, ) self._route_prefix = self._version.route_prefix except Exception: raise RuntimeError(traceback.format_exc()) from None def _on_request_cancelled( self, request_metadata: RequestMetadata, e: asyncio.CancelledError ): """Recursively cancel child requests. This includes all requests that are pending assignment, and gRPC requests that have already been assigned. """ # Cancel child requests pending assignment requests_pending_assignment = ( ray.serve.context._get_requests_pending_assignment( request_metadata.internal_request_id ) ) for task in requests_pending_assignment.values(): task.cancel() # Cancel child requests that have already been assigned. # This is for gRPC requests and direct ingress requests. in_flight_requests = _get_in_flight_requests( request_metadata.internal_request_id ) for replica_result in in_flight_requests.values(): replica_result.cancel() def _on_request_failed(self, request_metadata: RequestMetadata, e: Exception): if ray.util.pdb._is_ray_debugger_post_mortem_enabled(): ray.util.pdb._post_mortem() @asynccontextmanager async def _start_request(self, request_metadata: RequestMetadata): reserved_slot_token = request_metadata._reserved_slot_token if reserved_slot_token: if reserved_slot_token not in self._reserved_slots: raise RuntimeError( "Request tried to consume an unknown reserved slot " f"{reserved_slot_token}." ) self._reserved_slots.remove(reserved_slot_token) else: await self._semaphore.acquire() try: try: self._metrics_manager.inc_num_ongoing_requests(request_metadata) yield finally: self._metrics_manager.dec_num_ongoing_requests(request_metadata) finally: self._semaphore.release() @contextmanager def _track_queued_request(self) -> Generator[Callable[[], None], None, None]: """Count this request against max_queued_requests while it waits for a slot. A direct-ingress request is queued from admission until it acquires an ongoing-request slot. Use this as a `with` block around that wait. Entering adds the request to the count. There are two ways it is removed. 1. The caller invokes the yielded callback once the slot is acquired, so a running request does not keep occupying a queue slot. 2. The block exit removes it when the slot was never acquired, for example a cancellation while the request is still queued. The callback is idempotent, so both happening is safe and a cancelled request cannot leak its count and wedge backpressure. """ self._num_queued_requests += 1 released = False def release() -> None: nonlocal released if not released: released = True self._num_queued_requests -= 1 try: yield release finally: release() async def _drain_ongoing_requests(self, min_draining_period_s: float = 0.0): """Wait until the minimum draining period has elapsed and no ongoing requests remain. The minimum draining period gives load balancers time to deregister this replica; a request admitted during it becomes ongoing and is waited for like any other. """ wait_loop_period_s = self._deployment_config.graceful_shutdown_wait_loop_s deadline = time.monotonic() + min_draining_period_s while True: await asyncio.sleep(wait_loop_period_s) num_ongoing_requests = self.get_num_ongoing_requests() min_period_remaining_s = deadline - time.monotonic() if num_ongoing_requests > 0 or min_period_remaining_s > 0: logger.info( f"Waiting for an additional {wait_loop_period_s}s to shut down: " f"{num_ongoing_requests} ongoing requests, " f"{max(0.0, min_period_remaining_s):.1f}s minimum draining " f"period remaining." ) else: logger.info( "Drained ongoing requests; shutting down servers.", extra={"log_to_stderr": False}, ) break async def shutdown(self): try: self._user_callable_wrapper.stop_user_loop_watchdog() await self._user_callable_wrapper.call_destructor() except: # noqa: E722 # We catch a blanket exception since the constructor may still be # running, so instance variables used by the destructor may not exist. if self._user_callable_initialized: logger.exception( "__del__ ran before replica finished initializing, and " "raised an exception." ) else: logger.exception("__del__ raised an exception.") await self._metrics_manager.shutdown() async def perform_graceful_shutdown(self): self._shutting_down = True # Shutdown budget, mirroring the controller's force-kill deadline (see # `DeploymentReplica.stop`). Each step below consumes from it; once # exhausted, steps degrade to an abrupt close. shutdown_budget_s = self._deployment_config.graceful_shutdown_timeout_s shutdown_deadline = time.monotonic() + shutdown_budget_s def remaining_grace_s() -> float: return max(0.0, shutdown_deadline - time.monotonic()) # If the replica was never initialized it never served traffic, so we # can skip the drain entirely. if self._user_callable_initialized: # In direct ingress mode, hold the replica open at least # RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S so load balancers can # deregister it; the drain also waits for in-flight requests. min_draining_period_s = ( RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S if RAY_SERVE_ENABLE_DIRECT_INGRESS and self._ingress else 0.0 ) await self._drain_ongoing_requests(min_draining_period_s) # Requests can still arrive after the drain (stale routers, keep-alive # connections). Quiesce before reporting shutdown complete: reject new # handle-path requests and shut every server down gracefully, so the # controller's subsequent `ray.kill` can't sever an in-flight request. self._quiescing = True if self._direct_ingress_http_server is not None: # Stop accepting, close idle keep-alives, finish in-flight # requests, then the serve task exits. self._direct_ingress_http_server.should_exit = True try: # On timeout, `wait_for` cancels the server task (abrupt close). await asyncio.wait_for( self._direct_ingress_http_server_task, timeout=remaining_grace_s(), ) except asyncio.TimeoutError: logger.warning( "Direct ingress HTTP server didn't shut down gracefully " "within the graceful shutdown budget; closing it abruptly." ) except Exception: logger.exception("Error shutting down the direct ingress HTTP server.") # Always cancel the task so the listener is guaranteed to stop even if # the graceful close above errored (mirrors the gRPC path below). if self._direct_ingress_http_server_task: self._direct_ingress_http_server_task.cancel() if self._direct_ingress_grpc_server is not None: # Stop accepting new RPCs; wait for in-flight ones up to the # remaining grace. try: await self._direct_ingress_grpc_server.stop(remaining_grace_s()) except Exception: logger.exception("Error shutting down the direct ingress gRPC server.") if self._direct_ingress_grpc_server_task: self._direct_ingress_grpc_server_task.cancel() # Stop the inter-deployment gRPC server (handle-path traffic) so # `ray.kill` can't sever an executing handle request. Gate on the # port: `self._server` is constructed in `__init__` but only started # in `_on_initialized`. if self._internal_grpc_port is not None: try: await self._server.stop(remaining_grace_s()) except Exception: logger.exception( "Error shutting down the inter-deployment gRPC server." ) await self.shutdown() logger.info( "Graceful shutdown complete; replica exiting.", extra={"log_to_stderr": False}, ) async def check_health(self): try: # Runs the user-defined check_health on the user code loop if defined. # Otherwise, if the background watchdog has detected the user loop is # unresponsive, raises immediately. If the watchdog is disabled # (MAX_FAIL=0) or the fail counter hasn't reached the threshold yet, # returns None. f = self._user_callable_wrapper.call_user_health_check() if f is not None: await f self._healthy = True except Exception as e: logger.warning("Replica health check failed.") self._healthy = False raise e from None async def record_routing_stats(self) -> Dict[str, Any]: try: f = self._user_callable_wrapper.call_user_record_routing_stats() if f is not None: return await f return {} except Exception as e: logger.warning("Replica record routing stats failed.") raise e from None @property def max_queued_requests(self) -> int: return self._deployment_config.max_queued_requests async def _maybe_start_direct_ingress_servers(self): if not RAY_SERVE_ENABLE_DIRECT_INGRESS: return if not self._ingress and not self._is_ingress_request_router: return async def allocate_and_start_server(start_server_fn, protocol): """Attempt to allocate a port and start the server with retries.""" is_port_in_use = False for _ in range(RAY_SERVE_DIRECT_INGRESS_PORT_RETRY_COUNT): port = await self._controller_handle.allocate_replica_port.remote( self._node_id, self._replica_id.unique_id, protocol ) logger.info(f"Allocated port {port} for {protocol}") try: server_task_and_server = await start_server_fn(port) logger.info( f"Successfully started {protocol} server on port {port}" ) return port, server_task_and_server except RuntimeError as e: logger.warning( f"Failed to start {protocol} server on port {port}: {e}. Retrying..." ) # `start_asgi_http_server` raises a RuntimeError with the original OSError as the cause. if isinstance(e.__cause__, OSError) and e.__cause__.errno in ( errno.EADDRINUSE, errno.EADDRNOTAVAIL, ): is_port_in_use = True else: is_port_in_use = False # setting block_port to True because we are concluding that the port is # in use by another service on the same node. Blocking port here is a small # optimization to avoid trying to start the server on a the same port # multiple times by other replicas. await self._controller_handle.release_replica_port.remote( self._node_id, self._replica_id.unique_id, port, protocol, block_port=True, ) err_msg = f"Failed to allocate and start {protocol} server after retries" if is_port_in_use: err_msg = f""" Failed to start {protocol} server: port already in use. Suggestion: Ensure that the Ray Serve direct ingress port ranges do not overlap with the Ray worker port range (min_worker_port to max_worker_port). """ raise RuntimeError(err_msg) if self._ingress: self._http_options, self._grpc_options, resolved_proxy_location = ray.get( [ self._controller_handle.get_http_config.remote(), self._controller_handle.get_grpc_config.remote(), self._controller_handle.get_proxy_location.remote(), ] ) else: self._http_options, resolved_proxy_location = ray.get( [ self._controller_handle.get_http_config.remote(), self._controller_handle.get_proxy_location.remote(), ] ) self._grpc_options = None grpc_enabled = self._ingress and is_grpc_enabled(self._grpc_options) # HTTP ingress is enabled unless the resolved proxy placement is Disabled. http_enabled = resolved_proxy_location != ProxyLocation.Disabled # Allocate and start HTTP server if http_enabled: async def start_http_server(port): options = configure_http_middlewares( configure_http_options_with_defaults( HTTPOptions(**{**self._http_options.model_dump(), "port": port}) ) ) return await start_asgi_http_server( self._direct_ingress_asgi, options, event_loop=self._event_loop, enable_so_reuseport=False, ) ( self._http_port, ( self._direct_ingress_http_server_task, self._direct_ingress_http_server, ), ) = await allocate_and_start_server( start_server_fn=start_http_server, protocol=RequestProtocol.HTTP, ) # Allocate and start gRPC server for ingress replicas if enabled. # Ingress request router replicas only need HTTP for /internal/route. if grpc_enabled: self._metrics_manager.enable_grpc_ingress_metrics() async def start_grpc_server_fn(port): options = gRPCOptions( **{**self._grpc_options.model_dump(), "port": port} ) return await start_grpc_server( self._direct_ingress_service_handler_factory, options, event_loop=self._event_loop, enable_so_reuseport=False, ) ( self._grpc_port, ( self._direct_ingress_grpc_server_task, self._direct_ingress_grpc_server, ), ) = await allocate_and_start_server( start_server_fn=start_grpc_server_fn, protocol=RequestProtocol.GRPC, ) started = [] if http_enabled: started.append(f"HTTP server on port {self._http_port}") if grpc_enabled: started.append(f"gRPC server on port {self._grpc_port}") if started: logger.info(f"Started {' and '.join(started)}") @contextmanager def _tracing_context(self, request_metadata: RequestMetadata): if not is_tracing_enabled(): yield return call_method = request_metadata.call_method trace_context = extract_propagated_context(request_metadata.tracing_context) trace_manager = TraceContextManager( trace_name=f"replica_handle_request {self._deployment_id.name} {call_method}", trace_context=trace_context, ) with trace_manager: if is_span_recording(): trace_attributes = { "request_id": request_metadata.request_id, "replica_id": self._replica_id.unique_id, "deployment": self._deployment_id.name, "app": self._deployment_id.app_name, "call_method": request_metadata.call_method, "route": request_metadata.route, "multiplexed_model_id": request_metadata.multiplexed_model_id, "is_streaming": request_metadata.is_streaming, } set_span_attributes(trace_attributes) yield @contextmanager def _wrap_request( self, request_metadata: RequestMetadata, ray_trace_ctx: Optional[Any] = None ) -> Generator[StatusCodeCallback, None, None]: """Context manager that wraps user method calls. 1) Sets the request context var with appropriate metadata. 2) Records the access log message (if not disabled). 3) Records per-request metrics via the metrics manager. """ with self._tracing_context(request_metadata): ray.serve.context._serve_request_context.set( ray.serve.context._RequestContext( route=request_metadata.route, request_id=request_metadata.request_id, _internal_request_id=request_metadata.internal_request_id, app_name=self._deployment_id.app_name, multiplexed_model_id=request_metadata.multiplexed_model_id, session_id=request_metadata.session_id, grpc_context=request_metadata.grpc_context, _client=request_metadata._client, cancel_on_parent_request_cancel=self._ingress and RAY_SERVE_ENABLE_DIRECT_INGRESS, _ray_trace_ctx=ray_trace_ctx, ) ) with self._handle_errors_and_metrics( request_metadata ) as status_code_callback: yield status_code_callback @_wrap_grpc_call async def HandleRequest( self, context: grpc.aio.ServicerContext, request_metadata: RequestMetadata, *request_args, **request_kwargs, ): result = await self.handle_request( request_metadata, *request_args, **request_kwargs ) if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) return result @_wrap_grpc_call async def HandleRequestStreaming( self, context: grpc.aio.ServicerContext, request_metadata: RequestMetadata, *request_args, **request_kwargs, ): async for result in self.handle_request_streaming( request_metadata, *request_args, **request_kwargs ): if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) yield result @_wrap_grpc_call async def HandleRequestWithRejection( self, context: grpc.aio.ServicerContext, request_metadata: RequestMetadata, *request_args, **request_kwargs, ): """gRPC entrypoint for all unary requests with strict max_ongoing_requests enforcement This generator yields a system message indicating if the request was accepted, then the actual response. If an exception occurred while processing the request, whether it's a user exception or an error intentionally raised by Serve, it will be returned as a gRPC response instead of raised directly. """ result_gen = self.handle_request_with_rejection( request_metadata, *request_args, **request_kwargs ) queue_len_info: ReplicaQueueLengthInfo = await result_gen.__anext__() await context.send_initial_metadata( [ ("accepted", str(int(queue_len_info.accepted))), ("num_ongoing_requests", str(queue_len_info.num_ongoing_requests)), ] ) if not queue_len_info.accepted: # NOTE(edoakes): in gRPC, it's not guaranteed that the initial metadata sent # by the server will be delivered for a stream with no messages. Therefore, # we send a dummy message here to ensure it is populated in every case. return b"" result = await result_gen.__anext__() # Consume the result generator to ensure all request operations are completed. async for _ in result_gen: pass if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) return result @_wrap_grpc_call async def HandleRequestWithRejectionStreaming( self, context: grpc.aio.ServicerContext, request_metadata: RequestMetadata, *request_args, **request_kwargs, ) -> AsyncGenerator[Any, None]: """gRPC entrypoint for all streaming requests with strict max_ongoing_requests enforcement This generator yields a system message indicating if the request was accepted, then the actual response(s). If an exception occurred while processing the request, whether it's a user exception or an error intentionally raised by Serve, it will be returned as a gRPC response instead of raised directly. """ result_gen = self.handle_request_with_rejection( request_metadata, *request_args, **request_kwargs ) queue_len_info: ReplicaQueueLengthInfo = await result_gen.__anext__() await context.send_initial_metadata( [ ("accepted", str(int(queue_len_info.accepted))), ("num_ongoing_requests", str(queue_len_info.num_ongoing_requests)), ] ) if not queue_len_info.accepted: # NOTE(edoakes): in gRPC, it's not guaranteed that the initial metadata sent # by the server will be delivered for a stream with no messages. Therefore, # we send a dummy message here to ensure it is populated in every case. yield b"" return async for result in result_gen: if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) yield result async def _dataplane_health_check(self) -> Tuple[bool, str]: healthy, message = True, HEALTHY_MESSAGE if self._shutting_down: healthy = False message = "DRAINING" elif not self._healthy: healthy = False message = "UNHEALTHY" return healthy, message def get_grpc_tracing_context(self, context: RayServegRPCContext): """Populate tracing context for gRPC requests. This method extracts the "traceparent" and "tracestate" metadata from the request headers and sets the tracing context from it. """ if not is_tracing_enabled(): return tracing_ctx = {} traceparent = context.traceparent() if traceparent is not None: tracing_ctx["traceparent"] = traceparent tracestate = context.tracestate() if tracestate is not None: tracing_ctx["tracestate"] = tracestate return tracing_ctx async def _gen_direct_ingress_grpc_response( self, service_method: str, context: grpc._cython.cygrpc._ServicerContext, *, request_input: Any, is_streaming: bool, ) -> AsyncGenerator[bytes, None]: """Shared generator for the four direct-ingress gRPC user-request handlers. Yields serialized response message bytes (zero or more) and sets the final gRPC status on `context` as a side effect. App-mismatch (NOT_FOUND) and backpressure (RESOURCE_EXHAUSTED) short-circuit by yielding nothing. The two axes that distinguish the four cardinalities are passed in: - input axis: `request_input` is the deserialized request proto (unary input) or a `gRPCInputStream` over the native request iterator (client/bidi streaming input). - output axis: `is_streaming` selects `call_user_generator` (many results) vs. `call_user_method` (a single result, wrapped as a 1-item generator). """ start_time = time.time() c = RayServegRPCContext(context) request_id = c.request_id() or generate_request_id() c.set_trailing_metadata([("request_id", request_id)]) # If the request targets a different application, return NOT_FOUND. # If no application is specified, serve this replica's app. requested_app = c.application() if requested_app and requested_app != self._deployment_id.app_name: status = ResponseStatus( code=grpc.StatusCode.NOT_FOUND, message=( f"Application '{requested_app}' not found. Ping " "/ray.serve.RayServeAPIService/ListApplications for available " "applications." ), is_error=True, ) set_grpc_code_and_details(context, status) self._metrics_manager.record_ingress_request_metrics( protocol=RequestProtocol.GRPC, method=service_method, route="", app_name="", deployment_name="", latency_ms=(time.time() - start_time) * 1000.0, is_error=True, status_code=grpc.StatusCode.NOT_FOUND.name, ) return request_metadata = RequestMetadata( request_id=request_id, internal_request_id=generate_request_id(), call_method=service_method.split("/")[-1], _grpc_service_method=service_method, _request_protocol=RequestProtocol.GRPC, grpc_context=c, app_name=self._deployment_id.app_name, # TODO(edoakes): populate this. multiplexed_model_id="", route=self._deployment_id.app_name, tracing_context=self.get_grpc_tracing_context(c), is_streaming=is_streaming, is_direct_ingress=True, _client=format_grpc_peer_address(context.peer()), ) if not self._can_accept_request(request_metadata): status = ResponseStatus( code=grpc.StatusCode.RESOURCE_EXHAUSTED, message="Request dropped due to backpressure", ) set_grpc_code_and_details(context, status) return method_info = self._user_callable_wrapper.get_user_method_info( request_metadata.call_method ) request_args = (request_input,) request_kwargs = ( {GRPC_CONTEXT_ARG_NAME: request_metadata.grpc_context} if method_info.takes_grpc_context_kwarg else {} ) if is_streaming: result_gen = self._user_callable_wrapper.call_user_generator( request_metadata, request_args, request_kwargs ) else: async def call_unary(): yield await self._user_callable_wrapper.call_user_method( request_metadata, request_args, request_kwargs ) result_gen = call_unary() with ( self._wrap_request(request_metadata) as status_code_callback, self._track_queued_request() as release_queue_slot, ): async with self._start_request(request_metadata): # Acquired an ongoing-request slot, so it's running, not queued. release_queue_slot() # Use the generic disconnect/timeout detecting wrapper. replica_response_generator = ReplicaResponseGenerator( result_gen, timeout_s=self._grpc_options.request_timeout_s, ) status = ResponseStatus(code=grpc.StatusCode.OK) try: async for result in replica_response_generator: yield result.SerializeToString() # Apply any user-set code/details/trailing metadata once the # request completes successfully (sent as HTTP/2 trailers). c._set_on_grpc_context(context) except BaseException as e: # For gRPC requests, wrap exception with user-set status code. e = self._maybe_wrap_grpc_exception(e, request_metadata) status = get_grpc_response_status( e, self._grpc_options.request_timeout_s, request_metadata.request_id, ) raise e finally: # Closing `result_gen` runs `call_user_generator`'s `finally`, # which cancels the unit running the user method (its task, # or the inline generator). This is a noop if the generator is # already exhausted. await result_gen.aclose() # Record the status code for both success and error paths so # ingress metrics are emitted for successful gRPC requests. status_code_callback(status.code.name) set_grpc_code_and_details(context, status) async def _maybe_handle_builtin_grpc_service( self, service_method: str, context: grpc._cython.cygrpc._ServicerContext, ) -> Optional[bytes]: """Handle the built-in RayServeAPIService unary-unary methods. `Healthz` and `ListApplications` are health-check-style endpoints that do not dispatch to user code; both run the dataplane health check, set the gRPC status, and record ingress metrics identically -- only the response message differs. Returns the serialized response bytes if `service_method` is one of them, otherwise None (the request targets a user-defined method). """ if service_method not in ( "/ray.serve.RayServeAPIService/Healthz", "/ray.serve.RayServeAPIService/ListApplications", ): return None start_time = time.time() healthy, message = await self._dataplane_health_check() code = grpc.StatusCode.OK if healthy else grpc.StatusCode.UNAVAILABLE context.set_code(code) context.set_details(message) self._metrics_manager.record_ingress_request_metrics( protocol=RequestProtocol.GRPC, method=service_method, route=self._deployment_id.app_name, app_name=self._deployment_id.app_name, deployment_name=self._deployment_id.name, latency_ms=(time.time() - start_time) * 1000.0, is_error=not healthy, status_code=code.name, ) if service_method == "/ray.serve.RayServeAPIService/Healthz": return HealthzResponse(message=message).SerializeToString() # NOTE(edoakes): ListApplications may be used for health checking. It # returns only the app name this replica is serving. return ListApplicationsResponse( application_names=[self._deployment_id.app_name] ).SerializeToString() async def _direct_ingress_unary_unary( self, service_method: str, request_proto: Any, context: grpc._cython.cygrpc._ServicerContext, ) -> bytes: builtin_response = await self._maybe_handle_builtin_grpc_service( service_method, context ) if builtin_response is not None: return builtin_response response_generator = self._gen_direct_ingress_grpc_response( service_method, context, request_input=request_proto, is_streaming=False, ) # Fully consume the generator (so finalizers run) and return the response bytes, # or empty bytes if none were produced (returning `None` to gRPC would cause # serialization errors). result = b"" async for message in response_generator: result = message return result async def _direct_ingress_unary_stream( self, service_method: str, request_proto: Any, context: grpc._cython.cygrpc._ServicerContext, ): response_generator = self._gen_direct_ingress_grpc_response( service_method, context, request_input=request_proto, is_streaming=True, ) async for message in response_generator: yield message def _make_grpc_input_stream( self, request_iterator: Any ) -> Tuple[gRPCInputStream, Optional[gRPCDIReceiveStream]]: """Build the gRPCInputStream the user method iterates over. The native request iterator is bound to the replica's server event loop, so a `gRPCDIReceiveStream` drains it on the server loop and forwards messages to the user method. The gRPCInputStream is built with a `cancel_event` that is set when the client disconnects/errors mid-stream, so the user method sees `is_cancelled()` and a graceful end (matching the proxy path) rather than a raw gRPC error. Args: request_iterator: The native gRPC request iterator (bound to the server event loop) for this client/bidirectional streaming request. Returns: (input_stream, receive_stream) where receive_stream is the bridge draining the native iterator, to be torn down after the request completes. """ cancel_event = asyncio.Event() receive_stream = gRPCDIReceiveStream( request_iterator, self._user_callable_wrapper.event_loop, cancel_event=cancel_event, ) receive_stream.start() return ( gRPCInputStream(receive_stream, cancel_event=cancel_event), receive_stream, ) async def _direct_ingress_stream_unary( self, service_method: str, request_iterator: Any, context: grpc._cython.cygrpc._ServicerContext, ) -> bytes: input_stream, receive_proxy = self._make_grpc_input_stream(request_iterator) try: response_generator = self._gen_direct_ingress_grpc_response( service_method, context, request_input=input_stream, is_streaming=False, ) # Fully consume the generator (so finalizers run) and return the response # bytes, or empty bytes if none were produced (returning `None` to gRPC # would cause serialization errors). result = b"" async for message in response_generator: result = message return result finally: if receive_proxy is not None: receive_proxy.cancel() async def _direct_ingress_stream_stream( self, service_method: str, request_iterator: Any, context: grpc._cython.cygrpc._ServicerContext, ): input_stream, receive_proxy = self._make_grpc_input_stream(request_iterator) try: response_generator = self._gen_direct_ingress_grpc_response( service_method, context, request_input=input_stream, is_streaming=True, ) async for message in response_generator: yield message finally: if receive_proxy is not None: receive_proxy.cancel() def _direct_ingress_service_handler_factory( self, service_method: str, streaming_type: gRPCStreamingType ) -> Callable: if streaming_type == gRPCStreamingType.UNARY_STREAM: async def handler(*args, **kwargs): async for result in self._direct_ingress_unary_stream( service_method, *args, **kwargs ): yield result elif streaming_type == gRPCStreamingType.UNARY_UNARY: async def handler(*args, **kwargs): return await self._direct_ingress_unary_unary( service_method, *args, **kwargs ) elif streaming_type == gRPCStreamingType.STREAM_UNARY: async def handler(*args, **kwargs): return await self._direct_ingress_stream_unary( service_method, *args, **kwargs ) elif streaming_type == gRPCStreamingType.STREAM_STREAM: async def handler(*args, **kwargs): async for result in self._direct_ingress_stream_stream( service_method, *args, **kwargs ): yield result else: raise ValueError(f"Unsupported streaming type: {streaming_type}") return handler def get_asgi_tracing_context(self, headers: List[Tuple[bytes, bytes]]): """Extract tracing context from ASGI request headers. This method extracts both "traceparent" and "tracestate" headers from the request headers to maintain proper trace context propagation. """ if not is_tracing_enabled(): return None tracing_ctx = None for key, value in headers: key_str = key.decode() if key_str in ("traceparent", "tracestate"): tracing_ctx = tracing_ctx or {} tracing_ctx[key_str] = value.decode() return tracing_ctx def _determine_http_route(self, scope: Scope) -> str: # Default to route prefix for consistency with non-DI mode route = self._route_prefix if self._user_callable_asgi_app is not None: try: matched_route = get_asgi_route_name(self._user_callable_asgi_app, scope) if matched_route is not None: route = matched_route except Exception: # If route matching fails, keep the route prefix pass return route def _parse_request_timeout(self, headers: Dict[bytes, bytes]) -> Optional[float]: """Gets the desired request timeout from the headers. If the header is missing or invalid, returns the default request timeout from HttpOptions. If the header is non-positive, timeout is disabled. """ return parse_request_timeout_header( headers, self._http_options.request_timeout_s ) async def _direct_ingress_asgi( self, scope: Scope, receive: Receive, send: Send, ): # NOTE(edoakes): it's important to only start the replica server after the # constructor runs because we are using SO_REUSEPORT. We don't want a new # replica to start handling connections until it's ready to serve traffic. # # This can be loosened to listen on the port but fail health checks once we no # longer rely on SO_REUSEPORT. assert ( self._user_callable_initialized ), "Replica server should only be started *after* the replica is initialized." if self._route_prefix and self._route_prefix != "/": scope["root_path"] = self._route_prefix start_time = time.time() method = scope.get("method", "WS").upper() route = scope.get("path", "") # Handle health check or routes request. if route in ["/-/healthz", "/-/routes"]: healthy, message = await self._dataplane_health_check() status_code = 200 if healthy else 503 if route == "/-/routes" and healthy: # routes endpoint returns only the route prefix andapp name the replica is serving. message = { self._route_prefix: self._deployment_id.app_name, } for msg in convert_object_to_asgi_messages( message, status_code=status_code, ): await send(msg) latency_ms = (time.time() - start_time) * 1000.0 self._metrics_manager.record_ingress_request_metrics( protocol=RequestProtocol.HTTP, method=method, route=route, app_name=self._deployment_id.app_name, deployment_name=self._deployment_id.name, latency_ms=latency_ms, is_error=not healthy, status_code=str(status_code), ) return # If the HTTP path does not match the deployment route prefix, # it is invalid and we should not serve it. Ingress request router # peer deployments (e.g. LLMRouter) have no route prefix; fall # back to "" so any path (including the empty-path ASGI edge case) # matches and downstream user code dispatches. route_prefix = self._route_prefix or "" if not route.startswith(route_prefix): status_code = 404 for msg in convert_object_to_asgi_messages( f"Path '{route}' not found. " "Ping http://.../-/routes for available routes.", status_code=status_code, ): await send(msg) latency_ms = (time.time() - start_time) * 1000.0 self._metrics_manager.record_ingress_request_metrics( protocol=RequestProtocol.HTTP, method=method, route="", app_name="", deployment_name="", latency_ms=latency_ms, is_error=True, status_code=str(status_code), ) return headers = dict(scope["headers"]) request_id = ( # RequestIdMiddleware populates the request ID in the headers if it isn't provided. headers.get(SERVE_HTTP_REQUEST_ID_HEADER.encode("utf-8"), b"").decode( "utf-8" ) or generate_request_id() ) request_disconnect_disabled = parse_disconnect_disabled_header(headers) request_timeout_s = self._parse_request_timeout(headers) session_id = parse_session_id_header(headers) request_metadata = RequestMetadata( request_id=request_id, internal_request_id=generate_request_id(), call_method="__call__", route=self._determine_http_route(scope), app_name=self._deployment_id.app_name, # TODO(edoakes): populate the multiplexed model ID. multiplexed_model_id="", session_id=session_id, is_streaming=True, _request_protocol=RequestProtocol.HTTP, tracing_context=self.get_asgi_tracing_context(scope["headers"]), _http_method=scope.get("method", "WS").upper(), is_direct_ingress=True, _client=format_client_address(scope.get("client")), ) if not self._can_accept_request(request_metadata): # NOTE(abrar): its possible that we drop more requests than actual max_queued_requests # because between incrementing and decrementing the queued requests, we yield to the event loop. for msg in convert_object_to_asgi_messages( "Request dropped due to backpressure", status_code=503, ): await send(msg) return # Optimization: we can avoid creating an async receive task if the client # has disabled handling disconnects for this request. if request_disconnect_disabled: receive_proxy = receive receive_task = None else: receive_proxy = ASGIDIReceiveProxy( scope, receive, self._user_callable_wrapper.event_loop ) receive_task = receive_proxy.fetch_until_disconnect_task() response_started = False response_finished = False first_message_peeked = False with ( self._wrap_request(request_metadata) as status_code_callback, self._track_queued_request() as release_queue_slot, ): async def send_user_message(msg: Dict): nonlocal response_started nonlocal response_finished nonlocal first_message_peeked if not first_message_peeked: first_message_peeked = True if msg["type"] == "http.response.start": status_code_callback(str(msg["status"])) await send(msg) response_started = True # more_body: "Signifies if there is additional content to come (as part # of a Response Body message). If False, and the server is not expecting # Response Trailers, response will be taken as complete and closed. # Optional; if missing defaults to False." # https://asgi.readthedocs.io/en/latest/specs/www.html#response-body-send-event if msg["type"] == "http.response.body" and not msg.get( "more_body", False ): response_finished = True async def call_asgi(): async with self._start_request(request_metadata): # Acquired an ongoing-request slot, so it's running, not queued. release_queue_slot() if ( not self._user_callable_wrapper._run_user_code_in_separate_thread ): user_method_info = ( self._user_callable_wrapper.get_user_method_info( request_metadata.call_method ) ) # `_call_http_entrypoint` will have already called # `send_user_message`, so the ASGI messages will have # already been sent back to the client. await self._user_callable_wrapper._call_http_entrypoint( user_method_info, scope, receive_proxy, send_user_message ) else: async for asgi_messages in self._user_callable_wrapper.call_http_entrypoint( request_metadata, status_code_callback, scope, receive_proxy ): for message in asgi_messages: await send_user_message(message) # Optimization: if Serve doesn't need to handle disconnects and # timeouts for this request, we can avoid event loop overhead by # directly awaiting the user code. if receive_task is None and request_timeout_s is None: return await call_asgi() # Otherwise, we'd always need the call_asgi() task. request_task = asyncio.create_task(call_asgi()) tasks = [request_task] if receive_task is not None: tasks.append(receive_task) done, _ = await asyncio.wait( tasks, timeout=request_timeout_s, return_when=asyncio.FIRST_COMPLETED, ) # NOTE(zcin): it's possible that the request task has finished sending # all ASGI messages, but the task is suspended and before it can fully # complete, the client has sent a disconnect message after the request # is completed. That is why we check for `response_finished` here. if request_task in done or response_finished: if receive_task is not None: receive_task.cancel() await request_task elif receive_task in done: request_task.cancel() status_code_callback("499") if not response_started: msg = ( f"Client for request {request_id} disconnected, " "cancelling request." ) await send_http_response(msg, 499, send) raise asyncio.CancelledError else: request_task.cancel() if receive_task is not None: receive_task.cancel() status_code_callback("408") if not response_started: msg = ( f"Request {request_id} timed out after " f"{self._http_options.request_timeout_s}s." ) await send_http_response(msg, 408, send) raise asyncio.CancelledError def _get_inflight_direct_ingress_task_counts_for_testing(self) -> Dict[str, int]: """Return counts of in-flight direct-ingress asyncio tasks. Used for testing.""" counts = {"request_tasks": 0, "receive_tasks": 0} for task in asyncio.all_tasks(self._event_loop): if task.done(): continue name = getattr(task.get_coro(), "__name__", "") if name == "call_asgi": counts["request_tasks"] += 1 elif name == "_fetch_until_disconnect": counts["receive_tasks"] += 1 return counts async def send_http_response(message, status_code, send): for msg in convert_object_to_asgi_messages( message, status_code=status_code, ): await send(msg) class ReplicaActor: """Actor definition for replicas of Ray Serve deployments. This class defines the interface that the controller and deployment handles (i.e., from proxies and other replicas) use to interact with a replica. All interaction with the user-provided callable is done via the `UserCallableWrapper` class. """ async def __init__( self, replica_id: ReplicaID, serialized_deployment_def: bytes, serialized_init_args: bytes, serialized_init_kwargs: bytes, deployment_config_proto_bytes: bytes, version: DeploymentVersion, ingress: bool, route_prefix: str, is_ingress_request_router: bool = False, ): deployment_config = DeploymentConfig.from_proto_bytes( deployment_config_proto_bytes ) deployment_def = cloudpickle.loads(serialized_deployment_def) if isinstance(deployment_def, str): deployment_def = _load_deployment_def_from_import_path(deployment_def) self._replica_impl: Replica = create_replica_impl( replica_id=replica_id, deployment_def=deployment_def, init_args=cloudpickle.loads(serialized_init_args), init_kwargs=cloudpickle.loads(serialized_init_kwargs), deployment_config=deployment_config, version=version, ingress=ingress, route_prefix=route_prefix, is_ingress_request_router=is_ingress_request_router, ) def push_proxy_handle(self, handle: ActorHandle): # NOTE(edoakes): it's important to call a method on the proxy handle to # initialize its state in the C++ core worker. handle.pong.remote() def get_num_ongoing_requests(self) -> int: """Fetch the number of ongoing requests at this replica (queue length). This runs on a separate thread (using a Ray concurrency group) so it will not be blocked by user code. """ return self._replica_impl.get_num_ongoing_requests() async def reserve_slot( self, request_metadata: RequestMetadata, slot_token: str ) -> Tuple[bool, int]: """Reserve capacity for a future choose_replica/dispatch request.""" return await self._replica_impl.reserve_slot(request_metadata, slot_token) def release_slot(self, slot_token: str) -> Tuple[bool, int]: """Release capacity reserved by choose_replica().""" return self._replica_impl.release_slot(slot_token) async def is_allocated(self) -> str: """poke the replica to check whether it's alive. When calling this method on an ActorHandle, it will complete as soon as the actor has started running. We use this mechanism to detect when a replica has been allocated a worker slot. At this time, the replica can transition from PENDING_ALLOCATION to PENDING_INITIALIZATION startup state. Returns: The PID, actor ID, node ID, node IP, and log filepath id of the replica. """ return ( os.getpid(), ray.get_runtime_context().get_actor_id(), ray.get_runtime_context().get_worker_id(), ray.get_runtime_context().get_node_id(), ray.util.get_node_ip_address(), ray.util.get_node_instance_id(), get_component_logger_file_path(), ) async def was_initialized(self) -> bool: """Whether this replica's user callable has finished initializing. Used by the controller during recovery to detect actors that were created but never received their initial ``initialize_and_get_metadata(rank=...)`` call (e.g., because the previous controller crashed mid-startup). Such an actor has neither a rank nor a fully-initialized user callable, and recovering it would silently complete its initialization with ``rank=None``, breaking rank tracking. The controller can call this method first and skip / kill the actor when it returns False. """ return self._replica_impl._user_callable_initialized def list_outbound_deployments(self) -> Optional[List[DeploymentID]]: return self._replica_impl.list_outbound_deployments() async def _get_inflight_direct_ingress_task_counts_for_testing( self, ) -> Dict[str, int]: return self._replica_impl._get_inflight_direct_ingress_task_counts_for_testing() async def initialize_and_get_metadata( self, deployment_config: DeploymentConfig = None, rank: ReplicaRank = None, gang_context: GangContext = None, ) -> ReplicaMetadata: """Handles initializing the replica. Returns: 5-tuple containing 1. DeploymentConfig of the replica 2. DeploymentVersion of the replica 3. Initialization duration in seconds 4. Port 5. FastAPI `docs_path`, if relevant (i.e. this is an ingress deployment integrated with FastAPI). """ # Unused `_after` argument is for scheduling: passing an ObjectRef # allows delaying this call until after the `_after` call has returned. await self._replica_impl.initialize(deployment_config, rank, gang_context) return self._replica_impl.get_metadata() async def check_health(self): await self._replica_impl.check_health() async def record_routing_stats(self) -> Dict[str, Any]: return await self._replica_impl.record_routing_stats() async def reconfigure( self, deployment_config, rank: ReplicaRank, route_prefix: Optional[str] = None ) -> ReplicaMetadata: await self._replica_impl.reconfigure(deployment_config, rank, route_prefix) return self._replica_impl.get_metadata() def _preprocess_request_args( self, pickled_request_metadata: bytes, request_args: Tuple[Any], ) -> Tuple[RequestMetadata, Tuple[Any]]: request_metadata = pickle.loads(pickled_request_metadata) if request_metadata.is_http_request or request_metadata.is_grpc_request: request_args = (pickle.loads(request_args[0]),) return request_metadata, request_args async def handle_request( self, pickled_request_metadata: bytes, *request_args, **request_kwargs, ) -> Tuple[bytes, Any]: """Entrypoint for `stream=False` calls.""" request_metadata, request_args = self._preprocess_request_args( pickled_request_metadata, request_args ) result = await self._replica_impl.handle_request( request_metadata, *request_args, **request_kwargs ) if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) return result async def handle_request_streaming( self, pickled_request_metadata: bytes, *request_args, **request_kwargs, ) -> AsyncGenerator[Any, None]: """Generator that is the entrypoint for all `stream=True` handle calls.""" request_metadata, request_args = self._preprocess_request_args( pickled_request_metadata, request_args ) async for result in self._replica_impl.handle_request_streaming( request_metadata, *request_args, **request_kwargs ): if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) yield result async def handle_request_with_rejection( self, pickled_request_metadata: bytes, *request_args, **request_kwargs, ) -> AsyncGenerator[Any, None]: """Entrypoint for all requests with strict max_ongoing_requests enforcement. The first response from this generator is always a system message indicating if the request was accepted (the replica has capacity for the request) or rejected (the replica is already at max_ongoing_requests). For non-streaming requests, there will only be one more message, the unary result of the user request handler. For streaming requests, the subsequent messages will be the results of the user request handler (which must be a generator). """ request_metadata, request_args = self._preprocess_request_args( pickled_request_metadata, request_args ) async for result in self._replica_impl.handle_request_with_rejection( request_metadata, *request_args, **request_kwargs ): if isinstance(result, ReplicaQueueLengthInfo): yield pickle.dumps(result) else: if request_metadata.is_grpc_request: result = (request_metadata.grpc_context, result.SerializeToString()) yield result async def handle_request_from_java( self, proto_request_metadata: bytes, *request_args, **request_kwargs, ) -> Any: from ray.serve.generated.serve_pb2 import ( RequestMetadata as RequestMetadataProto, ) proto = RequestMetadataProto.FromString(proto_request_metadata) request_metadata: RequestMetadata = RequestMetadata( request_id=proto.request_id, internal_request_id=proto.internal_request_id, call_method=proto.call_method, multiplexed_model_id=proto.multiplexed_model_id, route=proto.route, ) return await self._replica_impl.handle_request( request_metadata, *request_args, **request_kwargs ) async def perform_graceful_shutdown(self): await self._replica_impl.perform_graceful_shutdown() @dataclass class UserMethodInfo: """Wrapper for a user method and its relevant metadata.""" callable: Callable name: str is_asgi_app: bool takes_any_args: bool takes_grpc_context_kwarg: bool @classmethod def from_callable(cls, c: Callable, *, is_asgi_app: bool) -> "UserMethodInfo": params = inspect.signature(c).parameters return cls( callable=c, name=c.__name__, is_asgi_app=is_asgi_app, takes_any_args=len(params) > 0, takes_grpc_context_kwarg=GRPC_CONTEXT_ARG_NAME in params, ) class UserCallableWrapper: """Wraps a user-provided callable that is used to handle requests to a replica.""" service_unavailable_exceptions = (BackPressureError, DeploymentUnavailableError) def __init__( self, deployment_def: Callable, init_args: Tuple, init_kwargs: Dict, *, deployment_id: DeploymentID, run_sync_methods_in_threadpool: bool, run_user_code_in_separate_thread: bool, local_testing_mode: bool, deployment_config: DeploymentConfig, actor_id: str, ray_actor_options: Optional[Dict] = None, ): if not (inspect.isfunction(deployment_def) or inspect.isclass(deployment_def)): raise TypeError( "deployment_def must be a function or class. Instead, its type was " f"{type(deployment_def)}." ) self._deployment_def = deployment_def self._init_args = init_args self._init_kwargs = init_kwargs self._is_function = inspect.isfunction(deployment_def) self._deployment_id = deployment_id self._local_testing_mode = local_testing_mode self._destructor_called = False self._run_sync_methods_in_threadpool = run_sync_methods_in_threadpool self._run_user_code_in_separate_thread = run_user_code_in_separate_thread self._warned_about_sync_method_change = False self._cached_user_method_info: Dict[str, UserMethodInfo] = {} # This is for performance optimization https://docs.python.org/3/howto/logging.html#optimization self._is_enabled_for_debug = logger.isEnabledFor(logging.DEBUG) # Will be populated in `initialize_callable`. self._callable = None self._user_health_check: Optional[Callable] = None self._user_loop_probe_consecutive_fail_count: int = 0 self._user_loop_probe_task: Optional[asyncio.Task] = None self._deployment_config = deployment_config self._ray_actor_options = ray_actor_options or {} self._user_code_threadpool: Optional[ concurrent.futures.ThreadPoolExecutor ] = None if self._run_user_code_in_separate_thread: # All interactions with user code run on this loop to avoid blocking the # replica's main event loop. self._user_code_event_loop: asyncio.AbstractEventLoop = ( asyncio.new_event_loop() ) # Start event loop monitoring for the user code event loop. # We create the monitor here but start it inside the thread function # so the task is created on the correct thread. self._user_code_loop_monitor = EventLoopMonitor( component=EventLoopMonitor.COMPONENT_REPLICA, loop_type=EventLoopMonitor.LOOP_TYPE_USER_CODE, actor_id=actor_id, extra_tags={ "deployment": self._deployment_id.name, "application": self._deployment_id.app_name, }, ) def _run_user_code_event_loop(): # Required so that calls to get the current running event loop work # properly in user code. asyncio.set_event_loop(self._user_code_event_loop) self._configure_user_code_threadpool() # Start monitoring before run_forever so the task is scheduled. self._user_code_loop_monitor.start(self._user_code_event_loop) self._user_code_event_loop.run_forever() self._user_code_event_loop_thread = threading.Thread( daemon=True, target=_run_user_code_event_loop, ) self._user_code_event_loop_thread.start() else: self._user_code_event_loop = asyncio.get_running_loop() self._user_code_loop_monitor = None self._configure_user_code_threadpool() @property def event_loop(self) -> asyncio.AbstractEventLoop: return self._user_code_event_loop def _user_loop_watchdog_enabled(self) -> bool: """Whether we run the optional background user-loop probe (and may fast-fail HC).""" return ( self._run_user_code_in_separate_thread and self._user_health_check is None and USER_HEALTH_CHECK_PROBE_MAX_FAIL > 0 ) def start_user_loop_watchdog(self, main_loop: asyncio.AbstractEventLoop) -> None: """Periodically probe the user code loop from the replica main loop (opt-in). With user code on a separate thread, a wedged asyncio loop can queue health probes indefinitely. If ``RAY_SERVE_USER_HEALTH_CHECK_PROBE_MAX_FAIL`` > 0, repeated timeouts fail ``check_health`` without waiting for the controller RPC timeout. Applies even when ``RAY_SERVE_RUN_SYNC_IN_THREADPOOL=1``: async work on that loop stays observable regardless of sync thread-pool idle/busy mix. """ if not self._user_loop_watchdog_enabled(): return if self._user_loop_probe_task is not None: return self._user_loop_probe_task = main_loop.create_task(self._user_loop_probe_loop()) def stop_user_loop_watchdog(self) -> None: if self._user_loop_probe_task is not None: self._user_loop_probe_task.cancel() self._user_loop_probe_task = None # Reset so a subsequent start_user_loop_watchdog() doesn't begin with a stale # fail count that would immediately trip call_user_health_check(). self._user_loop_probe_consecutive_fail_count = 0 async def _user_loop_probe_loop(self) -> None: while True: await asyncio.sleep(USER_HEALTH_CHECK_PROBE_INTERVAL_S) fut = asyncio.run_coroutine_threadsafe( asyncio.sleep(0), self._user_code_event_loop ) try: await asyncio.wait_for( asyncio.wrap_future(fut), timeout=USER_HEALTH_CHECK_PROBE_TIMEOUT_S, ) except asyncio.TimeoutError: self._user_loop_probe_consecutive_fail_count += 1 logger.warning( "User event loop probe timed out " f"(failed {self._user_loop_probe_consecutive_fail_count}/" f"{USER_HEALTH_CHECK_PROBE_MAX_FAIL} before health check fails). " f"deployment={self._deployment_id.name} " f"app={self._deployment_id.app_name}", ) continue except Exception as e: self._user_loop_probe_consecutive_fail_count += 1 logger.warning( "User event loop probe failed: " f"{e} ({self._user_loop_probe_consecutive_fail_count}/" f"{USER_HEALTH_CHECK_PROBE_MAX_FAIL}). " f"deployment={self._deployment_id.name} " f"app={self._deployment_id.app_name}", ) continue self._user_loop_probe_consecutive_fail_count = 0 def _get_user_code_threadpool_max_workers(self) -> Optional[int]: num_cpus = self._ray_actor_options.get("num_cpus") if num_cpus is None: return None # Mirror ThreadPoolExecutor default behavior while respecting num_cpus. return min(32, max(1, int(math.ceil(num_cpus))) + 4) def _configure_user_code_threadpool(self) -> None: max_workers = self._get_user_code_threadpool_max_workers() if max_workers is None: return self._user_code_threadpool = concurrent.futures.ThreadPoolExecutor( max_workers=max_workers ) self._user_code_event_loop.set_default_executor(self._user_code_threadpool) def _run_user_code(f: Callable) -> Callable: """Decorator to run a coroutine method on the user code event loop. The method will be modified to be a sync function that returns a `asyncio.Future` if user code is running in a separate event loop. Otherwise, it will return the coroutine directly. """ assert inspect.iscoroutinefunction( f ), "_run_user_code can only be used on coroutine functions." @functools.wraps(f) def wrapper(self, *args, **kwargs) -> Any: coro = f(self, *args, **kwargs) if self._run_user_code_in_separate_thread: fut = asyncio.run_coroutine_threadsafe(coro, self._user_code_event_loop) if self._local_testing_mode: return fut return asyncio.wrap_future(fut) else: return coro return wrapper @_run_user_code async def set_sync_method_threadpool_limit(self, limit: int): # NOTE(edoakes): the limit is thread local, so this must # be run on the user code event loop. to_thread.current_default_thread_limiter().total_tokens = limit def get_user_method_info(self, method_name: str) -> UserMethodInfo: """Get UserMethodInfo for the provided call method name. This method is cached to avoid repeated expensive calls to `inspect.signature`. """ if method_name in self._cached_user_method_info: return self._cached_user_method_info[method_name] if self._is_function: user_method = self._callable elif hasattr(self._callable, method_name): user_method = getattr(self._callable, method_name) else: # Filter to methods that don't start with '__' prefix. def callable_method_filter(attr): if attr.startswith("__"): return False elif not callable(getattr(self._callable, attr)): return False return True methods = list(filter(callable_method_filter, dir(self._callable))) raise RayServeException( f"Tried to call a method '{method_name}' " "that does not exist. Available methods: " f"{methods}." ) info = UserMethodInfo.from_callable( user_method, is_asgi_app=isinstance(self._callable, ASGIAppReplicaWrapper), ) self._cached_user_method_info[method_name] = info return info async def _send_user_result_over_asgi( self, result: Any, asgi_args: ASGIArgs, ): """Handle the result from user code and send it over the ASGI interface. If the result is already a Response type, it is sent directly. Otherwise, it is converted to a custom Response type that handles serialization for common Python objects. """ scope, receive, send = asgi_args.to_args_tuple() if isinstance(result, starlette.responses.Response): await result(scope, receive, send) else: await Response(result).send(scope, receive, send) async def _call_func_or_gen( self, callable: Callable, *, args: Optional[Tuple[Any]] = None, kwargs: Optional[Dict[str, Any]] = None, is_streaming: bool = False, generator_result_callback: Optional[Callable] = None, run_sync_methods_in_threadpool_override: Optional[bool] = None, ) -> Tuple[Any, bool]: """Call the callable with the provided arguments. This is a convenience wrapper that will work for `def`, `async def`, generator, and async generator functions. Returns the result and a boolean indicating if the result was a sync generator that has already been consumed. """ sync_gen_consumed = False args = args if args is not None else tuple() kwargs = kwargs if kwargs is not None else dict() run_sync_in_threadpool = ( self._run_sync_methods_in_threadpool if run_sync_methods_in_threadpool_override is None else run_sync_methods_in_threadpool_override ) is_sync_method = ( inspect.isfunction(callable) or inspect.ismethod(callable) ) and not ( inspect.iscoroutinefunction(callable) or inspect.isasyncgenfunction(callable) ) if is_sync_method and run_sync_in_threadpool: is_generator = inspect.isgeneratorfunction(callable) if is_generator: sync_gen_consumed = True if not is_streaming: # TODO(edoakes): make this check less redundant with the one in # _handle_user_method_result. raise TypeError( f"Method '{callable.__name__}' returned a generator. " "You must use `handle.options(stream=True)` to call " "generators on a deployment." ) def run_callable(): result = callable(*args, **kwargs) if is_generator: for r in result: generator_result_callback(r) result = None return result # NOTE(edoakes): we use anyio.to_thread here because it's what Starlette # uses (and therefore FastAPI too). The max size of the threadpool is # set to max_ongoing_requests in the replica wrapper. # anyio.to_thread propagates ContextVars to the worker thread automatically. result = await to_thread.run_sync(run_callable) else: if ( is_sync_method and not self._warned_about_sync_method_change and run_sync_methods_in_threadpool_override is None ): self._warned_about_sync_method_change = True warnings.warn( RAY_SERVE_RUN_SYNC_IN_THREADPOOL_WARNING.format( method_name=callable.__name__, ) ) result = callable(*args, **kwargs) if inspect.iscoroutine(result): result = await result return result, sync_gen_consumed @property def user_callable(self) -> Optional[Callable]: return self._callable async def _initialize_asgi_callable(self) -> None: self._callable: ASGIAppReplicaWrapper build_asgi_app = getattr(self._callable, SERVE_BUILD_ASGI_APP_METHOD, None) is_late_bound = not hasattr(self._callable, "_asgi_app") if is_late_bound and build_asgi_app is None: raise TypeError( f"ASGI app was not provided to the wrapper and " f"`{SERVE_BUILD_ASGI_APP_METHOD}` is not defined on the deployment." ) if build_asgi_app is not None: app, _ = await self._call_func_or_gen( build_asgi_app, run_sync_methods_in_threadpool_override=False, ) if app is None: raise TypeError( f"`{SERVE_BUILD_ASGI_APP_METHOD}` must return an ASGI app." ) self._callable._set_asgi_app(app) app: ASGIApp = self._callable.app if hasattr(app, "add_exception_handler"): # The reason we need to do this is because BackPressureError is a serve # internal exception and FastAPI doesn't know how to handle it, so it # treats it as a 500 error. With same reasoning, we are not handling # TimeoutError because it's a generic exception the FastAPI knows how # to handle. See https://www.starlette.io/exceptions/ def handle_exception(_: Request, exc: Exception): return self.handle_exception(exc) for exc in self.service_unavailable_exceptions: app.add_exception_handler(exc, handle_exception) await self._callable._run_asgi_lifespan_startup() @_run_user_code async def initialize_callable(self) -> Optional[ASGIApp]: """Initialize the user callable. If the callable is an ASGI app wrapper (e.g., using @serve.ingress), returns the ASGI app object, which may be used *read only* by the caller. """ if self._callable is not None: raise RuntimeError("initialize_callable should only be called once.") # This closure initializes user code and finalizes replica # startup. By splitting the initialization step like this, # we can already access this actor before the user code # has finished initializing. # The supervising state manager can then wait # for allocation of this replica by using the `is_allocated` # method. After that, it calls `reconfigure` to trigger # user code initialization. logger.info( "Started initializing replica.", extra={"log_to_stderr": False}, ) if self._is_function: self._callable = self._deployment_def else: # This allows deployments to define an async __init__ # method (mostly used for testing). self._callable = self._deployment_def.__new__(self._deployment_def) await self._call_func_or_gen( self._callable.__init__, args=self._init_args, kwargs=self._init_kwargs, # Always run the constructor on the main user code thread. run_sync_methods_in_threadpool_override=False, ) if isinstance(self._callable, ASGIAppReplicaWrapper): await self._initialize_asgi_callable() if isinstance(self._callable, TaskConsumerWrapper): self._callable.initialize_callable( self._deployment_config.max_ongoing_requests ) ServeUsageTag.NUM_REPLICAS_USING_ASYNCHRONOUS_INFERENCE.record("1") self._user_health_check = getattr(self._callable, HEALTH_CHECK_METHOD, None) self._user_record_routing_stats = getattr( self._callable, REQUEST_ROUTING_STATS_METHOD, None ) self._user_record_replica_metadata = getattr( self._callable, RECORD_REPLICA_METADATA_METHOD, None ) self._user_autoscaling_stats = getattr( self._callable, "record_autoscaling_stats", None ) logger.info( "Finished initializing replica.", extra={"log_to_stderr": False}, ) return ( self._callable.app if isinstance(self._callable, ASGIAppReplicaWrapper) else None ) def _raise_if_not_initialized(self, method_name: str): if self._callable is None: raise RuntimeError( f"`initialize_callable` must be called before `{method_name}`." ) def call_user_health_check(self) -> Optional[concurrent.futures.Future]: self._raise_if_not_initialized("call_user_health_check") # If the user provided a health check, call it on the user code thread. If user # code blocks the event loop the health check may time out. # # When there is no user-defined health check, health is determined by the # optional watchdog fail counter. if ( self._user_loop_watchdog_enabled() and self._user_loop_probe_consecutive_fail_count >= USER_HEALTH_CHECK_PROBE_MAX_FAIL ): raise RuntimeError( "User event loop unresponsive: probe failed " f"{self._user_loop_probe_consecutive_fail_count} consecutive times " f"(limit {USER_HEALTH_CHECK_PROBE_MAX_FAIL})." ) if self._user_health_check is not None: return self._call_user_health_check() return None @property def has_user_routing_stats_method(self) -> bool: """Whether the user has defined a record_routing_stats method.""" return self._user_record_routing_stats is not None def call_user_record_routing_stats(self) -> Optional[concurrent.futures.Future]: self._raise_if_not_initialized("call_user_record_routing_stats") if self._user_record_routing_stats is not None: return self._call_user_record_routing_stats() return None @property def has_user_replica_metadata_method(self) -> bool: """Whether the user has defined a record_replica_metadata method.""" return self._user_record_replica_metadata is not None def call_user_record_replica_metadata( self, ) -> Optional[concurrent.futures.Future]: self._raise_if_not_initialized("call_user_record_replica_metadata") if self._user_record_replica_metadata is not None: return self._call_user_record_replica_metadata() return None def call_record_autoscaling_stats(self) -> Optional[concurrent.futures.Future]: self._raise_if_not_initialized("call_record_autoscaling_stats") if self._user_autoscaling_stats is not None: return self._call_user_autoscaling_stats() return None @_run_user_code async def _call_user_health_check(self): await self._call_func_or_gen(self._user_health_check) @_run_user_code async def _call_user_record_routing_stats(self) -> Dict[str, Any]: result, _ = await self._call_func_or_gen(self._user_record_routing_stats) return result @_run_user_code async def _call_user_record_replica_metadata(self) -> Dict[str, Any]: result, _ = await self._call_func_or_gen(self._user_record_replica_metadata) return result @_run_user_code async def _call_user_autoscaling_stats(self) -> Dict[str, Union[int, float]]: result, _ = await self._call_func_or_gen(self._user_autoscaling_stats) return result @_run_user_code async def call_reconfigure(self, user_config: Optional[Any], rank: ReplicaRank): self._raise_if_not_initialized("call_reconfigure") # NOTE(edoakes): there is the possibility of a race condition in user code if # they don't have any form of concurrency control between `reconfigure` and # other methods. See https://github.com/ray-project/ray/pull/42159. # NOTE(abrar): The only way to subscribe to rank changes is to provide some user config. # We can relax this in the future as more use cases arise for rank. I am reluctant to # introduce behavior change for a feature we might not need. user_subscribed_to_rank = False if not self._is_function and hasattr(self._callable, RECONFIGURE_METHOD): reconfigure_method = getattr(self._callable, RECONFIGURE_METHOD) params = inspect.signature(reconfigure_method).parameters user_subscribed_to_rank = "rank" in params if user_config is not None or user_subscribed_to_rank: if self._is_function: raise ValueError( "deployment_def must be a class to use user_config or rank" ) elif not hasattr(self._callable, RECONFIGURE_METHOD): raise RayServeException( "user_config or rank specified but deployment " + self._deployment_id + " missing " + RECONFIGURE_METHOD + " method" ) kwargs = {} if user_subscribed_to_rank: # For backwards compatibility, only pass rank if it is an argument to the reconfigure method. kwargs["rank"] = rank await self._call_func_or_gen( getattr(self._callable, RECONFIGURE_METHOD), args=(user_config,), kwargs=kwargs, ) async def _handle_user_method_result( self, result: Any, user_method_info: UserMethodInfo, *, is_streaming: bool, is_http_request: bool, sync_gen_consumed: bool, generator_result_callback: Optional[Callable], asgi_args: Optional[ASGIArgs], ) -> Any: """Postprocess the result of a user method. User methods can be regular unary functions or return a sync or async generator. This method will raise an exception if the result is not of the expected type (e.g., non-generator for streaming requests or generator for unary requests). Generator outputs will be written to the `generator_result_callback`. Note that HTTP requests are an exception: they are *always* streaming requests, but for ASGI apps (like FastAPI), the actual method will be a regular function implementing the ASGI `__call__` protocol. """ result_is_gen = inspect.isgenerator(result) result_is_async_gen = inspect.isasyncgen(result) if is_streaming: if result_is_gen: for r in result: generator_result_callback(r) elif result_is_async_gen: async for r in result: generator_result_callback(r) elif is_http_request and not user_method_info.is_asgi_app: # For the FastAPI codepath, the response has already been sent over # ASGI, but for the vanilla deployment codepath we need to send it. await self._send_user_result_over_asgi(result, asgi_args) elif not is_http_request and not sync_gen_consumed: # If a unary method is called with stream=True for anything EXCEPT # an HTTP request, raise an error. # HTTP requests are always streaming regardless of if the method # returns a generator, because it's provided the result queue as its # ASGI `send` interface to stream back results. raise TypeError( f"Called method '{user_method_info.name}' with " "`handle.options(stream=True)` but it did not return a " "generator." ) else: assert ( not is_http_request ), "All HTTP requests go through the streaming codepath." if result_is_gen or result_is_async_gen: raise TypeError( f"Method '{user_method_info.name}' returned a generator. " "You must use `handle.options(stream=True)` to call " "generators on a deployment." ) return result async def call_http_entrypoint( self, request_metadata: RequestMetadata, status_code_callback: StatusCodeCallback, scope: Scope, receive: Receive, ) -> Any: result_queue = MessageQueue() user_method_info = self.get_user_method_info(request_metadata.call_method) if self._run_user_code_in_separate_thread: # `asyncio.Event`s are not thread safe, so `call_soon_threadsafe` must be # used to interact with the result queue from the user callable thread. system_event_loop = asyncio.get_running_loop() async def enqueue(item: Any): system_event_loop.call_soon_threadsafe(result_queue.put_nowait, item) call_future = self._call_http_entrypoint( user_method_info, scope, receive, enqueue ) else: async def enqueue(item: Any): result_queue.put_nowait(item) call_future = asyncio.create_task( self._call_http_entrypoint(user_method_info, scope, receive, enqueue) ) first_message_peeked = False async for messages in result_queue.fetch_messages_from_queue(call_future): # HTTP (ASGI) messages are only consumed by the proxy so batch them # and use vanilla pickle (we know it's safe because these messages # only contain primitive Python types). # Peek the first ASGI message to determine the status code. if not first_message_peeked: msg = messages[0] first_message_peeked = True if msg["type"] == "http.response.start": # HTTP responses begin with exactly one # "http.response.start" message containing the "status" # field. Other response types like WebSockets may not. status_code_callback(str(msg["status"])) yield messages @_run_user_code async def _call_http_entrypoint( self, user_method_info: UserMethodInfo, scope: Scope, receive: Receive, send: Send, ) -> Any: """Call an HTTP entrypoint. `send` is used to communicate the results of streaming responses. Raises any exception raised by the user code so it can be propagated as a `RayTaskError`. """ self._raise_if_not_initialized("_call_http_entrypoint") if self._is_enabled_for_debug: logger.debug( f"Started executing request to method '{user_method_info.name}'.", extra={"log_to_stderr": False, "serve_access_log": True}, ) if user_method_info.is_asgi_app: request_args = (scope, receive, send) elif not user_method_info.takes_any_args: # Edge case to support empty HTTP handlers: don't pass the Request # argument if the callable has no parameters. request_args = tuple() else: # Non-FastAPI HTTP handlers take only the starlette `Request`. request_args = (starlette.requests.Request(scope, receive, send),) receive_task = None try: if hasattr(receive, "fetch_until_disconnect"): receive_task = asyncio.create_task(receive.fetch_until_disconnect()) result, sync_gen_consumed = await self._call_func_or_gen( user_method_info.callable, args=request_args, kwargs={}, is_streaming=True, generator_result_callback=send, ) final_result = await self._handle_user_method_result( result, user_method_info, is_streaming=True, is_http_request=True, sync_gen_consumed=sync_gen_consumed, generator_result_callback=send, asgi_args=ASGIArgs(scope, receive, send), ) if receive_task is not None and not receive_task.done(): receive_task.cancel() return final_result except Exception as e: if not user_method_info.is_asgi_app: response = self.handle_exception(e) await self._send_user_result_over_asgi( response, ASGIArgs(scope, receive, send) ) if receive_task is not None and not receive_task.done(): receive_task.cancel() raise except asyncio.CancelledError: if receive_task is not None and not receive_task.done(): # Do NOT cancel the receive task if the request has been # cancelled, but the call is a batched call. This is # because we cannot guarantee cancelling the batched # call, so in the case that the call continues executing # we should continue fetching data from the client. if not hasattr(user_method_info.callable, "set_max_batch_size"): receive_task.cancel() raise async def call_user_generator( self, request_metadata: RequestMetadata, request_args: Tuple[Any], request_kwargs: Dict[str, Any], ) -> AsyncGenerator[Any, None]: """Calls a user method for a streaming call and yields its results. The user method is called in an asyncio `Task` and places its results on a `result_queue`. This method pulls and yields from the `result_queue`. If this generator is closed before the user method finishes (e.g. the client disconnected mid-stream), the unit running the user code is cancelled so a `CancelledError` is raised inside the user generator: the `call_future` task in separate-thread mode, or (inline) the user generator itself in same-loop mode. """ if not self._run_user_code_in_separate_thread: gen = await self._call_user_generator( request_metadata, request_args, request_kwargs ) try: async for result in gen: yield result finally: # User code runs inline; closing the wrapper injects a # CancelledError into the user generator (see `_call_generator_async`). await gen.aclose() else: result_queue = MessageQueue() # `asyncio.Event`s are not thread safe, so `call_soon_threadsafe` must be # used to interact with the result queue from the user callable thread. system_event_loop = asyncio.get_running_loop() def _enqueue_thread_safe(item: Any): system_event_loop.call_soon_threadsafe(result_queue.put_nowait, item) call_future = self._call_user_generator( request_metadata, request_args, request_kwargs, enqueue=_enqueue_thread_safe, ) try: async for messages in result_queue.fetch_messages_from_queue( call_future ): for msg in messages: yield msg finally: # Cancel the user-code task so a CancelledError is raised inside the # user method. No-op if it already completed. if not call_future.done(): call_future.cancel() @_run_user_code async def _call_user_generator( self, request_metadata: RequestMetadata, request_args: Tuple[Any], request_kwargs: Dict[str, Any], *, enqueue: Optional[Callable] = None, ) -> Optional[AsyncGenerator[Any, None]]: """Call a user generator. The `generator_result_callback` is used to communicate the results of generator methods. Raises any exception raised by the user code so it can be propagated as a `RayTaskError`. """ self._raise_if_not_initialized("_call_user_generator") request_args = request_args if request_args is not None else tuple() request_kwargs = request_kwargs if request_kwargs is not None else dict() user_method_info = self.get_user_method_info(request_metadata.call_method) callable = user_method_info.callable is_sync_method = ( inspect.isfunction(callable) or inspect.ismethod(callable) ) and not ( inspect.iscoroutinefunction(callable) or inspect.isasyncgenfunction(callable) ) if self._is_enabled_for_debug: logger.debug( f"Started executing request to method '{user_method_info.name}'.", extra={"log_to_stderr": False, "serve_access_log": True}, ) async def _call_generator_async() -> AsyncGenerator[Any, None]: gen = callable(*request_args, **request_kwargs) if inspect.iscoroutine(gen): gen = await gen try: if inspect.isgenerator(gen): for result in gen: yield result elif inspect.isasyncgen(gen): async for result in gen: yield result else: raise TypeError( f"Called method '{user_method_info.name}' with " "`handle.options(stream=True)` but it did not return a " "generator." ) finally: # If this wrapper is closed before the user generator finishes # (e.g. the client disconnected), inject a CancelledError into the # user's async generator so its cancellation handling runs. (Sync # generators have no await points to cancel.) No-op if it already # finished, in which case athrow raises StopAsyncIteration. if inspect.isasyncgen(gen): try: await gen.athrow(asyncio.CancelledError()) except (StopAsyncIteration, asyncio.CancelledError): pass def _call_generator_sync(): gen = callable(*request_args, **request_kwargs) if inspect.isgenerator(gen): for result in gen: enqueue(result) else: raise TypeError( f"Called method '{user_method_info.name}' with " "`handle.options(stream=True)` but it did not return a generator." ) if enqueue and is_sync_method and self._run_sync_methods_in_threadpool: await to_thread.run_sync(_call_generator_sync) elif enqueue: async def gen_coro_wrapper(): async for result in _call_generator_async(): enqueue(result) await gen_coro_wrapper() else: return _call_generator_async() @_run_user_code async def call_user_method( self, request_metadata: RequestMetadata, request_args: Tuple[Any], request_kwargs: Dict[str, Any], ) -> Any: """Call a (unary) user method. Raises any exception raised by the user code so it can be propagated as a `RayTaskError`. """ self._raise_if_not_initialized("call_user_method") if self._is_enabled_for_debug: logger.debug( f"Started executing request to method '{request_metadata.call_method}'.", extra={"log_to_stderr": False, "serve_access_log": True}, ) user_method_info = self.get_user_method_info(request_metadata.call_method) result, _ = await self._call_func_or_gen( user_method_info.callable, args=request_args, kwargs=request_kwargs, is_streaming=False, ) if inspect.isgenerator(result) or inspect.isasyncgen(result): raise TypeError( f"Method '{user_method_info.name}' returned a generator. " "You must use `handle.options(stream=True)` to call " "generators on a deployment." ) return result def handle_exception(self, exc: Exception): if isinstance(exc, self.service_unavailable_exceptions): return starlette.responses.Response(exc.message, status_code=503) else: return starlette.responses.Response( "Internal Server Error", status_code=500 ) @_run_user_code async def call_destructor(self): """Explicitly call the `__del__` method of the user callable. Calling this multiple times has no effect; only the first call will actually call the destructor. """ if self._callable is None: logger.debug( "This replica has not yet started running user code. " "Skipping __del__." ) return # Only run the destructor once. This is safe because there is no `await` between # checking the flag here and flipping it to `True` below. if self._destructor_called: return self._destructor_called = True try: if hasattr(self._callable, "__del__"): # Make sure to accept `async def __del__(self)` as well. await self._call_func_or_gen( self._callable.__del__, # Always run the destructor on the main user callable thread. run_sync_methods_in_threadpool_override=False, ) if hasattr(self._callable, "__serve_multiplex_wrapper"): await getattr(self._callable, "__serve_multiplex_wrapper").shutdown() except Exception as e: logger.exception(f"Exception during graceful shutdown of replica: {e}") finally: if self._user_code_threadpool is not None: self._user_code_threadpool.shutdown(wait=False)