diff --git a/tensorrt_llm/executor/postproc_worker.py b/tensorrt_llm/executor/postproc_worker.py index a81fbb7fb05e..cbaeb33bab2f 100644 --- a/tensorrt_llm/executor/postproc_worker.py +++ b/tensorrt_llm/executor/postproc_worker.py @@ -11,6 +11,7 @@ from ..bindings import executor as tllm from ..llmapi.tokenizer import TransformersTokenizer, load_hf_tokenizer from ..llmapi.utils import print_traceback_on_error +from ..logger import logger from ..sampling_params import SamplingParams from .ipc import ZeroMqQueue from .utils import ErrorResponse, is_llm_response @@ -69,7 +70,6 @@ class Output(NamedTuple): client_id: int res: Any is_final: bool - error: str = "" metrics: Optional[dict[str, float]] = None request_perf_metrics: Any = None disaggregated_params: Any = None @@ -193,24 +193,36 @@ async def handle_single_input(inp: PostprocWorker.Input, batch.append(inp.rsp) self._records.pop(client_id, None) return - is_final = inp.rsp.result.is_final if is_llm_response( - inp.rsp) else True - res, metrics, perf_metrics, disaggregated_params = await self._handle_input( - inp) - record = self._records.get(client_id) - should_abort = record._aborted if record else False - batch.append( - PostprocWorker.Output( - client_id=client_id, - res=res, - is_final=is_final, - metrics=metrics, - request_perf_metrics=perf_metrics, - disaggregated_params=disaggregated_params, - should_abort=should_abort, - )) - if is_final: - self._records.pop(client_id) + try: + is_final = inp.rsp.result.is_final if is_llm_response( + inp.rsp) else True + res, metrics, perf_metrics, disaggregated_params = await self._handle_input( + inp) + record = self._records.get(client_id) + should_abort = record._aborted if record else False + batch.append( + PostprocWorker.Output( + client_id=client_id, + res=res, + is_final=is_final, + metrics=metrics, + request_perf_metrics=perf_metrics, + disaggregated_params=disaggregated_params, + should_abort=should_abort, + )) + if is_final: + self._records.pop(client_id) + except Exception as e: + logger.error( + f"Postprocessing error for client {client_id}: {e}\n" + f"{traceback.format_exc()}") + batch.append( + ErrorResponse( + client_id=client_id, + error_msg=f"Postprocessing error: {e}", + request_id=getattr(inp.rsp, 'request_id', -1), + )) + self._records.pop(client_id, None) while not self._to_stop.is_set(): batch = [] diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index e798c2c6b567..8ab75a0e81cb 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -172,6 +172,7 @@ def __init__(self, self.id = id self.sampling_params = sampling_params self.postproc_params = postproc_params + self._error_msg: Optional[str] = None self._disaggregated_params = None self.decoding_iter = 0 self.cached_tokens = 0 @@ -259,6 +260,11 @@ def disaggregated_params(self) -> Optional[DisaggregatedParams]: """Returns the disaggregated params.""" return self._disaggregated_params + @property + def error(self) -> Optional[str]: + """Return the error message if this result completed with an error.""" + return self._error_msg + def _handle_sequence(self, finish_reasons, response_tensors, @@ -446,15 +452,14 @@ def _handle_response(self, if response.should_abort and not self._aborted: self.abort() - if response.error: - if self._background_error_handler is not None and ( - handler := self._background_error_handler()): - handler(response.error) elif is_llm_response(response): if response.has_error(): + self._error_msg = response.error_msg + self._done = True if self._background_error_handler is not None and ( handler := self._background_error_handler()): handler(response.error_msg) + return # Never fall through to response.result response_result = response.result if hasattr(response_result, "_result") and isinstance( @@ -546,6 +551,7 @@ def _handle_response(self, handler := self._background_error_handler()): handler() elif isinstance(response, ErrorResponse): + self._error_msg = response.error_msg self._done = True if self._background_error_handler is not None and ( handler := self._background_error_handler()): diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 49fd606a1360..cb0c7d29c7c0 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -64,6 +64,7 @@ class RequestOutput(DetokenizedGenerationResultBase, GenerationResult): context_logits (torch.Tensor, optional): The logits on the prompt token ids. disaggregated_params (DisaggregatedParams, optional): Parameters for disaggregated serving, including multimodal embedding handles. finished (bool): Whether the whole request is finished. + error (str, optional): The error message if this result completed with an error. """ def __init__(self) -> None: diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index b1c9bbb0b851..0371f2da7e03 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -16,7 +16,7 @@ import asyncio import traceback from abc import ABC, abstractmethod -from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Type +from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Type import aiohttp @@ -99,6 +99,7 @@ def __init__( max_retries: int = 1, retry_interval_sec: int = 1, session: Optional[aiohttp.ClientSession] = None, + disagg_id_generator: Optional[Callable[[], int]] = None, ): self._router = router self._role = role @@ -115,6 +116,7 @@ def __init__( ) self._max_retries = max_retries self._retry_interval_sec = retry_interval_sec + self._disagg_id_generator = disagg_id_generator async def _send_request( self, @@ -161,9 +163,14 @@ async def _post_with_retry( request: UCompletionRequest, hooks: Optional[ResponseHooks] = None, ) -> AsyncGenerator[Any, None]: - json_data = request.model_dump(exclude_unset=True, mode="json") is_stream = request.stream for attempt in range(self._max_retries + 1): + # Regenerate disagg_request_id on retry to avoid ID collision on workers + if attempt > 0 and self._disagg_id_generator is not None: + dp = getattr(request, "disaggregated_params", None) + if dp is not None and getattr(dp, "disagg_request_id", None) is not None: + dp.disagg_request_id = self._disagg_id_generator() + json_data = request.model_dump(exclude_unset=True, mode="json") try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() @@ -183,7 +190,15 @@ async def _post_with_retry( yield line # don't finish the request here since the response generator is not done yet else: - http_response.raise_for_status() + if http_response.status >= 400: + error_body = await http_response.text() + raise aiohttp.ClientResponseError( + http_response.request_info, + http_response.history, + status=http_response.status, + message=f"{http_response.reason}: {error_body[:2048]}", + headers=http_response.headers, + ) response_dict = await http_response.json() # yield here since python forbids return statements in async generators yield response_dict diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index ee1edbe5d00c..67d4a418e86c 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,7 +33,8 @@ from tensorrt_llm.llmapi import tracing from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig, MetadataServerConfig, ServerRole, - get_ctx_gen_server_addrs) + get_ctx_gen_server_addrs, + get_global_disagg_request_id) from tensorrt_llm.logger import logger from tensorrt_llm.serve.cluster_storage import (HttpClusterStorageServer, create_cluster_storage) @@ -141,7 +142,10 @@ async def validation_exception_handler(_, exc): self.register_routes() def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: - client = OpenAIHttpClient(router, role, self._req_timeout_secs, max_retries) + node_id = self._config.node_id + client = OpenAIHttpClient( + router, role, self._req_timeout_secs, max_retries, + disagg_id_generator=lambda: get_global_disagg_request_id(node_id)) self._perf_metrics_collector.add_client(client) return client diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e77d9e07e178..7ab886580608 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -380,13 +380,23 @@ async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None: raise ValueError( f"Context server returned {len(ctx_response.choices)} choices, expecting 1." ) - if ctx_response.choices[0].disaggregated_params is None: - raise ValueError("Context server did not return disaggregated params") - if ctx_response.choices[0].disaggregated_params.ctx_request_id is None: - raise ValueError("Invalid disaggregated params in context phase response.") - if ctx_response.choices[0].disaggregated_params.disagg_request_id is None: + choice = ctx_response.choices[0] + if choice.disaggregated_params is None: raise ValueError( - "Invalid disaggregated params in context phase response. disagg_request_id is None" + f"Context server did not return disaggregated params." + f" finish_reason={choice.finish_reason!r}" + ) + if choice.disaggregated_params.ctx_request_id is None: + raise ValueError( + f"Invalid disaggregated params: ctx_request_id is None." + f" finish_reason={choice.finish_reason!r}," + f" disagg_request_id={choice.disaggregated_params.disagg_request_id!r}" + ) + if choice.disaggregated_params.disagg_request_id is None: + raise ValueError( + f"Invalid disaggregated params: disagg_request_id is None." + f" finish_reason={choice.finish_reason!r}," + f" ctx_request_id={choice.disaggregated_params.ctx_request_id!r}" ) return ctx_response diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index e248ac96b1fa..58cb6a857021 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -939,6 +939,8 @@ async def _create_chat_response( disaggregated_params: Optional[LlmDisaggregatedParams] = None ) -> ChatCompletionResponse: await promise.aresult() + if promise.error is not None: + raise RuntimeError(f"Generation failed: {promise.error}") if self.postproc_worker_enabled: chat_response = promise.outputs[0]._postprocess_result else: @@ -1265,6 +1267,8 @@ async def completion_response( postproc_params: Optional[PostprocParams] ) -> CompletionResponse: response = await promise + if response.error is not None: + raise RuntimeError(f"Generation failed: {response.error}") if not self.postproc_worker_enabled: post_processor, args = postproc_params.post_processor, postproc_params.postproc_args pp_result = post_processor(response, args) diff --git a/tests/unittest/api_stability/references/request_output.yaml b/tests/unittest/api_stability/references/request_output.yaml index e0a75bd5c25c..5ef2255bfcb1 100644 --- a/tests/unittest/api_stability/references/request_output.yaml +++ b/tests/unittest/api_stability/references/request_output.yaml @@ -32,4 +32,7 @@ methods: annotation: Optional[dict[str, float]] default: None return_annotation: None -properties: {} +properties: + error: + annotation: Optional[str] + default: inspect._empty diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index 698344da0311..ee97510580b3 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import aiohttp import pytest @@ -194,9 +194,9 @@ async def test_request_with_custom_server( mock_response = self.dummy_response() mock_http_response = AsyncMock() + mock_http_response.status = 200 mock_http_response.headers = {"Content-Type": "application/json"} mock_http_response.json = AsyncMock(return_value=mock_response.model_dump()) - mock_http_response.raise_for_status = Mock() mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) mock_http_response.__aexit__ = AsyncMock() @@ -231,9 +231,9 @@ async def test_request_with_retry( mock_response = self.dummy_response() mock_http_response = AsyncMock() + mock_http_response.status = 200 mock_http_response.headers = {"Content-Type": "application/json"} mock_http_response.json = AsyncMock(return_value=mock_response.model_dump()) - mock_http_response.raise_for_status = Mock() mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) mock_http_response.__aexit__ = AsyncMock() @@ -269,3 +269,146 @@ async def test_invalid_request_type(self, openai_client): """Test handling of invalid request type.""" with pytest.raises(ValueError, match="Invalid request type"): await openai_client.send_request("invalid_request") + + +class TestHttpErrorBodyPreservation: + """Test that HTTP 4xx/5xx errors include the response body (TRTLLM-11123).""" + + def _mock_http_error(self, status, body): + r = AsyncMock() + r.status = status + r.reason = "Bad Request" if status == 400 else "Internal Server Error" + r.text = AsyncMock(return_value=body) + r.headers = {"Content-Type": "application/json"} + r.request_info = MagicMock() + r.history = () + r.__aenter__ = AsyncMock(return_value=r) + r.__aexit__ = AsyncMock(return_value=False) + return r + + def _make_client(self, session, **kwargs): + from prometheus_client.registry import REGISTRY + + REGISTRY._names_to_collectors = {} + REGISTRY._collector_to_names = {} + router = AsyncMock(spec=Router) + router.servers = ["localhost:8000"] + router.get_next_server = AsyncMock(return_value=("localhost:8000", None)) + router.finish_request = AsyncMock() + return OpenAIHttpClient( + router=router, + role=ServerRole.CONTEXT, + timeout_secs=10, + max_retries=0, + session=session, + **kwargs, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "status,body", + [ + (400, '{"error":"missing field X"}'), + (500, "internal failure detail"), + ], + ) + async def test_error_body_in_exception(self, status, body): + session = AsyncMock(spec=aiohttp.ClientSession) + session.post.return_value = self._mock_http_error(status, body) + client = self._make_client(session) + req = CompletionRequest( + model="m", + prompt="hi", + stream=False, + disaggregated_params=DisaggregatedParams(request_type="context_only", ctx_request_id=1), + ) + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + await client.send_request(req) + assert body[:20] in str(exc_info.value.message) + + +class TestDisaggIdRegenOnRetry: + """Test that disagg_request_id is regenerated on retry (TRTLLM-11123).""" + + def _ok_response(self): + return CompletionResponse( + model="m", + usage=UsageInfo(prompt_tokens=1, completion_tokens=1), + choices=[CompletionResponseChoice(index=0, text="ok")], + ).model_dump() + + def _mock_http_ok(self, json_val): + r = AsyncMock() + r.status = 200 + r.headers = {"Content-Type": "application/json"} + r.json = AsyncMock(return_value=json_val) + r.__aenter__ = AsyncMock(return_value=r) + r.__aexit__ = AsyncMock() + return r + + def _make_client(self, session, **kwargs): + from prometheus_client.registry import REGISTRY + + REGISTRY._names_to_collectors = {} + REGISTRY._collector_to_names = {} + router = AsyncMock(spec=Router) + router.servers = ["localhost:8000"] + router.get_next_server = AsyncMock(return_value=("localhost:8000", None)) + router.finish_request = AsyncMock() + return OpenAIHttpClient( + router=router, + role=ServerRole.CONTEXT, + timeout_secs=10, + max_retries=2, + retry_interval_sec=0, + session=session, + **kwargs, + ) + + @pytest.mark.asyncio + async def test_retry_regenerates_disagg_id(self): + session = AsyncMock(spec=aiohttp.ClientSession) + ids = iter(range(1000, 2000)) + client = self._make_client(session, disagg_id_generator=lambda: next(ids)) + + session.post.side_effect = [ + aiohttp.ClientError("transient"), + self._mock_http_ok(self._ok_response()), + ] + req = CompletionRequest( + model="m", + prompt="hi", + stream=False, + disaggregated_params=DisaggregatedParams( + request_type="context_only", disagg_request_id=42 + ), + ) + + with patch("asyncio.sleep", new_callable=AsyncMock): + resp = await client.send_request(req) + + assert req.disaggregated_params.disagg_request_id != 42 + assert isinstance(resp, CompletionResponse) + + @pytest.mark.asyncio + async def test_no_generator_keeps_original_id(self): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(session) # no disagg_id_generator + + session.post.side_effect = [ + aiohttp.ClientError("transient"), + self._mock_http_ok(self._ok_response()), + ] + req = CompletionRequest( + model="m", + prompt="hi", + stream=False, + disaggregated_params=DisaggregatedParams( + request_type="context_only", disagg_request_id=42 + ), + ) + + with patch("asyncio.sleep", new_callable=AsyncMock): + await client.send_request(req) + + assert req.disaggregated_params.disagg_request_id == 42 diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index bcbef10b9e25..8e5c8e51a57e 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -159,6 +159,42 @@ async def _delayed_gen_response(*_args, **_kwargs): ) +class TestVerifyCtxResponseDiagnostics: + """Test enriched error messages in _verify_ctx_response (TRTLLM-11123).""" + + @pytest.mark.asyncio + async def test_missing_disagg_params_includes_finish_reason(self): + svc = _make_service("context_first") + resp = _make_completion_response("", finish_reason="error", disagg_request_id=1) + resp.choices[0].disaggregated_params = None + with pytest.raises(ValueError, match="finish_reason='error'"): + await svc._verify_ctx_response(resp) + + @pytest.mark.asyncio + async def test_missing_ctx_request_id_includes_disagg_id(self): + svc = _make_service("context_first") + resp = _make_completion_response("", finish_reason="length", disagg_request_id=999) + resp.choices[0].disaggregated_params.ctx_request_id = None + with pytest.raises(ValueError, match=r"ctx_request_id is None.*999"): + await svc._verify_ctx_response(resp) + + @pytest.mark.asyncio + async def test_missing_disagg_request_id_includes_ctx_id(self): + svc = _make_service("context_first") + resp = _make_completion_response("", finish_reason="stop", disagg_request_id=555) + resp.choices[0].disaggregated_params.disagg_request_id = None + resp.choices[0].disaggregated_params.ctx_request_id = 555 + with pytest.raises(ValueError, match=r"disagg_request_id is None.*555"): + await svc._verify_ctx_response(resp) + + @pytest.mark.asyncio + async def test_valid_response_passes(self): + svc = _make_service("context_first") + resp = _make_completion_response("ok", finish_reason="stop", disagg_request_id=42) + result = await svc._verify_ctx_response(resp) + assert result is resp + + class TestFirstGenLogProbsSerializeRoundtrip: """Roundtrip tests for _serialize/_deserialize_first_gen_log_probs."""