mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(router)!: keep a deployment's own model_id over the request body's (#49045)
* fix(router): keep a deployment's own model_id over the request body's A request body model_id overrode the deployment's pinned model_id, so a client could move SageMaker chat and embedding calls to another inference component. The router now drops the request's model_id when the deployment sets one. * fix(router): let deployments opt back in to a client model_id and stop client component headers on SageMaker chat Drop the request's model_id after the router defaults are applied, so a default_litellm_params model_id cannot win over the deployment either, and apply the same precedence on the sync text_completion path. A deployment that lists model_id in configurable_clientside_auth_params keeps the client's value. SageMaker chat drops any case variant of a client's X-Amzn-SageMaker-Inference-Component header when the deployment sets a component * fix(router): keep a SageMaker deployment's component when it pins it through extra_headers * fix(router): read the SageMaker component header as a pin only on SageMaker deployments * fix(sagemaker): send the inference component only as a header, never as a body model_id SageMaker picks the component from X-Amzn-SageMaker-Inference-Component, so the chat container body and legacy parameters no longer carry model_id. Nova endpoints accept a fixed body field list, so a pinned Nova deployment could reject the extra field. * fix(health): probe header-pinned SageMaker deployments at their component Health checks bypass the router, so a legacy sagemaker/ deployment pinned only through its inference component header sent no component on /health while routed requests now send it. Also update the two SageMaker integration tests that still asserted the old precedence and the dropped body field. * fix(health): apply the router's model_id pin rule to health probes Health probes now share the router's pin function, so a deployment that lets clients set model_id gets no header pin and its own model_id wins without reading the header. Header keys that are not strings are skipped when reading the inference component header instead of failing the call. * fix(sagemaker): keep header keys that aren't strings and report a string opt-in in health checks The chat transformation now keeps extra_headers keys that aren't strings, so such a deployment fails the way it did before instead of answering 500, and the model_id pin checks configurable_clientside_auth_params with isinstance, so a string there reaches the health probe's own error capture again. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
simonp9
GitHub
parent
373f1752d3
commit
d1758cda91
12 changed files with 414 additions and 13 deletions
|
|
@ -28,7 +28,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.utils import LlmProviders
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from ..common_utils import AWSEventStreamDecoder, SagemakerError
|
||||
from ..common_utils import INFERENCE_COMPONENT_HEADER, OBJECT_KEYED_MAPPING, AWSEventStreamDecoder, SagemakerError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
|
@ -59,7 +59,12 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
|
|||
inference_component_name: Final = optional_params.get("model_id")
|
||||
if not isinstance(inference_component_name, str):
|
||||
return headers
|
||||
return {**headers, "X-Amzn-SageMaker-Inference-Component": inference_component_name}
|
||||
other_headers: Final = {
|
||||
name: value
|
||||
for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items()
|
||||
if not isinstance(name, str) or name.lower() != INFERENCE_COMPONENT_HEADER.lower()
|
||||
}
|
||||
return {**other_headers, INFERENCE_COMPONENT_HEADER: inference_component_name}
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
|
|
@ -76,9 +81,8 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
|
|||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
container_request: Final = {
|
||||
key: value for key, value in request.items() if key not in self.aws_authentication_params
|
||||
}
|
||||
header_only_params: Final = frozenset((*self.aws_authentication_params, "model_id"))
|
||||
container_request: Final = {key: value for key, value in request.items() if key not in header_only_params}
|
||||
served_model_name: Final = litellm_params.get("hf_model_name")
|
||||
if not isinstance(served_model_name, str):
|
||||
return container_request
|
||||
|
|
|
|||
|
|
@ -14,9 +14,10 @@ from typing_extensions import ReadOnly, Required, TypedDict
|
|||
from litellm import verbose_logger
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
from litellm.types.utils import StreamingChatCompletionChunk
|
||||
from litellm.types.utils import LlmProviders, StreamingChatCompletionChunk
|
||||
|
||||
STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
OBJECT_KEYED_MAPPING: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
|
|
@ -167,6 +168,36 @@ def get_sagemaker_response_stream_shape():
|
|||
return _load_sagemaker_response_stream_shape()
|
||||
|
||||
|
||||
INFERENCE_COMPONENT_HEADER: Final = "X-Amzn-SageMaker-Inference-Component"
|
||||
|
||||
|
||||
_SAGEMAKER_PROVIDERS: Final = frozenset(
|
||||
provider.value for provider in (LlmProviders.SAGEMAKER, LlmProviders.SAGEMAKER_CHAT, LlmProviders.SAGEMAKER_NOVA)
|
||||
)
|
||||
|
||||
|
||||
def _deployment_provider(litellm_params: Mapping[str, object]) -> str | None:
|
||||
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
||||
if isinstance(custom_llm_provider, str) and custom_llm_provider:
|
||||
return custom_llm_provider
|
||||
model: Final = litellm_params.get("model")
|
||||
return model.split("/", 1)[0] if isinstance(model, str) else None
|
||||
|
||||
|
||||
def deployment_inference_component(litellm_params: Mapping[str, object]) -> str | None:
|
||||
headers: Final = litellm_params.get("extra_headers")
|
||||
if _deployment_provider(litellm_params) not in _SAGEMAKER_PROVIDERS or not isinstance(headers, Mapping):
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
value
|
||||
for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items()
|
||||
if isinstance(name, str) and name.lower() == INFERENCE_COMPONENT_HEADER.lower() and isinstance(value, str)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
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}})`"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.utils import ModelResponse, Usage
|
||||
from litellm.utils import token_counter
|
||||
|
||||
from ..common_utils import SagemakerError
|
||||
from ..common_utils import STR_KEYED_MAPPING, SagemakerError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
|
@ -153,7 +153,9 @@ class SagemakerConfig(BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
inference_params: Final = optional_params.copy()
|
||||
inference_params: Final = {
|
||||
key: value for key, value in STR_KEYED_MAPPING.validate_python(optional_params).items() if key != "model_id"
|
||||
}
|
||||
stream: Final = inference_params.pop("stream", False)
|
||||
data: Final[dict] = {"parameters": inference_params}
|
||||
if stream is True:
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import (
|
|||
classify_strategy_router_model,
|
||||
strategy_router_dependencies,
|
||||
)
|
||||
from litellm.router_utils.common_utils import is_model_capability_field
|
||||
from litellm.router_utils.common_utils import deployment_model_id_pin, is_model_capability_field
|
||||
from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params
|
||||
|
||||
# Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy
|
||||
|
|
@ -855,6 +855,8 @@ def update_litellm_params_for_health_check(
|
|||
`deployment_model` (default: the incoming `litellm_params` model), so a `health_check_model` or a `model`
|
||||
in `health_check_params` is judged on its own. Keys already set through `health_check_params` are kept,
|
||||
and `modalities` set there without "audio" turns the audio fields off
|
||||
- sends a SageMaker deployment's `X-Amzn-SageMaker-Inference-Component` extra header as `model_id` when the
|
||||
deployment sets none, so legacy `sagemaker/` probes reach that component the way routed requests do
|
||||
"""
|
||||
incoming_model: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model")
|
||||
base_model_owner: Final = deployment_model or (incoming_model if isinstance(incoming_model, str) else None)
|
||||
|
|
@ -935,9 +937,11 @@ def update_litellm_params_for_health_check(
|
|||
audio_output_params: Final = _health_check_audio_output_params(
|
||||
_MODEL_INFO_ADAPTER.validate_python(model_info), probe_params, mode, base_model_owner
|
||||
)
|
||||
model_id_pin: Final = deployment_model_id_pin(probe_params)
|
||||
model_id_params: Final = {} if model_id_pin is None else {"model_id": model_id_pin}
|
||||
return {
|
||||
key: value
|
||||
for key, value in {**audio_output_params, **probe_params}.items()
|
||||
for key, value in {**audio_output_params, **model_id_params, **probe_params}.items()
|
||||
if not is_model_capability_field(key)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -255,6 +255,7 @@ from litellm.router_utils.clientside_credential_handler import (
|
|||
from litellm.router_utils.common_utils import (
|
||||
MID_STREAM_CONTINUATION_KWARG,
|
||||
as_retry_skipped_deployment_ids,
|
||||
deployment_model_id_pin,
|
||||
deployments_the_caller_may_use,
|
||||
filter_team_based_models,
|
||||
filter_web_search_deployments,
|
||||
|
|
@ -5316,6 +5317,9 @@ class Router(DeprecatedRouterMembers):
|
|||
):
|
||||
kwargs[STREAM_USAGE_INJECTED_FIELD] = True
|
||||
self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name)
|
||||
pinned_model_id: Final = deployment_model_id_pin(deployment_context)
|
||||
if pinned_model_id is not None:
|
||||
kwargs["model_id"] = pinned_model_id # rebind-ok: this method's contract is filling the caller's kwargs
|
||||
hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs)
|
||||
if _is_fallback_hop(hop_kwargs):
|
||||
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
|
||||
|
|
@ -6242,6 +6246,9 @@ class Router(DeprecatedRouterMembers):
|
|||
kwargs[k] = v
|
||||
elif k == "metadata":
|
||||
kwargs[k].update(v)
|
||||
pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
|
||||
if pinned_model_id is not None:
|
||||
kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults
|
||||
|
||||
# call via litellm.completion()
|
||||
return litellm.text_completion(
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from litellm.exceptions import BadRequestError
|
|||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
|
||||
from litellm.llms.sagemaker.common_utils import deployment_inference_component
|
||||
from litellm.types.router import (
|
||||
CredentialLiteLLMParams,
|
||||
Deployment,
|
||||
|
|
@ -522,6 +523,16 @@ PROVIDER_SPECIFIC_CAPABILITY_FIELDS: Final = frozenset(ProviderSpecificModelInfo
|
|||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def deployment_model_id_pin(deployment_litellm_params: Mapping[str, object]) -> object | None:
|
||||
clientside_params: Final = deployment_litellm_params.get("configurable_clientside_auth_params")
|
||||
if isinstance(clientside_params, list | tuple) and "model_id" in clientside_params:
|
||||
return None
|
||||
own_model_id: Final = deployment_litellm_params.get("model_id")
|
||||
if own_model_id is not None:
|
||||
return own_model_id
|
||||
return deployment_inference_component(deployment_litellm_params)
|
||||
|
||||
|
||||
def is_model_capability_field(name: str) -> bool:
|
||||
return name in PROVIDER_SPECIFIC_CAPABILITY_FIELDS or (
|
||||
name.startswith("supports_") and name in GetModelCostMap.loaded_model_cost_map_fields()
|
||||
|
|
|
|||
|
|
@ -243,7 +243,7 @@ def test_sagemaker_chat_signs_the_inference_component_header_and_sends_hf_model_
|
|||
)
|
||||
assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";")
|
||||
assert JSON_OBJECT.validate_json(request.body) == {
|
||||
**container_body(wire.url, PROMPT, stream=False, model_id=_INFERENCE_COMPONENT),
|
||||
**container_body(wire.url, PROMPT, stream=False),
|
||||
"model": SERVED_MODEL,
|
||||
}, request.body
|
||||
|
||||
|
|
|
|||
|
|
@ -668,12 +668,12 @@ def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body(
|
|||
assert _body(request) == {"input": [text]}, request.body
|
||||
|
||||
|
||||
def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None:
|
||||
def test_a_caller_sent_model_id_cannot_override_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 request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
|
||||
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -343,6 +343,26 @@ def test_body_carries_no_aws_credentials(
|
|||
assert client.request_headers["X-Amz-Security-Token"] == "test-session-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["sagemaker_chat/my-endpoint", "sagemaker_nova/my-endpoint"])
|
||||
def test_inference_component_travels_only_in_the_header(monkeypatch: pytest.MonkeyPatch, model: str):
|
||||
client = _invoke_sagemaker_chat(monkeypatch, model=model, model_id="my-inference-component")
|
||||
|
||||
assert client.request_headers[_INFERENCE_COMPONENT_HEADER] == "my-inference-component"
|
||||
assert "model_id" not in client.request_body
|
||||
|
||||
|
||||
def test_inference_component_header_keeps_header_keys_that_are_not_strings():
|
||||
headers = SagemakerChatConfig().validate_environment(
|
||||
headers={7: "other", "x-amzn-sagemaker-inference-component": "client-component"},
|
||||
model="my-endpoint",
|
||||
messages=[],
|
||||
optional_params={"model_id": "my-inference-component"},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert headers == {7: "other", _INFERENCE_COMPONENT_HEADER: "my-inference-component"}
|
||||
|
||||
|
||||
def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch):
|
||||
"""Without `hf_model_name` the body must keep the model it has today."""
|
||||
client = _invoke_sagemaker_chat(monkeypatch)
|
||||
|
|
|
|||
|
|
@ -372,6 +372,27 @@ def test_completion_sagemaker_non_stream(
|
|||
assert response.choices[0].message.content == "SageMaker reply"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("httpx_transport")
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_completion_sagemaker_sends_the_inference_component_only_as_a_header(respx_mock: respx.MockRouter) -> None:
|
||||
url: Final = "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations"
|
||||
route: Final = respx_mock.post(url).mock(
|
||||
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
|
||||
)
|
||||
litellm.completion(
|
||||
model="sagemaker/my-endpoint",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
temperature=0.2,
|
||||
model_id="my-inference-component",
|
||||
aws_region_name="us-west-2",
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
request: Final = route.calls[0].request
|
||||
assert request.headers["X-Amzn-SageMaker-Inference-Component"] == "my-inference-component"
|
||||
assert json.loads(request.content) == {"inputs": "hi", "parameters": {"temperature": 0.2}}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("httpx_transport")
|
||||
@pytest.mark.usefixtures("fake_provider_credentials")
|
||||
def test_completion_sagemaker_non_stream_with_aws_params(respx_mock: respx.MockRouter) -> None:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
from typing import Final
|
||||
|
|
@ -1564,3 +1565,84 @@ async def test_a_failing_health_probe_leaves_the_deployment_serving_requests(mon
|
|||
|
||||
assert "error" in probe
|
||||
assert reply.choices[0].message.content == "ok"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("component_params", "probed_component"),
|
||||
[
|
||||
({"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}}, "header-component"),
|
||||
(
|
||||
{"model_id": "own-component", "extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}},
|
||||
"own-component",
|
||||
),
|
||||
(
|
||||
{"model_id": "own-component", "extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "x"}},
|
||||
"own-component",
|
||||
),
|
||||
(
|
||||
{"extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "header-component"}},
|
||||
"header-component",
|
||||
),
|
||||
(
|
||||
{
|
||||
"configurable_clientside_auth_params": ["model_id"],
|
||||
"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"},
|
||||
},
|
||||
None,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"header_pinned",
|
||||
"own_model_id_beats_the_header",
|
||||
"own_model_id_beside_a_non_string_header_key",
|
||||
"header_pinned_beside_a_non_string_header_key",
|
||||
"client_configurable_model_id_takes_no_header_pin",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_model_health_check_probes_a_legacy_sagemaker_embedding_at_the_deployments_component(
|
||||
component_params: dict[str, object], probed_component: str | None
|
||||
) -> None:
|
||||
embedding_reply: Final = {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}]}
|
||||
invoke_endpoint: Final = MagicMock(return_value={"Body": io.BytesIO(json.dumps(embedding_reply).encode())})
|
||||
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
|
||||
deployment: Final = {
|
||||
"litellm_params": {
|
||||
"model": "sagemaker/openai/embed-endpoint",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "test-key",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
**component_params,
|
||||
},
|
||||
"model_info": {"mode": "embedding"},
|
||||
}
|
||||
|
||||
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
|
||||
result: Final = await hc_module._run_model_health_check(deployment)
|
||||
|
||||
assert "error" not in result, result
|
||||
assert invoke_endpoint.call_count == 1
|
||||
assert invoke_endpoint.call_args.kwargs.get("InferenceComponentName") == probed_component
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_model_health_check_reports_why_a_string_configurable_clientside_auth_params_fails() -> None:
|
||||
invoke_endpoint: Final = MagicMock()
|
||||
deployment: Final = {
|
||||
"litellm_params": {
|
||||
"model": "sagemaker/openai/embed-endpoint",
|
||||
"model_id": "own-component",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "test-key",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
},
|
||||
"model_info": {"mode": "embedding", "health_check_params": {"configurable_clientside_auth_params": "model_id"}},
|
||||
}
|
||||
|
||||
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
|
||||
|
||||
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
|
||||
result: Final = await hc_module._run_model_health_check(deployment)
|
||||
|
||||
assert "configurable_clientside_auth_params" in result["error"]
|
||||
invoke_endpoint.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import copy
|
|||
import functools
|
||||
import gc
|
||||
import itertools
|
||||
import io
|
||||
import json
|
||||
import linecache
|
||||
import logging
|
||||
|
|
@ -31523,6 +31524,224 @@ def test_router_keeps_supported_fallback_shapes(
|
|||
assert router.fallbacks == expected_fallbacks
|
||||
|
||||
|
||||
_SAGEMAKER_ROUTER_CREDENTIALS: Final = MappingProxyType(
|
||||
{"aws_region_name": "us-east-1", "aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"}
|
||||
)
|
||||
|
||||
_PINNED_COMPONENT: Final = MappingProxyType({"model_id": "deployment-component"})
|
||||
_HEADER_PINNED_COMPONENT: Final = MappingProxyType(
|
||||
{"extra_headers": {"x-amzn-sagemaker-inference-component": "deployment-component"}}
|
||||
)
|
||||
|
||||
_INFERENCE_COMPONENT_PRECEDENCE: Final = pytest.mark.parametrize(
|
||||
("deployment_params", "sent_component"),
|
||||
[
|
||||
(_PINNED_COMPONENT, "deployment-component"),
|
||||
(_HEADER_PINNED_COMPONENT, "deployment-component"),
|
||||
(MappingProxyType({}), "client-component"),
|
||||
(
|
||||
MappingProxyType({**_PINNED_COMPONENT, "configurable_clientside_auth_params": ["model_id"]}),
|
||||
"client-component",
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"deployment_pins_the_component",
|
||||
"deployment_pins_the_component_with_a_header",
|
||||
"deployment_leaves_the_component_to_the_client",
|
||||
"deployment_lets_the_client_override_its_component",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _sagemaker_router(
|
||||
model: str,
|
||||
deployment_params: Mapping[str, object],
|
||||
mode: str,
|
||||
router_defaults: Mapping[str, str] = MappingProxyType({}),
|
||||
model_name: str = "sagemaker-model",
|
||||
) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model, **_SAGEMAKER_ROUTER_CREDENTIALS, **deployment_params},
|
||||
"model_info": {"mode": mode},
|
||||
}
|
||||
],
|
||||
default_litellm_params=dict(router_defaults),
|
||||
)
|
||||
|
||||
|
||||
_CHAT_COMPLETION_REPLY: Final = MappingProxyType(
|
||||
{
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "served-model",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _sagemaker_chat_route(respx_mock: respx.MockRouter) -> respx.Route:
|
||||
return respx_mock.post(
|
||||
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/chat-endpoint/invocations"
|
||||
).respond(json=dict(_CHAT_COMPLETION_REPLY))
|
||||
|
||||
|
||||
@_INFERENCE_COMPONENT_PRECEDENCE
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_model_id_cannot_move_a_sagemaker_chat_call_off_the_deployments_component(
|
||||
respx_mock: respx.MockRouter,
|
||||
_router_httpx_transport: None,
|
||||
deployment_params: Mapping[str, object],
|
||||
sent_component: str,
|
||||
) -> None:
|
||||
route: Final = _sagemaker_chat_route(respx_mock)
|
||||
router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_params, "chat")
|
||||
|
||||
await router.acompletion(
|
||||
model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component"
|
||||
)
|
||||
|
||||
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component
|
||||
|
||||
|
||||
@_INFERENCE_COMPONENT_PRECEDENCE
|
||||
def test_a_client_model_id_cannot_move_a_sync_sagemaker_text_completion_off_the_deployments_component(
|
||||
respx_mock: respx.MockRouter, deployment_params: Mapping[str, object], sent_component: str
|
||||
) -> None:
|
||||
route: Final = _sagemaker_chat_route(respx_mock)
|
||||
router: Final = _sagemaker_router(
|
||||
"sagemaker_chat/chat-endpoint", deployment_params, "completion", model_name="sagemaker_chat/chat-endpoint"
|
||||
)
|
||||
|
||||
router.text_completion(model="sagemaker_chat/chat-endpoint", prompt="hi", model_id="client-component")
|
||||
|
||||
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"client_header",
|
||||
[
|
||||
MappingProxyType({"extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"}}),
|
||||
MappingProxyType({"headers": {"x-amzn-sagemaker-inference-component": "client-component"}}),
|
||||
MappingProxyType(
|
||||
{
|
||||
"provider_specific_header": {
|
||||
"custom_llm_provider": "sagemaker_chat",
|
||||
"extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"},
|
||||
}
|
||||
}
|
||||
),
|
||||
],
|
||||
ids=["extra_headers", "headers", "provider_specific_header"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"deployment_pin", [_PINNED_COMPONENT, _HEADER_PINNED_COMPONENT], ids=["model_id_pin", "header_pin"]
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_inference_component_header_cannot_move_a_pinned_sagemaker_chat_call(
|
||||
respx_mock: respx.MockRouter,
|
||||
_router_httpx_transport: None,
|
||||
client_header: Mapping[str, object],
|
||||
deployment_pin: Mapping[str, object],
|
||||
) -> None:
|
||||
route: Final = _sagemaker_chat_route(respx_mock)
|
||||
router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_pin, "chat")
|
||||
|
||||
await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **client_header)
|
||||
|
||||
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_component",
|
||||
[MappingProxyType({}), MappingProxyType({"model_id": "deployment-component"})],
|
||||
ids=["request_without_model_id", "request_with_the_pinned_model_id"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_router_default_model_id_cannot_move_a_sagemaker_call_off_the_deployments_component(
|
||||
respx_mock: respx.MockRouter, _router_httpx_transport: None, request_component: Mapping[str, str]
|
||||
) -> None:
|
||||
route: Final = _sagemaker_chat_route(respx_mock)
|
||||
router: Final = _sagemaker_router(
|
||||
"sagemaker_chat/chat-endpoint",
|
||||
_PINNED_COMPONENT,
|
||||
"chat",
|
||||
router_defaults=MappingProxyType({"model_id": "default-component"}),
|
||||
)
|
||||
|
||||
await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **request_component)
|
||||
|
||||
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
|
||||
|
||||
|
||||
@_INFERENCE_COMPONENT_PRECEDENCE
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_deployments_component(
|
||||
deployment_params: Mapping[str, object], sent_component: str
|
||||
) -> None:
|
||||
component_embeddings: Final = MappingProxyType({"deployment-component": [0.1], "client-component": [0.2]})
|
||||
|
||||
def invoke_endpoint(**request: str | bytes) -> Mapping[str, io.BytesIO]:
|
||||
served: Final = component_embeddings[str(request["InferenceComponentName"])]
|
||||
return {
|
||||
"Body": io.BytesIO(
|
||||
json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode()
|
||||
)
|
||||
}
|
||||
|
||||
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
|
||||
router: Final = _sagemaker_router("sagemaker/openai/embed-endpoint", deployment_params, "embedding")
|
||||
|
||||
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
|
||||
response: Final = await router.aembedding(model="sagemaker-model", input=["hello"], model_id="client-component")
|
||||
|
||||
assert response.data[0]["embedding"] == component_embeddings[sent_component]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_custom_llm_provider(
|
||||
respx_mock: respx.MockRouter, _router_httpx_transport: None
|
||||
) -> None:
|
||||
route: Final = _sagemaker_chat_route(respx_mock)
|
||||
router: Final = _sagemaker_router(
|
||||
"chat-endpoint", MappingProxyType({**_HEADER_PINNED_COMPONENT, "custom_llm_provider": "sagemaker_chat"}), "chat"
|
||||
)
|
||||
|
||||
await router.acompletion(
|
||||
model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component"
|
||||
)
|
||||
|
||||
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone(
|
||||
respx_mock: respx.MockRouter, _router_httpx_transport: None
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY))
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gateway-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/served-model",
|
||||
"api_base": "https://gateway.local/v1",
|
||||
"api_key": "test-key",
|
||||
"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "gateway-component"},
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value")
|
||||
|
||||
assert json.loads(route.calls.last.request.content)["model_id"] == "client-value"
|
||||
|
||||
|
||||
def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog: pytest.LogCaptureFixture) -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
router: Final = Router(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue