wehub-resource-sync
303 lines
9.4 KiB
Python
303 lines
9.4 KiB
Python
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
|