# Copyright (c) 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. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # yapf: disable import asyncio import errno import json import os import traceback from abc import ABC, abstractmethod from contextlib import aclosing from typing import ( Any, AsyncGenerator, Awaitable, Callable, List, Optional, Protocol, Tuple, Type, TypeVar, ) import aiohttp import msgspec from pydantic import BaseModel from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.disaggregated_params import DisaggScheduleStyle from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.logger import logger from tensorrt_llm.serve.disagg_auth import ( build_internal_disagg_auth_headers, build_subagent_affinity_headers, request_requires_internal_disagg_auth, ) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, CompletionRequest, CompletionResponse, ResponsesRequest, ResponsesResponse, UCompletionRequest, UCompletionResponse, ) from tensorrt_llm.serve.perf_metrics import ( _PERF_METRICS_HEADER_BUDGET_BYTES, RETURN_METRICS_HEADER, SSE_METRICS_EVENT, ClientMetricsCollector, adjusted_clock_from_headers, build_metrics_record_from_headers, ) from tensorrt_llm.serve.responses_utils import ( ResponseHooks, UCompletionResponseOrGenerator, get_steady_clock_now_in_seconds, ) from tensorrt_llm.serve.router import Router # yapf: enable # The forwarded orchestrator->worker request body is encoded as msgspec # msgpack and flagged with the X-TRTLLM-Msgpack header. _msgpack_encoder = msgspec.msgpack.Encoder() MSGPACK_HEADERS = {"Content-Type": "application/json", "X-TRTLLM-Msgpack": "1"} # A Responses stream opens with these events before the model emits anything. _RESPONSES_OPENING_EVENTS = frozenset((b"response.created", b"response.in_progress")) def _has_responses_output_event(chunk: bytes) -> bool: """Whether an SSE chunk of a Responses stream carries an event past its opening.""" return any( line.startswith(b"event: ") and line[len(b"event: ") :].strip() not in _RESPONSES_OPENING_EVENTS for line in chunk.split(b"\n") ) # post_json's request and response are ordinary pydantic models, but not the # UCompletion* union: it carries auxiliary payloads such as count-tokens. The # protocol names the one method actually called, and the TypeVar ties the # returned object to the type the caller asked for, so a mismatch is a type # error rather than an Any that silently propagates. class _SerializableRequest(Protocol): def model_dump( self, *, mode: str = ..., exclude_unset: bool = ..., by_alias: bool = ... ) -> dict: ... def _encode_request_body(request: _SerializableRequest) -> bytes: """Encode a request forwarded to a worker. Fields go out under their wire names, e.g. ``schema``, which the model holds as ``schema_``. """ return _msgpack_encoder.encode( request.model_dump(mode="json", exclude_unset=True, by_alias=True) ) _ResponseT = TypeVar("_ResponseT", bound=BaseModel) def _metrics_phase(role: ServerRole) -> str: return "ctx" if role is ServerRole.CONTEXT else "gen" class OpenAIClient(ABC): async def send_request( self, request: UCompletionRequest, server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: if isinstance(request, CompletionRequest): return await self._send_request( "v1/completions", request, CompletionResponse, server, hooks, req_id ) elif isinstance(request, ChatCompletionRequest): return await self._send_request( "v1/chat/completions", request, ChatCompletionResponse, server, hooks, req_id, ) elif isinstance(request, ResponsesRequest): return await self._send_request( "v1/responses", request, ResponsesResponse, server, hooks, req_id ) else: raise ValueError(f"Invalid request type: {type(request)}") @abstractmethod async def _send_request( self, endpoint: str, request: UCompletionRequest, response_type: Type[UCompletionResponse], server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: """Send a request to the server and return the response and the body generator. The request is finished (in routers) when the generator is exhausted or there is an error. """ ... @abstractmethod async def check_ready(self) -> Tuple[List[str], List[str]]: """Return the list of ready servers and the list of unready servers.""" ... @abstractmethod async def post_json( self, endpoint: str, request: _SerializableRequest, response_type: Type[_ResponseT], server: str, ) -> _ResponseT: """Post a non-streaming auxiliary request to a selected worker. Unlike _send_request this carries no disaggregation state and does not touch the router's in-flight accounting: it is for side endpoints such as token counting, which need a worker's tokenizer but no KV transfer. """ ... @abstractmethod async def get_json( self, endpoint: str, response_type: Type[_ResponseT], server: str, ) -> _ResponseT: """GET a worker endpoint, outside routing state, like post_json.""" ... async def shutdown(self) -> None: ... @abstractmethod async def _finish_request( self, request: UCompletionRequest, success: bool = True, req_id: Optional[int] = None, ) -> None: """Finish the request in the router. ``success`` lets the router distinguish completed vs failed requests (e.g. for routed-block tracking). """ ... def _server_base_url(server: str) -> str: """Return the base URL of a router server entry, with or without a scheme.""" if server.startswith(("http://", "https://")): return server.rstrip("/") return f"http://{server}".rstrip("/") class OpenAIHttpClient(OpenAIClient): def __init__( self, router: Router, role: ServerRole, timeout_secs: int = 180, max_retries: int = 1, retry_interval_sec: int = 1, session: Optional[aiohttp.ClientSession] = None, disagg_id_generator: Optional[Callable[[], Awaitable[int]]] = None, request_perf_metrics: bool = False, internal_disagg_auth_key: Optional[str] = None, bounded_replay_on_generation: bool = False, ): self._router = router self._role = role self._metrics_collector = ClientMetricsCollector(role) self._session = session or aiohttp.ClientSession( connector=aiohttp.TCPConnector( limit=0, limit_per_host=0, force_close=False, # Set keepalive_timeout below the server-side keepalive timeout to avoid reusing stale connections. keepalive_timeout=1, ), timeout=aiohttp.ClientTimeout(total=timeout_secs), max_field_size=_PERF_METRICS_HEADER_BUDGET_BYTES, ) self._no_retry = os.getenv("TRTLLM_DISAGG_NO_RETRY", "0") == "1" self._max_retries = 0 if self._no_retry else max_retries if self._no_retry: logger.info( "Disaggregated HTTP retry is DISABLED by " f"TRTLLM_DISAGG_NO_RETRY=1 for role={role.name}" ) self._retry_interval_sec = retry_interval_sec self._disagg_id_generator = disagg_id_generator self._request_perf_metrics = request_perf_metrics self._internal_disagg_auth_key = internal_disagg_auth_key self._bounded_replay_on_generation = bounded_replay_on_generation def _get_request_headers(self, request: UCompletionRequest) -> dict[str, str]: headers = build_subagent_affinity_headers( self._internal_disagg_auth_key, request, self._role ) if self._role == ServerRole.GENERATION and request_requires_internal_disagg_auth(request): headers.update( build_internal_disagg_auth_headers(self._internal_disagg_auth_key, request) ) return headers async def post_json( self, endpoint: str, request: _SerializableRequest, response_type: Type[_ResponseT], server: str, ) -> _ResponseT: url = f"{_server_base_url(server)}/{endpoint}" # A count-tokens body carries the whole conversation, so it is as large # as a chat body and encoding it costs the orchestrator's event loop the # same way -- hence the same msgpack transport as _send_request. The # worker decodes it because _MsgspecRoute is the app's route_class, so # it covers every route rather than just the completion ones. body = _encode_request_body(request) headers = dict(MSGPACK_HEADERS) async with self._session.post(url, data=body, headers=headers) as response: if response.status >= 400: error_body = await response.text() raise aiohttp.ClientResponseError( response.request_info, response.history, status=response.status, message=f"{response.reason}: {error_body[:2048]}", headers=response.headers, ) return response_type(**await response.json()) async def get_json( self, endpoint: str, response_type: Type[_ResponseT], server: str, ) -> _ResponseT: url = f"{_server_base_url(server)}/{endpoint}" async with self._session.get(url) as response: if response.status >= 400: error_body = await response.text() raise aiohttp.ClientResponseError( response.request_info, response.history, status=response.status, message=f"{response.reason}: {error_body[:2048]}", headers=response.headers, ) return response_type(**await response.json()) async def _send_request( self, endpoint: str, request: UCompletionRequest, response_type: Type[UCompletionResponse], server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: if server is None: if req_id is None: server, _ = await self._router.get_next_server(request) else: server, _ = await self._router.get_next_server(request, req_id=req_id) # A retry below re-issues disagg_request_id, but the router keyed this # request's reservation by the id it was routed with. Pin that id here so # every later renew/finish still addresses the original reservation -- # otherwise the coordinator never sees a release for it and the placement # load only drains when the expiration task fires. Only the context role # needs this: the coordinator keys generation requests by ctx_request_id, # which the retry loop never rewrites. if ( req_id is None and self._role == ServerRole.CONTEXT and self._disagg_id_generator is not None ): pinned_dp = request.disaggregated_params if pinned_dp is not None: req_id = pinned_dp.disagg_request_id url = f"{_server_base_url(server)}/{endpoint}" # disaggregated_params is None when conditional_disagg bypasses ctx. _dp = request.disaggregated_params _ctx_rid = _dp.ctx_request_id if _dp is not None else None logger.debug(f"Sending {self._role} request {_ctx_rid} to {url}") self._metrics_collector.total_requests.inc() if request.stream: resp_generator = self._stream_with_cleanup(server, url, request, hooks, req_id) # Enter its cleanup scope before returning: aclose() on an unstarted # async generator would otherwise leave the router reservation live. await resp_generator.__anext__() return resp_generator success = False try: response = None async with aclosing( self._post_with_retry(server, url, request, hooks, req_id) ) as response_generator: async for resp_json in response_generator: response = response_type(**resp_json) if hooks: if self._role == ServerRole.CONTEXT: hooks.on_ctx_resp(server, response) else: hooks.on_first_token(server, request) hooks.on_resp_done(server, request, response) success = True return response finally: if not success: self._metrics_collector.error_requests.inc() await self._finish_request(request, success=success, req_id=req_id) async def _stream_with_cleanup( self, server: str, url: str, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> AsyncGenerator[Any, None]: success = False try: # Consumed by _send_request, before the caller sees this generator. yield None async with aclosing( self._post_with_retry(server, url, request, hooks, req_id) ) as response_generator: async for chunk in response_generator: yield chunk success = True finally: # Own the logical request, including errors before response headers # and cancellation. Individual HTTP attempts must not release it. if not success: self._metrics_collector.error_requests.inc() await self._finish_request(request, success=success, req_id=req_id) async def _post_with_retry( self, server: str, url: str, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> AsyncGenerator[Any, None]: is_stream = request.stream disagg_params = request.disaggregated_params is_remote_tail_handoff = disagg_params is not None and ( ( self._role == ServerRole.GENERATION and disagg_params.request_type == "generation_only" and disagg_params.remote_tail_start is not None ) or ( self._bounded_replay_on_generation and disagg_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST and ( ( self._role == ServerRole.CONTEXT and disagg_params.request_type == "context_only" ) or ( self._role == ServerRole.GENERATION and disagg_params.request_type == "generation_only" ) ) ) ) # Remote-tail retries cannot safely replay or cancel the paired transfer. # Generation-first dispatch precedes the context response carrying the # tail boundary, so identify that handoff from the deployment config. max_retries = 0 if is_remote_tail_handoff else self._max_retries # Loop range must cover the transient-TCP extended budget (up to 5) # so the conditional raise inside the except block can actually decide # to keep retrying. Non-transient errors still raise on the first # attempt that reaches self._max_retries. _TRANSIENT_TCP_BUDGET = 0 if self._no_retry or is_remote_tail_handoff else 5 loop_max = max(max_retries, _TRANSIENT_TCP_BUDGET) + 1 for attempt in range(loop_max): if attempt > 0: await self._router.renew_request(request, req_id=req_id) # 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 = await self._disagg_id_generator() if hooks: hooks.on_disagg_request_id(dp.disagg_request_id) # Serialize once on the orchestrator's single event-loop thread. # Content-Type stays application/json so FastAPI still routes the # body through Request.json(); the header picks the decoder. body = _encode_request_body(request) headers = dict(MSGPACK_HEADERS) if self._request_perf_metrics: headers[RETURN_METRICS_HEADER] = "1" headers.update(self._get_request_headers(request)) try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() async with self._session.post( url, data=body, headers=headers, ) as http_response: response_clock = AdjustedSteadyClock() if self._request_perf_metrics: destination_time = get_steady_clock_now_in_seconds() response_clock = adjusted_clock_from_headers( http_response.headers, start_time, destination_time ) content_type = http_response.headers.get("Content-Type", "") if self._request_perf_metrics: role = _metrics_phase(self._role) disagg_params = getattr(request, "disaggregated_params", None) request_id = "" if disagg_params is not None: request_id = str( disagg_params.disagg_request_id or disagg_params.ctx_request_id or "" ) response_metrics = build_metrics_record_from_headers( http_response.headers, role, request_id=request_id, adjusted_clock=response_clock, ) if hooks and response_metrics: hooks.on_perf_metrics( server, role, response_metrics, ) if not is_stream and "text/event-stream" in content_type: raise ValueError( "Received an event-stream although request stream was False" ) if is_stream: # do NOT return generator directly here or the response will go # out of scope and get destroyed async with aclosing( self._response_generator( request, http_response, start_time, server, hooks, req_id, response_clock, ) ) as response_generator: async for line in response_generator: lines_yielded += 1 yield line else: 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 self._metrics_collector.complete_latency_seconds.observe( get_steady_clock_now_in_seconds() - start_time ) break # break and skip retries if the whole response is processed without exception except (aiohttp.ClientError, OSError) as e: if lines_yielded > 0: logger.error( f"Client error to {url}: {e} - cannot retry since {lines_yielded} lines were yielded", traceback.format_exc(), ) raise # ServerDisconnectedError and connection resets can result from # transient TCP races, such as client/server keepalive races at # burst start. Give them an extended retry budget to recover # while preserving the original fail-fast for other errors. # aiohttp can wrap an OS connection reset in ClientOSError # rather than a ConnectionResetError subclass. is_transient_tcp = isinstance( e, (aiohttp.ServerDisconnectedError, ConnectionResetError), ) or (isinstance(e, aiohttp.ClientOSError) and e.errno == errno.ECONNRESET) effective_max = max_retries if is_transient_tcp: effective_max = max(max_retries, _TRANSIENT_TCP_BUDGET) if attempt >= effective_max: logger.error( f"Client error to {url}: {e} - last retry {attempt} of {effective_max}" "failed", traceback.format_exc(), ) raise logger.error( f"{self._role} client error to {url}: {e} - retry {attempt} of {effective_max}", traceback.format_exc(), ) await self._router.renew_request(request, req_id=req_id) await asyncio.sleep(self._retry_interval_sec) self._metrics_collector.retry_requests.inc() except Exception as e: logger.error( f"Unexpected error while processing {self._role} request to {url}: {e}" ) raise async def _response_generator( self, request: UCompletionRequest, http_response: aiohttp.ClientResponse, start_time: float, server: str, hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, response_clock: Optional[AdjustedSteadyClock] = None, ) -> AsyncGenerator[Any, None]: assert request.stream, "Request is not streaming" assert "text/event-stream" in http_response.headers.get("Content-Type", ""), ( "Response is not streaming" ) try: last_token_time = start_time chunk_count = 0 # The first token of a Responses stream is the first frame after # its opening events, which the worker sends before generating. awaiting_output = isinstance(request, ResponsesRequest) marker = f"event: {SSE_METRICS_EVENT}\n".encode() event_delimiter = b"\n\n" pending = b"" metrics_event = b"" async for chunk in http_response.content.iter_any(): now_time = get_steady_clock_now_in_seconds() if awaiting_output and chunk: awaiting_output = not _has_responses_output_event(chunk) if chunk and not awaiting_output: if chunk_count == 0 and hooks: hooks.on_first_token(server, request) latency = now_time - last_token_time metric = ( self._metrics_collector.first_token_latency_seconds if chunk_count == 0 else self._metrics_collector.per_token_latency_seconds ) metric.observe(latency) chunk_count += 1 last_token_time = now_time if not self._request_perf_metrics: yield chunk await asyncio.sleep(0) continue if metrics_event: metrics_event += chunk continue pending += chunk marker_index = pending.find(marker) if marker_index >= 0: if marker_index: yield pending[:marker_index] metrics_event = pending[marker_index:] pending = b"" continue event_end = pending.rfind(event_delimiter) if event_end >= 0: emit_size = event_end + len(event_delimiter) yield pending[:emit_size] pending = pending[emit_size:] await asyncio.sleep(0) if pending: yield pending if hooks and metrics_event: data_prefix = b"data: " data_index = metrics_event.find(data_prefix) if data_index >= 0: data = metrics_event[data_index + len(data_prefix) :].split(b"\n", 1)[0] try: headers = json.loads(data) metrics = build_metrics_record_from_headers( headers, _metrics_phase(self._role), adjusted_clock=response_clock, ) except (TypeError, ValueError) as error: logger.warning("Ignoring malformed perf metrics event: %s", error) metrics = None if metrics: hooks.on_perf_metrics(server, _metrics_phase(self._role), metrics) if hooks: hooks.on_resp_done(server, request, None) self._metrics_collector.complete_latency_seconds.observe( get_steady_clock_now_in_seconds() - start_time ) except aiohttp.ClientError as e: # a client error is expected when the response stream is done if the connector has close=True logger.error(f"{self._role} client {server} error: {e}") raise async def _finish_request( self, request: UCompletionRequest, success: bool = True, req_id: Optional[int] = None, ) -> None: self._metrics_collector.completed_requests.inc() if req_id is None: await self._router.finish_request(request, self._session, success=success) else: await self._router.finish_request( request, self._session, success=success, req_id=req_id ) async def shutdown(self) -> None: await self._session.close() async def check_ready(self) -> Tuple[List[str], List[str]]: ready_servers, unready_servers = await OpenAIHttpClient.check_ready_for_servers( self._session, self._router.servers ) if ready_servers: await self._router.prepare_servers(ready_servers) return ready_servers, unready_servers @staticmethod async def check_ready_for_servers( session: aiohttp.ClientSession, servers: List[str] ) -> Tuple[List[str], List[str]]: async def check_server_ready(server: str) -> bool: try: url = f"{_server_base_url(server)}/health" async with session.get(url) as response: return response.status == 200 except Exception: return False servers_ready = await asyncio.gather(*[check_server_ready(server) for server in servers]) return [server for server, ready in zip(servers, servers_ready) if ready], [ server for server, ready in zip(servers, servers_ready) if not ready ]