mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
feat(sagemaker)!: inference components and OpenAI-compatible embeddings (#44985)
* feat(sagemaker): inference components and OpenAI-compatible embeddings
Send model_id on a SageMaker embedding deployment as the invoke call's
InferenceComponentName, sync and async, and keep it out of the request body.
Add the sagemaker/openai/<endpoint> route, which sends the OpenAI embeddings
request shape ({"input": [...]} plus extra_body fields) and reads
data[].embedding and usage from the response for containers that serve the
OpenAI embeddings API. Plain sagemaker/<endpoint> deployments keep the
Hugging Face shape
* refactor(sagemaker): type validate_environment messages and cover base64 embeddings
* fix(sagemaker): name model_id in inference component errors and reject a non-string model_id on embeddings
* fix(sagemaker): fall back to the error text when SageMaker sends a null message
A ClientError whose Error.Message is null (SageMaker's InternalFailure shape) reached the
inference component hint as None and crashed with a TypeError; the message now falls back
to the error's own text on the chat and embedding paths
* test(sagemaker): audit matrix for inference-component embeddings
* fix(sagemaker): read the error status and text through one shape-checked helper
botocore's ConnectionClosedError carries response=None, so the except blocks that
read e.response.get(...) crashed with an AttributeError and the proxy answered 400
"Invalid request format" for a dropped connection. client_error_details validates
the botocore error shape with pydantic and falls back to 500 plus the error's own
text for any other shape (a missing response, an httpx Response, a plain exception).
* fix(sagemaker): drop the inference component hint from chat and derive the SDK wire expectation from the request
The chat except blocks never see SageMaker's text: HTTPHandler.post raises from
raise_for_status before the handler's own SagemakerError line, so the hint check ran
on httpx's generic "Client error '400 Bad Request' for url" line and could not match on
any chat path. The hint now names litellm.embedding only and the chat blocks keep the
status and text they had, without the dead check. The two OpenAI SDK wire tests pinned
the SDK's default encoding_format; they now expect whatever the SDK sent.
* refactor(sagemaker): meet main's per-rule lint and type ceilings
Main started gating every lint rule on its merge-base count and added LIT015
(frozen pydantic models) after this branch was cut, so the merge with main
pushed LIT001, LIT011, LIT015, BLE001 and reportArgumentType past their
ceilings. The response models are frozen with tuple fields, the OpenAI
embedding response is built with model_copy instead of writing into the
caller's object, the invoke_endpoint arguments are a TypedDict, and
optional_params and extra_body are validated into Mapping[str, object] before
they are read
* refactor(sagemaker): pick the embedding config in one place and stop pinning upstream error text
The openai/, Voyage and Cohere checks lived in both get_model_config and
get_supported_openai_params, so they now share one lookup. Tests compare error
messages against the error's own text, or against botocore's and json's own
wording, instead of copying that wording, and assert whole messages where the
text is ours
* fix(caching): key direct calls on model_id so inference components don't share cached embeddings
* test(sagemaker): name the fake-boto3 SDK test class for what it covers
* test(sagemaker): kill and stop chaos workers through the owned-proxy harness
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b892bc0c10
commit
4fae3c64bb
12 changed files with 2115 additions and 45 deletions
|
|
@ -445,7 +445,7 @@ class Cache:
|
|||
|
||||
1. If caching groups are set, then return the caching group as the model https://docs.litellm.ai/docs/routing#caching-across-model-groups
|
||||
2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router()
|
||||
3. Else use the `model` passed in kwargs
|
||||
3. Else use the `model` passed in kwargs, plus `model_id` when set (e.g. a SageMaker inference component)
|
||||
"""
|
||||
litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
|
||||
metadata_sources: Final[tuple[dict, ...]] = (
|
||||
|
|
@ -460,7 +460,10 @@ class Cache:
|
|||
caching_group: Final = next(
|
||||
(group for source in metadata_sources if (group := self._get_caching_group(source, model_group))), None
|
||||
)
|
||||
return caching_group or model_group or kwargs["model"]
|
||||
model_id_suffix: Final = (
|
||||
f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else ""
|
||||
)
|
||||
return caching_group or model_group or f"{kwargs['model']}{model_id_suffix}"
|
||||
|
||||
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
|
||||
caching_groups: Final[list | None] = metadata.get("caching_groups", [])
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, RootModel, TypeAdapter, with_config
|
||||
from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, ValidationError, with_config
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -13,7 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
from litellm.types.utils import StreamingChatCompletionChunk
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
|
|
@ -54,7 +54,9 @@ def _load_sagemaker_response_stream_shape():
|
|||
from botocore.model import ServiceModel
|
||||
|
||||
loader: Final = Loader()
|
||||
service_dict: Final = _JSON_OBJECT.validate_python(loader.load_service_model("sagemaker-runtime", "service-2"))
|
||||
service_dict: Final = STR_KEYED_MAPPING.validate_python(
|
||||
loader.load_service_model("sagemaker-runtime", "service-2")
|
||||
)
|
||||
return ServiceModel(service_dict).shape_for("InvokeEndpointWithResponseStreamOutput")
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
|
|
@ -76,6 +78,38 @@ def get_sagemaker_response_stream_shape():
|
|||
return _load_sagemaker_response_stream_shape()
|
||||
|
||||
|
||||
def with_inference_component_hint(error_message: str) -> str:
|
||||
if "Inference Component Name header is required" in error_message:
|
||||
return f"{error_message}\n pass in via `litellm.embedding(..., model_id={{InferenceComponentName}})`"
|
||||
if "Inference Component Name header is not allowed" in error_message:
|
||||
return f"{error_message}\n remove `model_id` from this deployment, the endpoint has no inference components"
|
||||
return error_message
|
||||
|
||||
|
||||
class _ClientErrorBody(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
message: str | None = Field(default=None, alias="Message")
|
||||
|
||||
|
||||
class _ClientErrorMetadata(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
http_status_code: int = Field(default=500, alias="HTTPStatusCode")
|
||||
|
||||
|
||||
class _ClientErrorResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
error: _ClientErrorBody = Field(default_factory=_ClientErrorBody, alias="Error")
|
||||
response_metadata: _ClientErrorMetadata = Field(default_factory=_ClientErrorMetadata, alias="ResponseMetadata")
|
||||
|
||||
|
||||
def client_error_details(error: Exception) -> tuple[int, str]:
|
||||
try:
|
||||
response: Final = _ClientErrorResponse.model_validate(getattr(error, "response", None))
|
||||
except ValidationError:
|
||||
return 500, str(error)
|
||||
return response.response_metadata.http_status_code, response.error.message or str(error)
|
||||
|
||||
|
||||
class SagemakerError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from copy import deepcopy
|
|||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -24,12 +25,26 @@ from litellm.utils import (
|
|||
get_secret,
|
||||
)
|
||||
|
||||
from ..common_utils import AWSEventStreamDecoder, SagemakerError
|
||||
from ..common_utils import (
|
||||
STR_KEYED_MAPPING,
|
||||
AWSEventStreamDecoder,
|
||||
SagemakerError,
|
||||
client_error_details,
|
||||
with_inference_component_hint,
|
||||
)
|
||||
from ..embedding.transformation import SagemakerEmbeddingConfig
|
||||
from .transformation import SagemakerConfig
|
||||
|
||||
sagemaker_config: Final = SagemakerConfig()
|
||||
|
||||
|
||||
class _InvokeEndpointArgs(TypedDict):
|
||||
EndpointName: ReadOnly[str]
|
||||
ContentType: ReadOnly[str]
|
||||
Body: ReadOnly[bytes]
|
||||
CustomAttributes: ReadOnly[str]
|
||||
|
||||
|
||||
"""
|
||||
SAGEMAKER AUTH Keys/Vars
|
||||
os.environ['AWS_ACCESS_KEY_ID'] = ""
|
||||
|
|
@ -288,10 +303,7 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
raise e
|
||||
except Exception as e:
|
||||
verbose_logger.error("Sagemaker error %s", str(e))
|
||||
status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500)
|
||||
error_message = getattr(e, "response", {}).get("Error", {}).get("Message", str(e))
|
||||
if "Inference Component Name header is required" in error_message:
|
||||
error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`"
|
||||
status_code, error_message = client_error_details(e)
|
||||
raise SagemakerError(status_code=status_code, message=error_message)
|
||||
|
||||
return sagemaker_config.transform_response(
|
||||
|
|
@ -512,10 +524,7 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
)
|
||||
raise e
|
||||
except Exception as e:
|
||||
error_message = f"{e}"
|
||||
if "Inference Component Name header is required" in error_message:
|
||||
error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`"
|
||||
raise SagemakerError(status_code=500, message=error_message)
|
||||
raise SagemakerError(status_code=500, message=str(e))
|
||||
return sagemaker_config.transform_response(
|
||||
model=model,
|
||||
raw_response=response,
|
||||
|
|
@ -574,16 +583,29 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
#### EMBEDDING LOGIC
|
||||
# Transform request based on model type
|
||||
provider_config: Final = SagemakerEmbeddingConfig.get_model_config(model)
|
||||
request_data: Final = provider_config.transform_embedding_request(model, input, optional_params, {})
|
||||
endpoint_name: Final = SagemakerEmbeddingConfig.get_endpoint_name(model)
|
||||
embedding_params: Final = STR_KEYED_MAPPING.validate_python(optional_params)
|
||||
inference_component_name: Final = embedding_params.get("model_id")
|
||||
if inference_component_name is not None and not isinstance(inference_component_name, str):
|
||||
raise SagemakerError(status_code=400, message="model_id must be the inference component name as a string")
|
||||
body_params: Final = {k: v for k, v in embedding_params.items() if k != "model_id"}
|
||||
request_data: Final = provider_config.transform_embedding_request(endpoint_name, input, body_params, {})
|
||||
data: Final = json.dumps(request_data).encode("utf-8")
|
||||
invoke_args: Final = _InvokeEndpointArgs(
|
||||
EndpointName=endpoint_name,
|
||||
ContentType="application/json",
|
||||
Body=data,
|
||||
CustomAttributes="accept_eula=true",
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
request_str: Final = f"""
|
||||
response = client.invoke_endpoint(
|
||||
EndpointName={model},
|
||||
EndpointName={endpoint_name},
|
||||
ContentType="application/json",
|
||||
Body=f"{data!r}", # Use !r for safe representation
|
||||
CustomAttributes="accept_eula=true",
|
||||
InferenceComponentName={inference_component_name},
|
||||
)"""
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
|
|
@ -592,16 +614,14 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
)
|
||||
## EMBEDDING CALL
|
||||
try:
|
||||
response = client.invoke_endpoint(
|
||||
EndpointName=model,
|
||||
ContentType="application/json",
|
||||
Body=data,
|
||||
CustomAttributes="accept_eula=true",
|
||||
response = (
|
||||
client.invoke_endpoint(**invoke_args, InferenceComponentName=inference_component_name)
|
||||
if inference_component_name
|
||||
else client.invoke_endpoint(**invoke_args)
|
||||
)
|
||||
except Exception as e:
|
||||
status_code: Final = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500)
|
||||
error_message: Final = getattr(e, "response", {}).get("Error", {}).get("Message", str(e))
|
||||
raise SagemakerError(status_code=status_code, message=error_message)
|
||||
status_code, error_message = client_error_details(e)
|
||||
raise SagemakerError(status_code=status_code, message=with_inference_component_hint(error_message))
|
||||
|
||||
response = json.loads(response["Body"].read().decode("utf8"))
|
||||
## LOGGING
|
||||
|
|
@ -628,7 +648,7 @@ class SagemakerLLM(BaseAWSLLM):
|
|||
|
||||
# Use the request_data that was already transformed above
|
||||
return provider_config.transform_embedding_response(
|
||||
model=model,
|
||||
model=endpoint_name,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
|
|
|
|||
146
litellm/llms/sagemaker/embedding/openai_transformation.py
Normal file
146
litellm/llms/sagemaker/embedding/openai_transformation.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
"""
|
||||
Translate from OpenAI's `/v1/embeddings` to SageMaker's `/invoke`
|
||||
|
||||
For endpoints whose container already serves the OpenAI embeddings API, opted into per
|
||||
deployment with the `openai/` route (`sagemaker/openai/<endpoint-name>`): the request
|
||||
goes out as `{"input": [...]}` plus the caller's extra body fields, and the response is
|
||||
read as the OpenAI embeddings object (`data[].embedding`, `usage`).
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import AllEmbeddingInputValues
|
||||
|
||||
from httpx._models import Headers, Response
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
|
||||
from ..common_utils import STR_KEYED_MAPPING, SagemakerError
|
||||
|
||||
SUPPORTED_OPENAI_PARAMS: Final = ("dimensions", "encoding_format", "user")
|
||||
|
||||
|
||||
class OpenAIEmbeddingUsage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
prompt_tokens: int = 0
|
||||
total_tokens: int | None = None
|
||||
|
||||
|
||||
class OpenAIEmbeddingRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
embedding: tuple[float, ...] | str
|
||||
index: int | None = None
|
||||
|
||||
|
||||
class OpenAIEmbeddingsPayload(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
data: tuple[OpenAIEmbeddingRow, ...]
|
||||
model: str | None = None
|
||||
usage: OpenAIEmbeddingUsage | None = None
|
||||
|
||||
|
||||
class SagemakerOpenAIEmbeddingConfig(BaseEmbeddingConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig returns a list
|
||||
return list(SUPPORTED_OPENAI_PARAMS)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: callers keep adding to the mapped params
|
||||
mapped: Final = {k: v for k, v in non_default_params.items() if k in SUPPORTED_OPENAI_PARAMS}
|
||||
return {**optional_params, **mapped}
|
||||
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: dict[str, str] | Headers, # mutable-ok: SagemakerError takes the same header types
|
||||
) -> BaseLLMException:
|
||||
return SagemakerError(message=error_message, status_code=status_code, headers=headers)
|
||||
|
||||
def transform_embedding_request(
|
||||
self,
|
||||
model: str,
|
||||
input: "AllEmbeddingInputValues",
|
||||
optional_params: Mapping[str, object],
|
||||
headers: Mapping[str, str],
|
||||
) -> dict[str, object]: # mutable-ok: BaseEmbeddingConfig returns the request body as a dict
|
||||
extra_body: Final = optional_params.get("extra_body")
|
||||
extra_fields: Final[Mapping[str, object]] = (
|
||||
STR_KEYED_MAPPING.validate_python(extra_body) if isinstance(extra_body, dict) else {}
|
||||
)
|
||||
top_level_params: Final = {k: v for k, v in optional_params.items() if k != "extra_body"}
|
||||
return {"input": input, **top_level_params, **extra_fields}
|
||||
|
||||
def transform_embedding_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: Response,
|
||||
model_response: "EmbeddingResponse",
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
api_key: str | None = None,
|
||||
request_data: Mapping[str, object] | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> "EmbeddingResponse":
|
||||
try:
|
||||
response_data: Final = raw_response.json()
|
||||
except ValueError as e:
|
||||
raise SagemakerError(
|
||||
message=f"Failed to parse response: {e}",
|
||||
status_code=raw_response.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
payload: Final = OpenAIEmbeddingsPayload.model_validate(response_data)
|
||||
except ValidationError as e:
|
||||
raise SagemakerError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"Unexpected response format. Expected an OpenAI embeddings object with a 'data' list of "
|
||||
f"{{'embedding': [...]}} rows: {e}"
|
||||
),
|
||||
)
|
||||
|
||||
prompt_tokens: Final = payload.usage.prompt_tokens if payload.usage else 0
|
||||
total_tokens: Final = (
|
||||
payload.usage.total_tokens if payload.usage and payload.usage.total_tokens is not None else prompt_tokens
|
||||
)
|
||||
|
||||
return model_response.model_copy(
|
||||
update={
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"object": "embedding",
|
||||
"index": row.index if row.index is not None else idx,
|
||||
"embedding": row.embedding if isinstance(row.embedding, str) else list(row.embedding),
|
||||
}
|
||||
for idx, row in enumerate(payload.data)
|
||||
],
|
||||
"model": payload.model or model,
|
||||
"usage": Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=total_tokens),
|
||||
}
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: BaseConfig returns the request headers as a dict
|
||||
return {"Content-Type": "application/json"}
|
||||
|
|
@ -19,6 +19,20 @@ from litellm.types.utils import EmbeddingResponse, Usage
|
|||
|
||||
from ..common_utils import SagemakerError
|
||||
from .cohere_transformation import SagemakerCohereEmbeddingConfig
|
||||
from .openai_transformation import SagemakerOpenAIEmbeddingConfig
|
||||
|
||||
OPENAI_COMPATIBLE_ROUTE: Final = "openai/"
|
||||
|
||||
|
||||
def _model_specific_config(model: str) -> BaseEmbeddingConfig | None:
|
||||
if model.startswith(OPENAI_COMPATIBLE_ROUTE):
|
||||
return SagemakerOpenAIEmbeddingConfig()
|
||||
model_lower: Final = model.lower()
|
||||
if "voyage" in model_lower:
|
||||
return VoyageEmbeddingConfig()
|
||||
if "cohere" in model_lower:
|
||||
return SagemakerCohereEmbeddingConfig()
|
||||
return None
|
||||
|
||||
|
||||
class SagemakerEmbeddingConfig(BaseEmbeddingConfig):
|
||||
|
|
@ -35,25 +49,22 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig):
|
|||
Factory method to get the appropriate embedding config based on model type
|
||||
|
||||
Args:
|
||||
model: The model name
|
||||
model: The model name, optionally prefixed with the `openai/` route for
|
||||
endpoints whose container serves the OpenAI embeddings API
|
||||
|
||||
Returns:
|
||||
Appropriate embedding config instance
|
||||
"""
|
||||
model_lower: Final = model.lower()
|
||||
if "voyage" in model_lower:
|
||||
return VoyageEmbeddingConfig()
|
||||
if "cohere" in model_lower:
|
||||
return SagemakerCohereEmbeddingConfig()
|
||||
return cls()
|
||||
specific_config: Final = _model_specific_config(model)
|
||||
return cls() if specific_config is None else specific_config
|
||||
|
||||
@staticmethod
|
||||
def get_endpoint_name(model: str) -> str:
|
||||
return model.removeprefix(OPENAI_COMPATIBLE_ROUTE)
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
model_lower: Final = model.lower()
|
||||
if "voyage" in model_lower:
|
||||
return VoyageEmbeddingConfig().get_supported_openai_params(model)
|
||||
if "cohere" in model_lower:
|
||||
return SagemakerCohereEmbeddingConfig().get_supported_openai_params(model)
|
||||
return []
|
||||
specific_config: Final = _model_specific_config(model)
|
||||
return [] if specific_config is None else specific_config.get_supported_openai_params(model)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
|
|
|
|||
210
tests/integration/_support/sagemaker_embedding.py
Normal file
210
tests/integration/_support/sagemaker_embedding.py
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
import base64
|
||||
import json
|
||||
import struct
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
ACCESS_KEY: Final = "AKIAINTEGRATION000004"
|
||||
SECRET_KEY: Final = "synthetic-secret-key-for-testing"
|
||||
REGION: Final = "us-east-1"
|
||||
COMPONENT: Final = "integration-embed-component"
|
||||
SERVED_MODEL: Final = "integration-org/served-embedding-model"
|
||||
VOYAGE_SERVED_MODEL: Final = "voyage-served-model"
|
||||
ENDPOINT_PREFIX: Final = "integration-embed-"
|
||||
COMPONENT_HEADER: Final = "x-amzn-sagemaker-inference-component"
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
NOT_JSON_BODY: Final = b"<html>upstream gateway error</html>"
|
||||
EXPECTED_SHAPE: Final = "Expected an OpenAI embeddings object with a 'data' list of {'embedding': [...]} rows"
|
||||
NON_STRING_MODEL_ID: Final = "model_id must be the inference component name as a string"
|
||||
IC_REQUIRED: Final = (
|
||||
"Inference Component Name header is required for endpoints to which you plan to deploy inference components. "
|
||||
"Please include Inference Component Name header or consider using SageMaker models."
|
||||
)
|
||||
IC_NOT_ALLOWED: Final = (
|
||||
"Inference Component Name header is not allowed for endpoints to which you dont plan to deploy inference "
|
||||
"components. Please remove the Inference Component Name header and try again."
|
||||
)
|
||||
PASS_IN_HINT: Final = "pass in via `litellm.embedding(..., model_id={InferenceComponentName})`"
|
||||
REMOVE_HINT: Final = "remove `model_id` from this deployment, the endpoint has no inference components"
|
||||
|
||||
|
||||
def endpoint(scenario: str) -> str:
|
||||
return f"{ENDPOINT_PREFIX}{scenario}"
|
||||
|
||||
|
||||
def openai_route(scenario: str) -> str:
|
||||
return f"sagemaker/openai/{endpoint(scenario)}"
|
||||
|
||||
|
||||
def plain_route(scenario: str) -> str:
|
||||
return f"sagemaker/{endpoint(scenario)}"
|
||||
|
||||
|
||||
def invocations_target(name: str) -> str:
|
||||
return f"/endpoints/{name}/invocations"
|
||||
|
||||
|
||||
def vector(index: int) -> list[float]:
|
||||
return [0.25, -0.5, 1.0 + index]
|
||||
|
||||
|
||||
def inputs(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]:
|
||||
carried: Final = next((body[key] for key in ("input", "inputs", "texts") if key in body), [])
|
||||
return tuple(carried) if isinstance(carried, list) else (carried,)
|
||||
|
||||
|
||||
def embedding_value(index: int, encoding_format: JsonValue) -> JsonValue:
|
||||
floats: Final = vector(index)
|
||||
if encoding_format == "base64":
|
||||
return base64.b64encode(struct.pack(f"<{len(floats)}f", *floats)).decode()
|
||||
return floats
|
||||
|
||||
|
||||
def openai_rows(count: int, encoding_format: JsonValue, indexed: bool) -> list[JsonValue]:
|
||||
return [
|
||||
{
|
||||
"object": "embedding",
|
||||
**({"index": index} if indexed else {}),
|
||||
"embedding": embedding_value(index, encoding_format),
|
||||
}
|
||||
for index in range(count)
|
||||
]
|
||||
|
||||
|
||||
def _openai_usage(scenario: str) -> dict[str, JsonValue]:
|
||||
if scenario == "nousage":
|
||||
return {}
|
||||
if scenario == "promptonly":
|
||||
return {"usage": {"prompt_tokens": 5}}
|
||||
return {"usage": {"prompt_tokens": 7, "total_tokens": 7}}
|
||||
|
||||
|
||||
def openai_payload(body: Mapping[str, JsonValue], scenario: str = "openai") -> dict[str, JsonValue]:
|
||||
return {
|
||||
"object": "list",
|
||||
"data": openai_rows(len(inputs(body)), body.get("encoding_format"), indexed=scenario != "noindex"),
|
||||
**({} if scenario == "nomodel" else {"model": SERVED_MODEL}),
|
||||
**_openai_usage(scenario),
|
||||
}
|
||||
|
||||
|
||||
def json_reply(payload: JsonValue) -> Reply:
|
||||
return Reply(body=json.dumps(payload).encode())
|
||||
|
||||
|
||||
def aws_error(status: int, code: str, message: str, **fields: JsonValue) -> Reply:
|
||||
return Reply(
|
||||
status=status,
|
||||
body=json.dumps({"__type": code, "Message": message, **fields}).encode(),
|
||||
headers={"x-amzn-ErrorType": code},
|
||||
)
|
||||
|
||||
|
||||
def endpoint_name_of(request: Request) -> str:
|
||||
parts: Final = request.target.split("/")
|
||||
return unquote(parts[2]) if len(parts) > 2 else ""
|
||||
|
||||
|
||||
def scenario_of(name: str) -> str:
|
||||
if not name.startswith(ENDPOINT_PREFIX):
|
||||
return ""
|
||||
return name.removeprefix(ENDPOINT_PREFIX).split("-")[0]
|
||||
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
name: Final = endpoint_name_of(request)
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
match scenario_of(name):
|
||||
case "openai" | "nousage" | "nomodel" | "noindex" | "promptonly":
|
||||
return json_reply(openai_payload(body, scenario_of(name)))
|
||||
case "hf":
|
||||
return json_reply([vector(index) for index in range(len(inputs(body)))])
|
||||
case "icrequired":
|
||||
if COMPONENT_HEADER not in request.headers:
|
||||
return aws_error(400, "ValidationError", IC_REQUIRED)
|
||||
return json_reply(openai_payload(body))
|
||||
case "noic":
|
||||
if COMPONENT_HEADER in request.headers:
|
||||
return aws_error(400, "ValidationError", IC_NOT_ALLOWED)
|
||||
return json_reply([vector(index) for index in range(len(inputs(body)))])
|
||||
case "hfdict":
|
||||
return json_reply({"embedding": [vector(index) for index in range(len(inputs(body)))]})
|
||||
case "voyage":
|
||||
return json_reply(
|
||||
{
|
||||
"object": "list",
|
||||
"data": openai_rows(len(inputs(body)), None, indexed=True),
|
||||
"model": VOYAGE_SERVED_MODEL,
|
||||
"usage": {"total_tokens": 9},
|
||||
}
|
||||
)
|
||||
case "cohere":
|
||||
return json_reply(
|
||||
{
|
||||
"id": "integration-cohere-response",
|
||||
"texts": list(inputs(body)),
|
||||
"embeddings": [vector(index) for index in range(len(inputs(body)))],
|
||||
"meta": {"api_version": {"version": "1"}, "billed_units": {"input_tokens": 6}},
|
||||
"response_type": "embeddings_floats",
|
||||
}
|
||||
)
|
||||
case "datadict":
|
||||
return json_reply({"object": "list", "data": {"embedding": vector(0)}})
|
||||
case "hfshape":
|
||||
return json_reply([vector(0)])
|
||||
case "notjson":
|
||||
return Reply(body=NOT_JSON_BODY, content_type="text/html")
|
||||
case "err400":
|
||||
return aws_error(
|
||||
400, "ValidationError", "1 validation error detected: Value at 'body' failed to satisfy constraint"
|
||||
)
|
||||
case "err424":
|
||||
return aws_error(
|
||||
424,
|
||||
"ModelError",
|
||||
'Received client error (400) from model with message "{\\"error\\":\\"unexpected field inputs\\"}".',
|
||||
OriginalStatusCode=400,
|
||||
OriginalMessage='{"error":"unexpected field inputs"}',
|
||||
LogStreamArn="arn:aws:logs:us-east-1:000000000000:log-group:/aws/sagemaker/Endpoints/integration",
|
||||
)
|
||||
case "err429":
|
||||
return aws_error(429, "ModelNotReadyException", "Model is not ready for inference yet")
|
||||
case "err503":
|
||||
return aws_error(503, "ServiceUnavailable", "The endpoint is scaling, try again later")
|
||||
case "err500":
|
||||
return aws_error(500, "InternalFailure", "An internal failure occurred")
|
||||
case "drop":
|
||||
return Reply(drop_connection=True)
|
||||
case _:
|
||||
return aws_error(400, "ValidationError", f"Endpoint {name} of account 000000000000 not found.")
|
||||
|
||||
|
||||
def litellm_params(model: str, **parameters: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"aws_access_key_id": ACCESS_KEY,
|
||||
"aws_secret_access_key": SECRET_KEY,
|
||||
"aws_region_name": REGION,
|
||||
**parameters,
|
||||
}
|
||||
|
||||
|
||||
def deployment(
|
||||
name: str, model: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue
|
||||
) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": litellm_params(model, **parameters),
|
||||
**({} if model_info is None else {"model_info": dict(model_info)}),
|
||||
}
|
||||
|
||||
|
||||
def openai_deployment(
|
||||
name: str, scenario: str, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue
|
||||
) -> dict[str, JsonValue]:
|
||||
defaults: Final[dict[str, JsonValue]] = {"model_id": COMPONENT, "extra_body": {"model": SERVED_MODEL}}
|
||||
return deployment(name, openai_route(scenario), model_info=model_info, **{**defaults, **parameters})
|
||||
302
tests/integration/providers/test_sagemaker_embedding_chaos.py
Normal file
302
tests/integration/providers/test_sagemaker_embedding_chaos.py
Normal file
|
|
@ -0,0 +1,302 @@
|
|||
import asyncio
|
||||
import re
|
||||
import socket
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy, owned_proxy_process
|
||||
from integration._support.sagemaker_embedding import (
|
||||
COMPONENT,
|
||||
COMPONENT_HEADER,
|
||||
JSON_OBJECT,
|
||||
SERVED_MODEL,
|
||||
json_reply,
|
||||
openai_deployment,
|
||||
openai_payload,
|
||||
respond,
|
||||
)
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark = pytest.mark.timeout(300)
|
||||
|
||||
_MODEL: Final = "sm-openai-chaos"
|
||||
_ALIASES: Final = (
|
||||
"/v1/embeddings",
|
||||
"/embeddings",
|
||||
f"/openai/deployments/{_MODEL}/embeddings",
|
||||
f"/engines/{_MODEL}/embeddings",
|
||||
)
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
_ROWS_BY_CALL: Final = (
|
||||
"SELECT litellm_call_id, status FROM \"LiteLLM_SpendLogs\" WHERE litellm_call_id = ANY(string_to_array(%s, ','))"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
text: str
|
||||
status: int
|
||||
body: str
|
||||
call_id: str
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
return probe.getsockname()[1]
|
||||
|
||||
|
||||
def _config(directory: Path) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["model_list"] = [openai_deployment(_MODEL, "openai", model_info={"mode": "embedding"})]
|
||||
path: Final = directory / "sagemaker-embedding-chaos.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _runtime(peer_url: str) -> dict[str, str]:
|
||||
return {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": peer_url, "AWS_MAX_ATTEMPTS": "1"}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig() -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as gateway:
|
||||
yield gateway
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def peer_port() -> int:
|
||||
return _free_port()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy(rig: Gateway, peer_port: int, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("sagemaker-embedding-chaos")
|
||||
runtime: Final = _runtime(f"http://127.0.0.1:{peer_port}")
|
||||
with owned_proxy(rig, directory, runtime, config=_config(directory), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _texts(count: int) -> tuple[str, ...]:
|
||||
return tuple(f"integration embedding {uuid.uuid4().hex}" for _ in range(count))
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, index: int, text: str) -> _Served:
|
||||
path: Final = _ALIASES[index % len(_ALIASES)]
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"input": text,
|
||||
"cache": {"no-cache": True},
|
||||
**({"model": _MODEL} if path.endswith("/v1/embeddings") or path == "/embeddings" else {}),
|
||||
}
|
||||
response: Final = await client.post(path, json=body, headers={"Authorization": f"Bearer {key}"})
|
||||
return _Served(text, response.status_code, response.text, response.headers.get("x-litellm-call-id", ""))
|
||||
|
||||
|
||||
async def _burst(
|
||||
base_url: str, key: str, texts: tuple[str, ...], *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_send(client, key, index, text) for index, text in enumerate(texts)),
|
||||
return_exceptions=tolerate_transport_errors,
|
||||
)
|
||||
for result in results:
|
||||
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
def _assert_embedded(served: _Served) -> None:
|
||||
assert served.status == 200, served.body
|
||||
payload: Final = JSON_OBJECT.validate_json(served.body)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], served.body
|
||||
|
||||
|
||||
def _assert_peer_saw(received: tuple[Request, ...], texts: tuple[str, ...]) -> None:
|
||||
assert {request.headers.get(COMPONENT_HEADER) for request in received} == {COMPONENT}, received
|
||||
bodies: Final = [JSON_OBJECT.validate_json(request.body) for request in received]
|
||||
assert sorted(string_value(body["input"]) for body in bodies) == sorted(texts), bodies
|
||||
assert {body["model"] for body in bodies} == {SERVED_MODEL}, bodies
|
||||
|
||||
|
||||
def _spend_status_by_call(served: tuple[_Served, ...]) -> dict[str, JsonValue]:
|
||||
wanted: Final = sorted(item.call_id for item in served)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(_ROWS_BY_CALL, (",".join(wanted),)), lambda found: len(found) >= len(wanted), seconds=70
|
||||
)
|
||||
assert sorted(string_value(row["litellm_call_id"]) for row in rows) == wanted, rows
|
||||
return {string_value(row["litellm_call_id"]): row["status"] for row in rows}
|
||||
|
||||
|
||||
def _unhealthy_count(gateway: Gateway) -> JsonValue:
|
||||
response: Final = gateway.request("GET", "/health", params={"model": _MODEL})
|
||||
return JSON_OBJECT.validate_json(response.content).get("unhealthy_count")
|
||||
|
||||
|
||||
async def test_concurrent_burst_across_every_alias_reaches_the_component_and_is_logged_once_per_call(
|
||||
proxy: Gateway, peer_port: int
|
||||
) -> None:
|
||||
texts: Final = _texts(24)
|
||||
with wire_server(respond, port=peer_port) as wire:
|
||||
served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts)
|
||||
assert len(served) == 24
|
||||
for item in served:
|
||||
_assert_embedded(item)
|
||||
_assert_peer_saw(wire.drain(), texts)
|
||||
assert _spend_status_by_call(served) == {item.call_id: "success" for item in served}
|
||||
|
||||
|
||||
async def test_peer_outage_fails_each_call_once_reports_unhealthy_and_recovers_on_restart(
|
||||
proxy: Gateway, peer_port: int
|
||||
) -> None:
|
||||
before: Final = _texts(12)
|
||||
during: Final = _texts(4)
|
||||
after: Final = _texts(4)
|
||||
base_url: Final = str(proxy.client.base_url)
|
||||
with wire_server(respond, port=peer_port) as wire:
|
||||
served: Final = await _burst(base_url, proxy.key, before)
|
||||
assert len(served) == 12
|
||||
for item in served:
|
||||
_assert_embedded(item)
|
||||
_assert_peer_saw(wire.drain(), before)
|
||||
refused: Final = await _burst(base_url, proxy.key, during)
|
||||
assert len(refused) == 4
|
||||
for item in refused:
|
||||
assert item.status == 503, item.body
|
||||
assert "Could not connect to the endpoint URL" in item.body, item.body
|
||||
assert await asyncio.to_thread(eventually, lambda: _unhealthy_count(proxy), lambda count: count == 1, 60)
|
||||
with wire_server(respond, port=peer_port) as restarted:
|
||||
recovered: Final = await _burst(base_url, proxy.key, after)
|
||||
assert len(recovered) == 4
|
||||
for item in recovered:
|
||||
_assert_embedded(item)
|
||||
healthy: Final = [request for request in restarted.drain() if "test from litellm" not in request.body.decode()]
|
||||
_assert_peer_saw(tuple(healthy), after)
|
||||
assert _spend_status_by_call((*served, *refused, *recovered)) == {
|
||||
**{item.call_id: "success" for item in served},
|
||||
**{item.call_id: "failure" for item in refused},
|
||||
**{item.call_id: "success" for item in recovered},
|
||||
}
|
||||
|
||||
|
||||
def _slow(request: Request) -> Reply:
|
||||
body: Final = json_reply(openai_payload(JSON_OBJECT.validate_json(request.body))).body
|
||||
return Reply(chunks=(body[:8], body[8:]), pause_between_chunks=0.5)
|
||||
|
||||
|
||||
async def test_slow_peer_answers_every_caller_whole_without_duplicate_rows(proxy: Gateway, peer_port: int) -> None:
|
||||
texts: Final = _texts(10)
|
||||
with wire_server(_slow, port=peer_port) as wire:
|
||||
served: Final = await _burst(str(proxy.client.base_url), proxy.key, texts)
|
||||
assert len(served) == 10
|
||||
for item in served:
|
||||
_assert_embedded(item)
|
||||
_assert_peer_saw(wire.drain(), texts)
|
||||
assert _spend_status_by_call(served) == {item.call_id: "success" for item in served}
|
||||
|
||||
|
||||
def _live_workers(log: Path) -> tuple[int, ...]:
|
||||
started: Final = (int(pid) for pid in _STARTED_WORKER.findall(log.read_text()))
|
||||
return tuple(pid for pid in started if psutil.pid_exists(pid))
|
||||
|
||||
|
||||
def _open_peer_connections(pid: int, peer_url: str) -> int:
|
||||
port: Final = urlsplit(peer_url).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _HeldPeer:
|
||||
release: threading.Event
|
||||
held: SimpleQueue[str]
|
||||
|
||||
def respond(self, request: Request) -> Reply:
|
||||
self.held.put(string_value(JSON_OBJECT.validate_json(request.body)["input"]))
|
||||
assert self.release.wait(timeout=120), "The burst was never released"
|
||||
return respond(request)
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_component(rig: Gateway, tmp_path: Path) -> None:
|
||||
texts: Final = _texts(8)
|
||||
peer: Final = _HeldPeer(threading.Event(), SimpleQueue())
|
||||
with wire_server(peer.respond) as wire:
|
||||
with owned_proxy_process(rig, tmp_path, _runtime(wire.url), config=_config(tmp_path), workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
workers: Final = eventually(lambda: _live_workers(owned.log), lambda pids: len(pids) == 2, seconds=30)
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(str(candidate.client.base_url), candidate.key, texts, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120)
|
||||
held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == len(texts), held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
with owned.replacing_worker(victim_pid) as kill:
|
||||
kill.freeze()
|
||||
await asyncio.to_thread(kill)
|
||||
peer.release.set()
|
||||
served: Final = await burst
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
for item in served:
|
||||
_assert_embedded(item)
|
||||
(follow_up,) = _texts(1)
|
||||
(answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,))
|
||||
_assert_embedded(answered)
|
||||
_assert_peer_saw(wire.drain(), (*texts, follow_up))
|
||||
assert _spend_status_by_call((*served, answered)) == {
|
||||
item.call_id: "success" for item in (*served, answered)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_proxy_sigterm_mid_burst_finishes_in_flight_calls_and_a_relaunch_serves_again(
|
||||
rig: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
texts: Final = _texts(8)
|
||||
peer: Final = _HeldPeer(threading.Event(), SimpleQueue())
|
||||
with wire_server(peer.respond) as wire:
|
||||
first: Final = tmp_path / "first"
|
||||
first.mkdir()
|
||||
with owned_proxy_process(rig, first, _runtime(wire.url), config=_config(first), workers=2) as owned:
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(str(owned.gateway.client.base_url), owned.gateway.key, texts, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, peer.held.qsize, lambda size: size == len(texts), 120)
|
||||
owned.begin_shutdown()
|
||||
peer.release.set()
|
||||
served: Final = await burst
|
||||
exit_code: Final = await asyncio.to_thread(
|
||||
eventually, owned.process.poll, lambda code: code is not None, 120
|
||||
)
|
||||
assert exit_code is not None
|
||||
for item in served:
|
||||
_assert_embedded(item)
|
||||
_assert_peer_saw(wire.drain(), texts)
|
||||
assert len(served) == len(texts), (len(served), len(texts))
|
||||
second: Final = tmp_path / "second"
|
||||
second.mkdir()
|
||||
with owned_proxy_process(rig, second, _runtime(wire.url), config=_config(second), workers=2) as relaunched:
|
||||
(follow_up,) = _texts(1)
|
||||
(answered,) = await _burst(str(relaunched.gateway.client.base_url), relaunched.gateway.key, (follow_up,))
|
||||
_assert_embedded(answered)
|
||||
_assert_peer_saw(wire.drain(), (follow_up,))
|
||||
assert _spend_status_by_call((*served, answered)) == {
|
||||
item.call_id: "success" for item in (*served, answered)
|
||||
}
|
||||
742
tests/integration/providers/test_sagemaker_embedding_wire.py
Normal file
742
tests/integration/providers/test_sagemaker_embedding_wire.py
Normal file
|
|
@ -0,0 +1,742 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from botocore.exceptions import ConnectionClosedError
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.sagemaker_embedding import (
|
||||
ACCESS_KEY,
|
||||
COMPONENT,
|
||||
COMPONENT_HEADER,
|
||||
EXPECTED_SHAPE,
|
||||
IC_NOT_ALLOWED,
|
||||
IC_REQUIRED,
|
||||
JSON_OBJECT,
|
||||
NON_STRING_MODEL_ID,
|
||||
NOT_JSON_BODY,
|
||||
PASS_IN_HINT,
|
||||
REMOVE_HINT,
|
||||
SERVED_MODEL,
|
||||
deployment,
|
||||
embedding_value,
|
||||
endpoint,
|
||||
invocations_target,
|
||||
openai_deployment,
|
||||
openai_route,
|
||||
plain_route,
|
||||
respond,
|
||||
)
|
||||
from integration._support.wire import Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}
|
||||
_SPEND_ROW: Final = 'SELECT status, call_type, model_group FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s'
|
||||
|
||||
|
||||
def _model_list(health_endpoint: str) -> list[JsonValue]:
|
||||
return [
|
||||
openai_deployment("sm-openai", "openai"),
|
||||
deployment(
|
||||
"sm-openai-explicit",
|
||||
f"openai/{endpoint('openai')}",
|
||||
custom_llm_provider="sagemaker",
|
||||
model_id=COMPONENT,
|
||||
extra_body={"model": SERVED_MODEL},
|
||||
),
|
||||
deployment("sm-openai-bare", openai_route("openai")),
|
||||
deployment("sm-openai-no-model-id", openai_route("openai"), extra_body={"model": SERVED_MODEL}),
|
||||
openai_deployment("sm-openai-custom-key", "openai", input_type="query"),
|
||||
deployment(
|
||||
"sm-openai-pinned",
|
||||
openai_route("openai"),
|
||||
model_id=COMPONENT,
|
||||
extra_body={"model": SERVED_MODEL, "encoding_format": "float"},
|
||||
),
|
||||
openai_deployment("sm-openai-model-id-int", "openai", model_id=7),
|
||||
openai_deployment("sm-openai-model-id-list", "openai", model_id=[COMPONENT]),
|
||||
openai_deployment("sm-openai-model-id-empty", "openai", model_id=""),
|
||||
openai_deployment("sm-openai-model-id-null", "openai", model_id=None),
|
||||
openai_deployment("sm-openai-model-id-long", "openai", model_id="c" * 5120),
|
||||
openai_deployment("sm-openai-extra-body-string", "openai", extra_body="not-an-object"),
|
||||
openai_deployment("sm-openai-extra-body-empty", "openai", extra_body={}),
|
||||
openai_deployment("sm-openai-nousage", "nousage"),
|
||||
openai_deployment("sm-openai-nomodel", "nomodel"),
|
||||
openai_deployment("sm-openai-noindex", "noindex"),
|
||||
openai_deployment("sm-openai-promptonly", "promptonly"),
|
||||
openai_deployment("sm-openai-datadict", "datadict"),
|
||||
openai_deployment("sm-openai-hfshape", "hfshape"),
|
||||
openai_deployment("sm-openai-notjson", "notjson"),
|
||||
openai_deployment("sm-openai-err400", "err400"),
|
||||
openai_deployment("sm-openai-err424", "err424"),
|
||||
openai_deployment("sm-openai-err429", "err429"),
|
||||
openai_deployment("sm-openai-err503", "err503"),
|
||||
openai_deployment("sm-openai-err500", "err500"),
|
||||
openai_deployment("sm-openai-drop", "drop"),
|
||||
deployment("sm-openai-unknown", "sagemaker/openai/unknown-embed-endpoint", model_id=COMPONENT),
|
||||
openai_deployment("sm-openai-ic", "icrequired"),
|
||||
deployment("sm-openai-ic-missing", openai_route("icrequired"), extra_body={"model": SERVED_MODEL}),
|
||||
deployment("sm-plain-noic", plain_route("noic"), model_id=COMPONENT),
|
||||
deployment("sm-plain-noic-control", plain_route("noic")),
|
||||
deployment(
|
||||
"sm-openai-health",
|
||||
f"sagemaker/openai/{health_endpoint}",
|
||||
model_id=COMPONENT,
|
||||
extra_body={"model": SERVED_MODEL},
|
||||
model_info={"mode": "embedding"},
|
||||
),
|
||||
deployment("sm-plain", plain_route("hf"), model_id=COMPONENT),
|
||||
deployment("sm-plain-control", plain_route("hf")),
|
||||
deployment("sm-plain-drop-params", plain_route("hf"), model_id=COMPONENT, drop_params=True),
|
||||
deployment("sm-plain-dict", plain_route("hfdict"), model_id=COMPONENT),
|
||||
deployment("sm-plain-err424", plain_route("err424"), model_id=COMPONENT),
|
||||
deployment("sm-voyage", plain_route("voyage"), model_id=COMPONENT),
|
||||
deployment("sm-cohere", plain_route("cohere"), model_id=COMPONENT),
|
||||
]
|
||||
|
||||
|
||||
def _config(directory: Path, health_endpoint: str) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["model_list"] = _model_list(health_endpoint)
|
||||
path: Final = directory / "sagemaker-embedding.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig() -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as gateway:
|
||||
yield gateway
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wire() -> Iterator[Wire]:
|
||||
with wire_server(respond) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained_wire(wire: Wire) -> None:
|
||||
wire.drain()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def health_endpoint() -> str:
|
||||
return endpoint(f"openai-health-{uuid.uuid4().hex[:8]}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy(
|
||||
rig: Gateway, wire: Wire, health_endpoint: str, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("sagemaker-embedding")
|
||||
overrides: Final = {"AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME": wire.url, "AWS_MAX_ATTEMPTS": "2"}
|
||||
with owned_proxy(rig, directory, overrides, config=_config(directory, health_endpoint), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _text() -> str:
|
||||
return f"integration embedding {uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _embed(proxy: Gateway, model: str, text: JsonValue, **fields: JsonValue) -> httpx.Response:
|
||||
return proxy.request("POST", "/v1/embeddings", {"model": model, "input": text, "cache": _NO_CACHE, **fields})
|
||||
|
||||
|
||||
def _spend_row(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
call_id: Final = response.headers["x-litellm-call-id"]
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _usage(payload: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]:
|
||||
usage: Final = payload["usage"]
|
||||
assert isinstance(usage, dict), payload
|
||||
return tuple(usage[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens"))
|
||||
|
||||
|
||||
def _only_request(wire: Wire) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert len(received) == 1, [(request.method, request.target) for request in received]
|
||||
return received[0]
|
||||
|
||||
|
||||
def _body(request: Request) -> dict[str, JsonValue]:
|
||||
return JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def _forwarded_body(sent: httpx.Request, text: str) -> dict[str, JsonValue]:
|
||||
asked: Final = JSON_OBJECT.validate_json(sent.content)
|
||||
encoding: Final = {"encoding_format": asked["encoding_format"]} if "encoding_format" in asked else {}
|
||||
return {"input": [text], **encoding, "model": SERVED_MODEL}
|
||||
|
||||
|
||||
def _signed_headers(request: Request) -> tuple[str, ...]:
|
||||
authorization: Final = request.headers["authorization"]
|
||||
assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization
|
||||
signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders="))
|
||||
return tuple(signed.removeprefix("SignedHeaders=").split(";"))
|
||||
|
||||
|
||||
def _sdk_base_url(proxy: Gateway) -> str:
|
||||
return f"{str(proxy.client.base_url).rstrip('/')}/v1"
|
||||
|
||||
|
||||
def _sdk(proxy: Gateway) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60)
|
||||
|
||||
|
||||
def _async_sdk(proxy: Gateway) -> openai.AsyncOpenAI:
|
||||
return openai.AsyncOpenAI(base_url=_sdk_base_url(proxy), api_key=proxy.key, max_retries=0, timeout=60)
|
||||
|
||||
|
||||
def test_openai_route_signs_model_id_as_the_inference_component_header_and_sends_the_openai_body(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
texts: Final = [_text(), _text()]
|
||||
response: Final = _embed(proxy, "sm-openai", texts)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]},
|
||||
{"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]},
|
||||
], response.text
|
||||
assert payload["model"] == "sm-openai", response.text
|
||||
assert _usage(payload) == (7, 0, 7), response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert (request.method, request.target) == ("POST", invocations_target(endpoint("openai"))), request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert request.headers["x-amzn-sagemaker-custom-attributes"] == "accept_eula=true", dict(request.headers)
|
||||
assert request.headers["content-type"] == "application/json", dict(request.headers)
|
||||
assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"]
|
||||
assert _body(request) == {"input": texts, "model": SERVED_MODEL}, request.body
|
||||
row: Final = _spend_row(response)
|
||||
assert row == {"status": "success", "call_type": "aembedding", "model_group": "sm-openai"}, row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "model_in_body"),
|
||||
[
|
||||
("/v1/embeddings", True),
|
||||
("/embeddings", True),
|
||||
("/openai/deployments/sm-openai/embeddings", False),
|
||||
("/engines/sm-openai/embeddings", False),
|
||||
],
|
||||
)
|
||||
def test_every_embeddings_alias_reaches_the_inference_component(
|
||||
proxy: Gateway, wire: Wire, path: str, model_in_body: bool
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"input": text,
|
||||
"cache": _NO_CACHE,
|
||||
**({"model": "sm-openai"} if model_in_body else {}),
|
||||
}
|
||||
response: Final = proxy.request("POST", path, body)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": text, "model": SERVED_MODEL}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_openai_sdk_sync_client_reads_the_vectors_in_the_encoding_it_asked_for(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
raw: Final = _sdk(proxy).embeddings.with_raw_response.create(
|
||||
model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE}
|
||||
)
|
||||
created: Final = raw.parse()
|
||||
assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created
|
||||
assert created.model == "sm-openai", created
|
||||
assert (created.usage.prompt_tokens, created.usage.total_tokens) == (7, 7), created
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == _forwarded_body(raw.http_request, text), request.body
|
||||
call_id: Final = raw.headers["x-litellm-call-id"]
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
assert rows[0]["status"] == "success", rows
|
||||
|
||||
|
||||
async def test_openai_sdk_async_client_reads_the_vectors_in_the_encoding_it_asked_for(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
raw: Final = await _async_sdk(proxy).embeddings.with_raw_response.create(
|
||||
model="sm-openai", input=[text], extra_body={"cache": _NO_CACHE}
|
||||
)
|
||||
created: Final = raw.parse()
|
||||
assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created
|
||||
assert created.model == "sm-openai", created
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == _forwarded_body(raw.http_request, text), request.body
|
||||
call_id: Final = raw.headers["x-litellm-call-id"]
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_ROW, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
assert rows[0]["status"] == "success", rows
|
||||
|
||||
|
||||
def test_explicit_custom_llm_provider_spelling_takes_the_openai_route(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-explicit", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["model"] == "sm-openai-explicit", response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == f"/endpoints/{endpoint('openai')}/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_plain_route_signs_model_id_as_the_header_and_keeps_the_hf_body(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
assert payload["model"] == "sm-plain", response.text
|
||||
assert _usage(payload) == (3, 0, 3), response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == f"/endpoints/{endpoint('hf')}/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"]
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_plain_route_without_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain-control", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"]
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_plain_route_reads_the_wrapped_hf_embedding_key(proxy: Gateway, wire: Wire) -> None:
|
||||
texts: Final = [_text(), _text()]
|
||||
response: Final = _embed(proxy, "sm-plain-dict", texts)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]},
|
||||
{"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]},
|
||||
], response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"inputs": texts}, request.body
|
||||
|
||||
|
||||
def test_voyage_endpoint_keeps_the_voyage_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-voyage", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
assert payload["model"] == "sm-voyage", response.text
|
||||
assert _usage(payload) == (9, 0, 9), response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == f"/endpoints/{endpoint('voyage')}/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": endpoint("voyage")}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_cohere_endpoint_keeps_the_cohere_body_and_signs_model_id_as_the_header(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-cohere", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
assert _usage(payload) == (6, 0, 6), response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == f"/endpoints/{endpoint('cohere')}/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"texts": [text], "input_type": "search_document"}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_openai_route_forwards_dimensions_encoding_format_and_user_in_the_body(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(
|
||||
proxy, "sm-openai", [text], dimensions=3, encoding_format="float", user="integration-embedding-user"
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {
|
||||
"input": [text],
|
||||
"dimensions": 3,
|
||||
"encoding_format": "float",
|
||||
"user": "integration-embedding-user",
|
||||
"model": SERVED_MODEL,
|
||||
}, request.body
|
||||
|
||||
|
||||
def test_openai_route_without_model_id_or_extra_body_sends_the_bare_openai_body(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-bare", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert _body(request) == {"input": [text]}, request.body
|
||||
|
||||
|
||||
def test_openai_route_sends_a_custom_deployment_key_at_the_top_level_of_the_body(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-custom-key", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "input_type": "query", "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
def test_deployment_extra_body_wins_over_the_callers_encoding_format(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
created: Final = _sdk(proxy).embeddings.create(
|
||||
model="sm-openai-pinned", input=[text], extra_body={"cache": _NO_CACHE}
|
||||
)
|
||||
assert [row.embedding for row in created.data] == [[0.25, -0.5, 1.0]], created
|
||||
request: Final = _only_request(wire)
|
||||
assert _body(request) == {"input": [text], "encoding_format": "float", "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
def test_openai_route_returns_the_base64_string_a_raw_caller_asked_for_unchanged(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai", [text], encoding_format="base64")
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": embedding_value(0, "base64")}], (
|
||||
response.text
|
||||
)
|
||||
request: Final = _only_request(wire)
|
||||
assert _body(request) == {"input": [text], "encoding_format": "base64", "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
def test_openai_route_reports_zero_usage_when_the_container_sends_none(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-nousage", [_text()])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert _usage(payload) == (0, 0, 0), response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_openai_route_names_the_endpoint_when_the_container_sends_no_model(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-nomodel", [_text()])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["model"] == "sm-openai-nomodel", response.text
|
||||
assert _usage(payload) == (7, 0, 7), response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
|
||||
|
||||
def test_openai_route_numbers_rows_the_container_left_unindexed(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-noindex", [_text(), _text()])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]},
|
||||
{"object": "embedding", "index": 1, "embedding": [0.25, -0.5, 2.0]},
|
||||
], response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
|
||||
|
||||
def test_openai_route_takes_total_tokens_from_prompt_tokens_when_the_container_omits_it(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-promptonly", [_text()])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert _usage(payload) == (5, 0, 5), response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["sm-openai-datadict", "sm-openai-hfshape"])
|
||||
def test_openai_route_rejects_a_container_body_that_is_not_an_openai_embeddings_object(
|
||||
proxy: Gateway, wire: Wire, model: str
|
||||
) -> None:
|
||||
response: Final = _embed(proxy, model, [_text()])
|
||||
assert response.status_code == 503, response.text
|
||||
assert EXPECTED_SHAPE in response.text, response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_openai_route_reports_a_non_json_container_body_as_a_server_error(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-notjson", [_text()])
|
||||
with pytest.raises(json.JSONDecodeError) as decode_error:
|
||||
json.loads(NOT_JSON_BODY)
|
||||
assert response.status_code == 500, response.text
|
||||
assert str(decode_error.value) in response.text, response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_container_validation_error_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-err400", [_text()])
|
||||
assert response.status_code == 400, response.text
|
||||
assert "1 validation error detected" in response.text, response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_container_model_error_reaches_the_caller_with_the_original_message(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-err424", [_text()])
|
||||
assert response.status_code == 424, response.text
|
||||
assert "unexpected field inputs" in response.text, response.text
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_plain_route_body_rejected_by_an_openai_container_reaches_the_caller(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain-err424", [text])
|
||||
assert response.status_code == 424, response.text
|
||||
assert "unexpected field inputs" in response.text, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "status", "message"),
|
||||
[
|
||||
("sm-openai-err429", 429, "Model is not ready for inference yet"),
|
||||
("sm-openai-err503", 503, "The endpoint is scaling, try again later"),
|
||||
("sm-openai-err500", 503, "An internal failure occurred"),
|
||||
],
|
||||
)
|
||||
def test_container_outage_errors_reach_the_caller_and_every_attempt_carries_the_header(
|
||||
proxy: Gateway, wire: Wire, model: str, status: int, message: str
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, model, [text])
|
||||
assert response.status_code == status, response.text
|
||||
assert message in response.text, response.text
|
||||
attempts: Final = wire.drain()
|
||||
assert len(attempts) >= 1, attempts
|
||||
assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts
|
||||
assert {attempt.body for attempt in attempts} == {json.dumps({"input": [text], "model": SERVED_MODEL}).encode()}
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_dropped_connection_reaches_the_caller_as_a_server_error(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-drop", [_text()])
|
||||
assert response.status_code == 503, response.text
|
||||
assert ConnectionClosedError.fmt.partition(":")[0] in response.text, response.text
|
||||
attempts: Final = wire.drain()
|
||||
assert len(attempts) >= 1, attempts
|
||||
assert {attempt.headers[COMPONENT_HEADER] for attempt in attempts} == {COMPONENT}, attempts
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_inference_component_endpoint_answers_once_the_header_unlocks_it(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-ic", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, -0.5, 1.0]}], response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_inference_component_endpoint_without_model_id_tells_the_caller_how_to_pass_it(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-ic-missing", [text])
|
||||
assert response.status_code == 400, response.text
|
||||
assert IC_REQUIRED in response.text, response.text
|
||||
assert PASS_IN_HINT in response.text, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_endpoint_without_inference_components_tells_the_deployment_to_drop_model_id(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain-noic", [text])
|
||||
assert response.status_code == 400, response.text
|
||||
assert IC_NOT_ALLOWED in response.text, response.text
|
||||
assert REMOVE_HINT in response.text, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_endpoint_without_inference_components_serves_a_deployment_without_model_id(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain-noic-control", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_unknown_endpoint_reaches_the_caller_as_400(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-unknown", [_text()])
|
||||
assert response.status_code == 400, response.text
|
||||
assert "Endpoint unknown-embed-endpoint of account 000000000000 not found." in response.text, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == "/endpoints/unknown-embed-endpoint/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_plain_route_rejects_an_unsupported_openai_param_without_drop_params(proxy: Gateway, wire: Wire) -> None:
|
||||
response: Final = _embed(proxy, "sm-plain", [_text()], dimensions=3)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "does not support parameters: {'dimensions': 3}" in response.text, response.text
|
||||
assert "drop_params" in response.text, response.text
|
||||
assert wire.drain() == (), response.text
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
def test_plain_route_drops_an_unsupported_openai_param_under_drop_params(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-plain-drop-params", [text], dimensions=3)
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"inputs": [text]}, request.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["sm-openai-model-id-int", "sm-openai-model-id-list"])
|
||||
def test_a_non_string_model_id_fails_before_any_request_leaves_the_proxy(
|
||||
proxy: Gateway, wire: Wire, model: str
|
||||
) -> None:
|
||||
response: Final = _embed(proxy, model, [_text()])
|
||||
assert response.status_code == 400, response.text
|
||||
assert NON_STRING_MODEL_ID in response.text, response.text
|
||||
assert wire.drain() == (), response.text
|
||||
assert _spend_row(response)["status"] == "failure"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["sm-openai-model-id-empty", "sm-openai-model-id-null"])
|
||||
def test_an_empty_or_null_model_id_sends_no_inference_component_header(proxy: Gateway, wire: Wire, model: str) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, model, [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
def test_a_long_model_id_is_signed_and_forwarded_whole(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-model-id-long", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == "c" * 5120, len(request.headers[COMPONENT_HEADER])
|
||||
assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"]
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["sm-openai-extra-body-string", "sm-openai-extra-body-empty"])
|
||||
def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body(
|
||||
proxy: Gateway, wire: Wire, model: str
|
||||
) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, model, [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text]}, request.body
|
||||
|
||||
|
||||
def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai", [text], model_id="caller-chosen-component")
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
def test_a_caller_sent_model_id_is_forwarded_when_the_deployment_has_none(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai-bare", [text], model_id="caller-chosen-component")
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers)
|
||||
assert _body(request) == {"input": [text]}, request.body
|
||||
|
||||
|
||||
def test_a_caller_sent_non_string_model_id_fails_before_any_request_leaves_the_proxy(
|
||||
proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
response: Final = _embed(proxy, "sm-openai-bare", [_text()], model_id=["caller-chosen-component"])
|
||||
assert response.status_code == 400, response.text
|
||||
assert NON_STRING_MODEL_ID in response.text, response.text
|
||||
assert wire.drain() == (), response.text
|
||||
|
||||
|
||||
def test_a_caller_sent_extra_body_replaces_the_deployments_extra_body(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai", [text], extra_body={"truncate": "END"})
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "truncate": "END"}, request.body
|
||||
|
||||
|
||||
def test_a_failed_call_does_not_poison_the_next_one(proxy: Gateway, wire: Wire) -> None:
|
||||
failed: Final = _embed(proxy, "sm-openai-err503", [_text()])
|
||||
assert failed.status_code == 503, failed.text
|
||||
wire.drain()
|
||||
text: Final = _text()
|
||||
response: Final = _embed(proxy, "sm-openai", [text])
|
||||
assert response.status_code == 200, response.text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
assert _spend_row(response)["status"] == "success"
|
||||
|
||||
|
||||
def test_two_identical_uncached_requests_each_reach_the_container_and_log_once(proxy: Gateway, wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
first: Final = _embed(proxy, "sm-openai", [text])
|
||||
second: Final = _embed(proxy, "sm-openai", [text])
|
||||
assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text)
|
||||
assert first.headers["x-litellm-call-id"] != second.headers["x-litellm-call-id"], (first.headers, second.headers)
|
||||
requests: Final = wire.drain()
|
||||
assert [_body(request) for request in requests] == [{"input": [text], "model": SERVED_MODEL}] * 2, requests
|
||||
assert {request.headers[COMPONENT_HEADER] for request in requests} == {COMPONENT}, requests
|
||||
assert (_spend_row(first)["status"], _spend_row(second)["status"]) == ("success", "success")
|
||||
|
||||
|
||||
def test_health_check_reaches_the_container_with_the_inference_component_header(
|
||||
proxy: Gateway, wire: Wire, health_endpoint: str
|
||||
) -> None:
|
||||
report: Final = proxy.get("/health", {"model": "sm-openai-health"})
|
||||
assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == f"/endpoints/{health_endpoint}/invocations", request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
body: Final = _body(request)
|
||||
assert body["model"] == SERVED_MODEL, request.body
|
||||
assert isinstance(body["input"], list) and len(body["input"]) == 1, request.body
|
||||
171
tests/integration/sdk/test_sagemaker_embedding_sdk.py
Normal file
171
tests/integration/sdk/test_sagemaker_embedding_sdk.py
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.sagemaker_embedding import (
|
||||
ACCESS_KEY,
|
||||
COMPONENT,
|
||||
COMPONENT_HEADER,
|
||||
JSON_OBJECT,
|
||||
NON_STRING_MODEL_ID,
|
||||
REGION,
|
||||
SECRET_KEY,
|
||||
SERVED_MODEL,
|
||||
endpoint,
|
||||
invocations_target,
|
||||
litellm_params,
|
||||
openai_route,
|
||||
plain_route,
|
||||
respond,
|
||||
)
|
||||
from integration._support.wire import Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
|
||||
_AWS: Final[dict[str, str]] = {
|
||||
"aws_access_key_id": ACCESS_KEY,
|
||||
"aws_secret_access_key": SECRET_KEY,
|
||||
"aws_region_name": REGION,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wire() -> Iterator[Wire]:
|
||||
with wire_server(respond) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _runtime_at_the_peer(wire: Wire, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
wire.drain()
|
||||
monkeypatch.setenv("AWS_ENDPOINT_URL_SAGEMAKER_RUNTIME", wire.url)
|
||||
monkeypatch.setenv("AWS_MAX_ATTEMPTS", "1")
|
||||
|
||||
|
||||
def _text() -> str:
|
||||
return f"integration embedding {uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _only_request(wire: Wire) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert len(received) == 1, [(request.method, request.target) for request in received]
|
||||
return received[0]
|
||||
|
||||
|
||||
def _signed_headers(request: Request) -> tuple[str, ...]:
|
||||
authorization: Final = request.headers["authorization"]
|
||||
assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), authorization
|
||||
signed: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders="))
|
||||
return tuple(signed.removeprefix("SignedHeaders=").split(";"))
|
||||
|
||||
|
||||
def _assert_component_request(request: Request, scenario: str, body: dict[str, JsonValue]) -> None:
|
||||
assert (request.method, request.target) == ("POST", invocations_target(endpoint(scenario))), request
|
||||
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert COMPONENT_HEADER in _signed_headers(request), request.headers["authorization"]
|
||||
assert JSON_OBJECT.validate_json(request.body) == body, request.body
|
||||
|
||||
|
||||
def _vectors(response: litellm.EmbeddingResponse) -> list[JsonValue]:
|
||||
return [row["embedding"] for row in response.data]
|
||||
|
||||
|
||||
def test_embedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = litellm.embedding(
|
||||
model=openai_route("openai"), input=[text], model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS
|
||||
)
|
||||
assert response.model == SERVED_MODEL, response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (7, 7), response.usage
|
||||
_assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL})
|
||||
|
||||
|
||||
async def test_aembedding_signs_model_id_as_the_inference_component_and_reads_the_served_model(wire: Wire) -> None:
|
||||
texts: Final = [_text(), _text()]
|
||||
response: Final = await litellm.aembedding(
|
||||
model=openai_route("openai"), input=texts, model_id=COMPONENT, extra_body={"model": SERVED_MODEL}, **_AWS
|
||||
)
|
||||
assert response.model == SERVED_MODEL, response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0], [0.25, -0.5, 2.0]], response
|
||||
_assert_component_request(_only_request(wire), "openai", {"input": texts, "model": SERVED_MODEL})
|
||||
|
||||
|
||||
def test_embedding_names_the_endpoint_when_the_container_reports_no_model(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = litellm.embedding(model=openai_route("nomodel"), input=[text], model_id=COMPONENT, **_AWS)
|
||||
assert response.model == endpoint("nomodel"), response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
_assert_component_request(_only_request(wire), "nomodel", {"input": [text]})
|
||||
|
||||
|
||||
def test_plain_route_keeps_the_hf_body_and_signs_the_component_header(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = litellm.embedding(model=plain_route("hf"), input=[text], model_id=COMPONENT, **_AWS)
|
||||
assert response.model == endpoint("hf"), response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
assert response.usage.prompt_tokens == 3, response.usage
|
||||
_assert_component_request(_only_request(wire), "hf", {"inputs": [text]})
|
||||
|
||||
|
||||
def test_plain_route_without_model_id_sends_no_component_header(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = litellm.embedding(model=plain_route("hf"), input=[text], **_AWS)
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
request: Final = _only_request(wire)
|
||||
assert (request.method, request.target) == ("POST", invocations_target(endpoint("hf"))), request
|
||||
assert COMPONENT_HEADER not in request.headers, dict(request.headers)
|
||||
assert COMPONENT_HEADER not in _signed_headers(request), request.headers["authorization"]
|
||||
assert JSON_OBJECT.validate_json(request.body) == {"inputs": [text]}, request.body
|
||||
|
||||
|
||||
def test_explicit_provider_spelling_signs_model_id_as_the_inference_component(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
response: Final = litellm.embedding(
|
||||
model=f"openai/{endpoint('openai')}",
|
||||
custom_llm_provider="sagemaker",
|
||||
input=[text],
|
||||
model_id=COMPONENT,
|
||||
extra_body={"model": SERVED_MODEL},
|
||||
**_AWS,
|
||||
)
|
||||
assert response.model == SERVED_MODEL, response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
_assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL})
|
||||
|
||||
|
||||
async def test_router_aembedding_forwards_the_deployment_model_id_as_the_component(wire: Wire) -> None:
|
||||
text: Final = _text()
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "sm-openai-router",
|
||||
"litellm_params": litellm_params(
|
||||
openai_route("openai"), model_id=COMPONENT, extra_body={"model": SERVED_MODEL}
|
||||
),
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
response: Final = await router.aembedding(model="sm-openai-router", input=[text])
|
||||
assert response.model == SERVED_MODEL, response
|
||||
assert _vectors(response) == [[0.25, -0.5, 1.0]], response
|
||||
_assert_component_request(_only_request(wire), "openai", {"input": [text], "model": SERVED_MODEL})
|
||||
|
||||
|
||||
def test_embedding_rejects_a_non_string_model_id_before_calling_aws(wire: Wire) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as caught:
|
||||
litellm.embedding(model=openai_route("openai"), input=[_text()], model_id=[COMPONENT], **_AWS)
|
||||
assert caught.value.status_code == 400, caught.value
|
||||
assert NON_STRING_MODEL_ID in str(caught.value), caught.value
|
||||
assert wire.drain() == (), "AWS was called with a list-valued model_id"
|
||||
|
||||
|
||||
def test_embedding_surfaces_the_container_model_error_with_its_status(wire: Wire) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as caught:
|
||||
litellm.embedding(model=openai_route("err424"), input=[_text()], model_id=COMPONENT, **_AWS)
|
||||
assert caught.value.status_code == 424, caught.value
|
||||
assert "unexpected field inputs" in str(caught.value), caught.value
|
||||
assert _only_request(wire).headers[COMPONENT_HEADER] == COMPONENT
|
||||
|
|
@ -1548,3 +1548,23 @@ async def test_redis_caching_ttl_sadd(
|
|||
)
|
||||
|
||||
redis_client.expire.assert_awaited_once_with("cache-key", timedelta(seconds=120))
|
||||
|
||||
|
||||
def test_cache_key_separates_model_ids_unless_a_model_group_owns_the_key() -> None:
|
||||
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
variants: Final[tuple[dict[str, object], ...]] = (
|
||||
{"model_id": "component-a"},
|
||||
{"model_id": "component-b"},
|
||||
{},
|
||||
{"model_id": "component-a"},
|
||||
{"model_id": "component-a", "metadata": {"model_group": "embed"}},
|
||||
{"model_id": "component-b", "metadata": {"model_group": "embed"}},
|
||||
)
|
||||
key_a, key_b, bare, key_a_again, grouped_a, grouped_b = (
|
||||
cache.get_cache_key(model="sagemaker/openai/embed", input=["hello"], **extra) for extra in variants
|
||||
)
|
||||
|
||||
assert key_a != key_b
|
||||
assert key_a != bare
|
||||
assert key_a == key_a_again
|
||||
assert grouped_a == grouped_b
|
||||
|
|
|
|||
|
|
@ -4,7 +4,9 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder
|
||||
from botocore.exceptions import ClientError, ConnectionClosedError
|
||||
|
||||
from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder, client_error_details
|
||||
from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
|
||||
|
||||
|
||||
|
|
@ -32,9 +34,7 @@ def test_sagemaker_response_stream_shape_lazy_loads_once():
|
|||
import litellm.llms.sagemaker.common_utils as mod
|
||||
|
||||
sentinel = MagicMock()
|
||||
with patch.object(
|
||||
mod, "_load_sagemaker_response_stream_shape", return_value=sentinel
|
||||
) as mock_load:
|
||||
with patch.object(mod, "_load_sagemaker_response_stream_shape", return_value=sentinel) as mock_load:
|
||||
assert mod.get_sagemaker_response_stream_shape() is sentinel
|
||||
assert mod.get_sagemaker_response_stream_shape() is sentinel
|
||||
mock_load.assert_called_once()
|
||||
|
|
@ -94,9 +94,7 @@ def test_sagemaker_response_stream_shape_is_structure_shape():
|
|||
from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape
|
||||
|
||||
shape = get_sagemaker_response_stream_shape()
|
||||
assert (
|
||||
shape is not None
|
||||
), "get_sagemaker_response_stream_shape() is None — botocore may not be installed"
|
||||
assert shape is not None, "get_sagemaker_response_stream_shape() is None — botocore may not be installed"
|
||||
assert isinstance(shape, StructureShape)
|
||||
assert shape.name == "InvokeEndpointWithResponseStreamOutput"
|
||||
|
||||
|
|
@ -274,3 +272,38 @@ class TestSagemakerTransform:
|
|||
|
||||
# The function should properly map max_tokens if max_completion_tokens is not provided
|
||||
assert result == {"temperature": 0.7, "max_new_tokens": 200}
|
||||
|
||||
|
||||
NULL_MESSAGE_ERROR = ClientError(
|
||||
{"Error": {"Code": "InternalFailure", "Message": None}, "ResponseMetadata": {"HTTPStatusCode": 500}},
|
||||
"InvokeEndpoint",
|
||||
)
|
||||
DROPPED_CONNECTION_ERROR = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x")
|
||||
HTTPX_STATUS_ERROR = httpx.HTTPStatusError(
|
||||
"boom", request=httpx.Request("POST", "https://runtime.sagemaker.test/x"), response=httpx.Response(400)
|
||||
)
|
||||
PLAIN_ERROR = RuntimeError("plain failure")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "expected_status", "expected_message"),
|
||||
[
|
||||
(
|
||||
ClientError(
|
||||
{
|
||||
"Error": {"Code": "ValidationError", "Message": "bad input"},
|
||||
"ResponseMetadata": {"HTTPStatusCode": 400},
|
||||
},
|
||||
"InvokeEndpoint",
|
||||
),
|
||||
400,
|
||||
"bad input",
|
||||
),
|
||||
(NULL_MESSAGE_ERROR, 500, str(NULL_MESSAGE_ERROR)),
|
||||
(DROPPED_CONNECTION_ERROR, 500, str(DROPPED_CONNECTION_ERROR)),
|
||||
(HTTPX_STATUS_ERROR, 500, str(HTTPX_STATUS_ERROR)),
|
||||
(PLAIN_ERROR, 500, str(PLAIN_ERROR)),
|
||||
],
|
||||
)
|
||||
def test_client_error_details_reads_every_error_shape(error, expected_status, expected_message):
|
||||
assert client_error_details(error) == (expected_status, expected_message)
|
||||
|
|
|
|||
378
tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py
Normal file
378
tests/unit/llms/sagemaker/test_sagemaker_embedding_openai.py
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
"""
|
||||
Test cases for OpenAI-compatible SageMaker embedding endpoints and inference components
|
||||
|
||||
Covers the `sagemaker/openai/<endpoint>` route (request and response in the OpenAI
|
||||
embeddings shape) and `model_id` forwarded to `invoke_endpoint` as `InferenceComponentName`
|
||||
for every SageMaker embedding config.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError, ConnectionClosedError
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.llms.sagemaker.common_utils import SagemakerError
|
||||
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
|
||||
from litellm.llms.sagemaker.embedding.openai_transformation import (
|
||||
SagemakerOpenAIEmbeddingConfig,
|
||||
)
|
||||
from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig
|
||||
from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
OPENAI_RESPONSE = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]},
|
||||
{"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]},
|
||||
],
|
||||
"model": "served-embedding-model",
|
||||
"usage": {"prompt_tokens": 11, "total_tokens": 11},
|
||||
}
|
||||
|
||||
|
||||
def fake_sagemaker_session(response_body: object) -> tuple[MagicMock, MagicMock]:
|
||||
client = MagicMock()
|
||||
client.invoke_endpoint.return_value = {
|
||||
"Body": MagicMock(read=MagicMock(return_value=json.dumps(response_body).encode()))
|
||||
}
|
||||
session = MagicMock()
|
||||
session.client.return_value = client
|
||||
return session, client
|
||||
|
||||
|
||||
def invoke_kwargs(client: MagicMock) -> Mapping[str, object]:
|
||||
client.invoke_endpoint.assert_called_once()
|
||||
return client.invoke_endpoint.call_args.kwargs
|
||||
|
||||
|
||||
class TestSagemakerEmbeddingFactoryOpenAIRoute:
|
||||
def test_openai_route_selects_openai_config(self):
|
||||
config = SagemakerEmbeddingConfig.get_model_config("openai/my-embed-endpoint")
|
||||
|
||||
assert isinstance(config, SagemakerOpenAIEmbeddingConfig)
|
||||
assert SagemakerEmbeddingConfig().get_supported_openai_params("openai/my-embed-endpoint") == [
|
||||
"dimensions",
|
||||
"encoding_format",
|
||||
"user",
|
||||
]
|
||||
|
||||
def test_openai_route_wins_over_model_name_substrings(self):
|
||||
assert isinstance(
|
||||
SagemakerEmbeddingConfig.get_model_config("openai/voyage-endpoint"),
|
||||
SagemakerOpenAIEmbeddingConfig,
|
||||
)
|
||||
assert isinstance(
|
||||
SagemakerEmbeddingConfig.get_model_config("openai/cohere-endpoint"),
|
||||
SagemakerOpenAIEmbeddingConfig,
|
||||
)
|
||||
|
||||
def test_plain_endpoint_names_keep_their_config(self):
|
||||
assert type(SagemakerEmbeddingConfig.get_model_config("my-embed-endpoint")) is SagemakerEmbeddingConfig
|
||||
assert isinstance(SagemakerEmbeddingConfig.get_model_config("voyage-endpoint"), VoyageEmbeddingConfig)
|
||||
|
||||
def test_endpoint_name_strips_only_the_route(self):
|
||||
assert SagemakerEmbeddingConfig.get_endpoint_name("openai/my-embed-endpoint") == "my-embed-endpoint"
|
||||
assert SagemakerEmbeddingConfig.get_endpoint_name("my-embed-endpoint") == "my-embed-endpoint"
|
||||
assert SagemakerEmbeddingConfig.get_endpoint_name("my-openai-endpoint") == "my-openai-endpoint"
|
||||
|
||||
|
||||
class TestSagemakerOpenAIEmbeddingConfig:
|
||||
def setup_method(self):
|
||||
self.config = SagemakerOpenAIEmbeddingConfig()
|
||||
|
||||
def test_request_is_openai_shaped_with_extra_body_flattened(self):
|
||||
result = self.config.transform_embedding_request(
|
||||
model="my-embed-endpoint",
|
||||
input=["hello", "world"],
|
||||
optional_params={"truncate": True, "extra_body": {"model": "served-embedding-model"}},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result == {"input": ["hello", "world"], "truncate": True, "model": "served-embedding-model"}
|
||||
|
||||
def test_request_without_extra_body_sends_input_only(self):
|
||||
result = self.config.transform_embedding_request(
|
||||
model="my-embed-endpoint", input=["hello"], optional_params={}, headers={}
|
||||
)
|
||||
|
||||
assert result == {"input": ["hello"]}
|
||||
|
||||
def test_openai_params_pass_through(self):
|
||||
mapped = self.config.map_openai_params(
|
||||
non_default_params={"dimensions": 256, "encoding_format": "float", "user": "u1", "stream": True},
|
||||
optional_params={},
|
||||
model="my-embed-endpoint",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert mapped == {"dimensions": 256, "encoding_format": "float", "user": "u1"}
|
||||
|
||||
def test_response_reads_data_rows_and_usage(self):
|
||||
raw = httpx.Response(200, json=OPENAI_RESPONSE)
|
||||
|
||||
result = self.config.transform_embedding_response(
|
||||
model="my-embed-endpoint",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert result.object == "list"
|
||||
assert result.data == [
|
||||
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]},
|
||||
{"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]},
|
||||
]
|
||||
assert result.model == "served-embedding-model"
|
||||
assert result.usage is not None
|
||||
assert (result.usage.prompt_tokens, result.usage.total_tokens) == (11, 11)
|
||||
|
||||
def test_response_without_usage_or_model_falls_back(self):
|
||||
raw = httpx.Response(200, json={"data": [{"embedding": [0.9]}]})
|
||||
|
||||
result = self.config.transform_embedding_response(
|
||||
model="my-embed-endpoint",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert result.data == [{"object": "embedding", "index": 0, "embedding": [0.9]}]
|
||||
assert result.model == "my-embed-endpoint"
|
||||
assert result.usage is not None
|
||||
assert (result.usage.prompt_tokens, result.usage.total_tokens) == (0, 0)
|
||||
|
||||
def test_response_keeps_base64_embedding_strings(self):
|
||||
raw = httpx.Response(
|
||||
200, json={"data": [{"index": 0, "embedding": "AACAPwAAAEA="}], "usage": {"prompt_tokens": 2}}
|
||||
)
|
||||
|
||||
result = self.config.transform_embedding_response(
|
||||
model="my-embed-endpoint",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert result.data == [{"object": "embedding", "index": 0, "embedding": "AACAPwAAAEA="}]
|
||||
assert result.usage is not None
|
||||
assert (result.usage.prompt_tokens, result.usage.total_tokens) == (2, 2)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
{"embedding": [[0.1, 0.2]]},
|
||||
[[0.1, 0.2]],
|
||||
{"data": [{"index": 0}]},
|
||||
{"data": "not-a-list"},
|
||||
],
|
||||
)
|
||||
def test_response_rejects_non_openai_shapes(self, body):
|
||||
raw = httpx.Response(200, json=body)
|
||||
|
||||
with pytest.raises(SagemakerError) as exc_info:
|
||||
self.config.transform_embedding_response(
|
||||
model="my-embed-endpoint",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "'data'" in exc_info.value.message
|
||||
|
||||
|
||||
class TestSagemakerEmbeddingInferenceComponents:
|
||||
def setup_method(self):
|
||||
self.sagemaker_llm = SagemakerLLM()
|
||||
|
||||
def embed_with(self, session: MagicMock, model: str, optional_params: Mapping[str, object]) -> EmbeddingResponse:
|
||||
with patch("boto3.Session", return_value=session):
|
||||
return self.sagemaker_llm.embedding(
|
||||
model=model,
|
||||
input=["hello", "good morning"],
|
||||
model_response=EmbeddingResponse(),
|
||||
print_verbose=print,
|
||||
encoding=None,
|
||||
logging_obj=MagicMock(),
|
||||
optional_params={"aws_region_name": "us-east-1", **optional_params},
|
||||
)
|
||||
|
||||
def run_embedding(
|
||||
self, model: str, optional_params: Mapping[str, object], response_body: object
|
||||
) -> tuple[EmbeddingResponse, Mapping[str, object]]:
|
||||
session, client = fake_sagemaker_session(response_body)
|
||||
result = self.embed_with(session, model, optional_params)
|
||||
return result, invoke_kwargs(client)
|
||||
|
||||
def test_openai_route_forwards_component_and_keeps_it_out_of_the_body(self):
|
||||
result, kwargs = self.run_embedding(
|
||||
model="openai/my-embed-endpoint",
|
||||
optional_params={"model_id": "my-component", "extra_body": {"model": "served-embedding-model"}},
|
||||
response_body=OPENAI_RESPONSE,
|
||||
)
|
||||
|
||||
assert kwargs["EndpointName"] == "my-embed-endpoint"
|
||||
assert kwargs["InferenceComponentName"] == "my-component"
|
||||
assert json.loads(kwargs["Body"]) == {"input": ["hello", "good morning"], "model": "served-embedding-model"}
|
||||
assert [row["embedding"] for row in result.data] == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
assert result.usage is not None
|
||||
assert result.usage.prompt_tokens == 11
|
||||
|
||||
def test_hf_endpoint_forwards_component_and_keeps_its_body_shape(self):
|
||||
result, kwargs = self.run_embedding(
|
||||
model="my-embed-endpoint",
|
||||
optional_params={"model_id": "my-component", "normalize": True},
|
||||
response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]},
|
||||
)
|
||||
|
||||
assert kwargs["EndpointName"] == "my-embed-endpoint"
|
||||
assert kwargs["InferenceComponentName"] == "my-component"
|
||||
assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"], "normalize": True}
|
||||
assert [row["embedding"] for row in result.data] == [[0.1, 0.2], [0.3, 0.4]]
|
||||
|
||||
def test_no_component_without_model_id(self):
|
||||
_, kwargs = self.run_embedding(
|
||||
model="my-embed-endpoint",
|
||||
optional_params={},
|
||||
response_body={"embedding": [[0.1, 0.2], [0.3, 0.4]]},
|
||||
)
|
||||
|
||||
assert "InferenceComponentName" not in kwargs
|
||||
assert json.loads(kwargs["Body"]) == {"inputs": ["hello", "good morning"]}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sagemaker_message", "expected_hint"),
|
||||
[
|
||||
(
|
||||
"Inference Component Name header is not allowed for endpoints to which you dont plan to deploy"
|
||||
" inference components. Please remove the Inference Component Name header and try again.",
|
||||
"remove `model_id` from this deployment, the endpoint has no inference components",
|
||||
),
|
||||
(
|
||||
"Inference Component Name header is required for this endpoint.",
|
||||
"pass in via `litellm.embedding(..., model_id={InferenceComponentName})`",
|
||||
),
|
||||
('Received server error (500) from model with message "boom".', ""),
|
||||
],
|
||||
)
|
||||
def test_component_errors_keep_status_and_name_model_id(self, sagemaker_message: str, expected_hint: str):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
client.invoke_endpoint.side_effect = ClientError(
|
||||
{
|
||||
"Error": {"Code": "ValidationError", "Message": sagemaker_message},
|
||||
"ResponseMetadata": {"HTTPStatusCode": 400},
|
||||
},
|
||||
"InvokeEndpoint",
|
||||
)
|
||||
with pytest.raises(SagemakerError) as raised:
|
||||
self.embed_with(session, "my-embed-endpoint", {"model_id": "stray-component"})
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert raised.value.message == (
|
||||
sagemaker_message + "\n " + expected_hint if expected_hint else sagemaker_message
|
||||
)
|
||||
|
||||
def test_null_error_message_surfaces_the_error_code(self):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
error = ClientError(
|
||||
{
|
||||
"Error": {"Code": "InternalFailure", "Message": None},
|
||||
"ResponseMetadata": {"HTTPStatusCode": 500},
|
||||
},
|
||||
"InvokeEndpoint",
|
||||
)
|
||||
client.invoke_endpoint.side_effect = error
|
||||
with pytest.raises(SagemakerError) as raised:
|
||||
self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "x" * 2048})
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert raised.value.message == str(error)
|
||||
|
||||
def test_dropped_connection_surfaces_as_a_server_error(self):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
error = ConnectionClosedError(endpoint_url="https://runtime.sagemaker.test/x")
|
||||
client.invoke_endpoint.side_effect = error
|
||||
with pytest.raises(SagemakerError) as raised:
|
||||
self.embed_with(session, "openai/my-embed-endpoint", {"model_id": "my-component"})
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert raised.value.message == str(error)
|
||||
|
||||
def test_non_string_model_id_is_rejected_before_the_call(self):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
with pytest.raises(SagemakerError) as raised:
|
||||
self.embed_with(session, "openai/my-embed-endpoint", {"model_id": ["my-component"]})
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert raised.value.message == "model_id must be the inference component name as a string"
|
||||
client.invoke_endpoint.assert_not_called()
|
||||
|
||||
|
||||
SDK_KWARGS = {
|
||||
"model": "sagemaker/openai/my-embed-endpoint",
|
||||
"input": ["hello", "good morning"],
|
||||
"model_id": "my-component",
|
||||
"extra_body": {"model": "served-embedding-model"},
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "test-key",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
}
|
||||
|
||||
EXPECTED_BODY = {"input": ["hello", "good morning"], "model": "served-embedding-model"}
|
||||
|
||||
|
||||
class TestSagemakerOpenAIEmbeddingSdkCalls:
|
||||
def test_sdk_sync(self):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
with patch("boto3.Session", return_value=session):
|
||||
result = litellm.embedding(**SDK_KWARGS)
|
||||
|
||||
kwargs = invoke_kwargs(client)
|
||||
assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component")
|
||||
assert json.loads(kwargs["Body"]) == EXPECTED_BODY
|
||||
assert len(result.data) == 2
|
||||
assert result.usage is not None
|
||||
assert result.usage.total_tokens == 11
|
||||
|
||||
def test_sdk_async(self):
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
with patch("boto3.Session", return_value=session):
|
||||
result = asyncio.run(litellm.aembedding(**SDK_KWARGS))
|
||||
|
||||
kwargs = invoke_kwargs(client)
|
||||
assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component")
|
||||
assert json.loads(kwargs["Body"]) == EXPECTED_BODY
|
||||
assert len(result.data) == 2
|
||||
|
||||
def test_router_deployment_with_model_id(self):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "my-embeddings",
|
||||
"litellm_params": {
|
||||
"model": "sagemaker/openai/my-embed-endpoint",
|
||||
"model_id": "my-component",
|
||||
"extra_body": {"model": "served-embedding-model"},
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "test-key",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
},
|
||||
"model_info": {"mode": "embedding"},
|
||||
}
|
||||
]
|
||||
)
|
||||
session, client = fake_sagemaker_session(OPENAI_RESPONSE)
|
||||
with patch("boto3.Session", return_value=session):
|
||||
result = asyncio.run(router.aembedding(model="my-embeddings", input=["hello", "good morning"]))
|
||||
|
||||
kwargs = invoke_kwargs(client)
|
||||
assert (kwargs["EndpointName"], kwargs["InferenceComponentName"]) == ("my-embed-endpoint", "my-component")
|
||||
assert json.loads(kwargs["Body"]) == EXPECTED_BODY
|
||||
assert [row["index"] for row in result.data] == [0, 1]
|
||||
Loading…
Add table
Reference in a new issue