diff --git a/tests/integration/configuration/test_openapi_docs_route.py b/tests/integration/configuration/test_openapi_docs_route.py new file mode 100644 index 00000000000..ccb30c2f784 --- /dev/null +++ b/tests/integration/configuration/test_openapi_docs_route.py @@ -0,0 +1,165 @@ +import time +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from pydantic import JsonValue, TypeAdapter + +from litellm.proxy._lazy_features import LAZY_FEATURES + +DOCS_SECONDS: Final = 120 +COLD_DOCS: Final = 4 +CHATS: Final = 8 +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +TWO_OWNED_PROXIES_TIMEOUT: Final = pytest.mark.timeout(2 * (2 * graceful_stop_seconds() + 120)) +LAZY_ROUTES_FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES" +WARMED_FEATURE: Final = next(feature for feature in LAZY_FEATURES if feature.name == "policies") + + +def _timed_get(candidate: Gateway, path: str) -> tuple[httpx.Response, float]: + started: Final = time.perf_counter() + response: Final = candidate.request("GET", path) + return response, time.perf_counter() - started + + +def _liveliness_until(done: Future[tuple[httpx.Response, float]], candidate: Gateway) -> Iterator[tuple[int, float]]: + while not done.done(): + response, seconds = _timed_get(candidate, "/health/liveliness") + yield response.status_code, seconds + + +def _schema(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text[:500] + return JSON_OBJECT.validate_json(response.content) + + +def _servers(response: httpx.Response) -> list[JsonValue]: + return list(list_value(_schema(response).get("servers", []))) + + +def _chat(candidate: Gateway, model: str, key: str, content: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": content}], "cache": {"no-cache": True}}, + key=key, + ) + + +def _warmed_feature_operations(response: httpx.Response) -> dict[str, JsonValue]: + paths: Final = object_value(_schema(response)["paths"]) + return {path: operations for path, operations in paths.items() if WARMED_FEATURE.matches(path)} + + +@OWNED_PROXY_TIMEOUT +def test_liveliness_answers_while_a_cold_openapi_schema_builds(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(_timed_get, owned.gateway, "/openapi.json") + probes: Final = tuple(_liveliness_until(docs, owned.gateway)) + docs_response, docs_seconds = docs.result() + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + assert probes, "no liveliness probe ran while the docs request was pending" + assert {status for status, _ in probes} == {200}, probes + slowest: Final = max(seconds for _, seconds in probes) + assert slowest < docs_seconds / 3, ( + f"a liveliness probe waited {slowest:.2f}s of the {docs_seconds:.2f}s cold docs request" + ) + + +@OWNED_PROXY_TIMEOUT +def test_concurrent_cold_openapi_requests_share_the_schema_later_requests_get(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(COLD_DOCS) as pool, + ): + cold: Final = tuple(pool.map(lambda _: owned.gateway.request("GET", "/openapi.json"), range(COLD_DOCS))) + warm: Final = owned.gateway.request("GET", "/openapi.json") + assert [response.status_code for response in (*cold, warm)] == [200] * (COLD_DOCS + 1) + assert {response.content for response in cold} == {warm.content} + assert "/v1/chat/completions" in object_value(_schema(warm)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_chat_completions_overlapping_a_cold_openapi_build_all_succeed(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(CHATS + 1) as pool, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + first_chat: Final = _chat(owned.gateway, model, key, "before the docs build") + docs: Final = pool.submit(owned.gateway.request, "GET", "/openapi.json") + chats: Final = tuple( + pool.map(lambda index: _chat(owned.gateway, model, key, f"overlapping docs build {index}"), range(CHATS)) + ) + chats_returned_during_build: Final = not docs.done() + docs_response: Final = docs.result() + assert first_chat.status_code == 200, first_chat.text[:300] + assert chats_returned_during_build, "the chat calls returned only after the cold docs build finished" + assert [(response.status_code, response.text[:300]) for response in chats if response.status_code != 200] == [] + replies: Final = tuple(JSON_OBJECT.validate_json(response.content) for response in chats) + assert {reply["object"] for reply in replies} == {"chat.completion"} + assert all(list_value(reply["choices"]) for reply in replies), replies + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_docs_filtered_mode_serves_only_llm_routes(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {"DOCS_FILTERED": "True"}, client_timeout=DOCS_SECONDS) as owned: + paths: Final = object_value(_schema(owned.gateway.request("GET", "/openapi.json"))["paths"]) + assert "/v1/chat/completions" in paths + assert {"/key/generate", "/user/new", "/team/new"}.isdisjoint(paths), sorted(paths) + + +@OWNED_PROXY_TIMEOUT +def test_each_root_path_prefix_lists_itself_first_in_servers(gateway: Gateway, tmp_path: Path) -> None: + prefixes: Final = ("/tenant-a", "/tenant-b", "/tenant-a") + with owned_proxy_process( + gateway, tmp_path, {"SERVER_ROOT_PATHS": "/tenant-a,/tenant-b"}, client_timeout=DOCS_SECONDS + ) as owned: + served: Final = tuple( + (prefix, _servers(owned.gateway.request("GET", f"{prefix}/openapi.json"))) for prefix in prefixes + ) + assert served == tuple((prefix, [{"url": prefix}]) for prefix in prefixes) + + +@OWNED_PROXY_TIMEOUT +def test_a_scalar_server_root_path_is_listed_once_in_servers(gateway: Gateway, tmp_path: Path) -> None: + overrides: Final = {"SERVER_ROOT_PATH": "/gateway", "LITELLM_UI_PATH": str(tmp_path / "ui")} + with owned_proxy_process(gateway, tmp_path, overrides, client_timeout=DOCS_SECONDS) as owned: + servers: Final = _servers(owned.gateway.request("GET", "/gateway/openapi.json")) + assert servers == [{"url": "/gateway"}] + + +@TWO_OWNED_PROXIES_TIMEOUT +def test_a_lazy_feature_loaded_during_a_cold_openapi_build_shows_its_live_routes_in_the_next_schema( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy_process(gateway, tmp_path, {LAZY_ROUTES_FLAG: "true"}, client_timeout=DOCS_SECONDS) as eager: + live: Final = _warmed_feature_operations(eager.gateway.request("GET", "/openapi.json")) + with ( + owned_proxy_process( + gateway, tmp_path, {}, remove_environment=(LAZY_ROUTES_FLAG,), client_timeout=DOCS_SECONDS + ) as lazy, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(lazy.gateway.request, "GET", "/openapi.json") + warmed: Final = lazy.gateway.request("POST", f"/lazy/warm/{WARMED_FEATURE.name}") + warmed_during_build: Final = not docs.done() + during: Final = _warmed_feature_operations(docs.result()) + after: Final = _warmed_feature_operations(lazy.gateway.request("GET", "/openapi.json")) + assert warmed.status_code == 200, warmed.text + assert warmed_during_build, "the feature finished loading only after the cold docs build finished" + assert live, "the eager proxy published no route of the warmed feature" + assert during != live, "the cold docs build started after the feature had already loaded" + assert after == live