mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
fix(passthrough): dispatch the upstream error report at the preview budget and bound late relay teardown
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
12d0145a8c
commit
18cd44cc5d
3 changed files with 127 additions and 97 deletions
|
|
@ -1531,6 +1531,8 @@ MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_T
|
|||
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096
|
||||
PASSTHROUGH_UPSTREAM_ERROR_REPORT_TASK_NAME: Final = "passthrough-upstream-error-report"
|
||||
PASSTHROUGH_UPSTREAM_ERROR_REPORT_SHUTDOWN_WAIT_SECONDS: Final = 10.0
|
||||
PASSTHROUGH_UPSTREAM_ERROR_REDACT_MARGIN_CHARS: Final = 256
|
||||
SSE_LATE_RELAY_BACKGROUND_WAIT_SECONDS: Final = 10.0
|
||||
|
||||
# Headers to control callbacks
|
||||
X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.constants import (
|
|||
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
SSE_LATE_RELAY_BACKGROUND_WAIT_SECONDS,
|
||||
STREAM_SSE_DATA_PREFIX,
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES,
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS,
|
||||
|
|
@ -1141,7 +1142,9 @@ def ttft_keepalive_interval(request_data: Mapping[str, object], llm_router: Rout
|
|||
return resolve_ttft_keepalive_interval(deployments, litellm.sse_keepalive_ping_interval_seconds)
|
||||
|
||||
|
||||
async def _aclose_late_response(produced: Response) -> None:
|
||||
async def _aclose_late_response(
|
||||
produced: Response, *, background_wait_seconds: float = SSE_LATE_RELAY_BACKGROUND_WAIT_SECONDS
|
||||
) -> None:
|
||||
"""Run the cleanup Starlette would have run, for a response it never called.
|
||||
|
||||
Closing an already-closed async generator is a no-op, so this is safe to call
|
||||
|
|
@ -1162,10 +1165,16 @@ async def _aclose_late_response(produced: Response) -> None:
|
|||
except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended
|
||||
verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc)
|
||||
if produced.background is not None:
|
||||
try:
|
||||
await produced.background()
|
||||
except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended
|
||||
verbose_proxy_logger.debug("error running relayed response background task: %s", exc)
|
||||
with anyio.move_on_after(background_wait_seconds) as scope:
|
||||
try:
|
||||
await produced.background()
|
||||
except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended
|
||||
verbose_proxy_logger.debug("error running relayed response background task: %s", exc)
|
||||
if scope.cancel_called:
|
||||
verbose_proxy_logger.warning(
|
||||
"relayed response background task still running after %.0fs, leaving it to shutdown",
|
||||
background_wait_seconds,
|
||||
)
|
||||
|
||||
|
||||
async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]:
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm._uuid import uuid
|
|||
from litellm.constants import (
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS,
|
||||
PASSTHROUGH_UPSTREAM_ERROR_REDACT_MARGIN_CHARS,
|
||||
PASSTHROUGH_UPSTREAM_ERROR_REPORT_TASK_NAME,
|
||||
REDACTED_BY_LITELLM,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
|
|
@ -892,7 +893,7 @@ class _PreviewReportingStream(httpx.AsyncByteStream):
|
|||
self._aborted = False
|
||||
self._report_task: asyncio.Task[None] | None = None
|
||||
|
||||
def _dispatch(self) -> asyncio.Task[None]:
|
||||
def dispatch(self) -> asyncio.Task[None]:
|
||||
if self._report_task is None:
|
||||
self._report_task = asyncio.get_running_loop().create_task(
|
||||
self._run_report(), name=PASSTHROUGH_UPSTREAM_ERROR_REPORT_TASK_NAME
|
||||
|
|
@ -908,7 +909,7 @@ class _PreviewReportingStream(httpx.AsyncByteStream):
|
|||
await self._report(b"".join(self._collected))
|
||||
|
||||
async def report_collected(self) -> None:
|
||||
await asyncio.shield(self._dispatch())
|
||||
await asyncio.shield(self.dispatch())
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
total = 0 # rebind-ok: running byte count against the preview budget
|
||||
|
|
@ -919,6 +920,7 @@ class _PreviewReportingStream(httpx.AsyncByteStream):
|
|||
total += len(chunk)
|
||||
if total > PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS:
|
||||
self._budget_crossed = True
|
||||
self.dispatch()
|
||||
yield chunk
|
||||
self._completed = True
|
||||
except httpx.HTTPError as err:
|
||||
|
|
@ -932,10 +934,10 @@ class _PreviewReportingStream(httpx.AsyncByteStream):
|
|||
await self.report_collected()
|
||||
raise
|
||||
finally:
|
||||
self._dispatch()
|
||||
self.dispatch()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._dispatch()
|
||||
self.dispatch()
|
||||
await self._upstream.aclose()
|
||||
|
||||
|
||||
|
|
@ -952,13 +954,18 @@ def _passthrough_upstream_failure_reporter(
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
proxy_logging: "ProxyLogging",
|
||||
log_warning: Callable[..., None],
|
||||
redact: Callable[[str], str] = redact_secrets,
|
||||
) -> _ReportPreview:
|
||||
async def report(preview: bytes) -> None:
|
||||
preview_text: Final = preview.decode(response.encoding or "utf-8", errors="replace")
|
||||
preview_text: Final = preview[
|
||||
: (PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS + PASSTHROUGH_UPSTREAM_ERROR_REDACT_MARGIN_CHARS) * 4
|
||||
].decode(response.encoding or "utf-8", errors="replace")[
|
||||
: PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS + PASSTHROUGH_UPSTREAM_ERROR_REDACT_MARGIN_CHARS
|
||||
]
|
||||
upstream_error_body: Final = (
|
||||
REDACTED_BY_LITELLM
|
||||
if should_redact_message_logging(logging_obj.model_call_details)
|
||||
else _truncate_upstream_error_body(_sanitize_upstream_error_body(redact_secrets(preview_text)))
|
||||
else _truncate_upstream_error_body(_sanitize_upstream_error_body(redact(preview_text)))
|
||||
)
|
||||
log_warning(
|
||||
"pass_through_endpoint: upstream %s %s returned %s: %s",
|
||||
|
|
@ -1001,6 +1008,7 @@ def _passthrough_upstream_failure_reporter(
|
|||
class _UpstreamRelay:
|
||||
response: httpx.Response
|
||||
background: BackgroundTask | None
|
||||
dispatch: Callable[[], object] | None
|
||||
|
||||
|
||||
async def _log_passthrough_upstream_failure(
|
||||
|
|
@ -1010,7 +1018,7 @@ async def _log_passthrough_upstream_failure(
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> _UpstreamRelay:
|
||||
if response.status_code < 400:
|
||||
return _UpstreamRelay(response=response, background=None)
|
||||
return _UpstreamRelay(response=response, background=None, dispatch=None)
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
log_warning: Final = verbose_proxy_logger.warning
|
||||
|
|
@ -1019,7 +1027,7 @@ async def _log_passthrough_upstream_failure(
|
|||
)
|
||||
if response.is_stream_consumed:
|
||||
await report(response.content)
|
||||
return _UpstreamRelay(response=response, background=None)
|
||||
return _UpstreamRelay(response=response, background=None, dispatch=None)
|
||||
stream: Final = _PreviewReportingStream(
|
||||
upstream=response,
|
||||
report=report,
|
||||
|
|
@ -1034,6 +1042,7 @@ async def _log_passthrough_upstream_failure(
|
|||
extensions=response.extensions,
|
||||
),
|
||||
background=BackgroundTask(stream.report_collected),
|
||||
dispatch=stream.dispatch,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1477,53 +1486,58 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Call response headers hook for streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=relay.response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=relay.response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
try:
|
||||
# Call response headers hook for streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=relay.response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=relay.response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=relay.response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=relay.response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
upstream_status=relay.response.status_code,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
),
|
||||
upstream_status=relay.response.status_code,
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=relay.response.headers,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=relay.response.headers,
|
||||
),
|
||||
headers=_response_headers,
|
||||
status_code=relay.response.status_code,
|
||||
background=relay.background,
|
||||
)
|
||||
headers=_response_headers,
|
||||
status_code=relay.response.status_code,
|
||||
background=relay.background,
|
||||
)
|
||||
except BaseException: # noqa: BLE001 # the relay owns the report until the response takes it over
|
||||
if relay.dispatch is not None:
|
||||
relay.dispatch()
|
||||
raise
|
||||
|
||||
if state_raw_body is not None:
|
||||
# SigV4-signed callers (Bedrock) require the exact pre-signed bytes
|
||||
|
|
@ -1570,53 +1584,58 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Call response headers hook for detected streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=detected_relay.response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=detected_relay.response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
try:
|
||||
# Call response headers hook for detected streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=detected_relay.response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=detected_relay.response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
_response_headers.update(callback_headers)
|
||||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=detected_relay.response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=detected_relay.response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
upstream_status=detected_relay.response.status_code,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
),
|
||||
upstream_status=detected_relay.response.status_code,
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=detected_relay.response.headers,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=detected_relay.response.headers,
|
||||
),
|
||||
headers=_response_headers,
|
||||
status_code=detected_relay.response.status_code,
|
||||
background=detected_relay.background,
|
||||
)
|
||||
headers=_response_headers,
|
||||
status_code=detected_relay.response.status_code,
|
||||
background=detected_relay.background,
|
||||
)
|
||||
except BaseException: # noqa: BLE001 # the relay owns the report until the response takes it over
|
||||
if detected_relay.dispatch is not None:
|
||||
detected_relay.dispatch()
|
||||
raise
|
||||
|
||||
if not _should_buffer_passthrough_response(response):
|
||||
relay_custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue