diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 579309035f0..14f0e3ead62 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,13 +1,15 @@ +import codecs import functools import json +import re from collections.abc import AsyncIterator, Iterator, Mapping -from typing import Final +from types import MappingProxyType +from typing import Final, Protocol import httpx from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, Required, TypedDict -import litellm from litellm import verbose_logger from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk @@ -40,12 +42,90 @@ class _SageMakerMessagesChunkDataRoot(RootModel[_SageMakerMessagesChunkData]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) -def _parse_sagemaker_chunk_data(data: str) -> _SageMakerChunkData: - return _SageMakerChunkDataRoot.model_validate(json.loads(data)).root +_STREAM_FRAMING: Final = re.compile(r"(?:\s|data:|\[DONE\]|(?::|event:|id:|retry:)[^\n]*\n)*") +_JSON_DECODER: Final = json.JSONDecoder() +_RAW_DECODED: Final = TypeAdapter(tuple[object, int]) +_NO_RECORD: Final = object() +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_CLIENT_ERROR_TYPE_STATUS: Final = MappingProxyType( + { + "overloaded": 429, + "validation": 422, + "template_error": 422, + "missing_template_variable": 422, + "tool_error": 422, + } +) +_CLIENT_ERROR_STATUSES: Final = frozenset(_CLIENT_ERROR_TYPE_STATUS.values()) -def _parse_sagemaker_messages_chunk_data(data: str) -> _SageMakerMessagesChunkData: - return _SageMakerMessagesChunkDataRoot.model_validate(json.loads(data)).root +def _framing_end(text: str) -> int: + framing: Final = _STREAM_FRAMING.match(text) + return framing.end() if framing else 0 + + +class _StreamRecordBuffer: + """ + SageMaker cuts a container's stream into PayloadParts at arbitrary points, so one part can hold + several SSE events or JSON lines, or only part of one, down to half of a UTF-8 character + """ + + def __init__(self) -> None: + self._decoder: Final = codecs.getincrementaldecoder("utf-8")(errors="replace") + self._pending: str = "" + + def feed(self, data: bytes) -> tuple[object, ...]: + self._pending += self._decoder.decode(data) + return tuple(iter(self._pop_record, _NO_RECORD)) + + def close(self) -> "SagemakerError | str": + unparsed: Final = self._pending + self._decoder.decode(b"", final=True) + record_start: Final = unparsed[_framing_end(unparsed) :] + if record_start.startswith(("{", "[")): + return SagemakerError( + status_code=500, + message=f"SageMaker stream ended inside a record, {len(record_start)} characters into it", + ) + return unparsed + + def _pop_record(self) -> object: + start: Final = _framing_end(self._pending) + try: + record, end = _RAW_DECODED.validate_python(_JSON_DECODER.raw_decode(self._pending, start)) + except json.JSONDecodeError: + self._pending = self._pending[start:] + return _NO_RECORD + self._pending = self._pending[end:] + return record + + +def _fields(value: object) -> Mapping[str, object]: + return STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else _NO_FIELDS + + +def _client_error_status(value: object) -> int | None: + return value if isinstance(value, int) and value in _CLIENT_ERROR_STATUSES else None + + +def _streamed_error_status(record: Mapping[str, object]) -> int: + detail: Final = _fields(record.get("error")) + candidates: Final = ( + detail.get("code"), + detail.get("http_status_code"), + record.get("code"), + _CLIENT_ERROR_TYPE_STATUS.get(str(record.get("error_type"))), + ) + return next(filter(None, map(_client_error_status, candidates)), 500) + + +def _failed_generation(record: Mapping[str, object]) -> "SagemakerError | None": + error_message: Final = _fields(record.get("token")).get("error_msg") + if not error_message and _fields(record.get("details")).get("finish_reason") != "error": + return None + return SagemakerError( + status_code=500, + message=f"SageMaker endpoint streamed a failed generation: {error_message or 'no error message'}", + ) def _load_sagemaker_response_stream_shape(): @@ -110,6 +190,10 @@ def client_error_details(error: Exception) -> tuple[int, str]: return response.response_metadata.http_status_code, response.error.message or str(error) +class _EventStreamMessage(Protocol): + def to_response_dict(self) -> Mapping[str, object]: ... + + class SagemakerError(BaseLLMException): def __init__( self, @@ -120,6 +204,24 @@ class SagemakerError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +class _StreamExceptionBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, hide_input_in_errors=True) + message: str = Field(default="", alias="Message") + error_code: str | None = Field(default=None, alias="ErrorCode") + + +def _stream_exception_error(headers: Mapping[str, object], body: object) -> SagemakerError: + name: Final = headers.get(":exception-type") or headers.get(":error-code") or "an error" + try: + detail: Final = _StreamExceptionBody.model_validate_json(body if isinstance(body, bytes) else b"") + except ValidationError: + return SagemakerError( + status_code=500, message=f"SageMaker stream failed with {name}: {headers.get(':error-message', '')}" + ) + code: Final = f" ({detail.error_code})" if detail.error_code else "" + return SagemakerError(status_code=500, message=f"SageMaker stream failed with {name}{code}: {detail.message}") + + class AWSEventStreamDecoder: def __init__(self, model: str, is_messages_api: bool | None = None) -> None: from botocore.parsers import EventStreamJSONParser @@ -159,52 +261,55 @@ class AWSEventStreamDecoder: usage=None, ) + def _record_to_chunk(self, record: object) -> GChunk | StreamingChatCompletionChunk | SagemakerError | None: + fields: Final = _fields(record) + if fields.get("error") is not None or fields.get("object") == "error": + return SagemakerError( + status_code=_streamed_error_status(fields), + message=f"SageMaker endpoint streamed an error: {json.dumps(record)}", + ) + if failed_generation := _failed_generation(fields): + return failed_generation + try: + if self.is_messages_api: + return self._chunk_parser_messages_api( + chunk_data=_SageMakerMessagesChunkDataRoot.model_validate(record).root + ) + return self._chunk_parser(chunk_data=_SageMakerChunkDataRoot.model_validate(record).root) + except ValidationError: + verbose_logger.error( + "Skipping a SageMaker stream record that is not a chunk: %s with keys %s", + type(record).__name__, + sorted(fields), + ) + return None + + def _chunks(self, records: tuple[object, ...]) -> Iterator[GChunk | StreamingChatCompletionChunk]: + for outcome in map(self._record_to_chunk, records): + if isinstance(outcome, SagemakerError): + raise outcome + if outcome is not None: + yield outcome + def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[GChunk | StreamingChatCompletionChunk | None]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() - accumulated_json = "" + records: Final = _StreamRecordBuffer() for chunk in iterator: event_stream_buffer.add_data(chunk) for event in event_stream_buffer: - message = self._parse_message_from_event(event) - if message: - # remove data: prefix and "\n\n" at the end - message = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(message) or "" - message = message.replace("\n\n", "") + if payload := self._parse_message_from_event(event): + yield from self._chunks(records.feed(payload)) - # Accumulate JSON data - accumulated_json += message - - # Try to parse the accumulated JSON - try: - if self.is_messages_api: - yield self._chunk_parser_messages_api( - chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) - ) - else: - yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) - # Reset accumulated_json after successful parsing - accumulated_json = "" - except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event - continue - - # Handle any remaining data after the iterator is exhausted - if accumulated_json: - try: - if self.is_messages_api: - yield self._chunk_parser_messages_api( - chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) - ) - else: - yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) - except json.JSONDecodeError: - # Handle or log any unparseable data at the end - verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json) - yield None + unparsed: Final = records.close() + if isinstance(unparsed, SagemakerError): + raise unparsed + if unparsed: + verbose_logger.error("Warning: Unparseable JSON data remained, %d characters of it", len(unparsed)) + yield None async def aiter_bytes( self, iterator: AsyncIterator[bytes] @@ -213,58 +318,23 @@ class AWSEventStreamDecoder: from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() - accumulated_json = "" + records: Final = _StreamRecordBuffer() async for chunk in iterator: event_stream_buffer.add_data(chunk) for event in event_stream_buffer: - try: - message = self._parse_message_from_event(event) - if message: - verbose_logger.debug("sagemaker parsed chunk bytes %s", message) - # remove data: prefix and "\n\n" at the end - message = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(message) or "" - message = message.replace("\n\n", "") + if payload := self._parse_message_from_event(event): + for parsed_chunk in self._chunks(records.feed(payload)): + yield parsed_chunk - # Accumulate JSON data - accumulated_json += message + unparsed: Final = records.close() + if isinstance(unparsed, SagemakerError): + raise unparsed + if unparsed: + verbose_logger.error("Warning: Unparseable JSON data remained, %d characters of it", len(unparsed)) + yield None - # Try to parse the accumulated JSON - if self.is_messages_api: - yield self._chunk_parser_messages_api( - chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) - ) - else: - yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) - # Reset accumulated_json after successful parsing - accumulated_json = "" - except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event - continue - except UnicodeDecodeError as e: - verbose_logger.warning("UnicodeDecodeError: %s. Attempting to combine with next event.", e) - continue - except Exception as e: - verbose_logger.error("Error parsing message: %s. Attempting to combine with next event.", e) - continue - - # Handle any remaining data after the iterator is exhausted - if accumulated_json: - try: - if self.is_messages_api: - yield self._chunk_parser_messages_api( - chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) - ) - else: - yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) - except json.JSONDecodeError: - # Handle or log any unparseable data at the end - verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json) - yield None - except Exception as e: - verbose_logger.error("Final error parsing accumulated JSON: %s", e) - - def _parse_message_from_event(self, event) -> str | None: + def _parse_message_from_event(self, event: _EventStreamMessage) -> bytes | None: response_stream_shape: Final = get_sagemaker_response_stream_shape() if response_stream_shape is None: raise SagemakerError( @@ -276,18 +346,14 @@ class AWSEventStreamDecoder: ) response_dict: Final = event.to_response_dict() parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + body: Final = response_dict.get("body") if response_dict["status_code"] != 200: - raise ValueError(f"Bad response code, expected 200: {response_dict}") + raise _stream_exception_error(_fields(response_dict.get("headers")), body) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: return None - return chunk.get("bytes").decode() - else: - chunk = response_dict.get("body") - if not chunk: - return None - - return chunk.decode() + return chunk.get("bytes") + return body if isinstance(body, bytes) and body else None diff --git a/tests/integration/providers/_sagemaker_chat_container.py b/tests/integration/providers/_sagemaker_chat_container.py index 6a06b92e215..4cec8f8b6fe 100644 --- a/tests/integration/providers/_sagemaker_chat_container.py +++ b/tests/integration/providers/_sagemaker_chat_container.py @@ -16,6 +16,7 @@ SECRET_KEY: Final = "synthetic-secret-key-for-testing" REGION: Final = "us-east-1" PROMPT: Final = "synthetic inference component request" CONTENT: Final = ("sagemaker ", "wire control") +USAGE: Final[Mapping[str, int]] = MappingProxyType({"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}) NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _MARKER_PREFIX: Final = "marker-" @@ -69,27 +70,61 @@ def _event_header(name: str, value: str) -> bytes: return struct.pack("!B", len(name_bytes)) + name_bytes + b"\x07" + struct.pack("!H", len(value_bytes)) + value_bytes +def _frame(headers: bytes, body: bytes) -> bytes: + prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers)) + message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) + + def _payload_part(body: bytes) -> bytes: headers: Final = ( _event_header(":event-type", "PayloadPart") + _event_header(":content-type", "application/octet-stream") + _event_header(":message-type", "event") ) - prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers)) - message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body - return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) + return _frame(headers, body) + + +def model_stream_error(message: str, error_code: str) -> bytes: + headers: Final = ( + _event_header(":exception-type", "ModelStreamError") + + _event_header(":content-type", "application/json") + + _event_header(":message-type", "exception") + ) + return _frame(headers, json.dumps({"Message": message, "ErrorCode": error_code}).encode()) def container_reply(identity: str, stream: bool) -> Reply: if not stream: return Reply(body=_completion(identity)) - chunks: Final = ( + return event_stream_reply( + ( + _chunk(identity, {"role": "assistant", "content": CONTENT[0]}, None), + _chunk(identity, {"content": CONTENT[1]}, None), + _chunk(identity, {}, "stop", usage=dict(USAGE)), + ) + ) + + +def event_stream_reply(payloads: tuple[bytes, ...], *, then: tuple[bytes, ...] = ()) -> Reply: + return Reply(chunks=(*map(_payload_part, payloads), *then), content_type="application/vnd.amazon.eventstream") + + +def vllm_stream_events(identity: str) -> tuple[bytes, ...]: + usage_event: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": SERVED_MODEL, + "choices": [], + "usage": dict(USAGE), + } + return ( _chunk(identity, {"role": "assistant", "content": CONTENT[0]}, None), _chunk(identity, {"content": CONTENT[1]}, None), - _chunk(identity, {}, "stop", usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}), - ) - return Reply( - chunks=tuple(_payload_part(chunk) for chunk in chunks), content_type="application/vnd.amazon.eventstream" + _chunk(identity, {}, "stop"), + f"data: {json.dumps(usage_event, separators=(',', ':'))}\n\n".encode(), + b"data: [DONE]\n\n", ) diff --git a/tests/integration/providers/test_sagemaker_chat_wire.py b/tests/integration/providers/test_sagemaker_chat_wire.py index 0416404cec8..261787fec38 100644 --- a/tests/integration/providers/test_sagemaker_chat_wire.py +++ b/tests/integration/providers/test_sagemaker_chat_wire.py @@ -1,11 +1,14 @@ import asyncio +import itertools +import json import uuid -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias import httpx +import openai import pytest from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value from integration._support.database import read_rows @@ -26,14 +29,18 @@ from integration.providers._sagemaker_chat_container import ( REGION, SECRET_KEY, SERVED_MODEL, + USAGE, answered_choices, assert_signed, by_marker, container_body, container_reply, + event_stream_reply, marked_prompt, marker, + model_stream_error, only_request, + vllm_stream_events, ) from pydantic import JsonValue @@ -260,6 +267,77 @@ def test_sagemaker_chat_keeps_the_deployment_aws_credentials_out_of_the_containe assert SECRET_KEY not in string_value(row["stored"]), row +def _seven_byte_parts(stream: bytes) -> tuple[bytes, ...]: + return tuple(stream[start : start + 7] for start in range(0, len(stream), 7)) + + +def _event_pairs(events: tuple[bytes, ...]) -> tuple[bytes, ...]: + return tuple(b"".join(events[start : start + 2]) for start in range(0, len(events), 2)) + + +_PART_CUTTINGS: Final[Mapping[str, Callable[[tuple[bytes, ...]], tuple[bytes, ...]]]] = MappingProxyType( + { + "one-event-per-part": lambda events: events, + "usage-and-done-share-a-part": lambda events: (*events[:-2], events[-2] + events[-1]), + "event-pairs-share-a-part": _event_pairs, + "seven-byte-parts": lambda events: _seven_byte_parts(b"".join(events)), + "keep-alive-comments-between-events": lambda events: tuple( + itertools.chain.from_iterable((b":\n\n", event) for event in events) + ), + } +) + + +def _sse(record: Mapping[str, JsonValue]) -> bytes: + return f"data: {json.dumps(record)}\n\n".encode() + + +@dataclass(frozen=True, slots=True) +class _StreamedAnswer: + call_id: str + text: str + usages: tuple[tuple[int, int, int], ...] + + +async def _chat_openai_stream_with_usage(gateway: Gateway, model: str, prompt: str) -> _StreamedAnswer: + async with async_openai_client(gateway) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + max_tokens=16, + stream=True, + stream_options={"include_usage": True}, + extra_body=dict(NO_CACHE), + ) + chunks: Final = [chunk async for chunk in stream] + return _StreamedAnswer( + stream.response.headers["x-litellm-call-id"], + "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices), + tuple( + (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + for chunk in chunks + if chunk.usage + ), + ) + + +@pytest.mark.parametrize("cutting", tuple(_PART_CUTTINGS)) +def test_sagemaker_chat_vllm_stream_reaches_the_caller_whole_however_sagemaker_cuts_its_payload_parts( + gateway: Gateway, cutting: str +) -> None: + identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}" + parts: Final = _PART_CUTTINGS[cutting](vllm_stream_events(identity)) + expected_usage: Final = (USAGE["prompt_tokens"], USAGE["completion_tokens"], USAGE["total_tokens"]) + + with wire_server(lambda _request: event_stream_reply(parts)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, {}) + answer: Final = asyncio.run(_chat_openai_stream_with_usage(gateway, model, marked_prompt())) + only_request(wire) + assert (answer.text, answer.usages) == ("".join(CONTENT), (expected_usage,)), answer + stored: Final = JSON_OBJECT.validate_json(string_value(_spend_row(answer.call_id)["stored"])) + assert (stored["prompt_tokens"], stored["completion_tokens"], stored["total_tokens"]) == expected_usage, stored + + @dataclass(frozen=True, slots=True) class _Shape: provider: str @@ -388,6 +466,148 @@ def test_sagemaker_chat_container_error_reaches_the_caller_after_one_attempt_wit assert_signed(request, ACCESS_KEY, REGION, None) +def test_sagemaker_chat_container_validation_error_record_answers_422_after_one_attempt(gateway: Gateway) -> None: + failure: Final = _sse({"error": {"message": "synthetic prompt too long", "http_status_code": 422}}) + + with wire_server(lambda _request: event_stream_reply((failure,))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, {}) + response: Final = _chat(gateway, model, marked_prompt(), {"stream": True}) + only_request(wire) + assert response.status_code == 422, response.text + assert "synthetic prompt too long" in response.text, response.text + + +async def _texts_until_error( + stream: openai.AsyncStream[openai.types.chat.ChatCompletionChunk], +) -> AsyncIterator[str | openai.APIError]: + try: + async for chunk in stream: + yield "".join(choice.delta.content or "" for choice in chunk.choices) + except openai.APIError as error: + yield error + + +async def _chat_openai_stream_until_error( + gateway: Gateway, model: str, prompt: str, extra: Mapping[str, JsonValue] = MappingProxyType({}) +) -> tuple[str | openai.APIError, ...]: + async with async_openai_client(gateway) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + max_tokens=16, + stream=True, + extra_body={**NO_CACHE, **extra}, + ) + return tuple([item async for item in _texts_until_error(stream)]) + + +_MID_STREAM_FAILURES: Final[Mapping[str, tuple[Callable[[tuple[bytes, ...]], Reply], str]]] = MappingProxyType( + { + "error-record": ( + lambda events: event_stream_reply( + (*events[:2], _sse({"error": {"message": "synthetic decode failure", "code": 500}})) + ), + "synthetic decode failure", + ), + "model-stream-error": ( + lambda events: event_stream_reply( + events[:2], then=(model_stream_error("synthetic stream timeout", "ModelInvocationTimeExceeded"),) + ), + "SageMaker stream failed with ModelStreamError (ModelInvocationTimeExceeded): synthetic stream timeout", + ), + "stream-ends-inside-a-record": ( + lambda events: event_stream_reply((*events[:2], events[2][:30])), + "SageMaker stream ended inside a record", + ), + } +) + + +@pytest.mark.parametrize("failure", tuple(_MID_STREAM_FAILURES)) +def test_sagemaker_chat_stream_failure_after_the_first_chunks_reaches_the_caller_as_an_error( + gateway: Gateway, failure: str +) -> None: + reply, expected = _MID_STREAM_FAILURES[failure] + events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}") + + with wire_server(lambda _request: reply(events)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, {}) + outcome: Final = asyncio.run(_chat_openai_stream_until_error(gateway, model, marked_prompt())) + only_request(wire) + errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError)) + assert "".join(item for item in outcome if isinstance(item, str)) == "".join(CONTENT), outcome + assert len(errors) == 1 and expected in errors[0].message, outcome + + +_PRE_CONTENT_FAILURES: Final[Mapping[str, Reply]] = MappingProxyType( + { + "error-record": event_stream_reply((_sse({"error": {"message": "synthetic decode failure", "code": 500}}),)), + "model-stream-error": event_stream_reply( + (), then=(model_stream_error("synthetic stream timeout", "ModelInvocationTimeExceeded"),) + ), + } +) + + +@pytest.mark.parametrize("failure", tuple(_PRE_CONTENT_FAILURES)) +def test_sagemaker_chat_stream_failure_before_any_content_falls_back_to_the_next_model( + gateway: Gateway, failure: str +) -> None: + events: Final = vllm_stream_events(f"sagemaker-chat-{uuid.uuid4().hex}") + + with ( + wire_server(lambda _request: _PRE_CONTENT_FAILURES[failure]) as primary_wire, + wire_server(lambda _request: event_stream_reply(events)) as fallback_wire, + gateway.scenario() as scenario, + ): + primary: Final = _deployment(scenario, primary_wire, {}) + fallback: Final = _deployment(scenario, fallback_wire, {}) + outcome: Final = asyncio.run( + _chat_openai_stream_until_error(gateway, primary, marked_prompt(), {"fallbacks": [fallback]}) + ) + primary_attempts: Final = primary_wire.drain() + only_request(fallback_wire) + assert outcome and all(isinstance(item, str) for item in outcome), outcome + assert "".join(item for item in outcome if isinstance(item, str)) == "".join(CONTENT), outcome + assert len(primary_attempts) == 1, primary_attempts + + +def _lmi_line(record: Mapping[str, JsonValue]) -> bytes: + return f"{json.dumps(dict(record), ensure_ascii=False)}\n".encode() + + +_LMI_TOKENS: Final = ("sagemaker", " lmi tokens") +_LMI_ERROR: Final = "exception occurred during rolling batch inference" +# djl-serving 8a99e6e5: rolling_batch.py's error Token rendered by output_formatter._jsonlines_output_formatter +_LMI_FAILED_GENERATION: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "token": {"id": -1, "text": "", "log_prob": -1, "special_token": True, "error_msg": _LMI_ERROR}, + "generated_text": "".join(_LMI_TOKENS), + "details": {"finish_reason": "error"}, + } +) + + +def test_sagemaker_lmi_failed_generation_after_the_first_tokens_reaches_the_caller_as_an_error( + gateway: Gateway, +) -> None: + tokens: Final = tuple( + _lmi_line({"token": {"id": index, "text": text, "log_prob": -0.1}}) for index, text in enumerate(_LMI_TOKENS, 1) + ) + reply: Final = event_stream_reply((*tokens, _lmi_line(_LMI_FAILED_GENERATION))) + + with wire_server(lambda _request: reply) as wire, gateway.scenario() as scenario: + model: Final = _deployment( + scenario, wire, {"hf_model_name": "meta-llama/Llama-2-7b-chat-hf"}, provider="sagemaker" + ) + outcome: Final = asyncio.run(_chat_openai_stream_until_error(gateway, model, marked_prompt())) + only_request(wire) + errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError)) + assert "".join(item for item in outcome if isinstance(item, str)) == "".join(_LMI_TOKENS), outcome + assert len(errors) == 1, outcome + assert f"SageMaker endpoint streamed a failed generation: {_LMI_ERROR}" in errors[0].message, outcome + + @dataclass(frozen=True, slots=True) class _Call: endpoint: Endpoint diff --git a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py index 17fac935679..ac1c816f86e 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py @@ -1,4 +1,9 @@ import json +import logging +import struct +import zlib +from collections.abc import AsyncIterator, Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -6,8 +11,10 @@ import pytest from botocore.exceptions import ClientError, ConnectionClosedError -from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, client_error_details +from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, SagemakerError, client_error_details from litellm.llms.sagemaker.completion.transformation import SagemakerConfig +from litellm.types.utils import GenericStreamingChunk as GChunk +from litellm.types.utils import StreamingChatCompletionChunk # --------------------------------------------------------------------------- # @@ -232,6 +239,423 @@ async def test_aiter_bytes_valid_chunk_followed_by_unicode_error(): assert chunks[0]["text"] == "hello" # Verify the content of the valid chunk +def _event_header(name: str, value: str) -> bytes: + return struct.pack("!B", len(name)) + name.encode() + b"\x07" + struct.pack("!H", len(value)) + value.encode() + + +_PAYLOAD_PART_HEADERS: Final = ( + _event_header(":event-type", "PayloadPart") + + _event_header(":content-type", "application/octet-stream") + + _event_header(":message-type", "event") +) +_MODEL_STREAM_ERROR_HEADERS: Final = ( + _event_header(":exception-type", "ModelStreamError") + + _event_header(":content-type", "application/json") + + _event_header(":message-type", "exception") +) + + +def _frame(headers: bytes, body: str | bytes) -> bytes: + encoded: Final = body.encode() if isinstance(body, str) else body + prelude: Final = struct.pack("!II", 16 + len(headers) + len(encoded), len(headers)) + message: Final = prelude + struct.pack("!I", zlib.crc32(prelude)) + headers + encoded + return message + struct.pack("!I", zlib.crc32(message)) + + +def _payload_part(body: str | bytes) -> bytes: + return _frame(_PAYLOAD_PART_HEADERS, body) + + +def _cut(stream: str, size: int) -> tuple[bytes, ...]: + encoded: Final = stream.encode() + return tuple(encoded[start : start + size] for start in range(0, len(encoded), size)) + + +def _sse(record: Mapping[str, object]) -> str: + return f"data: {json.dumps(record, ensure_ascii=False)}\n\n" + + +def _vllm_chunk(choices: list[Mapping[str, object]], **extra: object) -> str: + return _sse( + {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": choices, **extra} + ) + + +_VLLM_EVENTS: Final = ( + _vllm_chunk([{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}]), + _vllm_chunk([{"index": 0, "delta": {"content": "Hél"}, "finish_reason": None}]), + _vllm_chunk([{"index": 0, "delta": {"content": "lo 👋"}, "finish_reason": "stop"}]), + _vllm_chunk([], usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}), + "data: [DONE]\n\n", +) +_VLLM_STREAM: Final = "".join(_VLLM_EVENTS) +_TGI_EVENTS: Final = tuple( + f"data:{json.dumps({'index': 0, 'token': {'text': text}}, ensure_ascii=False)}\n\n" for text in ("Hél", "lo 👋") +) +_TGI_STREAM: Final = "".join(_TGI_EVENTS) +_TGI_KEEP_ALIVE_STREAM: Final = ( + f"{_TGI_EVENTS[0]}:\n\n: keep-alive\n\nevent: message\nid: 2\nretry: 1000\n{_TGI_EVENTS[1]}" +) +_LMI_STREAM: Final = "".join( + f"{json.dumps({'token': {'text': text}}, ensure_ascii=False)}\n" for text in ("Hél", "lo 👋") +) + + +async def _stream_frames( + frames: tuple[bytes, ...], is_messages_api: bool, mode: str +) -> AsyncIterator[GChunk | StreamingChatCompletionChunk | None]: + decoder: Final = AWSEventStreamDecoder(model="m", is_messages_api=is_messages_api) + if mode == "sync": + for chunk in decoder.iter_bytes(iter(frames)): + yield chunk + return + + async def stream(): + for frame in frames: + yield frame + + async for chunk in decoder.aiter_bytes(stream()): + yield chunk + + +async def _decode( + parts: tuple[str | bytes, ...], is_messages_api: bool, mode: str +) -> list[GChunk | StreamingChatCompletionChunk | None]: + return [chunk async for chunk in _stream_frames(tuple(map(_payload_part, parts)), is_messages_api, mode)] + + +def _text(chunk: GChunk | StreamingChatCompletionChunk | None, is_messages_api: bool) -> str | None: + if chunk is None: + return None + return chunk.choices[0].delta.content if is_messages_api else chunk["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + "parts", + [ + pytest.param(_VLLM_EVENTS, id="one-event-per-part"), + pytest.param((*_VLLM_EVENTS[:3], _VLLM_EVENTS[3] + _VLLM_EVENTS[4]), id="usage-and-done-share-a-part"), + pytest.param( + (_VLLM_EVENTS[0] + _VLLM_EVENTS[1], _VLLM_EVENTS[2] + _VLLM_EVENTS[3], _VLLM_EVENTS[4]), + id="event-pairs-share-a-part", + ), + pytest.param((*_VLLM_EVENTS[:4], "data: [DO", "NE]\n\n"), id="done-split-across-parts"), + pytest.param((_VLLM_STREAM,), id="whole-stream-in-one-part"), + pytest.param(_cut(_VLLM_STREAM, 7), id="seven-byte-parts"), + pytest.param(_cut(_VLLM_STREAM, 1), id="one-byte-parts-split-every-utf8-character"), + ], +) +async def test_messages_api_stream_keeps_every_chunk_however_sagemaker_cuts_the_parts( + parts: tuple[str | bytes, ...], mode: str, caplog: pytest.LogCaptureFixture +): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + chunks: Final = await _decode(parts, is_messages_api=True, mode=mode) + + assert None not in chunks + assert "".join(chunk.choices[0].delta.content for chunk in chunks if chunk.choices) == "Héllo 👋" + assert [(chunk.usage.prompt_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage] == [(7, 10)] + assert caplog.records == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + "parts", + [ + pytest.param(_TGI_EVENTS, id="tgi-sse-one-event-per-part"), + pytest.param((_TGI_STREAM,), id="tgi-sse-whole-stream-in-one-part"), + pytest.param(_cut(_TGI_STREAM, 5), id="tgi-sse-five-byte-parts"), + pytest.param(_cut(_TGI_STREAM, 1), id="tgi-sse-one-byte-parts-split-every-utf8-character"), + pytest.param((_TGI_KEEP_ALIVE_STREAM,), id="tgi-sse-comments-and-fields-whole-stream-in-one-part"), + pytest.param(_cut(_TGI_KEEP_ALIVE_STREAM, 1), id="tgi-sse-comments-and-fields-one-byte-parts"), + pytest.param((_LMI_STREAM,), id="json-lines-whole-stream-in-one-part"), + pytest.param(_cut(_LMI_STREAM, 5), id="json-lines-five-byte-parts"), + pytest.param(tuple(_LMI_STREAM.splitlines()), id="undelimited-objects-one-per-part"), + ], +) +async def test_completion_stream_keeps_every_token_however_sagemaker_cuts_the_parts( + parts: tuple[str | bytes, ...], mode: str, caplog: pytest.LogCaptureFixture +): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + chunks: Final = await _decode(parts, is_messages_api=False, mode=mode) + + assert None not in chunks + assert "".join(chunk["text"] for chunk in chunks) == "Héllo 👋" + assert caplog.records == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + "last_part", + [ + pytest.param(_VLLM_EVENTS[2][:20], id="cut-between-characters"), + pytest.param( + _VLLM_EVENTS[2].encode()[: _VLLM_EVENTS[2].encode().index("👋".encode()) + 2], id="cut-inside-a-character" + ), + ], +) +async def test_stream_that_ends_inside_a_record_raises_after_the_complete_ones(mode: str, last_part: str | bytes): + stream: Final = _stream_frames(tuple(map(_payload_part, (*_VLLM_EVENTS[:2], last_part))), True, mode) + + assert [_text(await anext(stream), True) for _ in range(2)] == ["", "Hél"] + with pytest.raises(SagemakerError, match="ended inside a record") as exc_info: + await anext(stream) + assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_text_left_over_that_is_not_a_record_is_logged_once(mode: str, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + chunks: Final = await _decode((*_VLLM_EVENTS[:3], "data: upstream hiccup\n\n"), is_messages_api=True, mode=mode) + + assert chunks[-1] is None + assert "".join(chunk.choices[0].delta.content for chunk in chunks[:-1]) == "Héllo 👋" + assert [record.levelno for record in caplog.records] == [logging.ERROR] + assert "hiccup" not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("parts", "is_messages_api", "status_code"), + [ + pytest.param( + ( + _VLLM_EVENTS[1], + _sse( + { + "error": { + "message": "CUDA out of memory while decoding", + "type": "InternalServerError", + "code": 500, + } + } + ) + + "data: [DONE]\n\n", + ), + True, + 500, + id="vllm-500-error-event-shares-a-part-with-done", + ), + pytest.param( + ( + _VLLM_EVENTS[1], + _sse( + {"error": {"message": "CUDA out of memory while decoding", "type": "BadRequestError", "code": 400}} + ), + "data: [DONE]\n\n", + ), + True, + 500, + id="vllm-400-error-event-in-its-own-part", + ), + pytest.param( + ( + _VLLM_EVENTS[1], + _sse( + { + "object": "error", + "message": "CUDA out of memory while decoding", + "type": "BadRequestError", + "code": 400, + } + ), + ), + True, + 500, + id="legacy-vllm-flat-error-event", + ), + pytest.param( + ( + _VLLM_EVENTS[1], + _sse({"error": {"message": "CUDA out of memory while decoding", "http_status_code": 422}}), + ), + True, + 422, + id="tgi-messages-api-error-event", + ), + pytest.param( + ( + _VLLM_EVENTS[1], + _sse({"error": {"message": "CUDA out of memory while decoding", "http_status_code": 424}}), + ), + True, + 500, + id="tgi-messages-api-generation-error-event", + ), + pytest.param( + ( + _VLLM_EVENTS[1], + _sse( + {"error": {"message": "CUDA out of memory while decoding", "type": "RateLimitError", "code": 429}} + ), + ), + True, + 429, + id="vllm-429-error-event", + ), + pytest.param( + (_VLLM_EVENTS[1], _sse({"error": {"message": "CUDA out of memory while decoding", "code": "oom"}})), + True, + 500, + id="error-event-without-a-status", + ), + pytest.param( + ( + _TGI_EVENTS[0] + + f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'generation'})}\n\n", + ), + False, + 500, + id="tgi-generation-error-event-shares-a-part-with-a-token", + ), + pytest.param( + ( + _TGI_EVENTS[0], + f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'overloaded'})}\n\n", + ), + False, + 429, + id="tgi-overloaded-error-event", + ), + pytest.param( + ( + _TGI_EVENTS[0], + f"data:{json.dumps({'error': 'CUDA out of memory while decoding', 'error_type': 'incomplete_generation'})}\n\n", + ), + False, + 500, + id="tgi-incomplete-generation-error-event", + ), + ], +) +async def test_container_error_event_raises_with_its_message_and_status_after_the_chunks_before_it( + parts: tuple[str, ...], is_messages_api: bool, status_code: int, mode: str +): + stream: Final = _stream_frames(tuple(map(_payload_part, parts)), is_messages_api, mode) + + assert _text(await anext(stream), is_messages_api) == "Hél" + with pytest.raises(SagemakerError, match="CUDA out of memory while decoding") as exc_info: + await anext(stream) + assert exc_info.value.status_code == status_code + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("last_record", "message"), + [ + pytest.param( + { + "token": {"id": 2, "text": "", "log_prob": 0.0, "error_msg": "CUDA out of memory while decoding"}, + "generated_text": "Hél", + "details": {"finish_reason": "error"}, + }, + "SageMaker endpoint streamed a failed generation: CUDA out of memory while decoding", + id="token-error-message", + ), + pytest.param( + { + "token": {"id": 2, "text": "", "log_prob": 0.0}, + "generated_text": "Hél", + "details": {"finish_reason": "error"}, + }, + "SageMaker endpoint streamed a failed generation: no error message", + id="error-finish-without-a-message", + ), + ], +) +async def test_lmi_failed_generation_raises_after_the_chunks_before_it( + last_record: Mapping[str, object], message: str, mode: str +): + lines: Final = "".join( + f"{json.dumps(record, ensure_ascii=False)}\n" + for record in ({"token": {"id": 1, "text": "Hél", "log_prob": -0.1}}, last_record) + ) + stream: Final = _stream_frames((_payload_part(lines),), False, mode) + + assert _text(await anext(stream), False) == "Hél" + with pytest.raises(SagemakerError) as exc_info: + await anext(stream) + assert exc_info.value.status_code == 500 + assert exc_info.value.message == message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_model_stream_error_raises_after_the_chunks_before_it(mode: str): + frames: Final = ( + _payload_part(_VLLM_EVENTS[1]), + _frame( + _MODEL_STREAM_ERROR_HEADERS, + json.dumps({"Message": "container exited", "ErrorCode": "ModelInvocationTimeExceeded"}), + ), + _payload_part(_VLLM_EVENTS[2]), + ) + stream: Final = _stream_frames(frames, True, mode) + + assert _text(await anext(stream), True) == "Hél" + with pytest.raises(SagemakerError) as exc_info: + await anext(stream) + assert exc_info.value.status_code == 500 + assert exc_info.value.message == ( + "SageMaker stream failed with ModelStreamError (ModelInvocationTimeExceeded): container exited" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize( + ("parts", "is_messages_api"), + [ + pytest.param( + tuple( + _vllm_chunk([{"index": 0, "delta": {"content": text}, "finish_reason": None}], error=None) + for text in ("Hél", "lo 👋") + ), + True, + id="vllm-chunks-with-null-error", + ), + pytest.param( + tuple( + f"data:{json.dumps({'index': 0, 'token': {'text': text}, 'error': None}, ensure_ascii=False)}\n\n" + for text in ("Hél", "lo 👋") + ), + False, + id="tgi-tokens-with-null-error", + ), + ], +) +async def test_chunk_with_a_null_error_field_is_delivered( + parts: tuple[str, ...], is_messages_api: bool, mode: str, caplog: pytest.LogCaptureFixture +): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + chunks: Final = await _decode(parts, is_messages_api=is_messages_api, mode=mode) + + assert "".join(_text(chunk, is_messages_api) or "" for chunk in chunks) == "Héllo 👋" + assert caplog.records == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_record_that_is_not_a_chunk_is_logged_and_skipped(mode: str, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + chunks: Final = await _decode( + (_VLLM_EVENTS[1], _sse({"object": "keepalive", "detail": "still decoding"}), _VLLM_EVENTS[2]), + is_messages_api=True, + mode=mode, + ) + + assert "".join(chunk.choices[0].delta.content for chunk in chunks) == "Héllo 👋" + assert [record.levelno for record in caplog.records] == [logging.ERROR] + assert "['detail', 'object']" in caplog.records[0].getMessage() + assert "still decoding" not in caplog.records[0].getMessage() + + class TestSagemakerTransform: def setup_method(self): self.config = SagemakerConfig()