scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): keep a cost settings update from being undone by an in-flight config reload (#51273)

Browse source 
Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-11-12 00:54:10 +00:00 • committed by GitHub
parent 72319cc350
commit 194c984cff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 79 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
litellm/proxy/management_endpoints/cost_tracking_settings.py
View file

@ -253,6 +253,7 @@ async def update_cost_discount_config(
```
"""
from litellm.proxy.proxy_server import (
DB_RECONCILE_LOCK,
prisma_client,
proxy_config,
store_model_in_db,
@ -301,21 +302,13 @@ async def update_cost_discount_config(
)
try:
# Load existing config
config: Final = await proxy_config.get_config()
# Ensure litellm_settings exists
if "litellm_settings" not in config:
config["litellm_settings"] = {}
# Update cost_discount_config
config["litellm_settings"]["cost_discount_config"] = cost_discount_config
# Save the updated config to DB
await proxy_config.save_config(new_config=config)
# Update in-memory litellm.cost_discount_config
litellm.cost_discount_config = cost_discount_config
async with DB_RECONCILE_LOCK:
config: Final = await proxy_config.get_config()
if "litellm_settings" not in config:
config["litellm_settings"] = {}
config["litellm_settings"]["cost_discount_config"] = cost_discount_config
await proxy_config.save_config(new_config=config)
litellm.cost_discount_config = cost_discount_config
verbose_proxy_logger.info("Updated cost_discount_config: %s", cost_discount_config)
@ -397,6 +390,7 @@ async def update_cost_margin_config(
```
"""
from litellm.proxy.proxy_server import (
DB_RECONCILE_LOCK,
prisma_client,
proxy_config,
store_model_in_db,
@ -475,21 +469,13 @@ async def update_cost_margin_config(
)
try:
# Load existing config
config: Final = await proxy_config.get_config()
# Ensure litellm_settings exists
if "litellm_settings" not in config:
config["litellm_settings"] = {}
# Update cost_margin_config
config["litellm_settings"]["cost_margin_config"] = cost_margin_config
# Save the updated config to DB
await proxy_config.save_config(new_config=config)
# Update in-memory litellm.cost_margin_config
litellm.cost_margin_config = cost_margin_config
async with DB_RECONCILE_LOCK:
config: Final = await proxy_config.get_config()
if "litellm_settings" not in config:
config["litellm_settings"] = {}
config["litellm_settings"]["cost_margin_config"] = cost_margin_config
await proxy_config.save_config(new_config=config)
litellm.cost_margin_config = cost_margin_config
verbose_proxy_logger.info("Updated cost_margin_config: %s", cost_margin_config)
@ -534,6 +520,7 @@ async def update_block_requests_for_models_without_pricing(
request: BlockUnpricedModelsRequest,
) -> BlockUnpricedModelsResponse:
from litellm.proxy.proxy_server import (
DB_RECONCILE_LOCK,
prisma_client,
proxy_config,
store_model_in_db,
@ -552,13 +539,13 @@ async def update_block_requests_for_models_without_pricing(
)
try:
config = await proxy_config.get_config()
if "litellm_settings" not in config:
config["litellm_settings"] = {}
config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled
await proxy_config.save_config(new_config=config)
litellm.block_requests_for_models_without_pricing = request.enabled
async with DB_RECONCILE_LOCK:
config = await proxy_config.get_config()
if "litellm_settings" not in config:
config["litellm_settings"] = {}
config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled
await proxy_config.save_config(new_config=config)
litellm.block_requests_for_models_without_pricing = request.enabled
verbose_proxy_logger.info("Updated block_requests_for_models_without_pricing: %s", request.enabled)
return BlockUnpricedModelsResponse(enabled=request.enabled)

56
tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py
View file

@ -4,6 +4,7 @@ Tests for cost tracking settings management endpoints.
Tests the GET and PATCH endpoints for managing cost discount configuration.
"""
import asyncio
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
@ -15,7 +16,11 @@ from pydantic import ValidationError
import litellm
from litellm._internal_context import pinned_billing_time
from litellm.proxy._types import CostEstimateRequest
from litellm.proxy.management_endpoints.cost_tracking_settings import router
from litellm.proxy.management_endpoints.cost_tracking_settings import (
router,
update_cost_discount_config,
update_cost_margin_config,
)
from litellm.proxy.proxy_server import app
client = TestClient(app)
@ -762,6 +767,55 @@ class TestEstimateCostOnPremProvider:
class TestCostSettingsPatchAgainstConfigReload:
@pytest.mark.asyncio
@pytest.mark.parametrize(
("setting", "update", "stale", "fresh"),
[
(
"cost_margin_config",
update_cost_margin_config,
{"openai": {"percentage": 0.1, "fixed_amount": 0.001}},
{"global": 1},
),
("cost_discount_config", update_cost_discount_config, {"openai": 0.05}, {}),
],
)
async def test_patch_is_not_undone_by_a_reload_that_read_the_previous_row(self, setting, update, stale, fresh):
from litellm.proxy import proxy_server
reload_has_read_the_row = asyncio.Event()
release_reload = asyncio.Event()
async def reload_holding_the_previous_row() -> None:
async with proxy_server.DB_RECONCILE_LOCK:
reload_has_read_the_row.set()
await release_reload.wait()
proxy_server.ProxyConfig()._apply_litellm_settings_db_values({setting: stale})
mock_proxy_config = AsyncMock()
mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {setting: stale}})
mock_proxy_config.save_config = AsyncMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", asyncio.Lock()),
patch.object(litellm, setting, stale),
):
reload = asyncio.create_task(reload_holding_the_previous_row())
await reload_has_read_the_row.wait()
patched = asyncio.create_task(update(fresh, user_api_key_dict=MagicMock()))
for _ in range(20):
await asyncio.sleep(0)
release_reload.set()
await asyncio.gather(reload, patched)
assert patched.result()["values"] == fresh
assert getattr(litellm, setting) == fresh
class TestBlockRequestsForModelsWithoutPricing:
"""Test suite for the block_requests_for_models_without_pricing toggle endpoints"""