chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,302 @@
|
||||
import logging
|
||||
import pickle
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, AsyncIterator, List, Optional, Tuple, Union
|
||||
|
||||
import grpc
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from ray.serve._private.common import StreamingHTTPRequest, gRPCRequest
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
from ray.serve._private.logging_utils import format_grpc_peer_address
|
||||
from ray.serve._private.tracing_utils import (
|
||||
extract_propagated_context,
|
||||
is_tracing_enabled,
|
||||
set_trace_context,
|
||||
)
|
||||
from ray.serve._private.utils import DEFAULT
|
||||
from ray.serve.grpc_util import RayServegRPCContext
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class gRPCStreamingType(str, Enum):
|
||||
"""Enum representing the gRPC streaming type."""
|
||||
|
||||
UNARY_UNARY = "unary_unary" # Single request, single response
|
||||
UNARY_STREAM = (
|
||||
"unary_stream" # Single request, streaming response (server streaming)
|
||||
)
|
||||
STREAM_UNARY = (
|
||||
"stream_unary" # Streaming request, single response (client streaming)
|
||||
)
|
||||
STREAM_STREAM = (
|
||||
"stream_stream" # Streaming request, streaming response (bidirectional)
|
||||
)
|
||||
|
||||
|
||||
class ProxyRequest(ABC):
|
||||
"""Base ProxyRequest class to use in the common interface among proxies"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def request_type(self) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def method(self) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def route_path(self) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_route_request(self) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_health_request(self) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def client(self) -> str:
|
||||
return ""
|
||||
|
||||
@abstractmethod
|
||||
def populate_tracing_context(self):
|
||||
"""Implement this method to populate tracing context so the parent and
|
||||
child spans will be connected into a single trace."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class ASGIProxyRequest(ProxyRequest):
|
||||
"""ProxyRequest implementation to wrap ASGI scope, receive, and send."""
|
||||
|
||||
def __init__(self, scope: Scope, receive: Receive, send: Send):
|
||||
self.scope = scope
|
||||
self.receive = receive
|
||||
self.send = send
|
||||
|
||||
@property
|
||||
def request_type(self) -> str:
|
||||
return self.scope.get("type", "")
|
||||
|
||||
@property
|
||||
def method(self) -> str:
|
||||
# WebSocket messages don't have a 'method' field.
|
||||
return self.scope.get("method", "WS").upper()
|
||||
|
||||
@property
|
||||
def route_path(self) -> str:
|
||||
return self.scope.get("path", "")[len(self.root_path) :]
|
||||
|
||||
@property
|
||||
def is_route_request(self) -> bool:
|
||||
return self.route_path == "/-/routes"
|
||||
|
||||
@property
|
||||
def is_health_request(self) -> bool:
|
||||
return self.route_path == "/-/healthz"
|
||||
|
||||
@property
|
||||
def client(self) -> str:
|
||||
return self.scope.get("client", "")
|
||||
|
||||
@property
|
||||
def root_path(self) -> str:
|
||||
return self.scope.get("root_path", "")
|
||||
|
||||
@property
|
||||
def path(self) -> str:
|
||||
return self.scope.get("path", "")
|
||||
|
||||
@property
|
||||
def headers(self) -> List[Tuple[bytes, bytes]]:
|
||||
return self.scope.get("headers", [])
|
||||
|
||||
def set_path(self, path: str):
|
||||
self.scope["path"] = path
|
||||
|
||||
def set_root_path(self, root_path: str):
|
||||
self.scope["root_path"] = root_path
|
||||
|
||||
def serialized_replica_arg(self, proxy_actor_name: str) -> bytes:
|
||||
# NOTE(edoakes): it's important that the request is sent as raw bytes to
|
||||
# skip the Ray cloudpickle serialization codepath for performance.
|
||||
return pickle.dumps(
|
||||
StreamingHTTPRequest(
|
||||
asgi_scope=self.scope,
|
||||
proxy_actor_name=proxy_actor_name,
|
||||
)
|
||||
)
|
||||
|
||||
def populate_tracing_context(self):
|
||||
"""Populate tracing context for ASGI requests.
|
||||
|
||||
This method extracts the "traceparent" header from the request headers and sets
|
||||
the tracing context from it.
|
||||
"""
|
||||
if not is_tracing_enabled():
|
||||
return
|
||||
|
||||
for key, value in self.headers:
|
||||
if key.decode() == "traceparent":
|
||||
trace_context = extract_propagated_context(
|
||||
{key.decode(): value.decode()}
|
||||
)
|
||||
set_trace_context(trace_context)
|
||||
|
||||
|
||||
class gRPCProxyRequest(ProxyRequest):
|
||||
"""ProxyRequest implementation to wrap gRPC request protobuf and metadata."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request_proto: Any,
|
||||
context: grpc._cython.cygrpc._ServicerContext,
|
||||
service_method: str,
|
||||
stream: bool,
|
||||
*,
|
||||
streaming_type: gRPCStreamingType = None,
|
||||
request_iterator: Optional[AsyncIterator[Any]] = None,
|
||||
):
|
||||
self._request_proto = request_proto
|
||||
self._request_iterator = request_iterator
|
||||
self.context = context
|
||||
self.service_method = service_method
|
||||
self.stream = stream
|
||||
# Determine streaming type based on parameters
|
||||
if streaming_type is not None:
|
||||
self.streaming_type = streaming_type
|
||||
elif request_iterator is not None:
|
||||
# Has input stream
|
||||
self.streaming_type = (
|
||||
gRPCStreamingType.STREAM_STREAM
|
||||
if stream
|
||||
else gRPCStreamingType.STREAM_UNARY
|
||||
)
|
||||
else:
|
||||
# No input stream
|
||||
self.streaming_type = (
|
||||
gRPCStreamingType.UNARY_STREAM
|
||||
if stream
|
||||
else gRPCStreamingType.UNARY_UNARY
|
||||
)
|
||||
self.app_name = ""
|
||||
self.request_id = None
|
||||
self.method_name = "__call__"
|
||||
self.multiplexed_model_id = DEFAULT.VALUE
|
||||
self.session_id = DEFAULT.VALUE
|
||||
# ray_serve_grpc_context is a class implemented by us to be able to serialize
|
||||
# the object and pass it into the deployment.
|
||||
self.ray_serve_grpc_context = RayServegRPCContext(context)
|
||||
self.setup_variables()
|
||||
|
||||
def setup_variables(self):
|
||||
if not self.is_route_request and not self.is_health_request:
|
||||
service_method_split = self.service_method.split("/")
|
||||
self.method_name = service_method_split[-1]
|
||||
for key, value in self.context.invocation_metadata():
|
||||
if key == "application":
|
||||
self.app_name = value
|
||||
elif key == "request_id":
|
||||
self.request_id = value
|
||||
elif key == "multiplexed_model_id":
|
||||
self.multiplexed_model_id = value
|
||||
elif key == "session_id":
|
||||
self.session_id = value
|
||||
|
||||
@property
|
||||
def request_type(self) -> str:
|
||||
return "grpc"
|
||||
|
||||
@property
|
||||
def method(self) -> str:
|
||||
return self.service_method
|
||||
|
||||
@property
|
||||
def route_path(self) -> str:
|
||||
return self.app_name
|
||||
|
||||
@property
|
||||
def is_route_request(self) -> bool:
|
||||
return self.service_method == "/ray.serve.RayServeAPIService/ListApplications"
|
||||
|
||||
@property
|
||||
def is_health_request(self) -> bool:
|
||||
return self.service_method == "/ray.serve.RayServeAPIService/Healthz"
|
||||
|
||||
@property
|
||||
def client(self) -> str:
|
||||
return format_grpc_peer_address(self.context.peer())
|
||||
|
||||
@property
|
||||
def has_input_stream(self) -> bool:
|
||||
"""Returns True if this request has a streaming input (client/bidi streaming)."""
|
||||
return self.streaming_type in (
|
||||
gRPCStreamingType.STREAM_UNARY,
|
||||
gRPCStreamingType.STREAM_STREAM,
|
||||
)
|
||||
|
||||
@property
|
||||
def request_iterator(self) -> Optional[AsyncIterator[Any]]:
|
||||
"""Returns the request iterator for client/bidi streaming, or None."""
|
||||
return self._request_iterator
|
||||
|
||||
def send_request_id(self, request_id: str):
|
||||
# Setting the trailing metadata on the ray_serve_grpc_context object, so it's
|
||||
# not overriding the ones set from the user and will be sent back to the
|
||||
# client altogether.
|
||||
self.ray_serve_grpc_context.set_trailing_metadata([("request_id", request_id)])
|
||||
|
||||
def serialized_replica_arg(self) -> bytes:
|
||||
# NOTE(edoakes): it's important that the request is sent as raw bytes to
|
||||
# skip the Ray cloudpickle serialization codepath for performance.
|
||||
return pickle.dumps(gRPCRequest(user_request_proto=self._request_proto))
|
||||
|
||||
def populate_tracing_context(self):
|
||||
"""Populate tracing context for gRPC requests.
|
||||
|
||||
This method extracts the "traceparent" metadata from the request headers and
|
||||
sets the tracing context from it.
|
||||
"""
|
||||
if not is_tracing_enabled():
|
||||
return
|
||||
|
||||
for key, value in self.context.invocation_metadata():
|
||||
if key == "traceparent":
|
||||
trace_context = extract_propagated_context({key: value})
|
||||
set_trace_context(trace_context)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResponseStatus:
|
||||
code: Union[str, grpc.StatusCode] # Must be convertible to a string.
|
||||
is_error: bool = False
|
||||
message: str = ""
|
||||
|
||||
|
||||
# Yields protocol-specific messages followed by a final `ResponseStatus`.
|
||||
ResponseGenerator = AsyncIterator[Union[Any, ResponseStatus]]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HandlerMetadata:
|
||||
application_name: str = ""
|
||||
deployment_name: str = ""
|
||||
route: str = ""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResponseHandlerInfo:
|
||||
response_generator: ResponseGenerator
|
||||
metadata: HandlerMetadata
|
||||
should_record_access_log: bool
|
||||
should_increment_ongoing_requests: bool
|
||||
Reference in New Issue
Block a user