From dc96ade95681937dc830cdddc739aad2559b8a5d Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Fri, 27 Feb 2026 22:28:03 -0500 Subject: [PATCH] fix: preserve interval_hours in model cost map reload config (#22200) The upsert update branches for model_cost_map_reload_config were overwriting param_value with only the force_reload flag, dropping interval_hours. This caused scheduled reloads to self-destruct after their first execution. Co-authored-by: Claude Opus 4.6 (1M context) --- litellm/proxy/proxy_server.py | 26 ++- tests/test_litellm/proxy/test_proxy_server.py | 171 ++++++++++++++++++ 2 files changed, 191 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 54ad361c74..be68a26852 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4635,7 +4635,7 @@ class ProxyConfig: } ), }, - "update": {"param_value": safe_dumps({"force_reload": False})}, + "update": {"param_value": safe_dumps({"interval_hours": interval_hours, "force_reload": False})}, }, ) @@ -4736,7 +4736,7 @@ class ProxyConfig: } ), }, - "update": {"param_value": safe_dumps({"force_reload": False})}, + "update": {"param_value": safe_dumps({"interval_hours": interval_hours, "force_reload": False})}, }, ) @@ -12261,7 +12261,14 @@ async def reload_model_cost_map( current_time = datetime.utcnow() last_model_cost_map_reload = current_time.isoformat() - # Set force reload flag in database for other pods + # Set force reload flag in database for other pods, preserving existing interval_hours + existing_config = await prisma_client.db.litellm_config.find_unique( + where={"param_name": "model_cost_map_reload_config"} + ) + existing_interval = None + if existing_config and existing_config.param_value: + existing_interval = existing_config.param_value.get("interval_hours") + await prisma_client.db.litellm_config.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ @@ -12271,7 +12278,7 @@ async def reload_model_cost_map( {"interval_hours": None, "force_reload": True} ), }, - "update": {"param_value": safe_dumps({"force_reload": True})}, + "update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})}, }, ) @@ -12600,7 +12607,14 @@ async def reload_anthropic_beta_headers( current_time = datetime.utcnow() last_anthropic_beta_headers_reload = current_time.isoformat() - # Set force reload flag in database for other pods + # Set force reload flag in database for other pods, preserving existing interval_hours + existing_beta_config = await prisma_client.db.litellm_config.find_unique( + where={"param_name": "anthropic_beta_headers_reload_config"} + ) + existing_beta_interval = None + if existing_beta_config and existing_beta_config.param_value: + existing_beta_interval = existing_beta_config.param_value.get("interval_hours") + await prisma_client.db.litellm_config.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ @@ -12610,7 +12624,7 @@ async def reload_anthropic_beta_headers( {"interval_hours": None, "force_reload": True} ), }, - "update": {"param_value": safe_dumps({"force_reload": True})}, + "update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})}, }, ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5f54c151d8..b993d4c4cf 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2078,10 +2078,181 @@ class TestPriceDataReloadIntegration: param_value_json = call_args[1]["data"]["update"]["param_value"] param_value_dict = json.loads(param_value_json) assert param_value_dict["force_reload"] == False + assert param_value_dict.get("interval_hours") == 6 finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() + def test_distributed_reload_preserves_interval_hours(self): + """Test that _check_and_reload_model_cost_map preserves interval_hours after reload. + + Regression test: the update branch of the upsert was previously dropping + interval_hours, causing scheduled reloads to self-destruct after first execution. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + + # Set up config with interval_hours=24 and force_reload=True to trigger reload + mock_config = MagicMock() + mock_config.param_value = {"interval_hours": 24, "force_reload": True} + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + original_model_cost = litellm.model_cost.copy() + try: + with patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map: + mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + # Verify the upsert update branch preserves interval_hours + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == False + assert param_value_dict["interval_hours"] == 24, ( + "interval_hours must be preserved in the update branch; " + "dropping it causes the schedule to self-destruct" + ) + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_manual_reload_preserves_interval_hours(self): + """Test that manual reload via /reload/model_cost_map preserves existing interval_hours. + + Regression test: the manual reload endpoint was overwriting param_value with + only force_reload=True, dropping any existing interval_hours schedule. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.proxy_server import cleanup_router_config_variables + + cleanup_router_config_variables() + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + asyncio.run(initialize(config=config_fp, debug=True)) + + mock_auth = MagicMock() + mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + client = TestClient(app) + + original_model_cost = litellm.model_cost.copy() + try: + with patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map: + mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + # Simulate existing config with a schedule + mock_existing = MagicMock() + mock_existing.param_value = {"interval_hours": 12, "force_reload": False} + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_existing) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + response = client.post("/reload/model_cost_map") + assert response.status_code == 200 + + # Verify interval_hours was preserved in the upsert + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == True + assert param_value_dict["interval_hours"] == 12, ( + "interval_hours must be preserved when manual reload sets force_reload; " + "dropping it destroys any existing schedule" + ) + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_anthropic_beta_headers_reload_preserves_interval_hours(self): + """Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload. + + Regression test: the update branch of the upsert was dropping interval_hours, + identical to the model cost map bug. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + + # Set up config with interval_hours=12 and force_reload=True to trigger reload + mock_config = MagicMock() + mock_config.param_value = {"interval_hours": 12, "force_reload": True} + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + with patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload: + mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} + + asyncio.run(proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma)) + + # Verify the upsert update branch preserves interval_hours + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == False + assert param_value_dict["interval_hours"] == 12, ( + "interval_hours must be preserved in the update branch; " + "dropping it causes the schedule to self-destruct" + ) + + def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self): + """Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours. + + Regression test: the manual reload endpoint was overwriting param_value with + only force_reload=True, dropping any existing interval_hours schedule. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.proxy_server import cleanup_router_config_variables + + cleanup_router_config_variables() + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + asyncio.run(initialize(config=config_fp, debug=True)) + + mock_auth = MagicMock() + mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + client = TestClient(app) + + with patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload: + mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + # Simulate existing config with a schedule + mock_existing = MagicMock() + mock_existing.param_value = {"interval_hours": 8, "force_reload": False} + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_existing) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + response = client.post("/reload/anthropic_beta_headers") + assert response.status_code == 200 + + # Verify interval_hours was preserved in the upsert + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == True + assert param_value_dict["interval_hours"] == 8, ( + "interval_hours must be preserved when manual reload sets force_reload; " + "dropping it destroys any existing schedule" + ) + def test_config_file_parsing(self): """Test parsing of config file with reload settings""" config_content = """