wehub-resource-sync
259 lines
10 KiB
Python
259 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Verify PrefillMemoryExceededError maps to HTTP 400 in server.py.
|
|
|
|
Regression-arming test for the actual prefill-guard chain validated
|
|
end-to-end on 2026-05-15: the message string format matches what the
|
|
guard surfaces in production, so a refactor that changes either the
|
|
error body shape or the HTTP code will be caught here.
|
|
"""
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from omlx.exceptions import PrefillMemoryExceededError
|
|
|
|
|
|
def _build_test_app():
|
|
"""Build a minimal FastAPI app that re-uses the production handler."""
|
|
import omlx.server as srv
|
|
|
|
app = FastAPI()
|
|
app.add_exception_handler(
|
|
PrefillMemoryExceededError, srv.prefill_memory_exceeded_handler
|
|
)
|
|
|
|
@app.get("/v1/raise")
|
|
def raise_prefill_too_large():
|
|
raise PrefillMemoryExceededError(
|
|
message=(
|
|
"Prefill would require ~43.56 GB peak "
|
|
"(current 28.00 GB + KV+SDPA 15.56 GB) "
|
|
"but limit is 40.00 GB. "
|
|
"Reduce context length or increase --max-process-memory."
|
|
),
|
|
request_id="req-abc",
|
|
estimated_bytes=46_775_000_000,
|
|
limit_bytes=42_949_672_960,
|
|
)
|
|
|
|
@app.get("/health/raise")
|
|
def raise_prefill_too_large_health():
|
|
raise PrefillMemoryExceededError(
|
|
message="Prefill would require ~50 GB peak but limit is 40 GB.",
|
|
request_id="req-xyz",
|
|
)
|
|
|
|
return app
|
|
|
|
|
|
class TestPrefillMemoryHandler:
|
|
def test_returns_400(self):
|
|
with TestClient(_build_test_app()) as client:
|
|
resp = client.get("/v1/raise")
|
|
assert resp.status_code == 400
|
|
|
|
def test_api_route_uses_openai_error_body(self):
|
|
"""/v1/* routes get the OpenAI-style {"error": {"message": ...}} wrapper."""
|
|
with TestClient(_build_test_app()) as client:
|
|
resp = client.get("/v1/raise")
|
|
body = resp.json()
|
|
assert body["type"] == "error"
|
|
assert "error" in body
|
|
msg = body["error"]["message"]
|
|
# The guard's diagnostic format is part of the public contract — the
|
|
# CLI hint at the end tells the user exactly how to recover.
|
|
assert "Prefill would require" in msg
|
|
assert "KV+SDPA" in msg
|
|
assert "--max-process-memory" in msg
|
|
assert "Memory Guard to aggressive" in msg
|
|
assert "custom memory guard ceiling" in msg
|
|
assert body["error"]["code"] == "prefill_memory_exceeded"
|
|
assert body["error"]["omlx_code"] == "prefill_memory_exceeded"
|
|
|
|
def test_api_route_body_carries_estimated_and_limit_bytes(self):
|
|
"""Clients branch on the numeric ``estimated_bytes`` /
|
|
``limit_bytes`` fields rather than regex-matching the human
|
|
message (which is localized / format-prone). Regression for
|
|
the body-shape gap: prior to the fix on 2026-05-15 the handler
|
|
embedded these numbers only inside ``message`` and dropped the
|
|
structured fields, defeating the point of the typed exception
|
|
carrying them.
|
|
"""
|
|
with TestClient(_build_test_app()) as client:
|
|
resp = client.get("/v1/raise")
|
|
body = resp.json()
|
|
assert body["error"]["estimated_bytes"] == 46_775_000_000
|
|
assert body["error"]["limit_bytes"] == 42_949_672_960
|
|
|
|
def test_non_api_route_uses_plain_detail(self):
|
|
with TestClient(_build_test_app()) as client:
|
|
resp = client.get("/health/raise")
|
|
body = resp.json()
|
|
assert "detail" in body
|
|
assert "Prefill would require" in body["detail"]
|
|
assert body["omlx_code"] == "prefill_memory_exceeded"
|
|
|
|
|
|
class TestPostCommitPrefillMemorySurface:
|
|
@pytest.mark.asyncio
|
|
async def test_json_keepalive_emits_openai_error_body(self):
|
|
import json
|
|
|
|
import omlx.server as srv
|
|
|
|
class _Request:
|
|
async def is_disconnected(self):
|
|
return False
|
|
|
|
async def _raise_late():
|
|
raise PrefillMemoryExceededError(
|
|
message="Prefill context too large for available memory",
|
|
request_id="req-json",
|
|
estimated_bytes=123,
|
|
limit_bytes=100,
|
|
)
|
|
|
|
chunks = [
|
|
chunk
|
|
async for chunk in srv._with_json_keepalive(
|
|
_Request(), _raise_late(), disconnect_poll=0.001
|
|
)
|
|
]
|
|
body = json.loads("".join(chunks))
|
|
assert body["error"]["code"] == "prefill_memory_exceeded"
|
|
assert body["error"]["omlx_code"] == "prefill_memory_exceeded"
|
|
assert body["error"]["estimated_bytes"] == 123
|
|
assert body["error"]["limit_bytes"] == 100
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_keepalive_emits_openai_error_chunk(self):
|
|
import json
|
|
|
|
import omlx.server as srv
|
|
|
|
async def _gen():
|
|
raise PrefillMemoryExceededError(
|
|
message="Prefill context too large for available memory",
|
|
request_id="req-sse",
|
|
estimated_bytes=123,
|
|
limit_bytes=100,
|
|
)
|
|
yield ""
|
|
|
|
chunks = [
|
|
chunk
|
|
async for chunk in srv._with_sse_keepalive(_gen(), keepalive_chunk=None)
|
|
]
|
|
assert chunks[-1] == "data: [DONE]\n\n"
|
|
data = chunks[0].removeprefix("data: ").strip()
|
|
body = json.loads(data)
|
|
assert body["error"]["code"] == "prefill_memory_exceeded"
|
|
assert body["error"]["omlx_code"] == "prefill_memory_exceeded"
|
|
|
|
|
|
class TestResponsesEndpointReaches400:
|
|
"""End-to-end regression for ``/v1/responses``. The handler-shape tests
|
|
above use a synthetic ``/v1/raise`` route, which proves the handler
|
|
body but NOT the wiring of every prompt-bearing endpoint to the
|
|
preflight call. ``/v1/responses`` is the one route most-likely to
|
|
silently regress because it shares the StreamingResponse pattern
|
|
with ``/v1/chat/completions`` and reaches preflight via the same
|
|
code path. This test forces the preflight to raise and asserts
|
|
the route returns 400 instead of 200/500.
|
|
"""
|
|
|
|
def _make_app_with_failing_preflight(self):
|
|
"""Mount the real ``/v1/responses`` route with a mocked
|
|
engine_pool that returns an engine whose ``preflight_chat``
|
|
raises ``PrefillMemoryExceededError``. Hits the *production*
|
|
handler — not a synthesized stub — so a wiring regression is
|
|
caught.
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import omlx.server as srv
|
|
|
|
# Build an engine mock whose preflight_chat raises. The
|
|
# production handler awaits this BEFORE constructing
|
|
# StreamingResponse, so the raise propagates to the
|
|
# exception handler and the route can still emit 400.
|
|
async def _raising_preflight(*args, **kwargs):
|
|
raise PrefillMemoryExceededError(
|
|
message=(
|
|
"Prefill would require ~50 GB peak "
|
|
"(current 30 GB + KV+SDPA 20 GB) but limit "
|
|
"is 40 GB. Reduce context length or "
|
|
"increase --max-process-memory."
|
|
),
|
|
request_id="req-responses",
|
|
estimated_bytes=53_687_091_200,
|
|
limit_bytes=42_949_672_960,
|
|
)
|
|
|
|
engine = MagicMock()
|
|
engine.preflight_chat = AsyncMock(side_effect=_raising_preflight)
|
|
engine.start = AsyncMock()
|
|
# The handler calls ``count_chat_tokens`` and feeds the result
|
|
# into ``validate_context_window``; without a real int the
|
|
# comparison ``num_prompt_tokens > max_context`` raises before
|
|
# preflight ever runs.
|
|
engine.count_chat_tokens = MagicMock(return_value=128)
|
|
|
|
async def _get_engine_for_model(model_id, *, lease=None):
|
|
return engine
|
|
|
|
# Override the engine resolver and disable auth so the test
|
|
# talks to the real route.
|
|
srv.app.dependency_overrides[srv.verify_api_key] = lambda: True
|
|
srv.get_engine_for_model = _get_engine_for_model # type: ignore[assignment]
|
|
|
|
return srv.app
|
|
|
|
def test_v1_responses_returns_400_when_preflight_rejects(self):
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import omlx.server as srv
|
|
|
|
original_get_engine = srv.get_engine_for_model
|
|
original_overrides = dict(srv.app.dependency_overrides)
|
|
original_engine_pool = srv._server_state.engine_pool
|
|
try:
|
|
app = self._make_app_with_failing_preflight()
|
|
# Mock engine_pool so get_engine_pool() doesn't raise 503.
|
|
# get_entry returns None so the handler's preserve_thinking
|
|
# short-circuit doesn't fire.
|
|
from unittest.mock import AsyncMock
|
|
|
|
fake_pool = MagicMock()
|
|
fake_pool.get_entry = MagicMock(return_value=None)
|
|
fake_pool.preload_pinned_models = AsyncMock()
|
|
fake_pool.check_ttl_expirations = AsyncMock()
|
|
fake_pool.shutdown = AsyncMock()
|
|
srv._server_state.engine_pool = fake_pool
|
|
with TestClient(app, raise_server_exceptions=False) as client:
|
|
with (
|
|
patch.object(srv, "resolve_model_id", lambda name: name),
|
|
patch.object(srv, "validate_context_window", lambda *a, **k: None),
|
|
):
|
|
resp = client.post(
|
|
"/v1/responses",
|
|
json={
|
|
"model": "test-model",
|
|
"input": "Hello, world.",
|
|
"stream": False,
|
|
},
|
|
)
|
|
assert (
|
|
resp.status_code == 400
|
|
), f"expected 400, got {resp.status_code}: {resp.text}"
|
|
body = resp.json()
|
|
assert "error" in body, body
|
|
assert "Prefill would require" in body["error"]["message"]
|
|
assert "--max-process-memory" in body["error"]["message"]
|
|
finally:
|
|
srv.get_engine_for_model = original_get_engine
|
|
srv._server_state.engine_pool = original_engine_pool
|
|
srv.app.dependency_overrides.clear()
|
|
srv.app.dependency_overrides.update(original_overrides)
|