From c233776ac81e92f5984741bf45e4f4d11e9aa03e Mon Sep 17 00:00:00 2001 From: wefstratis Date: Mon, 14 Sep 2026 10:10:17 -0700 Subject: [PATCH] fix: handle llama.cpp context overflow in model adapter Signed-off-by: wefstratis --- responses_api_models/vllm_model/app.py | 37 ++++--- .../vllm_model/tests/test_app.py | 96 +++++++++++++++++++ 2 files changed, 122 insertions(+), 11 deletions(-) diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 7cf93b8c1c..ea42ac37f9 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -55,6 +55,30 @@ LOG = logging.getLogger("nemo_gym.vllm_model") + +def _is_context_length_error(error: ClientResponseError) -> bool: + """Recognize request overflow without swallowing runtime KV-cache failures.""" + if error.status != 400: + return False + + message = error.response_content.decode(errors="replace") + if "context length" in message or "max_tokens" in message: + return True + + # llama.cpp supplies a dedicated error type, unlike vLLM's BadRequestError. + try: + payload = json.loads(message) + except json.JSONDecodeError: + payload = None + if isinstance(payload, dict): + payload = payload.get("error", payload) + if isinstance(payload, dict) and payload.get("type") == "exceed_context_size_error": + return True + + # Retain compatibility when a proxy forwards only llama.cpp's message. + return "exceeds the available context size" in message or "is larger than the max context size" in message + + _TRANSPORT_LOG_CONTEXT_HEADERS = { "run_id": "x-nemo-gym-log-run-id", "adapter": "x-nemo-gym-log-adapter", @@ -827,12 +851,7 @@ async def chat_completions( 3. https://github.com/vllm-project/vllm/blob/685c99ee77b4818dcdd15b30fe0e0eff0d5d22ec/vllm/entrypoints/openai/serving_engine.py#L948 4. https://github.com/vllm-project/vllm/blob/685c99ee77b4818dcdd15b30fe0e0eff0d5d22ec/vllm/sampling_params.py#L463 """ - result_content_str = e.response_content.decode() - - is_out_of_context_length = e.status == 400 and ( - "context length" in result_content_str or "max_tokens" in result_content_str - ) - if is_out_of_context_length: + if _is_context_length_error(e): res = self._create_empty_chat_completion() res.choices[0].finish_reason = "length" return res @@ -1141,11 +1160,7 @@ async def _chat_completions_via_completions_api( try: completion_dict = await client.create_completion(**completion_body) except ClientResponseError as e: - result_content_str = e.response_content.decode() - is_out_of_context_length = e.status == 400 and ( - "context length" in result_content_str or "max_tokens" in result_content_str - ) - if is_out_of_context_length: + if _is_context_length_error(e): res = self._create_empty_chat_completion() res.choices[0].finish_reason = "length" return res diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index 927dc8a6ef..4198508bed 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/responses_api_models/vllm_model/tests/test_app.py @@ -18,6 +18,7 @@ from typing import Any, Union from unittest.mock import AsyncMock, MagicMock +from aiohttp.client_exceptions import ClientResponseError from fastapi.testclient import TestClient from pytest import MonkeyPatch, mark, raises @@ -3843,6 +3844,101 @@ def _make_completions_backend_model( ) +@mark.parametrize("use_completions_api", [False, True], ids=["chat-completions", "completions"]) +@mark.parametrize( + ("status", "error_content", "expect_length"), + [ + (400, '{"message": "This model\'s maximum context length is 262144 tokens."}', True), + (400, '{"message": "max_tokens must be at least 1, got -6677."}', True), + ( + 400, + json.dumps({"error": {"type": "exceed_context_size_error", "message": "request too long"}}), + True, + ), + ( + 400, + "request (268821 tokens) exceeds the available context size (262144 tokens), try increasing it", + True, + ), + (400, "input (268821 tokens) is larger than the max context size (262144 tokens)", True), + ( + 500, + json.dumps( + { + "error": { + "type": "exceed_context_size_error", + "message": "request (268821 tokens) exceeds the available context size (262144 tokens)", + } + } + ), + False, + ), + (500, '{"message": "context length allocation failed"}', False), + (500, "unified KV cache is full: context size exceeded", False), + (500, "speculative batch index is outside the sub-batch", False), + (401, '{"error": {"message": "Invalid API key"}}', False), + (400, '{"error": {"message": "Unknown model"}}', False), + (400, "Invalid context size configuration", False), + (400, "unified KV cache is full: context size exceeded", False), + ], + ids=[ + "vllm-context-length", + "vllm-max-tokens", + "llamacpp-error-type", + "llamacpp-request-text", + "llamacpp-input-text", + "http500-overflow", + "http500-context-length", + "http500-unified-kv", + "http500-speculative-batch", + "http401-auth", + "http400-unrelated", + "http400-context-config", + "http400-unified-kv", + ], +) +async def test_backend_context_overflow_handling( + monkeypatch: MonkeyPatch, + use_completions_api: bool, + status: int, + error_content: str, + expect_length: bool, +) -> None: + monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", MagicMock(return_value={})) + model = _make_completions_backend_model() if use_completions_api else TestApp()._setup_server(monkeypatch) + error = ClientResponseError(MagicMock(), (), status=status, message="backend request failed") + error.response_content = error_content.encode() + client = MagicMock(spec=NeMoGymAsyncOpenAI) + client.create_chat_completion = AsyncMock(side_effect=error) + client.create_completion = AsyncMock(side_effect=error) + model._clients = [client] + request = MagicMock() + request.session = {SESSION_ID_KEY: "context-overflow-test"} + request.headers = {} + body = NeMoGymChatCompletionCreateParamsNonStreaming( + messages=[NeMoGymChatCompletionUserMessageParam(role="user", content="hello")], + ) + + if expect_length: + result = await model.chat_completions(request, body) + assert result.object == "chat.completion" + assert result.model == model.config.model + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "length" + assert result.choices[0].message.role == "assistant" + assert result.choices[0].message.content is None + assert result.choices[0].message.tool_calls is None + else: + with raises(ClientResponseError) as exc_info: + await model.chat_completions(request, body) + assert exc_info.value is error + + used_method = client.create_completion if use_completions_api else client.create_chat_completion + unused_method = client.create_chat_completion if use_completions_api else client.create_completion + used_method.assert_awaited_once() + unused_method.assert_not_awaited() + + class TestCompletionsBackendRawRender: def test_single_user_message_string_content(self) -> None: model = _make_completions_backend_model()