wehub-resource-sync
246 lines
8.3 KiB
Python
246 lines
8.3 KiB
Python
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__]))
|