import asyncio import uuid from typing import Any, Dict, Optional, Tuple import aiohttp import httpx import pytest import starlette from aiohttp import ClientSession, TCPConnector from fastapi import FastAPI import ray from ray import serve from ray.serve._private.constants import ( SERVE_HTTP_REQUEST_ID_HEADER, SERVE_SESSION_ID, ) from ray.serve._private.test_utils import get_application_url from ray.serve._private.utils import generate_request_id def test_request_id_header_by_default(serve_instance): """Test that a request_id is generated by default and returned as a header.""" @serve.deployment class Model: def __call__(self): request_id = ray.serve.context._get_serve_request_context().request_id return request_id serve.run(Model.bind()) resp = httpx.get(f"{get_application_url()}") assert resp.status_code == 200 assert resp.text == resp.headers[SERVE_HTTP_REQUEST_ID_HEADER] def is_valid_uuid(num: str): try: uuid.UUID(num, version=4) return True except ValueError: return False assert is_valid_uuid(resp.text) class TestUserProvidedRequestIDHeader: def verify_result(self): for header_attr in ["X-Request-ID"]: resp = httpx.get( f"{get_application_url()}", headers={header_attr: "123-234"} ) assert resp.status_code == 200 assert resp.json() == 1 assert resp.headers[header_attr] == "123-234" def test_basic(self, serve_instance): @serve.deployment class Model: def __call__(self) -> int: request_id = ray.serve.context._get_serve_request_context().request_id assert request_id == "123-234" return 1 serve.run(Model.bind()) self.verify_result() def test_fastapi(self, serve_instance): app = FastAPI() @serve.deployment @serve.ingress(app) class Model: @app.get("/") def say_hi(self) -> int: request_id = ray.serve.context._get_serve_request_context().request_id assert request_id == "123-234" return 1 serve.run(Model.bind()) self.verify_result() def test_starlette_resp(self, serve_instance): @serve.deployment class Model: def __call__(self) -> int: request_id = ray.serve.context._get_serve_request_context().request_id assert request_id == "123-234" return starlette.responses.Response("1", media_type="application/json") serve.run(Model.bind()) self.verify_result() def test_set_request_id_headers_with_two_attributes(serve_instance): """Test that request id is set with X-Request-ID and RAY_SERVE_REQUEST_ID. x-request-id has higher priority. """ @serve.deployment class Model: def __call__(self): request_id = ray.serve.context._get_serve_request_context().request_id return request_id serve.run(Model.bind()) resp = httpx.get( get_application_url(), headers={ "X-Request-ID": "234", }, ) assert resp.status_code == 200 assert SERVE_HTTP_REQUEST_ID_HEADER in resp.headers assert resp.text == resp.headers[SERVE_HTTP_REQUEST_ID_HEADER] def test_reuse_request_id(serve_instance): """Test client re-uses request id. When multiple requests are submitted with the same request id at around the same time, the proxy should continue to track the correct request objects, setting the correct request id in the serve context, and return the original x-request-id request header as the response header. For more details, see https://github.com/ray-project/ray/issues/45723. """ app = FastAPI() @serve.deployment(num_replicas=3) @serve.ingress(app) class MyFastAPIDeployment: @app.post("/hello") def root(self, user_input: Dict[str, str]) -> Dict[str, str]: request_id = ray.serve.context._get_serve_request_context().request_id return { "app_name": user_input["app_name"], "serve_context_request_id": request_id, } serve.run(MyFastAPIDeployment.bind()) async def send_request( session: ClientSession, body: Dict[str, Any], request_id: Optional[str] ) -> Tuple[str, str]: headers = {SERVE_HTTP_REQUEST_ID_HEADER: request_id} url = "http://localhost:8000/hello" async with session.post(url=url, headers=headers, json=body) as response: result = await response.json() # Ensure the request object is tracked correctly. assert result["app_name"] == body["app_name"] # Ensure the request id from the serve context is set correctly. assert result["serve_context_request_id"] == request_id # Ensure the request id from the response header is returned correctly. assert response.headers[SERVE_HTTP_REQUEST_ID_HEADER] == request_id async def main(): """Sending 20 requests in parallel all with the same request id, but with different request body. """ bodies = [{"app_name": f"an_{generate_request_id()}"} for _ in range(20)] connector = TCPConnector(ssl=False) async with aiohttp.ClientSession(connector=connector) as session: request_id = f"rid_{generate_request_id()}" tasks = [ send_request(session, body, request_id=request_id) for body in bodies ] await asyncio.gather(*tasks) asyncio.run(main()) class TestSessionIdHeader: @pytest.fixture def session_id_app(self, serve_instance): @serve.deployment class Model: def __call__(self) -> str: return ray.serve.context._get_serve_request_context().session_id handle = serve.run(Model.bind()) return handle @pytest.mark.parametrize( "header_key", [ SERVE_SESSION_ID, # underscore form "x-session-id", # hyphenated form "X-Session-Id", # title-cased hyphenated form ], ) def test_header_forms(self, session_id_app, header_key): session_id = "sess_user_42" resp = httpx.get(get_application_url(), headers={header_key: session_id}) assert resp.status_code == 200 assert resp.text == session_id def test_handle_option(self, session_id_app): handle = session_id_app assert handle.options(session_id="from-handle").remote().result() == ( "from-handle" ) class TestMatchesSessionIdHeader: """Unit tests for the case/separator-tolerant comparison used by the proxy + Lua forwarder. Without this helper, an intermediate proxy that rewrites ``_``↔``-`` (nginx, AWS API Gateway, ...) would silently drop session affinity for half of all clients. """ @pytest.mark.parametrize( "stored,incoming,expected", [ # Default hyphenated form. ("x-session-id", "x-session-id", True), ("x-session-id", "X-Session-Id", True), ("x-session-id", "x_session_id", True), ("x-session-id", "x-Other-Id", False), # Operator stored the underscored form (e.g. legacy). ("x_session_id", "x-session-id", True), ("x_session_id", "x_session_id", True), # Custom configured value via env var (e.g. aiperf's correlation id). ("x-correlation-id", "x-correlation-id", True), ("x-correlation-id", "x_correlation_id", True), ("x-correlation-id", "X-Correlation-Id", True), ("x-correlation-id", "x-session-id", False), ], ) def test_match(self, monkeypatch, stored, incoming, expected): # SERVE_SESSION_ID is module-level constant; monkey-patch it on the # two modules that import it for the duration of this test. from ray.serve._private import constants, http_util monkeypatch.setattr(constants, "SERVE_SESSION_ID", stored) monkeypatch.setattr(http_util, "SERVE_SESSION_ID", stored) assert http_util._matches_session_id_header(incoming) is expected if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-s", __file__]))