scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-07 07:58:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

fix(proxy): keep model listing off the shared thread pool and out of credential lookups

Browse source 
Listings run on their own thread pool so many slow catalogs cannot starve
the default executor that chat and Responses calls share. File and batch
credential lookups match deployed names against a key's wildcard grants
instead of listing any provider catalog. The token counter keeps OpenRouter
deployments on their previous path, and a catch-all "*" over provider/*
lists bare ids that route back through it
This commit is contained in:
mateo-berri 2026-10-28 11:55:00 -07:00
parent 8d862c61d9
commit a5bbe70588
10 changed files with 210 additions and 108 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/auth/model_checks.py
View file

@ -1,11 +1,14 @@
# What is this?
## Common checks for /v1/models and `/model/info`
import asyncio
import contextvars
import copy
import functools
import itertools
import re
from collections.abc import Callable, Mapping, Sequence
from typing import Final
from concurrent.futures import ThreadPoolExecutor
from typing import Final, TypeVar
import litellm
from litellm._logging import verbose_proxy_logger
@ -25,6 +28,15 @@ _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)
_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)
_MODEL_LISTING_POOL: Final = ThreadPoolExecutor(thread_name_prefix="litellm-model-listing")
_Listed: Final = TypeVar("_Listed")
async def run_model_listing(listing: Callable[[], _Listed]) -> _Listed:
context: Final = contextvars.copy_context()
return await asyncio.get_running_loop().run_in_executor(_MODEL_LISTING_POOL, context.run, listing)
def _check_wildcard_routing(model: str) -> bool:
"""
@ -204,7 +216,7 @@ def get_complete_model_list(
include_model_access_groups: bool | None = False,
only_model_access_groups: bool | None = False,
team_id: str | None = None,
expand_deployed_wildcard_routes: bool = True,
expand_wildcard_routes: bool = True,
) -> list[str]:
"""Logic for returning complete model list for a given key + team pair"""
@ -250,7 +262,7 @@ def get_complete_model_list(
return_wildcard_routes=return_wildcard_routes,
llm_router=llm_router,
team_id=team_id,
expand_deployed_wildcard_routes=expand_deployed_wildcard_routes,
expand_wildcard_routes=expand_wildcard_routes,
)
complete_model_list: Final = unique_models + all_wildcard_models
@ -418,7 +430,11 @@ def get_known_models_from_wildcard(
wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1)
except ValueError: # safely fail
return []
deployment_pattern_naming_listed_ids: Final = _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params)
deployment_pattern_naming_listed_ids: Final = (
wildcard_model_to_expand
if wildcard_model_to_expand != wildcard_model
else _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params)
)
# Use provider from litellm_params when available, otherwise from wildcard prefix
# (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router)
@ -541,7 +557,7 @@ def _get_wildcard_models(
return_wildcard_routes: bool | None = False,
llm_router: Router | None = None,
team_id: str | None = None,
expand_deployed_wildcard_routes: bool = True,
expand_wildcard_routes: bool = True,
) -> list[str]:
models_to_remove: Final = set()
routes_to_model: Final = _routes_to_model(_cached_wildcard_routing(llm_router), team_id)
@ -552,10 +568,10 @@ def _get_wildcard_models(
all_wildcard_models.append(model)
models_to_remove.add(model)
if not expand_wildcard_routes:
continue
model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None
if model_list and not expand_deployed_wildcard_routes:
continue
if model_list:
for router_model in model_list:
all_wildcard_models.extend(

54
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -1,4 +1,5 @@
import base64
import itertools
import logging
import mimetypes
import re
@ -541,7 +542,8 @@ def get_team_provider_credentials(
listings on the team's own provider account/key instead of a shared
global one.
2. Fallback: any deployment the caller is granted access to for this
provider, expanding wildcard routes and the all-proxy-models sentinel.
provider, matching wildcard grants against deployed model names without
listing any provider's catalog, and expanding the all-proxy-models sentinel.
Skipped when ``include_accessible_deployments`` is False.
Credential lookup is scoped to both the team's allowlist and the key's own
@ -559,6 +561,7 @@ def get_team_provider_credentials(
return None
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list
from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models
team_id: Final = user_api_key_dict.team_id
@ -644,29 +647,34 @@ def get_team_provider_credentials(
grants_all_models: Final = SpecialModelNames.all_proxy_models.value in team_models
effective_team_models: Final = [] if grants_all_models else team_models
models_to_try: Final = list(
dict.fromkeys(
get_complete_model_list(
key_models=list(key_model_allowlist),
team_models=effective_team_models,
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=True,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=True,
team_id=team_id,
expand_deployed_wildcard_routes=False,
)
)
granted_models: Final = get_complete_model_list(
key_models=list(key_model_allowlist),
team_models=effective_team_models,
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=True,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=True,
team_id=team_id,
expand_wildcard_routes=False,
)
deployments_granted_by_wildcard: Final = (
model_name
for model_name in proxy_model_list
if model_matches_any_wildcard_pattern_in_list(model=model_name, allowed_model_list=granted_models)
)
return next(
(
credentials
for credentials in map(
_provider_credentials, dict.fromkeys(itertools.chain(granted_models, deployments_granted_by_wildcard))
)
if credentials is not None
),
None,
)
for model_name in models_to_try:
credentials = _provider_credentials(model_name)
if credentials is not None:
return credentials
return None
HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*")

55
litellm/proxy/proxy_server.py
View file

@ -193,6 +193,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
CountContextManagement,
LlmProviders,
ModelResponse,
ModelResponseStream,
StreamingChoices,
@ -436,6 +437,7 @@ from litellm.proxy.auth.model_checks import (
get_key_models,
get_mcp_server_ids,
get_team_models,
run_model_listing,
)
from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # legacy module exports
@ -12398,18 +12400,20 @@ async def model_list(
# Get complete model list including wildcard routes if requested
from litellm.proxy.auth.model_checks import get_complete_model_list
all_models = await asyncio.to_thread(
get_complete_model_list,
key_models=[],
team_models=[],
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=include_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups or False,
only_model_access_groups=only_model_access_groups or False,
all_models = await run_model_listing(
partial(
get_complete_model_list,
key_models=[],
team_models=[],
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=include_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups or False,
only_model_access_groups=only_model_access_groups or False,
)
)
expanded_undiscoverable_names: Final = undiscoverable_model_names(
@ -14649,6 +14653,8 @@ from litellm.responses.litellm_completion_transformation.count_request import (
with_provider_file_ids,
)
_PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: Final = frozenset({LlmProviders.OPENROUTER})
def _get_provider_token_counter(
deployment: dict, model_to_use: str
@ -14676,14 +14682,13 @@ def _get_provider_token_counter(
)
# Switch case pattern using existing get_provider_model_info
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
# Convert string provider to LlmProviders enum
llm_provider_enum: Final = LlmProviders(custom_llm_provider)
# Add more provider mappings as needed
if llm_provider_enum:
if llm_provider_enum not in _PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO:
provider_model_info: Final = ProviderConfigManager.get_provider_model_info(
model=full_model, provider=llm_provider_enum
)
@ -17752,7 +17757,7 @@ async def model_info_v1(
alias_models: Final = copy.deepcopy(llm_router.get_model_list_from_model_alias())
all_models.extend(alias_models)
all_models = await asyncio.to_thread(expand_wildcard_deployments_for_model_info, all_models, llm_router)
all_models = await run_model_listing(partial(expand_wildcard_deployments_for_model_info, all_models, llm_router))
allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups(
user_api_key_dict=user_api_key_dict,
@ -18087,15 +18092,17 @@ async def model_group_info(
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
all_models_str: Final = (
await asyncio.to_thread(
get_complete_model_list,
key_models=(),
team_models=(),
proxy_model_list=llm_router.get_model_names(),
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=False,
llm_router=llm_router,
await run_model_listing(
partial(
get_complete_model_list,
key_models=(),
team_models=(),
proxy_model_list=llm_router.get_model_names(),
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=False,
llm_router=llm_router,
)
)
if is_proxy_admin
else await get_available_models_for_user(

57
litellm/proxy/utils.py
View file

@ -9497,7 +9497,7 @@ async def _agent_access_group_visible_models(
) -> frozenset[str] | None:
"""Models an agent key may still list once its attached access groups cap it, ``None`` when
nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on."""
from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models
from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models, run_model_listing
if not user_api_key_dict.agent_id:
return None
@ -9517,18 +9517,20 @@ async def _agent_access_group_visible_models(
if not granted:
return frozenset()
return frozenset(
await asyncio.to_thread(
get_complete_model_list,
key_models=granted,
team_models=(),
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=return_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
team_id=team_id,
await run_model_listing(
partial(
get_complete_model_list,
key_models=granted,
team_models=(),
proxy_model_list=proxy_model_list,
user_model=None,
infer_model_from_keys=False,
return_wildcard_routes=return_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
team_id=team_id,
)
)
)
@ -9573,6 +9575,7 @@ async def get_available_models_for_user(
get_complete_model_list,
get_key_models,
get_team_models,
run_model_listing,
)
# Get proxy model list and access groups
@ -9653,19 +9656,21 @@ async def get_available_models_for_user(
)
# Get complete model list
complete_models: Final = await asyncio.to_thread(
get_complete_model_list,
key_models=granted_key_models,
team_models=granted_team_models,
proxy_model_list=proxy_model_list,
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=return_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
only_model_access_groups=only_model_access_groups,
team_id=effective_team_id,
complete_models: Final = await run_model_listing(
partial(
get_complete_model_list,
key_models=granted_key_models,
team_models=granted_team_models,
proxy_model_list=proxy_model_list,
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=return_wildcard_routes,
llm_router=llm_router,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
only_model_access_groups=only_model_access_groups,
team_id=effective_team_id,
)
)
all_models: Final = (
list(

1
tests/unit/proxy/auth/test_model_checks.py
View file

@ -1160,6 +1160,7 @@ def test_partial_bedrock_wildcard_lists_nothing_when_no_discovered_id_carries_it
),
("or-team/*", "openrouter/*", ["or-team/openai/gpt-5.4-mini", "or-team/anthropic/claude-haiku-5-5"]),
("openrouter/openai/*", "openrouter/openai/*", ["openrouter/openai/gpt-5.4-mini"]),
("*", "openrouter/*", ["openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"]),
],
)
@respx.mock

50
tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py
View file

@ -894,25 +894,49 @@ def test_a_provider_credential_lookup_never_lists_an_openrouter_wildcards_live_c
assert catalog_route.call_count == 0
def test_a_key_allowed_a_wider_wildcard_than_any_deployment_gets_the_matching_deployments_credentials(
@pytest.mark.parametrize(
("key_models", "deployments", "custom_llm_provider", "api_key"),
[
(
["anthropic/*"],
(("anthropic/claude-*", "anthropic/claude-*", "anthropic-narrow-key"),),
"anthropic",
"anthropic-narrow-key",
),
(["anthropic/claude-3*"], (("anthropic/*", "anthropic/*", "anthropic-wide-key"),), "anthropic", "anthropic-wide-key"),
(["openrouter/*"], (("openrouter/openai/*", "openrouter/openai/*", "sk-or-narrow"),), "openrouter", "sk-or-narrow"),
(
["team/*"],
(("team/*", "openrouter/*", "sk-or-team"), ("team/openai/*", "openai/*", "openai-team-key")),
"openai",
"openai-team-key",
),
],
)
def test_a_key_wildcard_finds_the_deployment_it_grants_without_listing_any_provider_catalog(
monkeypatch: pytest.MonkeyPatch,
key_models: list[str],
deployments: tuple[tuple[str, str, str], ...],
custom_llm_provider: str,
api_key: str,
):
monkeypatch.setattr(litellm, "check_provider_endpoint", False)
monkeypatch.setattr(litellm, "check_provider_endpoint", True)
router: Final = Router(
model_list=[
{
"model_name": "anthropic/claude-*",
"litellm_params": {"model": "anthropic/claude-*", "api_key": "anthropic-deployment-key"},
}
{"model_name": model_name, "litellm_params": {"model": model, "api_key": deployment_key}}
for model_name, model, deployment_key in deployments
]
)
credentials: Final = get_team_provider_credentials(
llm_router=router,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["anthropic/*"]),
custom_llm_provider="anthropic",
include_accessible_deployments=True,
)
with respx.mock(assert_all_called=False) as respx_mock:
catalog_requests: Final = respx_mock.route().mock(return_value=httpx.Response(200, json={"data": []}))
credentials: Final = get_team_provider_credentials(
llm_router=router,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=key_models),
custom_llm_provider=custom_llm_provider,
include_accessible_deployments=True,
)
assert credentials is not None
assert (credentials["custom_llm_provider"], credentials["api_key"]) == ("anthropic", "anthropic-deployment-key")
assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key)
assert catalog_requests.call_count == 0

12
tests/unit/proxy/proxy_server/test_routes_model_info.py
View file

@ -550,15 +550,12 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke
deployment = {
"model_name": "*",
"litellm_params": {"model": "openai/*"},
"litellm_params": {"model": "openai/*", "api_key": "sk-star"},
}
router = MagicMock()
router.get_model_access_groups = MagicMock(return_value={})
router.get_model_names = MagicMock(return_value=["*"])
router.get_model_list = MagicMock(return_value=[deployment])
router = litellm.Router(model_list=[deployment])
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment])
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(router.get_model_list(), router)
allowed_model_names = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test",
@ -575,7 +572,8 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke
allowed_model_names=allowed_model_names,
)
assert [model["model_name"] for model in result] == ["openai/gpt-4o"]
assert [model["model_name"] for model in result] == ["gpt-4o"]
assert router.wildcard_routed_models("gpt-4o", None) == frozenset({"openai/gpt-4o"})
# ---------------------------------------------------------------------------

13
tests/unit/proxy/proxy_server/test_routes_models.py
View file

@ -11,6 +11,7 @@ from __future__ import annotations
import asyncio
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Final
from unittest.mock import MagicMock
@ -543,7 +544,7 @@ class _CatalogHeldUntilTheLoopRuns:
@pytest.mark.parametrize("scope", [None, "expand"])
async def test_v1_models_fetches_a_live_openrouter_catalog_without_blocking_the_event_loop(
def test_v1_models_fetches_a_live_openrouter_catalog_without_holding_the_event_loop_or_its_default_thread_pool(
monkeypatch: pytest.MonkeyPatch, scope: str | None
):
monkeypatch.setattr(litellm, "check_provider_endpoint", True)
@ -568,8 +569,8 @@ async def test_v1_models_fetches_a_live_openrouter_catalog_without_blocking_the_
)
catalog: Final = _CatalogHeldUntilTheLoopRuns()
with respx.mock() as respx_mock:
catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer)
async def list_models_on_a_one_thread_default_pool() -> dict:
asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1))
listing: Final = asyncio.ensure_future(
proxy_server.model_list(
user_api_key_dict=UserAPIKeyAuth(
@ -579,7 +580,11 @@ async def test_v1_models_fetches_a_live_openrouter_catalog_without_blocking_the_
)
)
await catalog.release_from_the_loop()
response: Final = await listing
return await listing
with respx.mock() as respx_mock:
catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer)
response: Final = asyncio.run(list_models_on_a_one_thread_default_pool())
assert catalog_route.calls.last.response.status_code == 200
assert [model["id"] for model in response["data"]] == ["openrouter/openai/gpt-5.4-mini"]

33
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -175,6 +175,39 @@ async def test_token_counting_model_not_in_model_list():
assert response.model_used == "special-alias"
@pytest.mark.asyncio
async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_count_unchanged(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
litellm.proxy.proxy_server,
"llm_router",
Router(
model_list=[
{
"model_name": "or-gpt",
"litellm_params": {
"model": "openai/gpt-4o",
"custom_llm_provider": "openrouter",
"api_key": "sk-or-token-count",
},
}
]
),
)
request: Final = TokenCountRequest(
model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}]
)
via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True)
counted_locally: Final = await token_counter(request=request, call_endpoint=False)
assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == (
counted_locally.total_tokens,
counted_locally.model_used,
)
@pytest.mark.asyncio
async def test_gpt_token_counting():
"""

13
tests/unit/proxy/utils/helpers/test_model_access.py
View file

@ -1,5 +1,6 @@
import asyncio
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Final
from unittest.mock import MagicMock
@ -677,7 +678,7 @@ class _CatalogHeldUntilTheLoopRuns:
self.released.set()
async def test_get_available_models_for_user_fetches_a_live_catalog_without_blocking_the_event_loop(
def test_get_available_models_for_user_fetches_a_live_catalog_without_holding_the_event_loop_or_its_default_pool(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "check_provider_endpoint", True)
@ -695,8 +696,8 @@ async def test_get_available_models_for_user_fetches_a_live_catalog_without_bloc
]
)
with respx.mock() as respx_mock:
catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer)
async def list_models_on_a_one_thread_default_pool() -> list[str]:
asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1))
listing: Final = asyncio.ensure_future(
get_available_models_for_user(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-key", user_id="user-1"),
@ -706,7 +707,11 @@ async def test_get_available_models_for_user_fetches_a_live_catalog_without_bloc
)
)
await catalog.release_from_the_loop()
models: Final = await listing
return await listing
with respx.mock() as respx_mock:
catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer)
models: Final = asyncio.run(list_models_on_a_one_thread_default_pool())
assert catalog_route.calls.last.response.status_code == 200
assert models == ["openrouter/openai/gpt-5.4-mini"]