chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,252 @@
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.routing import Route
|
||||
from starlette.types import Scope
|
||||
|
||||
from ray.serve._private.common import ApplicationName, DeploymentID, EndpointInfo
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
from ray.serve._private.thirdparty.get_asgi_route_name import (
|
||||
RoutePattern,
|
||||
get_asgi_route_name,
|
||||
)
|
||||
from ray.serve.handle import DeploymentHandle
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
NO_ROUTES_MESSAGE = "Route table is not populated yet."
|
||||
NO_REPLICAS_MESSAGE = "No replicas are available yet."
|
||||
|
||||
|
||||
class ProxyRouter:
|
||||
"""Router interface for the proxy to use."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
get_handle: Callable[[str, str], DeploymentHandle],
|
||||
):
|
||||
# Function to get a handle given a name. Used to mock for testing.
|
||||
self._get_handle = get_handle
|
||||
# Contains a ServeHandle for each endpoint.
|
||||
self.handles: Dict[DeploymentID, DeploymentHandle] = dict()
|
||||
# Flipped to `True` once the route table has been updated at least once.
|
||||
# The proxy router is not ready for traffic until the route table is populated
|
||||
self._route_table_populated = False
|
||||
|
||||
# Info used for HTTP proxy
|
||||
# Routes sorted in order of decreasing length.
|
||||
self.sorted_routes: List[str] = list()
|
||||
# Endpoints associated with the routes.
|
||||
self.route_info: Dict[str, DeploymentID] = dict()
|
||||
# Map of application name to is_cross_language.
|
||||
self.app_to_is_cross_language: Dict[ApplicationName, bool] = dict()
|
||||
|
||||
# Info used for gRPC proxy
|
||||
# Endpoints info associated with endpoints.
|
||||
self.endpoints: Dict[DeploymentID, EndpointInfo] = dict()
|
||||
|
||||
# Map of route prefix to list of route patterns for that endpoint
|
||||
# Used to match incoming requests to ASGI route patterns for metrics
|
||||
# Route patterns are tuples of (methods, path) where methods can be None
|
||||
self.route_patterns: Dict[str, List[RoutePattern]] = dict()
|
||||
# Cache of mock Starlette apps for route pattern matching
|
||||
# Key: route prefix, Value: pre-built Starlette app with routes
|
||||
self._route_pattern_apps: Dict[str, Any] = dict()
|
||||
|
||||
def ready_for_traffic(self, is_head: bool) -> Tuple[bool, str]:
|
||||
"""Whether the proxy router is ready to serve traffic.
|
||||
|
||||
The first return value will be false if any of the following hold:
|
||||
- The route table has not been populated yet with a non-empty set of routes
|
||||
- The route table has been populated, but none of the handles
|
||||
have received running replicas yet AND it lives on a worker node.
|
||||
|
||||
Otherwise, the first return value will be true.
|
||||
"""
|
||||
|
||||
if not self._route_table_populated:
|
||||
return False, NO_ROUTES_MESSAGE
|
||||
|
||||
# NOTE(zcin): For the proxy on the head node, even if none of its handles have
|
||||
# been populated with running replicas yet, we MUST mark the proxy as ready for
|
||||
# traffic. This is to handle the case when all deployments have scaled to zero.
|
||||
# If the deployments (more precisely, ingress deployments) have all scaled down
|
||||
# to zero, at least one proxy needs to be able to receive incoming requests to
|
||||
# trigger upscale.
|
||||
if is_head:
|
||||
return True, ""
|
||||
|
||||
for handle in self.handles.values():
|
||||
if handle.running_replicas_populated():
|
||||
return True, ""
|
||||
|
||||
return False, NO_REPLICAS_MESSAGE
|
||||
|
||||
def update_routes(self, endpoints: Dict[DeploymentID, EndpointInfo]):
|
||||
logger.info(
|
||||
f"Got updated endpoints: {endpoints}.", extra={"log_to_stderr": True}
|
||||
)
|
||||
if endpoints:
|
||||
self._route_table_populated = True
|
||||
|
||||
self.endpoints = endpoints
|
||||
|
||||
existing_handles = set(self.handles.keys())
|
||||
routes = []
|
||||
route_info = {}
|
||||
app_to_is_cross_language = {}
|
||||
route_patterns = {}
|
||||
for endpoint, info in endpoints.items():
|
||||
routes.append(info.route)
|
||||
route_info[info.route] = endpoint
|
||||
app_to_is_cross_language[endpoint.app_name] = info.app_is_cross_language
|
||||
if info.route_patterns:
|
||||
route_patterns[info.route] = info.route_patterns
|
||||
if endpoint in self.handles:
|
||||
existing_handles.remove(endpoint)
|
||||
else:
|
||||
self.handles[endpoint] = self._get_handle(endpoint, info)
|
||||
|
||||
# Clean up any handles that are no longer used.
|
||||
if len(existing_handles) > 0:
|
||||
logger.info(
|
||||
f"Deleting {len(existing_handles)} unused handles.",
|
||||
extra={"log_to_stderr": False},
|
||||
)
|
||||
for endpoint in existing_handles:
|
||||
del self.handles[endpoint]
|
||||
|
||||
# Routes are sorted in order of decreasing length to enable longest
|
||||
# prefix matching.
|
||||
self.sorted_routes = sorted(routes, key=lambda x: len(x), reverse=True)
|
||||
self.route_info = route_info
|
||||
self.app_to_is_cross_language = app_to_is_cross_language
|
||||
self.route_patterns = route_patterns
|
||||
# Invalidate cached mock apps when route patterns change
|
||||
self._route_pattern_apps.clear()
|
||||
|
||||
def match_route(
|
||||
self, target_route: str
|
||||
) -> Optional[Tuple[str, DeploymentHandle, bool]]:
|
||||
"""Return the longest prefix match among existing routes for the route.
|
||||
Args:
|
||||
target_route: route to match against.
|
||||
Returns:
|
||||
(route, handle, is_cross_language) if found, else None.
|
||||
"""
|
||||
|
||||
for route in self.sorted_routes:
|
||||
if target_route.startswith(route):
|
||||
matched = False
|
||||
# If the route we matched on ends in a '/', then so does the
|
||||
# target route and this must be a match.
|
||||
if route.endswith("/"):
|
||||
matched = True
|
||||
# If the route we matched on doesn't end in a '/', we need to
|
||||
# do another check to ensure that either this is an exact match
|
||||
# or the next character in the target route is a '/'. This is
|
||||
# to guard against the scenario where we have '/route' as a
|
||||
# prefix and there's a request to '/routesuffix'. In this case,
|
||||
# it should *not* be a match.
|
||||
elif len(target_route) == len(route) or target_route[len(route)] == "/":
|
||||
matched = True
|
||||
|
||||
if matched:
|
||||
endpoint = self.route_info[route]
|
||||
return (
|
||||
route,
|
||||
self.handles[endpoint],
|
||||
self.app_to_is_cross_language[endpoint.app_name],
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
def get_handle_for_endpoint(
|
||||
self, target_app_name: str
|
||||
) -> Optional[Tuple[str, DeploymentHandle, bool]]:
|
||||
"""Return the handle that matches with endpoint.
|
||||
|
||||
Args:
|
||||
target_app_name: app_name to match against.
|
||||
Returns:
|
||||
(route, handle, is_cross_language) for the single app if there
|
||||
is only one, else find the app and handle for exact match. Else return None.
|
||||
"""
|
||||
for endpoint_tag, handle in self.handles.items():
|
||||
# If the target_app_name matches with the endpoint or if
|
||||
# there is only one endpoint.
|
||||
if target_app_name == endpoint_tag.app_name or len(self.handles) == 1:
|
||||
endpoint_info = self.endpoints[endpoint_tag]
|
||||
return (
|
||||
endpoint_info.route,
|
||||
handle,
|
||||
endpoint_info.app_is_cross_language,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
def match_route_pattern(self, route_prefix: str, asgi_scope: Scope) -> str:
|
||||
"""Match an incoming request to a specific route pattern.
|
||||
|
||||
This attempts to match the request path to a route pattern (e.g., /api/{user_id})
|
||||
rather than just the route prefix. This provides more granular metrics.
|
||||
|
||||
The mock Starlette app is cached per route_prefix for performance, avoiding
|
||||
the overhead of recreating the app and routes on every request.
|
||||
|
||||
Args:
|
||||
route_prefix: The matched route prefix from match_route()
|
||||
asgi_scope: The ASGI scope containing the request path and method
|
||||
|
||||
Returns:
|
||||
The matched route pattern if available, otherwise the route_prefix
|
||||
"""
|
||||
# If we don't have route patterns for this prefix, return the prefix
|
||||
if route_prefix not in self.route_patterns:
|
||||
return route_prefix
|
||||
|
||||
patterns = self.route_patterns[route_prefix]
|
||||
if not patterns:
|
||||
return route_prefix
|
||||
|
||||
# Get or create the cached mock app for this route_prefix
|
||||
mock_app = self._route_pattern_apps.get(route_prefix)
|
||||
if mock_app is None:
|
||||
try:
|
||||
# Create routes from patterns
|
||||
# We use a dummy endpoint since we only need pattern matching
|
||||
async def dummy_endpoint(request: Request):
|
||||
pass
|
||||
|
||||
routes = [
|
||||
Route(pattern.path, dummy_endpoint, methods=pattern.methods)
|
||||
for pattern in patterns
|
||||
]
|
||||
mock_app = Starlette(routes=routes)
|
||||
|
||||
# Cache the mock app for future requests
|
||||
self._route_pattern_apps[route_prefix] = mock_app
|
||||
except Exception:
|
||||
# If app creation fails, fall back to route prefix
|
||||
logger.debug(
|
||||
f"Failed to create mock app for route pattern matching: {route_prefix}",
|
||||
exc_info=True,
|
||||
)
|
||||
return route_prefix
|
||||
|
||||
# Use the cached mock app to match the route pattern
|
||||
try:
|
||||
matched = get_asgi_route_name(mock_app, asgi_scope)
|
||||
if matched:
|
||||
return matched
|
||||
except Exception:
|
||||
# If matching fails for any reason, fall back to route prefix
|
||||
logger.debug(
|
||||
f"Failed to match route pattern for {route_prefix}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
# Fall back to route prefix if no pattern matched
|
||||
return route_prefix
|
||||
Reference in New Issue
Block a user