mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
feat(router): keep a session on one kubernetes pod (#44393)
Requests with an explicit session id stick to one discovered pod via rendezvous hashing Requests without a session id keep round robin Pod routing records are included in logs, spend logs, and traces Resolves LIT-9168 Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
6fc60bc00d
commit
f85926a652
22 changed files with 1541 additions and 66 deletions
|
|
@ -47,6 +47,7 @@ KUBERNETES_POD_DISCOVERY_IDLE_EVICTION_SECONDS: Final = get_env_float(
|
|||
"KUBERNETES_POD_DISCOVERY_IDLE_EVICTION_SECONDS", 300.0
|
||||
)
|
||||
KUBERNETES_POD_DISCOVERY_DNS_TIMEOUT_SECONDS: Final = 2.0
|
||||
KUBERNETES_POD_ROUTING_KEY: Final = "kubernetes_pod_routing"
|
||||
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
|
||||
AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
|
||||
ROUTER_MAX_FALLBACKS: Final = get_env_int("ROUTER_MAX_FALLBACKS", 5)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY
|
||||
from litellm.integrations._types.open_inference import (
|
||||
OpenInferenceSpanKindValues,
|
||||
)
|
||||
|
|
@ -50,6 +51,7 @@ from litellm.types.utils import (
|
|||
Function,
|
||||
LLMResponseTypes,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingMetadata,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
|
||||
|
|
@ -2420,9 +2422,15 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
#############################################
|
||||
############ LLM CALL METADATA ##############
|
||||
#############################################
|
||||
metadata: Final = standard_logging_payload["metadata"]
|
||||
metadata: Final = cast( # cast-ok: payload metadata follows the standard logging contract
|
||||
StandardLoggingMetadata, standard_logging_payload["metadata"]
|
||||
)
|
||||
for key, value in metadata.items():
|
||||
self.safe_set_attribute(span=span, key=f"metadata.{key}", value=value)
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=f"metadata.{key}",
|
||||
value=safe_dumps(value) if key == KUBERNETES_POD_ROUTING_KEY and value is not None else value,
|
||||
)
|
||||
decision: Final = metadata.get("routing_decision")
|
||||
span.set_attributes(routing_decision_attributes(decision))
|
||||
|
||||
|
|
|
|||
|
|
@ -11,13 +11,14 @@ import httpx
|
|||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id
|
||||
from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
from litellm.types.utils import ModelResponseStream, StandardLoggingKubernetesPodRouting
|
||||
|
||||
Span = _Span | Any
|
||||
else:
|
||||
|
|
@ -385,6 +386,54 @@ def proxy_stamped_used_client_oauth_token(metadata: object, litellm_params: Mapp
|
|||
return metadata.get("used_client_oauth_token") if isinstance(metadata, Mapping) else None
|
||||
|
||||
|
||||
def proxy_stamped_kubernetes_pod_routing(
|
||||
metadata: object,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
) -> "StandardLoggingKubernetesPodRouting | None":
|
||||
litellm_metadata: Final = litellm_params.get("litellm_metadata") if litellm_params is not None else None
|
||||
litellm_routing: Final = _standard_logging_kubernetes_pod_routing(litellm_metadata)
|
||||
if litellm_routing is not None:
|
||||
return litellm_routing
|
||||
return _standard_logging_kubernetes_pod_routing(metadata)
|
||||
|
||||
|
||||
def _standard_logging_kubernetes_pod_routing(
|
||||
metadata: object,
|
||||
) -> "StandardLoggingKubernetesPodRouting | None":
|
||||
if not isinstance(metadata, Mapping):
|
||||
return None
|
||||
metadata_mapping: Final = cast( # cast-ok: routing metadata is a string-keyed mapping
|
||||
Mapping[str, object], metadata
|
||||
)
|
||||
routing: Final = metadata_mapping.get(KUBERNETES_POD_ROUTING_KEY)
|
||||
if not isinstance(routing, Mapping):
|
||||
return None
|
||||
routing_mapping: Final = cast( # cast-ok: pod routing records use string keys
|
||||
Mapping[str, object], routing
|
||||
)
|
||||
service_host: Final = routing_mapping.get("service_host")
|
||||
pod_ip: Final = routing_mapping.get("pod_ip")
|
||||
pod_count: Final = routing_mapping.get("pod_count")
|
||||
selection: Final = routing_mapping.get("selection")
|
||||
if (
|
||||
not isinstance(service_host, str)
|
||||
or not isinstance(pod_ip, str)
|
||||
or not isinstance(pod_count, int)
|
||||
or isinstance(pod_count, bool)
|
||||
or pod_count < 1
|
||||
):
|
||||
return None
|
||||
if selection in ("round_robin", "session_affinity", "session_affinity_retry"):
|
||||
routing_record: Final[StandardLoggingKubernetesPodRouting] = {
|
||||
"service_host": service_host,
|
||||
"pod_ip": pod_ip,
|
||||
"pod_count": pod_count,
|
||||
"selection": selection,
|
||||
}
|
||||
return routing_record
|
||||
return None
|
||||
|
||||
|
||||
def get_litellm_metadata_from_kwargs(kwargs: dict):
|
||||
"""
|
||||
Helper to get litellm metadata from all litellm request kwargs
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm.constants import (
|
|||
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
|
||||
EMPTY_MAPPING,
|
||||
INTERNAL_CALLBACK_NAMES,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
PROVIDER_REQUEST_ID_HEADERS,
|
||||
REDACTED_BY_LITELLM,
|
||||
)
|
||||
|
|
@ -75,6 +76,7 @@ from litellm.litellm_core_utils.classifier_logging import (
|
|||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_provider_response_headers_from_hidden_params,
|
||||
is_expected_client_error,
|
||||
proxy_stamped_kubernetes_pod_routing,
|
||||
proxy_stamped_used_client_oauth_token,
|
||||
reconstruct_model_name,
|
||||
set_response_cost_in_hidden_params,
|
||||
|
|
@ -298,7 +300,9 @@ else:
|
|||
_GENERIC_API_LOGGER_CLS: Final = GenericAPILogger
|
||||
_in_memory_loggers: Final[list[CustomLogger]] = []
|
||||
|
||||
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(("used_client_oauth_token", "usage_object"))
|
||||
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(
|
||||
("used_client_oauth_token", "usage_object", KUBERNETES_POD_ROUTING_KEY)
|
||||
)
|
||||
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = (
|
||||
frozenset(StandardLoggingMetadata.__annotations__.keys()) - _STANDARD_LOGGING_METADATA_RESOLVED_KEYS
|
||||
)
|
||||
|
|
@ -6275,7 +6279,7 @@ class StandardLoggingPayloadSetup:
|
|||
@staticmethod
|
||||
def get_standard_logging_metadata(
|
||||
metadata: Mapping[str, object] | None,
|
||||
litellm_params: dict | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
prompt_integration: str | None = None,
|
||||
applied_guardrails: list[str] | None = None,
|
||||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None,
|
||||
|
|
@ -6315,6 +6319,8 @@ class StandardLoggingPayloadSetup:
|
|||
prompt_integration=prompt_integration,
|
||||
)
|
||||
|
||||
kubernetes_pod_routing: Final = proxy_stamped_kubernetes_pod_routing(metadata, litellm_params)
|
||||
|
||||
# Initialize with default values
|
||||
clean_metadata = StandardLoggingMetadata(
|
||||
user_api_key_hash=None,
|
||||
|
|
@ -6351,6 +6357,7 @@ class StandardLoggingPayloadSetup:
|
|||
user_api_key_auth_metadata=None,
|
||||
team_alias=None,
|
||||
team_id=None,
|
||||
kubernetes_pod_routing=kubernetes_pod_routing,
|
||||
used_client_oauth_token=resolve_used_client_oauth_token(
|
||||
proxy_stamped_used_client_oauth_token(metadata, litellm_params),
|
||||
custom_llm_provider,
|
||||
|
|
@ -7418,6 +7425,7 @@ def get_standard_logging_metadata(
|
|||
user_api_key_auth_metadata=None,
|
||||
team_alias=None,
|
||||
team_id=None,
|
||||
kubernetes_pod_routing=None,
|
||||
used_client_oauth_token=None,
|
||||
)
|
||||
if isinstance(metadata, dict):
|
||||
|
|
@ -7491,6 +7499,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
|
|||
requester_ip_address="127.0.0.1",
|
||||
requester_metadata=None,
|
||||
user_api_key_end_user_id="test_end_user",
|
||||
kubernetes_pod_routing=None,
|
||||
)
|
||||
|
||||
hidden_params: Final = StandardLoggingHiddenParams(
|
||||
|
|
|
|||
|
|
@ -105,6 +105,7 @@ from litellm.types.utils import (
|
|||
ProviderField,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingKubernetesPodRouting,
|
||||
StandardLoggingMCPToolCall,
|
||||
StandardLoggingModelInformation,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
|
|
@ -3682,6 +3683,7 @@ class SpendLogsMetadata(TypedDict):
|
|||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None
|
||||
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None
|
||||
routing_decision: StandardLoggingRoutingDecision | None
|
||||
kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None]
|
||||
internal_call_origin: InternalCallOrigin | None
|
||||
litellm_roi_estimator: ReadOnly[NotRequired[bool | None]]
|
||||
guardrail_information: list[StandardLoggingGuardrailInformation] | None
|
||||
|
|
|
|||
|
|
@ -1672,11 +1672,9 @@ class LiteLLMProxyRequestSetup:
|
|||
metadata_from_headers["agent_id"] = agent_id_from_header
|
||||
verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header)
|
||||
|
||||
spend_log_session_id: Final = (
|
||||
caller_session_beside_trace: Final = (
|
||||
_caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name)
|
||||
if litellm.request_correlation_in_logs
|
||||
and chain_id is not None
|
||||
and chain_id == normalized_headers.get("x-litellm-trace-id")
|
||||
if chain_id is not None and chain_id == normalized_headers.get("x-litellm-trace-id")
|
||||
else None
|
||||
)
|
||||
if chain_id:
|
||||
|
|
@ -1738,8 +1736,8 @@ class LiteLLMProxyRequestSetup:
|
|||
|
||||
if isinstance(data[_metadata_variable_name], dict):
|
||||
data[_metadata_variable_name].update(
|
||||
{**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id}
|
||||
if spend_log_session_id
|
||||
{**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: caller_session_beside_trace}
|
||||
if caller_session_beside_trace
|
||||
else metadata_from_headers
|
||||
)
|
||||
return data
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.constants import (
|
|||
CLI_SESSION_KEY_PREFIX,
|
||||
EMPTY_MAPPING,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
LITELLM_PROXY_MASTER_KEY_ALIAS,
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
|
||||
|
|
@ -43,6 +44,7 @@ from litellm.litellm_core_utils.classifier_logging import (
|
|||
)
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_litellm_metadata_from_kwargs,
|
||||
proxy_stamped_kubernetes_pod_routing,
|
||||
proxy_stamped_used_client_oauth_token,
|
||||
reconstruct_model_name,
|
||||
)
|
||||
|
|
@ -76,6 +78,7 @@ from litellm.types.utils import (
|
|||
CostBreakdown,
|
||||
LlmProviders,
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingKubernetesPodRouting,
|
||||
StandardLoggingMCPToolCall,
|
||||
StandardLoggingModelInformation,
|
||||
StandardLoggingPayload,
|
||||
|
|
@ -178,6 +181,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset(
|
|||
"autorouter_savings_estimate",
|
||||
"autorouter_baseline_observation",
|
||||
"used_client_oauth_token",
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
"litellm_roi_estimator",
|
||||
"user_api_key_user_email",
|
||||
)
|
||||
|
|
@ -210,6 +214,7 @@ def _get_spend_logs_metadata(
|
|||
router_metadata: SpendLogsRouterMetadata | None = None,
|
||||
azure_spillover: AzureSpillover | None = None,
|
||||
used_client_oauth_token: bool | None = None,
|
||||
kubernetes_pod_routing: StandardLoggingKubernetesPodRouting | None = None,
|
||||
) -> SpendLogsMetadata:
|
||||
if metadata is None:
|
||||
return SpendLogsMetadata(
|
||||
|
|
@ -257,6 +262,7 @@ def _get_spend_logs_metadata(
|
|||
router_metadata=router_metadata,
|
||||
azure_spillover=azure_spillover,
|
||||
used_client_oauth_token=used_client_oauth_token,
|
||||
kubernetes_pod_routing=None,
|
||||
)
|
||||
verbose_proxy_logger.debug("getting payload for SpendLogs, available keys in metadata: %s", [*metadata])
|
||||
|
||||
|
|
@ -271,6 +277,7 @@ def _get_spend_logs_metadata(
|
|||
router_metadata=router_metadata,
|
||||
azure_spillover=azure_spillover,
|
||||
used_client_oauth_token=used_client_oauth_token,
|
||||
kubernetes_pod_routing=kubernetes_pod_routing,
|
||||
litellm_roi_estimator=metadata.get("litellm_roi_estimator") is True,
|
||||
user_api_key_user_email=_stored_user_email(metadata),
|
||||
)
|
||||
|
|
@ -557,6 +564,16 @@ def get_logging_payload(
|
|||
# standardize this function to be used across, s3, dynamoDB, langfuse logging
|
||||
litellm_params: Final = kwargs.get("litellm_params", {})
|
||||
metadata: Final = get_litellm_metadata_from_kwargs(kwargs)
|
||||
routing_litellm_params: Final[Mapping[str, object] | None] = (
|
||||
litellm_params if isinstance(litellm_params, Mapping) else None
|
||||
)
|
||||
routing_metadata: Final[object] = (
|
||||
routing_litellm_params.get("metadata") if routing_litellm_params is not None else None
|
||||
)
|
||||
kubernetes_pod_routing: Final = proxy_stamped_kubernetes_pod_routing(
|
||||
routing_metadata,
|
||||
routing_litellm_params,
|
||||
)
|
||||
completion_start_time: Final = kwargs.get("completion_start_time", end_time)
|
||||
call_type: Final = kwargs.get("call_type")
|
||||
cache_hit: Final = kwargs.get("cache_hit", False)
|
||||
|
|
@ -591,7 +608,9 @@ def get_logging_payload(
|
|||
id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata)
|
||||
standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None))
|
||||
|
||||
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
|
||||
end_user_id = get_end_user_id_for_cost_tracking(
|
||||
dict(litellm_params) if isinstance(litellm_params, Mapping) else litellm_params
|
||||
)
|
||||
|
||||
api_key = metadata.get("user_api_key", "")
|
||||
|
||||
|
|
@ -769,6 +788,7 @@ def get_logging_payload(
|
|||
used_client_oauth_token=resolve_used_client_oauth_token(
|
||||
proxy_stamped_used_client_oauth_token(litellm_params.get("metadata"), litellm_params), custom_llm_provider
|
||||
),
|
||||
kubernetes_pod_routing=kubernetes_pod_routing,
|
||||
azure_spillover=azure_spillover(
|
||||
response_headers=kwargs.get("response_headers")
|
||||
if isinstance(kwargs.get("response_headers"), Mapping)
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from litellm.constants import (
|
|||
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
|
||||
DEFAULT_MAX_LRU_CACHE_SIZE,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
MID_STREAM_FAILURE_LOG_WAIT_SECONDS,
|
||||
OUTPUT_TOKEN_CEILING_PARAMS,
|
||||
PROMPT_CACHE_PIN_CACHE_MAX_SIZE,
|
||||
|
|
@ -4606,6 +4607,7 @@ class Router:
|
|||
deployment_that_just_failed: Final = _deployment_that_just_failed(
|
||||
_FALLBACK_HOP_ADAPTER.validate_python(kwargs).get(metadata_variable_name)
|
||||
)
|
||||
routing: Final = deployment.get(KUBERNETES_POD_ROUTING_KEY)
|
||||
|
||||
kwargs.setdefault(metadata_variable_name, {}).update(
|
||||
{
|
||||
|
|
@ -4613,8 +4615,16 @@ class Router:
|
|||
"model_info": model_info,
|
||||
"api_base": deployment_api_base,
|
||||
"deployment_model_name": deployment_model_name,
|
||||
KUBERNETES_POD_ROUTING_KEY: routing,
|
||||
}
|
||||
)
|
||||
other_bucket: Final = "metadata" if metadata_variable_name == "litellm_metadata" else "litellm_metadata"
|
||||
other_bucket_metadata: Final = kwargs.get(other_bucket)
|
||||
if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata:
|
||||
kwargs[other_bucket] = { # rebind-ok: this method updates the router's request kwargs in place
|
||||
**other_bucket_metadata,
|
||||
KUBERNETES_POD_ROUTING_KEY: routing,
|
||||
}
|
||||
|
||||
# A retry/fallback reuses this same kwargs dict for the next deployment.
|
||||
# Refund and clear any reservation the previous deployment attempt left
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import asyncio
|
||||
import concurrent.futures
|
||||
import hashlib
|
||||
import inspect
|
||||
import ipaddress
|
||||
import socket
|
||||
import threading
|
||||
|
|
@ -13,6 +15,7 @@ from types import MappingProxyType
|
|||
from typing import (
|
||||
Concatenate,
|
||||
Final,
|
||||
Literal,
|
||||
ParamSpec,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
|
|
@ -24,9 +27,12 @@ import httpx
|
|||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
KUBERNETES_POD_DISCOVERY_DNS_TIMEOUT_SECONDS,
|
||||
KUBERNETES_POD_DISCOVERY_IDLE_EVICTION_SECONDS,
|
||||
KUBERNETES_POD_DISCOVERY_REFRESH_INTERVAL_SECONDS,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
)
|
||||
|
||||
_SocketAddress: TypeAlias = tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]
|
||||
|
|
@ -34,9 +40,22 @@ _AddressInfo: TypeAlias = tuple[socket.AddressFamily, socket.SocketKind, int, st
|
|||
_SyncLookup: TypeAlias = Future[list[_AddressInfo]]
|
||||
_AsyncLookup: TypeAlias = asyncio.Future[list[_AddressInfo]]
|
||||
_CacheKey: TypeAlias = tuple[str, int | None]
|
||||
_PodSelection: TypeAlias = tuple[str, int, Literal["round_robin", "session_affinity", "session_affinity_retry"]]
|
||||
_DeploymentT: Final = TypeVar("_DeploymentT")
|
||||
|
||||
|
||||
def _retry_count_from_metadata(metadata: object) -> int | None:
|
||||
if not isinstance(metadata, Mapping):
|
||||
return None
|
||||
metadata_mapping: Final = cast( # cast-ok: request metadata is a string-keyed mapping
|
||||
Mapping[str, object], metadata
|
||||
)
|
||||
retry_count: Final = metadata_mapping.get("request_retry_count")
|
||||
if type(retry_count) is int and retry_count >= 0:
|
||||
return retry_count
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PodSet:
|
||||
ips: tuple[str, ...]
|
||||
|
|
@ -99,16 +118,24 @@ class KubernetesPodDiscovery:
|
|||
self._refreshing = self._refreshing | {key}
|
||||
return None, True
|
||||
|
||||
def resolve_deployment(self, deployment: _DeploymentT) -> _DeploymentT:
|
||||
def resolve_deployment(
|
||||
self,
|
||||
deployment: _DeploymentT,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
) -> _DeploymentT:
|
||||
eligible: Final = self._eligible(deployment)
|
||||
if eligible is None:
|
||||
return deployment
|
||||
url, key, deployment_mapping = eligible
|
||||
host, port = key
|
||||
now: Final = self.clock()
|
||||
session_id: Final = self._session_id(request_kwargs)
|
||||
retry_count: Final = self._request_retry_count(request_kwargs)
|
||||
cached, should_refresh = self._begin_refresh(key, now)
|
||||
if cached is not None:
|
||||
return self._deployment_with_cached_ip(deployment, deployment_mapping, key, url, now)
|
||||
return self._deployment_with_cached_ip(
|
||||
deployment, deployment_mapping, key, url, now, session_id, retry_count
|
||||
)
|
||||
if not should_refresh:
|
||||
return deployment
|
||||
|
||||
|
|
@ -118,21 +145,36 @@ class KubernetesPodDiscovery:
|
|||
lookup.add_done_callback(lambda done: self._finish_refresh(key, done))
|
||||
try:
|
||||
return self._apply_lookup(
|
||||
deployment, deployment_mapping, key, url, self._await_lookup(key, host, lookup), now
|
||||
deployment,
|
||||
deployment_mapping,
|
||||
key,
|
||||
url,
|
||||
self._await_lookup(key, host, lookup),
|
||||
now,
|
||||
session_id,
|
||||
retry_count,
|
||||
)
|
||||
finally:
|
||||
self._end_refresh(key)
|
||||
|
||||
async def async_resolve_deployment(self, deployment: _DeploymentT) -> _DeploymentT:
|
||||
async def async_resolve_deployment(
|
||||
self,
|
||||
deployment: _DeploymentT,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
) -> _DeploymentT:
|
||||
eligible: Final = self._eligible(deployment)
|
||||
if eligible is None:
|
||||
return deployment
|
||||
url, key, deployment_mapping = eligible
|
||||
host, port = key
|
||||
now: Final = self.clock()
|
||||
session_id: Final = self._session_id(request_kwargs)
|
||||
retry_count: Final = self._request_retry_count(request_kwargs)
|
||||
cached, should_refresh = self._begin_refresh(key, now)
|
||||
if cached is not None:
|
||||
return self._deployment_with_cached_ip(deployment, deployment_mapping, key, url, now)
|
||||
return self._deployment_with_cached_ip(
|
||||
deployment, deployment_mapping, key, url, now, session_id, retry_count
|
||||
)
|
||||
if not should_refresh:
|
||||
return deployment
|
||||
|
||||
|
|
@ -142,10 +184,56 @@ class KubernetesPodDiscovery:
|
|||
lookup.add_done_callback(lambda done: self._finish_refresh(key, done))
|
||||
try:
|
||||
result: Final = await self._await_async_lookup(key, host, lookup)
|
||||
return self._apply_lookup(deployment, deployment_mapping, key, url, result, now)
|
||||
return self._apply_lookup(deployment, deployment_mapping, key, url, result, now, session_id, retry_count)
|
||||
finally:
|
||||
self._end_refresh(key)
|
||||
|
||||
@staticmethod
|
||||
def _session_id(request_kwargs: Mapping[str, object] | None) -> str | None:
|
||||
if request_kwargs is None:
|
||||
return None
|
||||
metadata_dicts: Final = tuple(
|
||||
cast(Mapping[str, object], metadata) # cast-ok: runtime dict check precedes metadata lookup
|
||||
for metadata in (request_kwargs.get("litellm_metadata"), request_kwargs.get("metadata"))
|
||||
if isinstance(metadata, dict)
|
||||
)
|
||||
if any(metadata.get(SESSION_ID_GENERATED_METADATA_KEY) for metadata in metadata_dicts):
|
||||
return None
|
||||
caller_session_id: Final = next(
|
||||
(
|
||||
caller_session
|
||||
for metadata in metadata_dicts
|
||||
if (caller_session := metadata.get(CALLER_SESSION_ID_METADATA_KEY)) is not None and caller_session != ""
|
||||
),
|
||||
None,
|
||||
)
|
||||
if caller_session_id is not None:
|
||||
return str(caller_session_id)
|
||||
top_level_session_id: Final = request_kwargs.get("litellm_session_id")
|
||||
if top_level_session_id is not None and top_level_session_id != "":
|
||||
return str(top_level_session_id)
|
||||
metadata_session_id: Final = next(
|
||||
(
|
||||
session_id
|
||||
for metadata in metadata_dicts
|
||||
if (session_id := metadata.get("session_id")) is not None and session_id != ""
|
||||
),
|
||||
None,
|
||||
)
|
||||
if metadata_session_id is not None:
|
||||
return str(metadata_session_id)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _request_retry_count(request_kwargs: Mapping[str, object] | None) -> int:
|
||||
if request_kwargs is None:
|
||||
return 0
|
||||
retry_counts: Final = tuple(
|
||||
_retry_count_from_metadata(request_kwargs.get(metadata_key))
|
||||
for metadata_key in ("litellm_metadata", "metadata")
|
||||
)
|
||||
return next((retry_count for retry_count in retry_counts if retry_count is not None), 0)
|
||||
|
||||
def _eligible(self, deployment: object) -> tuple[httpx.URL, _CacheKey, Mapping[str, object]] | None:
|
||||
if not isinstance(deployment, Mapping):
|
||||
return None
|
||||
|
|
@ -247,6 +335,8 @@ class KubernetesPodDiscovery:
|
|||
url: httpx.URL,
|
||||
result: Sequence[_AddressInfo] | OSError,
|
||||
now: float,
|
||||
session_id: str | None,
|
||||
retry_count: int,
|
||||
) -> _DeploymentT:
|
||||
if isinstance(result, OSError):
|
||||
if isinstance(result, socket.gaierror) and self._is_authoritative_no_pods(result):
|
||||
|
|
@ -254,12 +344,14 @@ class KubernetesPodDiscovery:
|
|||
return deployment
|
||||
self._stamp_failed_refresh(key, self.clock())
|
||||
verbose_router_logger.debug("Kubernetes pod discovery DNS refresh failed for %s: %s", key[0], result)
|
||||
return self._deployment_with_cached_ip(deployment, deployment_mapping, key, url, now)
|
||||
return self._deployment_with_cached_ip(
|
||||
deployment, deployment_mapping, key, url, now, session_id, retry_count
|
||||
)
|
||||
ips: Final = self._pod_ips(result)
|
||||
self._store(key, ips, self.clock())
|
||||
if not ips:
|
||||
return deployment
|
||||
return self._deployment_with_cached_ip(deployment, deployment_mapping, key, url, now)
|
||||
return self._deployment_with_cached_ip(deployment, deployment_mapping, key, url, now, session_id, retry_count)
|
||||
|
||||
@staticmethod
|
||||
def _is_ip_literal(host: str) -> bool:
|
||||
|
|
@ -305,20 +397,55 @@ class KubernetesPodDiscovery:
|
|||
)
|
||||
self._cache = MappingProxyType({**self._cache, key: stamped})
|
||||
|
||||
def _next_ip(self, key: _CacheKey, now: float) -> str | None:
|
||||
def _next_ip(
|
||||
self,
|
||||
key: _CacheKey,
|
||||
now: float,
|
||||
session_id: str | None,
|
||||
retry_count: int,
|
||||
) -> _PodSelection | None:
|
||||
session_hash_prefix: Final = (
|
||||
hashlib.blake2b(f"{session_id}\x00".encode(), digest_size=8) if session_id is not None else None
|
||||
)
|
||||
|
||||
def session_rank(address: str) -> int:
|
||||
if session_hash_prefix is None:
|
||||
return 0
|
||||
h: Final = session_hash_prefix.copy()
|
||||
h.update(address.encode())
|
||||
return int.from_bytes(h.digest(), "big")
|
||||
|
||||
with self._lock:
|
||||
cached: Final = self._cache.get(key)
|
||||
if cached is None or not cached.ips:
|
||||
return None
|
||||
ip: Final = cached.ips[cached.cursor]
|
||||
next_cursor: Final = (cached.cursor + 1) % len(cached.ips)
|
||||
pod_count: Final = len(cached.ips)
|
||||
ranked_ips: Final = (
|
||||
tuple(
|
||||
sorted(
|
||||
cached.ips,
|
||||
key=lambda address: (-session_rank(address), address),
|
||||
)
|
||||
)
|
||||
if session_id is not None
|
||||
else ()
|
||||
)
|
||||
ip: Final = cached.ips[cached.cursor] if session_id is None else ranked_ips[retry_count % pod_count]
|
||||
next_cursor: Final = (cached.cursor + 1) % len(cached.ips) if session_id is None else cached.cursor
|
||||
selection: Final = (
|
||||
"round_robin"
|
||||
if session_id is None
|
||||
else "session_affinity_retry"
|
||||
if retry_count > 0
|
||||
else "session_affinity"
|
||||
)
|
||||
self._cache = MappingProxyType(
|
||||
{
|
||||
**self._cache,
|
||||
key: replace(cached, cursor=next_cursor, last_used_at=max(cached.last_used_at, now)),
|
||||
}
|
||||
)
|
||||
return ip
|
||||
return ip, pod_count, selection
|
||||
|
||||
def _deployment_with_cached_ip(
|
||||
self,
|
||||
|
|
@ -327,10 +454,13 @@ class KubernetesPodDiscovery:
|
|||
key: _CacheKey,
|
||||
url: httpx.URL,
|
||||
now: float,
|
||||
session_id: str | None,
|
||||
retry_count: int,
|
||||
) -> _DeploymentT:
|
||||
ip: Final = self._next_ip(key, now)
|
||||
if ip is None:
|
||||
selection: Final = self._next_ip(key, now, session_id, retry_count)
|
||||
if selection is None:
|
||||
return deployment
|
||||
ip, pod_count, selection_kind = selection
|
||||
if self._proxy_bypasses_only_service_host(key, ip):
|
||||
self._warn_proxy_once(key[0])
|
||||
return deployment
|
||||
|
|
@ -345,6 +475,12 @@ class KubernetesPodDiscovery:
|
|||
{
|
||||
**deployment_mapping,
|
||||
"litellm_params": {**typed_params, "api_base": str(url.copy_with(host=ip))},
|
||||
KUBERNETES_POD_ROUTING_KEY: {
|
||||
"service_host": key[0],
|
||||
"pod_ip": ip,
|
||||
"pod_count": pod_count,
|
||||
"selection": selection_kind,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -402,9 +538,39 @@ _R: Final = TypeVar("_R")
|
|||
_S: Final = TypeVar("_S", bound=_HasPodDiscovery)
|
||||
|
||||
|
||||
def _request_kwargs_from_call(
|
||||
signature: inspect.Signature,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: forwards the call to Signature.bind_partial
|
||||
) -> Mapping[str, object] | None:
|
||||
try:
|
||||
bound_arguments: Final = cast( # cast-ok: Signature arguments use string parameter names
|
||||
Mapping[str, object],
|
||||
signature.bind_partial(*args, **kwargs).arguments,
|
||||
)
|
||||
request_kwargs: Final = bound_arguments.get(
|
||||
"request_kwargs",
|
||||
kwargs.get("request_kwargs"),
|
||||
)
|
||||
except TypeError:
|
||||
return _mapping_request_kwargs(kwargs.get("request_kwargs"))
|
||||
return _mapping_request_kwargs(request_kwargs)
|
||||
|
||||
|
||||
def _mapping_request_kwargs(request_kwargs: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(request_kwargs, Mapping):
|
||||
return None
|
||||
return cast( # cast-ok: router selector request kwargs use string keys
|
||||
Mapping[str, object],
|
||||
request_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def resolve_pods_after(
|
||||
fn: Callable[Concatenate[_S, _P], _R],
|
||||
) -> Callable[Concatenate[_S, _P], _R]:
|
||||
signature: Final = inspect.signature(fn)
|
||||
|
||||
@wraps(fn)
|
||||
def wrapped(
|
||||
self: _S,
|
||||
|
|
@ -412,7 +578,10 @@ def resolve_pods_after(
|
|||
**kwargs: _P.kwargs, # kwargs-ok: ParamSpec forwards the wrapped call signature
|
||||
) -> _R:
|
||||
deployment: Final = fn(self, *args, **kwargs)
|
||||
return self.kubernetes_pod_discovery.resolve_deployment(deployment)
|
||||
return self.kubernetes_pod_discovery.resolve_deployment(
|
||||
deployment,
|
||||
_request_kwargs_from_call(signature, self, *args, **kwargs),
|
||||
)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
|
@ -420,6 +589,8 @@ def resolve_pods_after(
|
|||
def async_resolve_pods_after(
|
||||
fn: Callable[Concatenate[_S, _P], Awaitable[_R]],
|
||||
) -> Callable[Concatenate[_S, _P], Coroutine[object, object, _R]]:
|
||||
signature: Final = inspect.signature(fn)
|
||||
|
||||
@wraps(fn)
|
||||
async def wrapped(
|
||||
self: _S,
|
||||
|
|
@ -427,7 +598,10 @@ def async_resolve_pods_after(
|
|||
**kwargs: _P.kwargs, # kwargs-ok: ParamSpec forwards the wrapped call signature
|
||||
) -> _R:
|
||||
deployment: Final = await fn(self, *args, **kwargs)
|
||||
return await self.kubernetes_pod_discovery.async_resolve_deployment(deployment)
|
||||
return await self.kubernetes_pod_discovery.async_resolve_deployment(
|
||||
deployment,
|
||||
_request_kwargs_from_call(signature, self, *args, **kwargs),
|
||||
)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
|
@ -436,13 +610,18 @@ def resolve_pods_after_bound(
|
|||
fn: Callable[_P, _R],
|
||||
discovery: KubernetesPodDiscovery,
|
||||
) -> Callable[_P, _R]:
|
||||
signature: Final = inspect.signature(fn)
|
||||
|
||||
@wraps(fn)
|
||||
def wrapped(
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs, # kwargs-ok: ParamSpec forwards the wrapped call signature
|
||||
) -> _R:
|
||||
deployment: Final = fn(*args, **kwargs)
|
||||
return discovery.resolve_deployment(deployment)
|
||||
return discovery.resolve_deployment(
|
||||
deployment,
|
||||
_request_kwargs_from_call(signature, *args, **kwargs),
|
||||
)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
|
@ -451,12 +630,17 @@ def async_resolve_pods_after_bound(
|
|||
fn: Callable[_P, Awaitable[_R]],
|
||||
discovery: KubernetesPodDiscovery,
|
||||
) -> Callable[_P, Coroutine[object, object, _R]]:
|
||||
signature: Final = inspect.signature(fn)
|
||||
|
||||
@wraps(fn)
|
||||
async def wrapped(
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs, # kwargs-ok: ParamSpec forwards the wrapped call signature
|
||||
) -> _R:
|
||||
deployment: Final = await fn(*args, **kwargs)
|
||||
return await discovery.async_resolve_deployment(deployment)
|
||||
return await discovery.async_resolve_deployment(
|
||||
deployment,
|
||||
_request_kwargs_from_call(signature, *args, **kwargs),
|
||||
)
|
||||
|
||||
return wrapped
|
||||
|
|
|
|||
|
|
@ -3242,6 +3242,13 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
tier_litellm_params: Mapping[str, object] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
|
||||
|
||||
class StandardLoggingKubernetesPodRouting(TypedDict):
|
||||
service_host: ReadOnly[str]
|
||||
pod_ip: ReadOnly[str]
|
||||
pod_count: ReadOnly[int]
|
||||
selection: ReadOnly[Literal["round_robin", "session_affinity", "session_affinity_retry"]]
|
||||
|
||||
|
||||
# Fields whose values quote the caller's prompt. Dropped when an operator turns message
|
||||
# logging off. Every other field aggregates the prompt without reproducing it and is kept,
|
||||
# so a redacted row stays explainable. `test_every_routing_decision_field_is_classified`
|
||||
|
|
@ -3308,6 +3315,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
|
|||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None
|
||||
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None
|
||||
routing_decision: StandardLoggingRoutingDecision | None
|
||||
kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None]
|
||||
applied_guardrails: list[str] | None
|
||||
usage_object: dict | None
|
||||
cold_storage_object_key: str | None # S3/GCS object key for cold storage retrieval
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
"user": "",
|
||||
"team_id": "",
|
||||
"organization_id": "",
|
||||
"metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
|
||||
"metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}",
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.00022500000000000002,
|
||||
"total_tokens": 30,
|
||||
|
|
|
|||
|
|
@ -51,7 +51,7 @@ from litellm.proxy import proxy_server
|
|||
from litellm.proxy._types import SpanAttributes
|
||||
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
|
||||
from collections.abc import AsyncIterator
|
||||
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY, LOGGING_WORKER_MAX_TIME_PER_COROUTINE
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Function
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
|
@ -363,6 +363,33 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada
|
|||
assert finished.attributes["metadata.routing_decision"] == str(decision)
|
||||
|
||||
|
||||
def test_v1_pod_routing_none_matches_other_none_metadata_attributes() -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
provider: Final = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
otel: Final = OpenTelemetry(tracer_provider=provider)
|
||||
with provider.get_tracer(__name__).start_as_current_span("inference") as span:
|
||||
otel.set_attributes(
|
||||
span,
|
||||
{
|
||||
"model": "answer",
|
||||
"standard_logging_object": {
|
||||
"id": "routing-call",
|
||||
"call_type": "acompletion",
|
||||
"metadata": {
|
||||
KUBERNETES_POD_ROUTING_KEY: None,
|
||||
"other_none": None,
|
||||
},
|
||||
},
|
||||
},
|
||||
{"model": "answer", "choices": [], "usage": {}},
|
||||
)
|
||||
(finished,) = exporter.get_finished_spans()
|
||||
routing_attribute: Final = finished.attributes["metadata.kubernetes_pod_routing"]
|
||||
assert routing_attribute == finished.attributes["metadata.other_none"]
|
||||
assert routing_attribute != "null"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decision", ["custom", 42, ["custom"]])
|
||||
def test_v1_malformed_routing_metadata_preserves_inference_attributes(decision: object) -> None:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
|
|
|
|||
|
|
@ -30,7 +30,12 @@ from litellm._internal_context import in_post_response_phase
|
|||
from litellm._logging import session_id_var, trace_id_var, verbose_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST
|
||||
from litellm.constants import (
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
|
||||
REDACTED_BY_LITELLM,
|
||||
SENTRY_PII_DENYLIST,
|
||||
)
|
||||
from litellm.cost_calculator import ocr_batch_cost
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
|
|
@ -5713,6 +5718,97 @@ def test_get_standard_logging_object_payload_takes_used_client_oauth_token_from_
|
|||
assert payload["metadata"]["used_client_oauth_token"] is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "litellm_metadata", "expected"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
}
|
||||
},
|
||||
{},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
},
|
||||
),
|
||||
(
|
||||
{},
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.2",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
}
|
||||
},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.2",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
},
|
||||
),
|
||||
(
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
}
|
||||
},
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.3",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity_retry",
|
||||
}
|
||||
},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.3",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity_retry",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_standard_logging_payload_resolves_kubernetes_pod_routing_from_metadata_buckets(
|
||||
logging_obj: LitellmLogging,
|
||||
metadata: dict[str, object],
|
||||
litellm_metadata: dict[str, object],
|
||||
expected: dict[str, object],
|
||||
) -> None:
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload
|
||||
|
||||
now: Final = datetime(2026, 1, 1)
|
||||
payload: Final = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o",
|
||||
"messages": [],
|
||||
"litellm_params": {"metadata": metadata, "litellm_metadata": litellm_metadata},
|
||||
},
|
||||
init_response_obj={},
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["metadata"][KUBERNETES_POD_ROUTING_KEY] == expected
|
||||
|
||||
|
||||
def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj):
|
||||
"""Access groups stamped at auth time reach the logging payload, so integrations see what a request billed."""
|
||||
from datetime import datetime
|
||||
|
|
|
|||
|
|
@ -3748,7 +3748,18 @@ class TestSpendLogsPayload:
|
|||
}
|
||||
)
|
||||
|
||||
differences = _compare_nested_dicts(payload, expected_payload, ignore_keys=ignored_keys)
|
||||
expected_payload_with_pod_routing: Final[SpendLogsPayload] = {
|
||||
**expected_payload,
|
||||
"metadata": expected_payload["metadata"].replace(
|
||||
'"routing_decision": null, ',
|
||||
'"routing_decision": null, "kubernetes_pod_routing": null, ',
|
||||
),
|
||||
}
|
||||
differences = _compare_nested_dicts(
|
||||
payload,
|
||||
expected_payload_with_pod_routing,
|
||||
ignore_keys=ignored_keys,
|
||||
)
|
||||
if differences:
|
||||
pytest.fail(f"Dictionary mismatch: {differences}")
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
import litellm
|
||||
import litellm.constants as litellm_constants
|
||||
import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils
|
||||
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo
|
||||
|
|
@ -76,6 +76,89 @@ def _get_additional_usage_values_for_usage(usage: litellm.Usage) -> dict:
|
|||
return metadata["additional_usage_values"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "litellm_metadata", "expected"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
}
|
||||
},
|
||||
{},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
},
|
||||
),
|
||||
(
|
||||
{},
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.2",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
}
|
||||
},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.2",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
},
|
||||
),
|
||||
(
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
}
|
||||
},
|
||||
{
|
||||
"kubernetes_pod_routing": {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.3",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity_retry",
|
||||
}
|
||||
},
|
||||
{
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.3",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity_retry",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_logging_payload_includes_kubernetes_pod_routing(
|
||||
metadata: dict[str, object],
|
||||
litellm_metadata: dict[str, object],
|
||||
expected: dict[str, object],
|
||||
) -> None:
|
||||
now: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_params": {"metadata": metadata, "litellm_metadata": litellm_metadata},
|
||||
},
|
||||
response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
assert json.loads(payload["metadata"])[KUBERNETES_POD_ROUTING_KEY] == expected
|
||||
assert _get_spend_logs_metadata(None)["kubernetes_pod_routing"] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("store_prompts,redact", [(True, False), (False, False), (True, True)])
|
||||
def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch, store_prompts, redact):
|
||||
from litellm.proxy import proxy_server
|
||||
|
|
|
|||
|
|
@ -3309,12 +3309,11 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain
|
|||
assert data["litellm_trace_id"] == "foo"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)])
|
||||
@pytest.mark.parametrize("correlation_in_logs", [True, False], ids=["correlation-on", "correlation-off"])
|
||||
@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"])
|
||||
def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only(
|
||||
def test_x_litellm_trace_id_keeps_the_body_session_beside_it_whatever_the_log_setting(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
correlation_in_logs: bool,
|
||||
expected_spend_log_session: str | None,
|
||||
session_in_metadata: bool,
|
||||
):
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs)
|
||||
|
|
@ -3325,7 +3324,7 @@ def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only(
|
|||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session
|
||||
assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == "convo-1"
|
||||
assert data["litellm_session_id"] == "req-trace-1"
|
||||
assert metadata["session_id"] == "req-trace-1"
|
||||
assert data["litellm_trace_id"] == "req-trace-1"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
|
|
@ -16,10 +17,22 @@ import respx
|
|||
from openai import AsyncOpenAI, OpenAI
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
)
|
||||
from litellm.router import CustomRoutingStrategyBase, Router
|
||||
from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery
|
||||
from litellm.router_utils.kubernetes_pod_discovery import (
|
||||
KubernetesPodDiscovery,
|
||||
async_resolve_pods_after,
|
||||
async_resolve_pods_after_bound,
|
||||
resolve_pods_after,
|
||||
resolve_pods_after_bound,
|
||||
)
|
||||
|
||||
_SERVICE_URL: Final = "http://vllm-headless.ns.svc.cluster.local:8000/v1"
|
||||
_THREE_POD_CHAT_PATTERN: Final = re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2|10\.0\.0\.3):8000/v1/chat/completions")
|
||||
_SocketAddress: TypeAlias = tuple[str, int] | tuple[str, int, int, int]
|
||||
_AddrInfo: TypeAlias = tuple[socket.AddressFamily, socket.SocketKind, int, str, _SocketAddress]
|
||||
_NO_POD_ERRNOS: Final = tuple(sorted((socket.EAI_NONAME, socket.EAI_NODATA)))
|
||||
|
|
@ -91,6 +104,14 @@ def _empty_proxy_environment() -> Mapping[str, str]:
|
|||
return {}
|
||||
|
||||
|
||||
def _patch_router_discovery(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*discoveries: KubernetesPodDiscovery,
|
||||
) -> None:
|
||||
discovery_iter: Final = iter(discoveries)
|
||||
monkeypatch.setattr(litellm.router, "KubernetesPodDiscovery", lambda: next(discovery_iter))
|
||||
|
||||
|
||||
def _stub_sync_dns(monkeypatch: pytest.MonkeyPatch, *ips: str) -> None:
|
||||
def getaddrinfo(
|
||||
host: str | None,
|
||||
|
|
@ -113,12 +134,88 @@ def _api_base(deployment: Mapping[str, object]) -> str:
|
|||
return api_base
|
||||
|
||||
|
||||
def _ranked_session_ips(session_id: str) -> tuple[str, ...]:
|
||||
ips: Final = ("10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
return tuple(
|
||||
sorted(
|
||||
ips,
|
||||
key=lambda ip: (
|
||||
-int.from_bytes(hashlib.blake2b(f"{session_id}\x00{ip}".encode(), digest_size=8).digest(), "big"),
|
||||
ip,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_session_hash_ranking_matches_full_hash_for_large_session_id(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
ips: Final = ("10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
_stub_sync_dns(monkeypatch, *ips)
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
|
||||
for session_id in ("s1", "session-2", "session-3"):
|
||||
expected: Final = _ranked_session_ips(session_id)
|
||||
actual: Final = tuple(
|
||||
httpx.URL(
|
||||
_api_base(
|
||||
discovery.resolve_deployment(
|
||||
_deployment(),
|
||||
{"metadata": {"session_id": session_id, "request_retry_count": retry_count}},
|
||||
)
|
||||
)
|
||||
).host
|
||||
for retry_count in range(len(ips))
|
||||
)
|
||||
assert actual == expected
|
||||
|
||||
large_ips: Final = tuple(f"10.1.0.{index}" for index in range(1, 101))
|
||||
_stub_sync_dns(monkeypatch, *large_ips)
|
||||
large_discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
large_session_id: Final = "x" * (1024 * 1024)
|
||||
|
||||
result: Final = large_discovery.resolve_deployment(
|
||||
_deployment(),
|
||||
{"metadata": {"session_id": large_session_id}},
|
||||
)
|
||||
|
||||
assert httpx.URL(_api_base(result)).host in large_ips
|
||||
|
||||
|
||||
def _request_body(content: bytes) -> Mapping[str, object]:
|
||||
body: Final = json.loads(content)
|
||||
assert isinstance(body, dict)
|
||||
return cast(Mapping[str, object], body)
|
||||
|
||||
|
||||
def _cache_async_client(router: Router) -> AsyncOpenAI:
|
||||
client: Final = AsyncOpenAI(api_key="fake", base_url=_SERVICE_URL)
|
||||
router.cache.set_cache(key="registered-id_async_client", value=client, local_only=True)
|
||||
return client
|
||||
|
||||
|
||||
def _cache_sync_client(router: Router) -> OpenAI:
|
||||
client: Final = OpenAI(api_key="fake", base_url=_SERVICE_URL)
|
||||
router.cache.set_cache(key="registered-id_client", value=client, local_only=True)
|
||||
return client
|
||||
|
||||
|
||||
async def _router_acompletion(router: Router, **request_kwargs: object) -> None:
|
||||
await router.acompletion(
|
||||
model="gpu-model",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _router_completion(router: Router, **request_kwargs: object) -> None:
|
||||
router.completion(
|
||||
model="gpu-model",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def test_sync_resolution_round_robins_sorted_ips_without_mutating_deployment(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
@ -1050,19 +1147,17 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag(
|
|||
monkeypatch.setattr(loop, "getaddrinfo", async_getaddrinfo)
|
||||
monkeypatch.setattr(socket, "getaddrinfo", sync_getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=re.compile(
|
||||
r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions"
|
||||
)
|
||||
).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE))
|
||||
discovery_enabled: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
|
||||
control: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery_enabled, control)
|
||||
discovery_enabled_router: Final = Router(model_list=[_deployment()])
|
||||
control_router: Final = Router(model_list=[_deployment(kubernetes_pod_discovery=None)])
|
||||
discovery_enabled_router.kubernetes_pod_discovery = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
cached_async_client: Final = AsyncOpenAI(api_key="fake", base_url=_SERVICE_URL)
|
||||
cached_sync_client: Final = OpenAI(api_key="fake", base_url=_SERVICE_URL)
|
||||
discovery_enabled_router.cache.set_cache(
|
||||
|
|
@ -1100,6 +1195,506 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag(
|
|||
assert next(sync_lookups) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_session_requests_stick_to_one_pod_for_sync_and_async(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
async_lookups: Final = count()
|
||||
sync_lookups: Final = count()
|
||||
|
||||
async def async_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
next(async_lookups)
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
def sync_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
next(sync_lookups)
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", async_getaddrinfo)
|
||||
monkeypatch.setattr(socket, "getaddrinfo", sync_getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
async_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
sync_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
|
||||
async_router: Final = Router(model_list=[_deployment()])
|
||||
sync_router: Final = Router(model_list=[_deployment()])
|
||||
async_client: Final = _cache_async_client(async_router)
|
||||
sync_client: Final = _cache_sync_client(sync_router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(async_router, metadata={"session_id": "s1"})
|
||||
for _ in range(6):
|
||||
_router_completion(sync_router, metadata={"session_id": "s1"})
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await async_client.close()
|
||||
sync_client.close()
|
||||
|
||||
assert len(set(hosts[:6])) == 1
|
||||
assert len(set(hosts[6:])) == 1
|
||||
assert next(async_lookups) == 1
|
||||
assert next(sync_lookups) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_session_mapping_repeats_across_pods(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
lookups: Final = count()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
next(lookups)
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
session_ids: Final = tuple(f"session-{index}" for index in range(30))
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for session_id in session_ids:
|
||||
await _router_acompletion(router, metadata={"session_id": session_id})
|
||||
await _router_acompletion(router)
|
||||
for session_id in session_ids:
|
||||
await _router_acompletion(router, metadata={"session_id": session_id})
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
first_pass: Final = hosts[: len(session_ids)]
|
||||
second_pass: Final = hosts[len(session_ids) + 1 :]
|
||||
assert len(set(first_pass)) > 1
|
||||
assert hosts[len(session_ids)] == "10.0.0.1"
|
||||
assert second_pass == first_pass
|
||||
assert next(lookups) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_session_id_uses_round_robin_pod_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(
|
||||
router,
|
||||
metadata={
|
||||
"session_id": "generated-session",
|
||||
SESSION_ID_GENERATED_METADATA_KEY: True,
|
||||
},
|
||||
)
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert hosts == ("10.0.0.1", "10.0.0.2", "10.0.0.3") * 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_session_id_uses_round_robin_pod_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(router, metadata={"session_id": ""})
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert hosts == ("10.0.0.1", "10.0.0.2", "10.0.0.3") * 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_requests_do_not_advance_round_robin_cursor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
await _router_acompletion(router)
|
||||
await _router_acompletion(router, metadata={"session_id": "s1"})
|
||||
await _router_acompletion(router, metadata={"session_id": "s1"})
|
||||
await _router_acompletion(router)
|
||||
await _router_acompletion(router, metadata={"session_id": "s1"})
|
||||
await _router_acompletion(router)
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert tuple(hosts[index] for index in (0, 3, 5)) == (
|
||||
"10.0.0.1",
|
||||
"10.0.0.2",
|
||||
"10.0.0.3",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_pod_membership_change_remaps_only_affected_sessions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
lookups: Final = count()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
lookup_number: Final = next(lookups)
|
||||
ips: Final = ("10.0.0.1", "10.0.0.2", "10.0.0.3") if lookup_number == 0 else ("10.0.0.2", "10.0.0.3")
|
||||
return _records(*ips, port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
session_ids: Final = tuple(f"session-{index}" for index in range(60))
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(
|
||||
refresh_interval_seconds=10,
|
||||
clock=_clock((0.0,) * 62 + (10.0,) * 10),
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for session_id in session_ids:
|
||||
await _router_acompletion(router, metadata={"session_id": session_id})
|
||||
initial_hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
removed_session_ids: Final = tuple(
|
||||
session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.1"
|
||||
)
|
||||
surviving_session_ids: Final = tuple(
|
||||
session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.2"
|
||||
)
|
||||
assert removed_session_ids
|
||||
assert len(surviving_session_ids) >= 2
|
||||
await _router_acompletion(router)
|
||||
await _router_acompletion(router, metadata={"session_id": surviving_session_ids[0]})
|
||||
await _router_acompletion(router, metadata={"session_id": removed_session_ids[0]})
|
||||
await _router_acompletion(router, metadata={"session_id": surviving_session_ids[1]})
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert hosts[60] == "10.0.0.1"
|
||||
assert hosts[61] == "10.0.0.2"
|
||||
assert hosts[62] in ("10.0.0.2", "10.0.0.3")
|
||||
assert hosts[63] == "10.0.0.2"
|
||||
assert next(lookups) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_routing_strategy_keeps_sync_and_async_sessions_on_one_pod(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
async_lookups: Final = count()
|
||||
sync_lookups: Final = count()
|
||||
|
||||
async def async_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
next(async_lookups)
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
def sync_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
next(sync_lookups)
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", async_getaddrinfo)
|
||||
monkeypatch.setattr(socket, "getaddrinfo", sync_getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
async_discovery: Final = KubernetesPodDiscovery(
|
||||
clock=lambda: 0.0,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
sync_discovery: Final = KubernetesPodDiscovery(
|
||||
clock=lambda: 0.0,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
|
||||
async_router: Final = Router(model_list=[_deployment()])
|
||||
sync_router: Final = Router(model_list=[_deployment()])
|
||||
async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router))
|
||||
sync_router.set_custom_routing_strategy(_ModelListRoutingStrategy(sync_router))
|
||||
async_client: Final = _cache_async_client(async_router)
|
||||
sync_client: Final = _cache_sync_client(sync_router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(async_router, metadata={"session_id": "custom-session"})
|
||||
for _ in range(6):
|
||||
_router_completion(sync_router, metadata={"session_id": "custom-session"})
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await async_client.close()
|
||||
sync_client.close()
|
||||
|
||||
assert len(set(hosts[:6])) == 1
|
||||
assert len(set(hosts[6:])) == 1
|
||||
assert next(async_lookups) == 1
|
||||
assert next(sync_lookups) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_top_level_litellm_session_id_keeps_requests_on_one_pod(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(router, litellm_session_id="s1")
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert len(set(hosts)) == 1
|
||||
|
||||
|
||||
def test_caller_session_beside_trace_header_keeps_the_pod_across_trace_ids(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
caller_session_id: Final = "chat-session-a"
|
||||
caller_pod: Final = _ranked_session_ips(caller_session_id)[0]
|
||||
trace_ids: Final = tuple(
|
||||
trace_id
|
||||
for trace_id in (f"trace-{turn}" for turn in range(20))
|
||||
if _ranked_session_ips(trace_id)[0] != caller_pod
|
||||
)[:3]
|
||||
assert len(trace_ids) == 3
|
||||
|
||||
pods: Final = tuple(
|
||||
httpx.URL(
|
||||
_api_base(
|
||||
discovery.resolve_deployment(
|
||||
_deployment(),
|
||||
{
|
||||
"litellm_session_id": trace_id,
|
||||
"metadata": {CALLER_SESSION_ID_METADATA_KEY: caller_session_id},
|
||||
},
|
||||
)
|
||||
)
|
||||
).host
|
||||
for trace_id in trace_ids
|
||||
)
|
||||
|
||||
assert pods == (caller_pod, caller_pod, caller_pod)
|
||||
|
||||
|
||||
def test_proxy_headers_with_trace_and_session_id_pin_the_pod_without_log_correlation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
||||
monkeypatch.setattr(litellm, "request_correlation_in_logs", False)
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
caller_session_id: Final = "chat-session-a"
|
||||
caller_pod: Final = _ranked_session_ips(caller_session_id)[0]
|
||||
trace_ids: Final = tuple(
|
||||
trace_id
|
||||
for trace_id in (f"trace-{turn}" for turn in range(20))
|
||||
if _ranked_session_ips(trace_id)[0] != caller_pod
|
||||
)[:3]
|
||||
assert len(trace_ids) == 3
|
||||
|
||||
def resolve_turn(trace_id: str) -> str:
|
||||
data: Final = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
{"x-litellm-trace-id": trace_id, "x-litellm-session-id": caller_session_id},
|
||||
{"model": "gpu-model", "metadata": {}},
|
||||
"metadata",
|
||||
)
|
||||
return httpx.URL(_api_base(discovery.resolve_deployment(_deployment(), data))).host
|
||||
|
||||
assert tuple(resolve_turn(trace_id) for trace_id in trace_ids) == (caller_pod, caller_pod, caller_pod)
|
||||
|
||||
|
||||
def test_top_level_litellm_session_id_precedes_metadata_session_id(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
top_level_session_id: Final = "s1"
|
||||
metadata_session_id: Final = "metadata-session"
|
||||
top_level_kwargs: Final = {"litellm_session_id": top_level_session_id}
|
||||
metadata_kwargs: Final = {"metadata": {"session_id": metadata_session_id}}
|
||||
combined_kwargs: Final = {
|
||||
"litellm_session_id": top_level_session_id,
|
||||
"metadata": {"session_id": metadata_session_id},
|
||||
}
|
||||
|
||||
assert _ranked_session_ips(top_level_session_id)[0] != _ranked_session_ips(metadata_session_id)[0]
|
||||
|
||||
top_level_pod: Final = httpx.URL(_api_base(discovery.resolve_deployment(_deployment(), top_level_kwargs))).host
|
||||
metadata_pod: Final = httpx.URL(_api_base(discovery.resolve_deployment(_deployment(), metadata_kwargs))).host
|
||||
combined_pod: Final = httpx.URL(_api_base(discovery.resolve_deployment(_deployment(), combined_kwargs))).host
|
||||
|
||||
assert top_level_pod != metadata_pod
|
||||
assert combined_pod == top_level_pod
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_metadata_ignores_top_level_litellm_session_id(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
|
||||
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
|
||||
)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()])
|
||||
client: Final = _cache_async_client(router)
|
||||
for _ in range(6):
|
||||
await _router_acompletion(
|
||||
router,
|
||||
litellm_session_id="s1",
|
||||
metadata={
|
||||
"session_id": "s1",
|
||||
SESSION_ID_GENERATED_METADATA_KEY: True,
|
||||
},
|
||||
)
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
await client.close()
|
||||
|
||||
assert hosts == ("10.0.0.1", "10.0.0.2", "10.0.0.3") * 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_requests(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -1134,23 +1729,23 @@ async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_req
|
|||
monkeypatch.setattr(loop, "getaddrinfo", async_getaddrinfo)
|
||||
monkeypatch.setattr(socket, "getaddrinfo", sync_getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=re.compile(
|
||||
r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions"
|
||||
)
|
||||
).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE))
|
||||
async_discovery: Final = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
sync_discovery: Final = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
|
||||
async_router: Final = Router(model_list=[_deployment()])
|
||||
sync_router: Final = Router(model_list=[_deployment()])
|
||||
async_router.kubernetes_pod_discovery = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
sync_router.kubernetes_pod_discovery = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router))
|
||||
sync_router.set_custom_routing_strategy(_ModelListRoutingStrategy(sync_router))
|
||||
for _ in range(4):
|
||||
|
|
@ -1202,16 +1797,13 @@ async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.M
|
|||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions")
|
||||
).mock(side_effect=response_for)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()], num_retries=1)
|
||||
router.kubernetes_pod_discovery = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
|
||||
response: Final = await router.acompletion(
|
||||
model="gpu-model",
|
||||
|
|
@ -1245,16 +1837,13 @@ def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPa
|
|||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions")
|
||||
).mock(side_effect=response_for)
|
||||
discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()], num_retries=1)
|
||||
router.kubernetes_pod_discovery = KubernetesPodDiscovery(
|
||||
clock=clock,
|
||||
proxy_environment=_empty_proxy_environment,
|
||||
)
|
||||
|
||||
response: Final = router.completion(
|
||||
model="gpu-model",
|
||||
|
|
@ -1264,3 +1853,240 @@ def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPa
|
|||
|
||||
assert response.choices[0].message.content == "ok"
|
||||
assert hosts == ("10.0.0.1", "10.0.0.2")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("litellm_retry_count", "metadata_retry_count", "expected_rank"),
|
||||
[
|
||||
(0, None, 0),
|
||||
(None, None, 0),
|
||||
(True, None, 0),
|
||||
("1", None, 0),
|
||||
(-1, None, 0),
|
||||
(1, None, 1),
|
||||
(True, 1, 1),
|
||||
(1, 0, 1),
|
||||
],
|
||||
)
|
||||
def test_session_retry_count_selects_ranked_pod(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
litellm_retry_count: object,
|
||||
metadata_retry_count: object,
|
||||
expected_rank: int,
|
||||
) -> None:
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
request_kwargs: Final = {
|
||||
"litellm_metadata": {
|
||||
"session_id": "s1",
|
||||
**({"request_retry_count": litellm_retry_count} if litellm_retry_count is not None else {}),
|
||||
},
|
||||
"metadata": ({"request_retry_count": metadata_retry_count} if metadata_retry_count is not None else {}),
|
||||
}
|
||||
|
||||
resolved: Final = discovery.resolve_deployment(_deployment(), request_kwargs)
|
||||
|
||||
assert httpx.URL(_api_base(resolved)).host == _ranked_session_ips("s1")[expected_rank]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_session_retry_uses_next_rendezvous_pod(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
attempts: Final = count()
|
||||
|
||||
async def getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
def response_for(request: httpx.Request) -> httpx.Response:
|
||||
attempt: Final = next(attempts)
|
||||
status_code: Final = 500 if attempt == 0 else 200
|
||||
body: Final = {"error": {"message": "retryable", "type": "server_error"}} if attempt == 0 else _CHAT_RESPONSE
|
||||
return httpx.Response(status_code, json=body, request=request)
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
session_id: Final = "retry-session"
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(side_effect=response_for)
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
_patch_router_discovery(monkeypatch, discovery)
|
||||
router: Final = Router(model_list=[_deployment()], num_retries=1)
|
||||
|
||||
response: Final = await router.acompletion(
|
||||
model="gpu-model",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
metadata={"session_id": session_id},
|
||||
)
|
||||
hosts: Final = tuple(call.request.url.host for call in route.calls)
|
||||
|
||||
assert response.choices[0].message.content == "ok"
|
||||
assert hosts == _ranked_session_ips(session_id)[:2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_positional_request_kwargs_reach_all_pod_resolver_decorators(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
||||
async def async_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
def sync_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
return _records("10.0.0.1", "10.0.0.2", "10.0.0.3", port=int(port or 0))
|
||||
|
||||
class _Selector:
|
||||
def __init__(self, discovery: KubernetesPodDiscovery) -> None:
|
||||
self.kubernetes_pod_discovery: Final = discovery
|
||||
|
||||
@resolve_pods_after
|
||||
def sync(self, deployment: dict[str, object], request_kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
return deployment
|
||||
|
||||
@async_resolve_pods_after
|
||||
async def async_call(
|
||||
self,
|
||||
deployment: dict[str, object],
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
return deployment
|
||||
|
||||
class _BoundSelector:
|
||||
def sync(self, deployment: dict[str, object], request_kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
return deployment
|
||||
|
||||
async def async_call(
|
||||
self,
|
||||
deployment: dict[str, object],
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
return deployment
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", async_getaddrinfo)
|
||||
monkeypatch.setattr(socket, "getaddrinfo", sync_getaddrinfo)
|
||||
session_id: Final = next(
|
||||
candidate
|
||||
for candidate in (f"positional-{index}" for index in range(100))
|
||||
if _ranked_session_ips(candidate)[0] != "10.0.0.1"
|
||||
)
|
||||
request_kwargs: Final = {"metadata": {"session_id": session_id}}
|
||||
sync_selector: Final = _Selector(KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment))
|
||||
async_selector: Final = _Selector(KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment))
|
||||
bound_selector: Final = _BoundSelector()
|
||||
bound_sync: Final = resolve_pods_after_bound(
|
||||
bound_selector.sync,
|
||||
KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment),
|
||||
)
|
||||
bound_async: Final = async_resolve_pods_after_bound(
|
||||
bound_selector.async_call,
|
||||
KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment),
|
||||
)
|
||||
|
||||
sync_result: Final = sync_selector.sync(_deployment(), request_kwargs)
|
||||
async_result: Final = await async_selector.async_call(_deployment(), request_kwargs)
|
||||
bound_sync_result: Final = bound_sync(_deployment(), request_kwargs)
|
||||
bound_async_result: Final = await bound_async(_deployment(), request_kwargs)
|
||||
|
||||
assert (
|
||||
tuple(
|
||||
httpx.URL(_api_base(result)).host
|
||||
for result in (sync_result, async_result, bound_sync_result, bound_async_result)
|
||||
)
|
||||
== (_ranked_session_ips(session_id)[0],) * 4
|
||||
)
|
||||
|
||||
|
||||
def test_direct_session_id_keyword_does_not_enable_session_affinity(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
request_kwargs: Final = {"session_id": "s1", "messages": [{"role": "user", "content": "hello"}]}
|
||||
|
||||
selected_ips: Final = tuple(
|
||||
httpx.URL(_api_base(discovery.resolve_deployment(_deployment(), request_kwargs))).host for _ in range(3)
|
||||
)
|
||||
|
||||
assert selected_ips == ("10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_kwargs", "selection"),
|
||||
[
|
||||
({}, "round_robin"),
|
||||
({"metadata": {"session_id": "s1"}}, "session_affinity"),
|
||||
({"metadata": {"session_id": "s1", "request_retry_count": 1}}, "session_affinity_retry"),
|
||||
],
|
||||
)
|
||||
def test_resolved_deployment_records_pod_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_kwargs: dict[str, object],
|
||||
selection: str,
|
||||
) -> None:
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1", "10.0.0.2", "10.0.0.3")
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
|
||||
|
||||
resolved: Final = discovery.resolve_deployment(_deployment(), request_kwargs)
|
||||
pod_ip: Final = httpx.URL(_api_base(resolved)).host
|
||||
|
||||
assert resolved[KUBERNETES_POD_ROUTING_KEY] == {
|
||||
"service_host": httpx.URL(_SERVICE_URL).host,
|
||||
"pod_ip": pod_ip,
|
||||
"pod_count": 3,
|
||||
"selection": selection,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fallback", ["https", "no_proxy", "empty_ips", "dns_failure"])
|
||||
def test_unresolved_deployment_does_not_record_pod_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
fallback: str,
|
||||
) -> None:
|
||||
if fallback == "empty_ips":
|
||||
_stub_sync_dns(monkeypatch)
|
||||
if fallback == "dns_failure":
|
||||
|
||||
def failing_getaddrinfo(
|
||||
host: str | None,
|
||||
port: str | int | None,
|
||||
family: int = 0,
|
||||
type: int = 0,
|
||||
proto: int = 0,
|
||||
flags: int = 0,
|
||||
) -> list[_AddrInfo]:
|
||||
raise socket.gaierror(socket.EAI_AGAIN, "temporary failure")
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", failing_getaddrinfo)
|
||||
if fallback == "no_proxy":
|
||||
_stub_sync_dns(monkeypatch, "10.0.0.1")
|
||||
|
||||
proxy_environment: Final = (
|
||||
_proxy_environment("vllm-headless.ns.svc.cluster.local") if fallback == "no_proxy" else _empty_proxy_environment
|
||||
)
|
||||
api_base: Final = "https://vllm-headless.ns.svc.cluster.local:8000/v1" if fallback == "https" else _SERVICE_URL
|
||||
discovery: Final = KubernetesPodDiscovery(proxy_environment=proxy_environment)
|
||||
deployment: Final = _deployment(api_base=api_base)
|
||||
|
||||
resolved: Final = discovery.resolve_deployment(deployment)
|
||||
|
||||
assert resolved is deployment
|
||||
assert KUBERNETES_POD_ROUTING_KEY not in resolved
|
||||
|
|
|
|||
|
|
@ -9978,6 +9978,94 @@ def test_update_kwargs_with_deployment_model_info_in_metadata():
|
|||
assert model_info["output_cost_per_token"] == 0.0015
|
||||
|
||||
|
||||
def test_update_kwargs_with_deployment_clears_pod_routing_on_non_discovery_fallback() -> None:
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY
|
||||
|
||||
routing_record: Final = {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
}
|
||||
discovery_deployment: Final = {
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"},
|
||||
"model_info": {"id": "discovery-id"},
|
||||
KUBERNETES_POD_ROUTING_KEY: routing_record,
|
||||
}
|
||||
fallback_deployment: Final = {
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"},
|
||||
"model_info": {"id": "fallback-id"},
|
||||
}
|
||||
router: Final = litellm.Router(model_list=[discovery_deployment])
|
||||
kwargs: Final = {"metadata": {}}
|
||||
|
||||
router._update_kwargs_with_deployment(deployment=discovery_deployment, kwargs=kwargs)
|
||||
assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record
|
||||
|
||||
router._update_kwargs_with_deployment(deployment=fallback_deployment, kwargs=kwargs)
|
||||
assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None
|
||||
|
||||
|
||||
def test_update_kwargs_with_deployment_synchronizes_existing_litellm_metadata_pod_routing() -> None:
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY
|
||||
|
||||
routing_record: Final = {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
}
|
||||
spoofed_record: Final = {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.9",
|
||||
"pod_count": 3,
|
||||
"selection": "round_robin",
|
||||
}
|
||||
deployment: Final = {
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"},
|
||||
"model_info": {"id": "discovery-id"},
|
||||
KUBERNETES_POD_ROUTING_KEY: routing_record,
|
||||
}
|
||||
original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: spoofed_record}
|
||||
kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata}
|
||||
router: Final = litellm.Router(model_list=[deployment])
|
||||
|
||||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion")
|
||||
|
||||
assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record
|
||||
assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record
|
||||
assert kwargs["litellm_metadata"] is not original_litellm_metadata
|
||||
assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == spoofed_record
|
||||
|
||||
|
||||
def test_update_kwargs_with_non_discovery_deployment_clears_other_pod_routing_bucket() -> None:
|
||||
from litellm.constants import KUBERNETES_POD_ROUTING_KEY
|
||||
|
||||
previous_record: Final = {
|
||||
"service_host": "vllm-headless.ns.svc.cluster.local",
|
||||
"pod_ip": "10.0.0.1",
|
||||
"pod_count": 3,
|
||||
"selection": "session_affinity",
|
||||
}
|
||||
deployment: Final = {
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"},
|
||||
"model_info": {"id": "fallback-id"},
|
||||
}
|
||||
original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: previous_record}
|
||||
kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata}
|
||||
router: Final = litellm.Router(model_list=[deployment])
|
||||
|
||||
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion")
|
||||
|
||||
assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None
|
||||
assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] is None
|
||||
assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == previous_record
|
||||
|
||||
|
||||
def test_combine_fallback_usage():
|
||||
"""Test that _combine_fallback_usage merges partial and fallback usage."""
|
||||
from litellm.router import Router
|
||||
|
|
|
|||
|
|
@ -33,6 +33,12 @@ export const CREDENTIAL_LABELS: Record<string, string> = {
|
|||
false: "Configured key",
|
||||
};
|
||||
|
||||
export const KUBERNETES_POD_ROUTING_LABELS = {
|
||||
round_robin: "Round robin",
|
||||
session_affinity: "Session affinity",
|
||||
session_affinity_retry: "Session affinity (retry)",
|
||||
} as const;
|
||||
|
||||
export const QUICK_SELECT_OPTIONS: { label: string; value: number; unit: string }[] = [
|
||||
{ label: "Last Minute", value: 1, unit: "minutes" },
|
||||
{ label: "Last 15 Minutes", value: 15, unit: "minutes" },
|
||||
|
|
|
|||
|
|
@ -431,6 +431,37 @@ describe("LogDetailContent", () => {
|
|||
expect(screen.queryByText("Configured key")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ selection: "round_robin", label: "Round robin" },
|
||||
{ selection: "session_affinity", label: "Session affinity" },
|
||||
{ selection: "session_affinity_retry", label: "Session affinity (retry)" },
|
||||
] as const)("shows the selected Kubernetes pod and routing mode", ({ selection, label }) => {
|
||||
render(
|
||||
<LogDetailContent
|
||||
logEntry={createLogEntry({
|
||||
metadata: {
|
||||
status: "success",
|
||||
kubernetes_pod_routing: {
|
||||
service_host: "vllm-headless.ns.svc.cluster.local",
|
||||
pod_ip: "10.0.0.2",
|
||||
pod_count: 3,
|
||||
selection,
|
||||
},
|
||||
},
|
||||
})}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Pod")).toBeInTheDocument();
|
||||
expect(screen.getByText(`10.0.0.2 · ${label} · 3 pods`)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("omits the Pod row when routing metadata is absent", () => {
|
||||
render(<LogDetailContent logEntry={createLogEntry({ metadata: { status: "success" } })} />);
|
||||
|
||||
expect(screen.queryByText("Pod")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display guardrail label when guardrail data exists", () => {
|
||||
render(
|
||||
<LogDetailContent
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ import {
|
|||
import { CostBreakdownViewer } from "./sections/CostBreakdownViewer";
|
||||
import { ConfigInfoMessage } from "./sections/ConfigInfoMessage";
|
||||
import { VectorStoreViewer } from "./sections/VectorStoreViewer";
|
||||
import { CREDENTIAL_LABELS } from "../constants";
|
||||
import { CREDENTIAL_LABELS, KUBERNETES_POD_ROUTING_LABELS } from "../constants";
|
||||
import { TruncatedValue } from "./TruncatedValue";
|
||||
import { TokenFlow } from "./TokenFlow";
|
||||
import { JsonViewer } from "./JsonViewer";
|
||||
|
|
@ -75,6 +75,7 @@ export function LogDetailContent({
|
|||
userEmail,
|
||||
}: LogDetailContentProps) {
|
||||
const metadata = logEntry.metadata || {};
|
||||
const podRouting = metadata.kubernetes_pod_routing;
|
||||
const hasError = metadata.status === "failure";
|
||||
const errorInfo = hasError ? metadata.error_information : null;
|
||||
const isClassifier =
|
||||
|
|
@ -160,6 +161,15 @@ export function LogDetailContent({
|
|||
<DescriptionItem label="API Base">
|
||||
<TruncatedValue value={logEntry.api_base} maxWidth={API_BASE_MAX_WIDTH} />
|
||||
</DescriptionItem>
|
||||
{podRouting && typeof podRouting === "object" && (
|
||||
<DescriptionItem label="Pod">
|
||||
{podRouting.pod_ip}
|
||||
{" · "}
|
||||
{KUBERNETES_POD_ROUTING_LABELS[podRouting.selection]}
|
||||
{" · "}
|
||||
{podRouting.pod_count} pods
|
||||
</DescriptionItem>
|
||||
)}
|
||||
{logEntry.requester_ip_address && (
|
||||
<DescriptionItem label="IP Address">{logEntry.requester_ip_address}</DescriptionItem>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,13 @@ export const LOGS_SORT_FIELD_MAP = {
|
|||
|
||||
export type LogsSortField = keyof typeof LOGS_SORT_FIELD_MAP;
|
||||
|
||||
export type KubernetesPodRouting = {
|
||||
service_host: string;
|
||||
pod_ip: string;
|
||||
pod_count: number;
|
||||
selection: "round_robin" | "session_affinity" | "session_affinity_retry";
|
||||
};
|
||||
|
||||
export type LogEntry = {
|
||||
request_id: string;
|
||||
litellm_call_id?: string | null;
|
||||
|
|
@ -29,7 +36,9 @@ export type LogEntry = {
|
|||
user?: string;
|
||||
end_user?: string;
|
||||
custom_llm_provider?: string;
|
||||
metadata?: Record<string, any>;
|
||||
metadata?: Record<string, any> & {
|
||||
kubernetes_pod_routing?: KubernetesPodRouting | null;
|
||||
};
|
||||
cache_hit: string;
|
||||
cache_key?: string;
|
||||
request_tags?: Record<string, any>;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue