scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(alerting): dedupe scheduled Slack spend reports across pods (#36489)

Browse source 
* fix(alerting): dedupe scheduled Slack spend reports across pods

Every pod ran its own weekly/monthly spend report jobs, prometheus
fallback stats cron, and daily report loop, so deployments with
multiple replicas or uvicorn workers received one copy per pod.

Gate each scheduled send behind the shared PodLockManager redis lock.
The lock is never released: its TTL (the full reporting window for the
weekly interval job, whose per-pod anchors drift by boot time and
jitter) doubles as a sent-this-window marker. acquire_lock returning
None (no redis wired) proceeds, preserving single-pod behavior.

Also generalize the pod lock could-not-acquire log line, which claimed
to be about spend tracking for every consumer.

Fixes #14809

* fix(alerting): harden spend report locks after adversarial review

Weekly lock TTL gets an hour haircut: with ttl equal to the interval,
the winner re-fires just before its own key expires, reacquires without
a TTL refresh, and the key then lapses in time for a trailing pod to
re-send. Job/lock ids move to litellm/constants.py per convention, and
spend_report_frequency now rejects non-positive day counts, which
previously coerced to an every-second schedule and would now compute a
negative lock TTL that silently never sends.

Adds the missing test coverage the review flagged: startup_event's
pod_lock_manager wiring (identity-asserted), the prometheus closure's
positive path, and the ungated immediate prometheus send pinned to
exactly one await.

* test(alerting): consolidate spend_report_frequency validator coverage

Drops a duplicate non-positive-days test and parametrizes the survivor
over the suffix half of the validator too

* fix(alerting): route the startup prometheus fallback send through the pod lock

Greptile caught that the boot-time send still ran once per pod when
PROMETHEUS_URL is set, the same duplication class this PR removes

* fix(alerting): make report lock acquisition non-reentrant

Greptile caught that a pod booting within an hour of the fallback stats
cron sent twice: the startup send takes the lock, then the cron fire
hits acquire_lock's reacquire branch, which returns True for the
holder. Window-marker gates now pass allow_reentrant=False so a live
lock blocks everyone including its holder; leader-election consumers
keep the reentrant default

* test(proxy): give spec'd ProxyLogging mocks a db_spend_update_writer

_initialize_slack_alerting_jobs now reads it for the pod lock manager,
and spec=ProxyLogging blocks instance-only attributes
This commit is contained in:
ryan-crabbe-berri 2026-08-11 12:41:11 -07:00 • committed by GitHub
parent 65fbeb5531
commit be71a8fdbf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 807 additions and 974 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -1478,6 +1478,10 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED
SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3))

27
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -17,7 +17,7 @@ import litellm.litellm_core_utils.litellm_logging
import litellm.types
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.constants import HOURS_IN_A_DAY
from litellm.constants import HOURS_IN_A_DAY, SLACK_DAILY_REPORT_LOCK_ID
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type
from litellm.integrations.SlackAlerting.hanging_request_check import (
@ -51,6 +51,7 @@ from .batching_handler import send_to_webhook, squash_payloads
from .utils import process_slack_alerting_variables
if TYPE_CHECKING:
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.router import Router as _Router
Router = _Router
@ -1576,7 +1577,11 @@ Model Info:
except Exception:
pass
async def _run_scheduler_helper(self, llm_router) -> bool:
async def _run_scheduler_helper(
self,
llm_router,
pod_lock_manager: "PodLockManager | None" = None,
) -> bool:
"""
Returns:
- True -> report sent
@ -1601,6 +1606,16 @@ Model Info:
interval_seconds: Final = self.alerting_args.daily_report_frequency
if current_time - report_sent >= interval_seconds:
if (
pod_lock_manager is not None
and (
await pod_lock_manager.acquire_lock(
cronjob_id=SLACK_DAILY_REPORT_LOCK_ID, ttl=interval_seconds, allow_reentrant=False
)
)
is False
):
return False
# Sneak in the reporting logic here
await self.send_daily_reports(router=llm_router)
# Also, don't forget to update the report_sent time after sending the report!
@ -1612,7 +1627,11 @@ Model Info:
return report_sent_bool
async def _run_scheduled_daily_report(self, llm_router: Any | None = None):
async def _run_scheduled_daily_report(
self,
llm_router: Any | None = None,
pod_lock_manager: "PodLockManager | None" = None,
):
"""
If 'daily_reports' enabled
@ -1625,7 +1644,7 @@ Model Info:
if "daily_reports" in self.alert_types:
while True:
await self._run_scheduler_helper(llm_router=llm_router)
await self._run_scheduler_helper(llm_router=llm_router, pod_lock_manager=pod_lock_manager)
interval = random.randint(
self.alerting_args.report_check_interval - 3,
self.alerting_args.report_check_interval + 3,

21
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

@ -43,6 +43,7 @@ end
self,
cronjob_id: str,
ttl: int | None = None,
allow_reentrant: bool = True,
) -> bool | None:
"""
Attempt to acquire the lock for a specific cron job using Redis.
@ -53,6 +54,10 @@ end
ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS.
Use a longer TTL for jobs that may take longer than the default 60s
(e.g. key rotation with many keys).
allow_reentrant: With the default True, a pod that already holds the lock
acquires it again (leader election semantics). Pass False when the live
lock marks work as already done for this window, so not even the holder
may redo it before the TTL expires.
"""
if self.redis_cache is None:
verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock")
@ -88,7 +93,7 @@ end
if current_value is not None:
if isinstance(current_value, bytes):
current_value = current_value.decode("utf-8")
if current_value == self.pod_id:
if current_value == self.pod_id and allow_reentrant:
verbose_proxy_logger.info(
"Pod %s already holds the Redis lock for cronjob_id=%s",
self.pod_id,
@ -96,14 +101,12 @@ end
)
self._emit_acquired_lock_event(cronjob_id, self.pod_id)
return True
else:
verbose_proxy_logger.info(
"Spend tracking - pod %s could not acquire lock for cronjob_id=%s, "
"held by pod %s. Spend updates in Redis will wait for the leader pod to commit.",
self.pod_id,
cronjob_id,
current_value,
)
verbose_proxy_logger.info(
"Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.",
self.pod_id,
cronjob_id,
current_value,
)
return False
except Exception as e:
verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e)

58
litellm/proxy/proxy_server.py
View file

@ -233,6 +233,8 @@ from litellm.constants import (
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_ADMIN_NAME,
LITELLM_PROXY_BUDGET_NAME,
MONTHLY_SPEND_REPORT_JOB_ID,
PROMETHEUS_FALLBACK_STATS_JOB_ID,
PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
PROXY_BATCH_POLLING_ENABLED,
PROXY_BATCH_POLLING_INTERVAL,
@ -240,6 +242,7 @@ from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
WEEKLY_SPEND_REPORT_JOB_ID,
)
from litellm.exceptions import RejectedRequestError
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -9043,41 +9046,76 @@ class ProxyStartupEvent:
spend_report_frequency: Final[str] = general_settings.get("spend_report_frequency", "7d") or "7d"
days: Final = int(spend_report_frequency[:-1])
if spend_report_frequency[-1].lower() != "d":
raise ValueError("spend_report_frequency must be specified in days, e.g., '1d', '7d'")
if spend_report_frequency[-1].lower() != "d" or days <= 0:
raise ValueError("spend_report_frequency must be a positive number of days, e.g., '1d', '7d'")
pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
weekly_lock_ttl: Final = duration_in_seconds(spend_report_frequency) - 3600
async def _scheduled_weekly_spend_report() -> None:
# TTL spans the whole reporting window: each pod's interval anchor is its own
# boot time + jitter, so a shorter lock would let a later pod re-send the report.
# Minus an hour so the next window's first firer finds a free key
if (
await pod_lock_manager.acquire_lock(
cronjob_id=WEEKLY_SPEND_REPORT_JOB_ID, ttl=weekly_lock_ttl, allow_reentrant=False
)
is False
):
return
await proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report(spend_report_frequency)
async def _scheduled_monthly_spend_report() -> None:
if (
await pod_lock_manager.acquire_lock(
cronjob_id=MONTHLY_SPEND_REPORT_JOB_ID, ttl=3600, allow_reentrant=False
)
is False
):
return
await proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report()
scheduler.add_job(
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report,
_scheduled_weekly_spend_report,
"interval",
days=days,
next_run_time=datetime.now() + timedelta(seconds=10 + random.randint(0, 300)),
args=[spend_report_frequency],
id="weekly_spend_report_job",
id=WEEKLY_SPEND_REPORT_JOB_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
scheduler.add_job(
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report,
_scheduled_monthly_spend_report,
"cron",
day=1,
id="monthly_spend_report_job",
id=MONTHLY_SPEND_REPORT_JOB_ID,
replace_existing=True,
)
if os.getenv("PROMETHEUS_URL"):
from zoneinfo import ZoneInfo
async def _scheduled_fallback_stats() -> None:
if (
await pod_lock_manager.acquire_lock(
cronjob_id=PROMETHEUS_FALLBACK_STATS_JOB_ID, ttl=3600, allow_reentrant=False
)
is False
):
return
await proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus()
scheduler.add_job(
proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus,
_scheduled_fallback_stats,
"cron",
hour=PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
minute=0,
timezone=ZoneInfo("America/Los_Angeles"),
id="prometheus_fallback_stats_job",
id=PROMETHEUS_FALLBACK_STATS_JOB_ID,
replace_existing=True,
)
await proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus()
await _scheduled_fallback_stats()
@classmethod
async def _setup_prisma_client(

5
litellm/proxy/utils.py
View file

@ -468,7 +468,10 @@ class ProxyLogging:
and not self.daily_report_started
):
asyncio.create_task(
self.slack_alerting_instance._run_scheduled_daily_report(llm_router=llm_router)
self.slack_alerting_instance._run_scheduled_daily_report(
llm_router=llm_router,
pod_lock_manager=self.db_spend_update_writer.pod_lock_manager,
)
) # RUN DAILY REPORT (if scheduled)
self.daily_report_started = True

170
tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py
View file

@ -1,17 +1,21 @@
import asyncio
import datetime
import json
import os
import sys
import time
import unittest
from typing import List, Optional, Tuple
from typing import Final, List, Optional, Tuple
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system-path
import pytest
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system-path
import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.proxy._types import CallInfo, Litellm_EntityType
from litellm.types.integrations.slack_alerting import SlackAlertingCacheKeys
class TestSlackAlerting(unittest.TestCase):
@ -20,37 +24,27 @@ class TestSlackAlerting(unittest.TestCase):
def test_get_percent_of_max_budget_left(self):
# Test case 1: When max_budget is None
user_info = CallInfo(
max_budget=None, spend=50.0, event_group=Litellm_EntityType.KEY
)
user_info = CallInfo(max_budget=None, spend=50.0, event_group=Litellm_EntityType.KEY)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 2: When max_budget is 0
user_info = CallInfo(
max_budget=0.0, spend=50.0, event_group=Litellm_EntityType.KEY
)
user_info = CallInfo(max_budget=0.0, spend=50.0, event_group=Litellm_EntityType.KEY)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 3: When spend is less than max_budget
user_info = CallInfo(
max_budget=100.0, spend=75.0, event_group=Litellm_EntityType.KEY
)
user_info = CallInfo(max_budget=100.0, spend=75.0, event_group=Litellm_EntityType.KEY)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.25)
# Test case 4: When spend equals max_budget
user_info = CallInfo(
max_budget=100.0, spend=100.0, event_group=Litellm_EntityType.KEY
)
user_info = CallInfo(max_budget=100.0, spend=100.0, event_group=Litellm_EntityType.KEY)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 5: When spend exceeds max_budget
user_info = CallInfo(
max_budget=100.0, spend=120.0, event_group=Litellm_EntityType.KEY
)
user_info = CallInfo(max_budget=100.0, spend=120.0, event_group=Litellm_EntityType.KEY)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, -0.2)
@ -189,7 +183,9 @@ class TestSlackAlerting(unittest.TestCase):
# Test the specific formatting logic we're interested in
alert_type_formatted = f"Alert type: `{alert_type.name}`\n"
formatted_message = f"{alert_type_formatted}\n Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
formatted_message = (
f"{alert_type_formatted}\n Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
)
# Verify alert_type is in the formatted message as expected
self.assertIn("Alert type: `llm_exceptions`", formatted_message)
@ -214,9 +210,7 @@ class TestSlackAlerting(unittest.TestCase):
json.dumps(outage_value)
# Verify the specific error message
self.assertIn(
"Object of type set is not JSON serializable", str(context.exception)
)
self.assertIn("Object of type set is not JSON serializable", str(context.exception))
def test_fixed_redis_serialization(self):
"""Test that our fix resolves the Redis serialization error."""
@ -245,3 +239,133 @@ class TestSlackAlerting(unittest.TestCase):
)
self.assertEqual(parsed_data["alerts"], [408])
self.assertEqual(parsed_data["provider_region_id"], "vertex_aius-east1")
_REPORT_SENT_KEY: Final = SlackAlertingCacheKeys.report_sent_key.value
_DAILY_REPORT_FREQUENCY: Final = 900
async def _slack_alerting_with_due_daily_report() -> SlackAlerting:
slack_alerting: Final = SlackAlerting(
internal_usage_cache=DualCache(),
alerting_args={"daily_report_frequency": _DAILY_REPORT_FREQUENCY},
)
await slack_alerting.internal_usage_cache.async_set_cache(
key=_REPORT_SENT_KEY,
value=time.time() - _DAILY_REPORT_FREQUENCY - 1,
)
slack_alerting.send_daily_reports = AsyncMock()
return slack_alerting
async def _read_report_sent(slack_alerting: SlackAlerting) -> float:
return await slack_alerting.internal_usage_cache.async_get_cache(
key=_REPORT_SENT_KEY,
parent_otel_span=None,
)
@pytest.mark.asyncio
async def test_daily_report_skipped_when_another_pod_holds_the_lock():
"""regression: issue #14809 - every pod sent its own copy of the daily report.
The losing pod must also leave report_sent untouched so the winner's window still counts.
"""
slack_alerting: Final = await _slack_alerting_with_due_daily_report()
report_sent_before: Final = await _read_report_sent(slack_alerting)
pod_lock_manager: Final = AsyncMock()
pod_lock_manager.acquire_lock.return_value = False
result: Final = await slack_alerting._run_scheduler_helper(
llm_router=MagicMock(),
pod_lock_manager=pod_lock_manager,
)
assert result is False
slack_alerting.send_daily_reports.assert_not_awaited()
assert await _read_report_sent(slack_alerting) == report_sent_before
pod_lock_manager.acquire_lock.assert_awaited_once_with(
cronjob_id="slack_daily_report",
ttl=_DAILY_REPORT_FREQUENCY,
allow_reentrant=False,
)
@pytest.mark.asyncio
async def test_daily_report_sent_by_the_pod_that_wins_the_lock():
slack_alerting: Final = await _slack_alerting_with_due_daily_report()
report_sent_before: Final = await _read_report_sent(slack_alerting)
llm_router: Final = MagicMock()
pod_lock_manager: Final = AsyncMock()
pod_lock_manager.acquire_lock.return_value = True
result: Final = await slack_alerting._run_scheduler_helper(
llm_router=llm_router,
pod_lock_manager=pod_lock_manager,
)
assert result is True
slack_alerting.send_daily_reports.assert_awaited_once_with(router=llm_router)
assert await _read_report_sent(slack_alerting) > report_sent_before
pod_lock_manager.acquire_lock.assert_awaited_once_with(
cronjob_id="slack_daily_report",
ttl=_DAILY_REPORT_FREQUENCY,
allow_reentrant=False,
)
@pytest.mark.parametrize("lock_state", ["no_pod_lock_manager", "no_redis_configured"])
@pytest.mark.asyncio
async def test_daily_report_still_sent_without_a_working_lock(lock_state: str):
"""Single-pod parity: a missing lock manager, or one whose acquire_lock returns None
because redis isn't configured, must not suppress the report."""
slack_alerting: Final = await _slack_alerting_with_due_daily_report()
report_sent_before: Final = await _read_report_sent(slack_alerting)
llm_router: Final = MagicMock()
pod_lock_manager: Final = (
None if lock_state == "no_pod_lock_manager" else AsyncMock(acquire_lock=AsyncMock(return_value=None))
)
result: Final = await slack_alerting._run_scheduler_helper(
llm_router=llm_router,
pod_lock_manager=pod_lock_manager,
)
assert result is True
slack_alerting.send_daily_reports.assert_awaited_once_with(router=llm_router)
assert await _read_report_sent(slack_alerting) > report_sent_before
@pytest.mark.asyncio
async def test_daily_report_lock_not_attempted_before_the_interval_elapses():
"""The lock is a per-window marker, so a pod must not burn it on a check that isn't due yet."""
slack_alerting: Final = await _slack_alerting_with_due_daily_report()
await slack_alerting.internal_usage_cache.async_set_cache(key=_REPORT_SENT_KEY, value=time.time())
pod_lock_manager: Final = AsyncMock()
pod_lock_manager.acquire_lock.return_value = True
result: Final = await slack_alerting._run_scheduler_helper(
llm_router=MagicMock(),
pod_lock_manager=pod_lock_manager,
)
assert result is False
pod_lock_manager.acquire_lock.assert_not_awaited()
slack_alerting.send_daily_reports.assert_not_awaited()
@pytest.mark.asyncio
async def test_scheduled_daily_report_threads_the_pod_lock_manager_through():
"""The loop in _run_scheduled_daily_report is where the lock manager reaches the gate."""
slack_alerting: Final = SlackAlerting(alert_types=["daily_reports"])
pod_lock_manager: Final = AsyncMock()
slack_alerting._run_scheduler_helper = AsyncMock(side_effect=asyncio.CancelledError)
with pytest.raises(asyncio.CancelledError):
await slack_alerting._run_scheduled_daily_report(
llm_router=MagicMock(),
pod_lock_manager=pod_lock_manager,
)
_, kwargs = slack_alerting._run_scheduler_helper.await_args
assert kwargs["pod_lock_manager"] is pod_lock_manager

35
tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py
View file

@ -7,9 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
@ -310,9 +308,7 @@ async def test_lock_takeover_race_condition(mock_redis):
@pytest.mark.asyncio
async def test_release_lock_uses_atomic_compare_delete_script_when_available(
pod_lock_manager, mock_redis
):
async def test_release_lock_uses_atomic_compare_delete_script_when_available(pod_lock_manager, mock_redis):
"""
Test that release_lock prefers atomic compare-and-delete Lua script when
redis cache exposes script registration.
@ -323,12 +319,8 @@ async def test_release_lock_uses_atomic_compare_delete_script_when_available(
await pod_lock_manager.release_lock(cronjob_id="test_job")
lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job")
mock_redis.async_register_script.assert_called_once_with(
PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT
)
script_callable.assert_called_once_with(
keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)]
)
mock_redis.async_register_script.assert_called_once_with(PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT)
script_callable.assert_called_once_with(keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)])
mock_redis.async_get_cache.assert_not_called()
mock_redis.async_delete_cache.assert_not_called()
@ -359,9 +351,7 @@ async def test_release_lock_lua_path_emits_released_event(pod_lock_manager, mock
with patch.object(pod_lock_manager, "_emit_released_lock_event") as mock_emit:
await pod_lock_manager.release_lock(cronjob_id="test_job")
mock_emit.assert_called_once_with(
cronjob_id="test_job", pod_id=pod_lock_manager.pod_id
)
mock_emit.assert_called_once_with(cronjob_id="test_job", pod_id=pod_lock_manager.pod_id)
class FakeRedisLockStore:
@ -437,9 +427,7 @@ async def test_release_lock_preserves_lock_held_by_other_pod():
@pytest.mark.asyncio
async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails(
pod_lock_manager, mock_redis
):
async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails(pod_lock_manager, mock_redis):
"""
Test that release_lock falls back to GET+DEL when Lua script execution
raises (e.g. Redis restart cleared loaded scripts).
@ -457,3 +445,14 @@ async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails(
mock_redis.async_delete_cache.assert_called_once_with(lock_key)
# Cached script handle should be reset so next call re-registers
assert pod_lock_manager._release_lock_script is None
@pytest.mark.asyncio
async def test_acquire_lock_own_lock_not_reentrant(pod_lock_manager, mock_redis):
"""With allow_reentrant=False a live lock means the window's work is done, so even
the holder gets False; the default stays reentrant for leader-election callers."""
mock_redis.async_set_cache.return_value = False
mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id
assert await pod_lock_manager.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False
assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True

184
tests/test_litellm/proxy/proxy_server/test_lifecycle.py
View file

@ -22,6 +22,7 @@ import inspect
import json
import logging
import os
from collections.abc import Awaitable, Callable
from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
@ -397,9 +398,7 @@ def test__redact_worker_config_for_logging_masks_nested_secret_fields():
"database_url": nested_db_url,
"database_extra_connection_params": {"password": nested_extra_pw},
"alert_to_webhook_url": {"budget_alerts": nested_webhook},
"pass_through_endpoints": [
{"path": "/up", "headers": {"Authorization": nested_bearer}}
],
"pass_through_endpoints": [{"path": "/up", "headers": {"Authorization": nested_bearer}}],
}
}
}
@ -451,16 +450,13 @@ def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch):
import litellm
sentinel_secret_mgr = object()
monkeypatch.setattr(
litellm, "secret_manager_client", sentinel_secret_mgr, raising=False
)
monkeypatch.setattr(litellm, "secret_manager_client", sentinel_secret_mgr, raising=False)
result = load_from_azure_key_vault(use_azure_key_vault=False)
observed = {
"return_value": result,
"secret_manager_unchanged": litellm.secret_manager_client
is sentinel_secret_mgr,
"secret_manager_unchanged": litellm.secret_manager_client is sentinel_secret_mgr,
"called_with": False,
}
assert normalize(observed) == {
@ -614,9 +610,7 @@ def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch):
observed = {
"called_arg": (
fake_get.call_args.args[0]
if fake_get.call_args.args
else fake_get.call_args.kwargs.get("model")
fake_get.call_args.args[0] if fake_get.call_args.args else fake_get.call_args.kwargs.get("model")
),
"returned_max_tokens": result.get("max_tokens"),
"returned_cost": result.get("input_cost_per_token"),
@ -663,9 +657,7 @@ def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch):
def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch):
"""Popen raising OSError must NOT propagate — function logs and returns."""
monkeypatch.setattr(
ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary"))
)
monkeypatch.setattr(ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary")))
result = run_ollama_serve()
assert result is None
@ -685,8 +677,7 @@ async def test_proxy_startup_event_is_async_context_manager_with_expected_signat
observed = {
"param_count": len(sig.parameters),
"has_app_param": "app" in sig.parameters,
"wrapped_is_async": inspect.iscoroutinefunction(wrapped)
or inspect.isasyncgenfunction(wrapped),
"wrapped_is_async": inspect.iscoroutinefunction(wrapped) or inspect.isasyncgenfunction(wrapped),
"has_asynccontextmanager_wrapper": wrapped is not None,
}
assert normalize(observed) == {
@ -777,3 +768,164 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
assert budget_check_pos < warn_pos < next_startup_section_pos, (
"DB-less budget warning must run after Prisma setup and the DB-backed budget block"
)
# ---------------------------------------------------------------------------
# _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809)
# ---------------------------------------------------------------------------
SlackAlertingJobs = dict[str, Callable[[], Awaitable[None]]]
def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> MagicMock:
proxy_logging_obj = MagicMock()
proxy_logging_obj.slack_alerting_instance.alerting = ["slack"]
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report = AsyncMock()
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report = AsyncMock()
proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus = AsyncMock()
pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
pod_lock_manager.acquire_lock = AsyncMock(return_value=acquire_lock_result)
pod_lock_manager.release_lock = AsyncMock()
return proxy_logging_obj
async def _init_slack_alerting_jobs(
acquire_lock_result: bool | None,
spend_report_frequency: str = "7d",
) -> tuple[SlackAlertingJobs, MagicMock]:
scheduler = MagicMock()
proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result)
await ProxyStartupEvent._initialize_slack_alerting_jobs(
scheduler=scheduler,
general_settings={"spend_report_frequency": spend_report_frequency},
proxy_logging_obj=proxy_logging_obj,
prisma_client=MagicMock(),
)
jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list}
return jobs, proxy_logging_obj
@pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"])
@pytest.mark.asyncio
async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str):
"""A non-positive window used to become an every-second APScheduler interval, and now also
computes a negative lock TTL that expires instantly and suppresses the report for good.
match= is load-bearing: drop the guard and "-1d" still raises, but from duration_in_seconds."""
with pytest.raises(ValueError, match="positive number of days"):
await _init_slack_alerting_jobs(
acquire_lock_result=True,
spend_report_frequency=spend_report_frequency,
)
@pytest.mark.asyncio
async def test_weekly_spend_report_skipped_when_another_pod_holds_the_lock():
"""regression: issue #14809 - every pod ran its own weekly spend report job."""
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
await jobs["weekly_spend_report_job"]()
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_not_awaited()
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
cronjob_id="weekly_spend_report_job",
ttl=7 * 86400 - 3600,
allow_reentrant=False,
)
@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_weekly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
"""None means redis isn't configured; a single-pod deploy must still report."""
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
await jobs["weekly_spend_report_job"]()
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("7d")
@pytest.mark.asyncio
async def test_weekly_spend_report_lock_ttl_tracks_the_configured_window():
"""TTL is the window less an hour: long enough that no second pod re-sends inside the
window, short enough that the lock is gone before the next one opens. A fixed TTL would
break one end or the other as soon as spend_report_frequency changes."""
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True, spend_report_frequency="1d")
await jobs["weekly_spend_report_job"]()
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
cronjob_id="weekly_spend_report_job",
ttl=86400 - 3600,
allow_reentrant=False,
)
proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("1d")
@pytest.mark.asyncio
async def test_monthly_spend_report_skipped_when_another_pod_holds_the_lock():
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
await jobs["monthly_spend_report_job"]()
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_not_awaited()
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
cronjob_id="monthly_spend_report_job",
ttl=3600,
allow_reentrant=False,
)
@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_monthly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
await jobs["monthly_spend_report_job"]()
proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_awaited_once_with()
@pytest.mark.asyncio
async def test_spend_report_locks_are_never_released():
"""The lock is a per-window marker, not a mutex: releasing it lets the next pod re-send."""
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True)
await jobs["weekly_spend_report_job"]()
await jobs["monthly_spend_report_job"]()
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()
@pytest.mark.asyncio
async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch):
"""The boot-time send goes through the same gate, so a losing pod sends nothing at all:
startup and the scheduled job both stay at zero."""
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
assert send_fallback_stats.await_count == 0
await jobs["prometheus_fallback_stats_job"]()
assert send_fallback_stats.await_count == 0
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_with(
cronjob_id="prometheus_fallback_stats_job",
ttl=3600,
allow_reentrant=False,
)
assert proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.await_count == 2
@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_prometheus_fallback_stats_job_runs_when_the_lock_is_free_or_absent(monkeypatch, acquire_lock_result):
monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
assert send_fallback_stats.await_count == 1
await jobs["prometheus_fallback_stats_job"]()
assert send_fallback_stats.await_count == 2

1241
tests/test_litellm/proxy/test_proxy_server.py
View file

File diff suppressed because it is too large Load diff

36
tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py
View file

@ -8,6 +8,7 @@ because they are direct dependents on the lifecycle state.
from __future__ import annotations
import asyncio
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock, patch
@ -138,6 +139,41 @@ def test_startup_event_propagates_init_callbacks_failure_raises(proxy_logging):
proxy_logging.startup_event(llm_router=None, redis_usage_cache=None)
@pytest.mark.asyncio
async def test_startup_event_hands_the_daily_report_this_pods_lock_manager(proxy_logging):
"""regression: issue #14809 - the daily report's dedupe lock only works if startup_event
passes the writer's pod_lock_manager down; dropping the argument silently restores the
every-pod-reports behavior."""
proxy_logging.slack_alerting_instance = MagicMock()
proxy_logging.slack_alerting_instance.alert_types = ["daily_reports"]
proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock()
proxy_logging._init_litellm_callbacks = MagicMock()
proxy_logging.update_values = MagicMock()
llm_router = MagicMock()
proxy_logging.startup_event(llm_router=llm_router, redis_usage_cache=None)
await asyncio.sleep(0)
call = proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_args
assert proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_count == 1
assert call.kwargs["pod_lock_manager"] is proxy_logging.db_spend_update_writer.pod_lock_manager
assert call.kwargs["llm_router"] is llm_router
@pytest.mark.asyncio
async def test_startup_event_skips_the_daily_report_when_it_is_not_an_alert_type(proxy_logging):
proxy_logging.slack_alerting_instance = MagicMock()
proxy_logging.slack_alerting_instance.alert_types = []
proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock()
proxy_logging._init_litellm_callbacks = MagicMock()
proxy_logging.update_values = MagicMock()
proxy_logging.startup_event(llm_router=None, redis_usage_cache=None)
await asyncio.sleep(0)
proxy_logging.slack_alerting_instance._run_scheduled_daily_report.assert_not_called()
# ---------------------------------------------------------------------------
# _add_proxy_hooks
# ---------------------------------------------------------------------------