mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): plan an earlier cron edit from now, and cover the job runner
Moving the cleanup cron earlier queued tomorrow's run because the next occurrence was planned after the pending run being cancelled. Plan from finished runs only. Register LiteLLM_ScheduledJobRun as a span table, keep the advisory lock query inside the tick's span, move the non-string cron startup test to the schedule parser, and add unit tests for the runner and job wiring.
This commit is contained in:
Jesus Velarde
parent
09f2926857
commit
309940cd22
6 changed files with 283 additions and 78 deletions
|
|
@ -327,6 +327,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_PromptTable",
|
||||
"LiteLLM_HealthCheckTable",
|
||||
"LiteLLM_SearchToolsTable",
|
||||
"LiteLLM_ScheduledJobRun",
|
||||
"LiteLLM_SSOConfig",
|
||||
"LiteLLM_ManagedVectorStoreIndexTable",
|
||||
"LiteLLM_CacheConfig",
|
||||
|
|
|
|||
|
|
@ -22,8 +22,6 @@ from litellm.proxy.db.db_span import db_span
|
|||
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import bounded_tx, trusted_sql
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_TABLE: Final = "LiteLLM_ScheduledJobRun"
|
||||
|
|
@ -246,7 +244,7 @@ _STATE_SQL: Final = f"""
|
|||
(SELECT max("finished_at") FROM "{_TABLE}"
|
||||
WHERE "job_name" = $1 AND "status" IN ('SUCCEEDED', 'FAILED')) AS "last_finished_at",
|
||||
(SELECT max("scheduled_for") FROM "{_TABLE}"
|
||||
WHERE "job_name" = $1 AND "status" <> 'CANCELLED') AS "last_scheduled_for"
|
||||
WHERE "job_name" = $1 AND "status" NOT IN ('CANCELLED', 'PENDING')) AS "last_scheduled_for"
|
||||
"""
|
||||
_PENDING_SQL: Final = f"""
|
||||
SELECT "id", "schedule", "scheduled_for" FROM "{_TABLE}"
|
||||
|
|
@ -334,7 +332,10 @@ class ScheduledJobRunner:
|
|||
async with db_span("scheduled_job_tick", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx:
|
||||
await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
|
||||
await tx.execute_raw(trusted_sql(f"SET LOCAL lock_timeout = {timeout_ms}"))
|
||||
if not await _try_group_lock(tx, self.job.lock_group):
|
||||
locked: Final = _LOCKED.validate_python(
|
||||
await tx.query_raw(trusted_sql(_GROUP_LOCK_SQL), f"litellm_scheduled_job:{self.job.lock_group}")
|
||||
)
|
||||
if not locked or not locked[0].locked:
|
||||
return None
|
||||
for reaped in _CLAIMED.validate_python(await tx.query_raw(trusted_sql(_REAP_SQL), self.job.lock_group)):
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -419,16 +420,6 @@ class ScheduledJobRunner:
|
|||
verbose_proxy_logger.info("Scheduled job %s run %s %s", self.job.name, run_id, status)
|
||||
|
||||
|
||||
async def _try_group_lock(tx: "Prisma", lock_group: str) -> bool:
|
||||
rows: Final = _LOCKED.validate_python(
|
||||
await tx.query_raw(
|
||||
trusted_sql('SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'),
|
||||
f"litellm_scheduled_job:{lock_group}",
|
||||
)
|
||||
)
|
||||
return bool(rows) and rows[0].locked
|
||||
|
||||
|
||||
class _Locked(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
|
|
@ -436,6 +427,7 @@ class _Locked(BaseModel):
|
|||
|
||||
|
||||
_LOCKED: Final = TypeAdapter(list[_Locked])
|
||||
_GROUP_LOCK_SQL: Final = 'SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'
|
||||
|
||||
|
||||
@lru_cache(maxsize=64)
|
||||
|
|
|
|||
|
|
@ -265,6 +265,22 @@ async def test_a_schedule_change_replaces_the_pending_run_planned_from_the_old_s
|
|||
assert job.started == 0
|
||||
|
||||
|
||||
async def test_moving_a_cron_earlier_plans_the_earlier_occurrence_instead_of_skipping_a_day() -> None:
|
||||
async with _scratch() as database:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
job: Final = _Job(schedule=f"0 {(now + timedelta(hours=2)).hour} * * *")
|
||||
runners: Final = _runners(job)
|
||||
await _tick_all(runners, database)
|
||||
(old,) = database.rows("\"status\" = 'PENDING'")
|
||||
|
||||
job.schedule_value = f"0 {(now + timedelta(hours=1)).hour} * * *"
|
||||
await _tick_all(runners, database)
|
||||
|
||||
(replacement,) = database.rows("\"status\" = 'PENDING'")
|
||||
assert replacement["scheduled_for"] < old["scheduled_for"]
|
||||
assert replacement["scheduled_for"] - datetime.now(timezone.utc) < timedelta(hours=1)
|
||||
|
||||
|
||||
async def test_a_run_cancelled_by_shutdown_is_recorded_as_failed_instead_of_holding_its_lease() -> None:
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job()
|
||||
|
|
|
|||
|
|
@ -1,21 +1,36 @@
|
|||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import pytest
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.proxy.db.scheduled_job_runs import (
|
||||
InvalidSchedule,
|
||||
JobFailed,
|
||||
JobResult,
|
||||
JobScheduler,
|
||||
JobState,
|
||||
JobSucceeded,
|
||||
NewRun,
|
||||
PendingRun,
|
||||
Schedule,
|
||||
ScheduledJob,
|
||||
ScheduledJobRunner,
|
||||
TickPlan,
|
||||
cron_schedule,
|
||||
interval_schedule,
|
||||
next_run_at,
|
||||
plan_tick,
|
||||
register_poll,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
NOW: Final = datetime(2026, 11, 4, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
|
|
@ -167,3 +182,183 @@ def test_cron_schedule_rejects_what_is_not_a_crontab(value: object) -> None:
|
|||
@pytest.mark.parametrize("value", ["soon", "99999999999d", "0s", 86400, None])
|
||||
def test_interval_schedule_rejects_what_is_not_a_positive_duration(value: object) -> None:
|
||||
assert isinstance(interval_schedule("maximum_spend_logs_retention_interval", value), InvalidSchedule)
|
||||
|
||||
|
||||
class _ScriptedDatabase:
|
||||
"""Answers the runner's statements the way Postgres would for one due run, and records every write."""
|
||||
|
||||
def __init__(self, *, locked: bool = True, finish_updates: int = 1, broken: bool = False) -> None:
|
||||
self.locked = locked
|
||||
self.finish_updates = finish_updates
|
||||
self.broken = broken
|
||||
self.writes: tuple[tuple[str, tuple[object, ...]], ...] = ()
|
||||
self.db = self
|
||||
|
||||
@asynccontextmanager
|
||||
async def tx(self, timeout: timedelta) -> AsyncIterator["_ScriptedDatabase"]:
|
||||
if self.broken:
|
||||
raise ConnectionError("database unreachable")
|
||||
yield self
|
||||
|
||||
async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
|
||||
if "pg_try_advisory_xact_lock" in sql:
|
||||
return [{"locked": self.locked}]
|
||||
if "SELECT now()" in sql:
|
||||
return [
|
||||
{
|
||||
"now": NOW,
|
||||
"group_running": False,
|
||||
"job_running": False,
|
||||
"last_finished_at": None,
|
||||
"last_scheduled_for": None,
|
||||
}
|
||||
]
|
||||
if "'RUNNING', \"claimed_by\"" in sql:
|
||||
self.writes = (*self.writes, ("claim", args))
|
||||
return [{"id": args[0], "job_name": "job"}]
|
||||
return []
|
||||
|
||||
async def execute_raw(self, sql: str, *args: object) -> int:
|
||||
if sql.lstrip().startswith("SET LOCAL"):
|
||||
return 0
|
||||
if 'SET "status" = $3' in sql:
|
||||
self.writes = (*self.writes, ("finish", args))
|
||||
self.locked = False
|
||||
return self.finish_updates
|
||||
if sql.lstrip().startswith("INSERT"):
|
||||
self.writes = (*self.writes, ("insert", args))
|
||||
return 1
|
||||
|
||||
def client(self) -> "PrismaClient":
|
||||
return cast("PrismaClient", self) # cast-ok: the runner only reaches the client through .db
|
||||
|
||||
|
||||
def _finishes(database: _ScriptedDatabase) -> tuple[tuple[object, ...], ...]:
|
||||
return tuple(args for kind, args in database.writes if kind == "finish")
|
||||
|
||||
|
||||
class _Job:
|
||||
def __init__(self, result: JobResult | BaseException = JobSucceeded("cleaned")) -> None:
|
||||
self.result = result
|
||||
self.runs = 0
|
||||
self.release = asyncio.Event()
|
||||
self.release.set()
|
||||
|
||||
async def run(self, _prisma_client: object) -> JobResult:
|
||||
self.runs += 1
|
||||
await self.release.wait()
|
||||
if isinstance(self.result, BaseException):
|
||||
raise self.result
|
||||
return self.result
|
||||
|
||||
def runner(self) -> ScheduledJobRunner:
|
||||
definition: Final = ScheduledJob(
|
||||
name="job",
|
||||
lock_group="group",
|
||||
schedule=lambda: _schedule("1h"),
|
||||
lease=lambda: timedelta(minutes=5),
|
||||
run=self.run,
|
||||
)
|
||||
return ScheduledJobRunner(definition, worker_id="pod-a", new_id=lambda: "run-1")
|
||||
|
||||
|
||||
async def test_a_due_run_is_claimed_by_this_pod_run_and_recorded_with_its_log() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job(JobSucceeded("cleaned"))
|
||||
await job.runner().tick(database.client())
|
||||
|
||||
assert job.runs == 1
|
||||
assert ("claim", ("run-1", "pod-a", 300.0)) in database.writes
|
||||
assert _finishes(database) == (("run-1", "pod-a", "SUCCEEDED", "cleaned"),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("result", "expected_log"),
|
||||
[
|
||||
pytest.param(JobFailed("aborted after 3 batch failures"), "aborted after 3 batch failures", id="reported"),
|
||||
pytest.param(RuntimeError("disk on fire"), "RuntimeError: disk on fire", id="raised"),
|
||||
],
|
||||
)
|
||||
async def test_a_failed_run_is_recorded_as_failed_with_why(
|
||||
result: JobResult | BaseException, expected_log: str
|
||||
) -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
await _Job(result).runner().tick(database.client())
|
||||
assert _finishes(database) == (("run-1", "pod-a", "FAILED", expected_log),)
|
||||
|
||||
|
||||
async def test_a_pod_that_does_not_get_the_group_lock_neither_plans_nor_runs() -> None:
|
||||
database: Final = _ScriptedDatabase(locked=False)
|
||||
job: Final = _Job()
|
||||
await job.runner().tick(database.client())
|
||||
assert job.runs == 0
|
||||
assert database.writes == ()
|
||||
|
||||
|
||||
async def test_an_unreachable_database_skips_the_tick_instead_of_raising_into_the_scheduler() -> None:
|
||||
job: Final = _Job()
|
||||
await job.runner().tick(_ScriptedDatabase(broken=True).client())
|
||||
assert job.runs == 0
|
||||
|
||||
|
||||
async def test_a_poll_landing_while_this_pod_still_runs_the_job_does_not_touch_the_database() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
runner: Final = job.runner()
|
||||
first: Final = asyncio.create_task(runner.tick(database.client()))
|
||||
while job.runs == 0:
|
||||
await asyncio.sleep(0)
|
||||
writes_during_run: Final = database.writes
|
||||
|
||||
await runner.tick(database.client())
|
||||
assert database.writes == writes_during_run
|
||||
|
||||
job.release.set()
|
||||
await first
|
||||
assert job.runs == 1
|
||||
|
||||
|
||||
async def test_a_run_cancelled_mid_flight_records_failed_and_still_cancels() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
tick: Final = asyncio.create_task(job.runner().tick(database.client()))
|
||||
while job.runs == 0:
|
||||
await asyncio.sleep(0)
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await tick
|
||||
((run_id, worker, status, log),) = _finishes(database)
|
||||
assert (run_id, worker, status) == ("run-1", "pod-a", "FAILED")
|
||||
assert "cancelled" in str(log)
|
||||
|
||||
|
||||
async def test_a_result_that_arrives_after_the_lease_expired_is_reported_not_recorded(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
database: Final = _ScriptedDatabase(finish_updates=0)
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
await _Job().runner().tick(database.client())
|
||||
assert any("after its lease expired" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
async def test_register_poll_adds_a_wanted_poll_once_and_removes_it_when_no_longer_wanted() -> None:
|
||||
scheduler: Final = AsyncIOScheduler()
|
||||
scheduler.start(paused=True)
|
||||
job_scheduler: Final = cast(JobScheduler, scheduler) # cast-ok: apscheduler ships no type information
|
||||
definition: Final = _Job().runner().job
|
||||
|
||||
def register(wanted: bool) -> None:
|
||||
register_poll(job_scheduler, None, "poll-id", definition, wanted=wanted, poll_seconds=30, misfire_grace_time=60)
|
||||
|
||||
try:
|
||||
register(True)
|
||||
first: Final = scheduler.get_job("poll-id")
|
||||
assert first is not None
|
||||
register(True)
|
||||
assert scheduler.get_job("poll-id") is first, "re-registering replaced the live poll and its runner"
|
||||
register(False)
|
||||
assert scheduler.get_job("poll-id") is None
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
|
|
|||
|
|
@ -6,10 +6,7 @@ import json
|
|||
import logging
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
|
@ -19,9 +16,6 @@ from typing import Final
|
|||
from unittest import mock
|
||||
from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch
|
||||
|
||||
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
||||
|
||||
import click
|
||||
import fastapi.routing
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -34,7 +28,6 @@ from starlette.datastructures import URL
|
|||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm._logging import turn_on_json, verbose_proxy_logger
|
||||
from litellm.caching.caching import RedisCache
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
|
@ -51,12 +44,13 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.login_throttle import LoginThrottle
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.banner import LITELLM_BANNER
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash
|
||||
from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
from litellm.proxy.common_utils.banner import LITELLM_BANNER
|
||||
from tests._master_key import MASTER_KEY
|
||||
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
||||
|
||||
PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234"
|
||||
|
||||
|
|
@ -1002,43 +996,6 @@ async def test_initialize_scheduled_jobs_registers_cleanup_when_retention_lives_
|
|||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_does_not_fall_back_to_the_interval_for_a_non_string_cron(monkeypatch):
|
||||
"""A truthy non-string cron is invalid, so startup must log it and register no cleanup job
|
||||
rather than silently pruning on the default interval the admin never configured."""
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_logging.db_spend_update_writer = MagicMock()
|
||||
settings = {"maximum_daily_tag_spend_retention_period": "30d", "maximum_spend_logs_cleanup_cron": 5}
|
||||
scheduler = AsyncIOScheduler()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.general_settings", settings),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings=settings,
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
assert scheduler.get_job("spend_log_cleanup_job") is None, "invalid cron fell back to the interval"
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
|
||||
"""
|
||||
|
|
@ -1701,7 +1658,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
|
|||
with open(config_path, "w") as f:
|
||||
yaml.dump(test_config, f)
|
||||
|
||||
print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {str(config_path)}")
|
||||
print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {config_path!s}")
|
||||
# Second setting of CONFIG_FILE_PATH to a different value
|
||||
monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path))
|
||||
print(f"config_path: {config_path}")
|
||||
|
|
@ -1720,7 +1677,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
|
|||
yaml.dump(empty_config, f)
|
||||
|
||||
monkeypatch.setenv("LITELLM_MASTER_KEY", test_env_master_key)
|
||||
print("test_env_master_key: {}".format(test_env_master_key))
|
||||
print(f"test_env_master_key: {test_env_master_key}")
|
||||
async with proxy_startup_event(app):
|
||||
from litellm.proxy.proxy_server import master_key
|
||||
|
||||
|
|
@ -2043,7 +2000,7 @@ def test_team_info_masking():
|
|||
all_teams_config=[team1_info],
|
||||
)
|
||||
|
||||
print("Got exception: {}".format(exc_info.value))
|
||||
print(f"Got exception: {exc_info.value}")
|
||||
assert "secret-test-key" not in str(exc_info.value)
|
||||
assert "public-test-key" not in str(exc_info.value)
|
||||
|
||||
|
|
@ -2055,7 +2012,6 @@ async def test_get_all_team_models():
|
|||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import LiteLLM_TeamTable
|
||||
from litellm.proxy.proxy_server import get_all_team_models
|
||||
|
||||
# Mock team data
|
||||
|
|
@ -4037,7 +3993,7 @@ async def test_load_config_warns_per_worker_login_counters_without_general_setti
|
|||
Redis must hear that its counters are per worker even when the config has no general_settings."""
|
||||
import logging
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.auth.login_throttle import warn_login_counters_are_per_worker
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -4197,7 +4153,6 @@ async def test_load_environment_variables_blocks_dangerous_keys():
|
|||
Test that _load_environment_variables rejects dangerous env var keys
|
||||
like PATH, LD_PRELOAD, PYTHONPATH, etc.
|
||||
"""
|
||||
import logging
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -4279,7 +4234,7 @@ async def test_write_config_to_file(monkeypatch):
|
|||
"""
|
||||
Do not write config to file if store_model_in_db is True
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -4327,7 +4282,7 @@ async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch):
|
|||
"""
|
||||
Test that config IS written to file when store_model_in_db is False
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -4424,7 +4379,7 @@ async def test_async_data_generator_midstream_error():
|
|||
try:
|
||||
async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data):
|
||||
yielded_data.append(data)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
# If there's an exception, that's also part of what we want to test
|
||||
pass
|
||||
|
||||
|
|
@ -4493,7 +4448,6 @@ async def test_chat_completion_result_no_nested_none_values():
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from fastapi import Request, Response
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -5560,7 +5514,7 @@ async def test_add_router_settings_from_db_config_merge_logic():
|
|||
This tests how router settings from config file and database are combined,
|
||||
including scenarios where nested dictionaries should be properly merged.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -6713,7 +6667,7 @@ async def test_init_sso_settings_in_db_error_handling():
|
|||
"""
|
||||
Test that _init_sso_settings_in_db handles database errors gracefully.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
@ -6843,7 +6797,6 @@ async def test_init_hashicorp_vault_config_override_retries_on_transport_error()
|
|||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise prisma.errors.ClientNotConnectedError()
|
||||
return None # No config in DB → function returns early after retry.
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock(side_effect=_flaky_find_unique)
|
||||
|
|
@ -6991,9 +6944,7 @@ def test_root_redirect_when_docs_url_not_root_and_redirect_url_set(monkeypatch):
|
|||
routes_to_remove = []
|
||||
for route in app.routes:
|
||||
if hasattr(route, "path") and route.path == "/":
|
||||
if hasattr(route, "methods") and "GET" in route.methods:
|
||||
routes_to_remove.append(route)
|
||||
elif not hasattr(route, "methods"): # Catch-all routes
|
||||
if hasattr(route, "methods") and "GET" in route.methods or not hasattr(route, "methods"):
|
||||
routes_to_remove.append(route)
|
||||
|
||||
for route in routes_to_remove:
|
||||
|
|
@ -10041,7 +9992,7 @@ async def test_reseed_locks_dict_is_bounded():
|
|||
long-lived deployments with high counter-key churn. Inserting more
|
||||
than the cap evicts the oldest entries.
|
||||
"""
|
||||
import litellm.constants as constants
|
||||
from litellm import constants
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
|
||||
orig_locks = SpendCounterReseed._locks.copy()
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
|
||||
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
|
||||
SPEND_LOG_RETENTION_JOB,
|
||||
SPEND_LOG_SIZE_CHECK_JOB,
|
||||
SizeCleanupState,
|
||||
SpendLogCleanup,
|
||||
SpendLogStorage,
|
||||
|
|
@ -33,7 +35,7 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
|
|||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
|
||||
SpendLogCleanupMetrics,
|
||||
)
|
||||
from litellm.proxy.db.scheduled_job_runs import InvalidSchedule
|
||||
from litellm.proxy.db.scheduled_job_runs import InvalidSchedule, JobFailed, JobSucceeded
|
||||
from tests.unit.proxy.db.fake_prisma_engine import engine_call
|
||||
|
||||
|
||||
|
|
@ -2178,3 +2180,51 @@ def test_the_retention_job_follows_the_cleanup_cron_when_set_and_the_retention_i
|
|||
schedule = retention_schedule(settings)
|
||||
assert not isinstance(schedule, InvalidSchedule)
|
||||
assert schedule.spec == expected_spec
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cron", [5, ["0 4 * * *"], "not a cron"])
|
||||
def test_an_invalid_cleanup_cron_is_rejected_instead_of_falling_back_to_the_retention_interval(cron: object):
|
||||
schedule = retention_schedule(
|
||||
{"maximum_spend_logs_cleanup_cron": cron, "maximum_spend_logs_retention_interval": "1h"}
|
||||
)
|
||||
assert isinstance(schedule, InvalidSchedule)
|
||||
|
||||
|
||||
async def test_the_retention_job_records_an_aborted_cleanup_as_failed(monkeypatch):
|
||||
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
||||
|
||||
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 1)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_retention_period": "7d"})
|
||||
mock_prisma_client = MagicMock()
|
||||
_wire_tx(mock_prisma_client.db)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
|
||||
|
||||
result = await SPEND_LOG_RETENTION_JOB.run(mock_prisma_client)
|
||||
|
||||
assert isinstance(result, JobFailed)
|
||||
assert "stop_reason=aborted" in result.log
|
||||
|
||||
|
||||
@pytest.mark.parametrize("job", [SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB], ids=["retention", "size"])
|
||||
async def test_a_cleanup_job_with_no_rule_configured_succeeds_without_touching_the_database(monkeypatch, job):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.execute_raw = AsyncMock()
|
||||
|
||||
result = await job.run(mock_prisma_client)
|
||||
|
||||
assert result == JobSucceeded("no retention or size rule is configured")
|
||||
mock_prisma_client.db.execute_raw.assert_not_awaited()
|
||||
|
||||
|
||||
def test_a_cleanup_lease_outlasts_the_configured_run_budget(monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "10m"})
|
||||
for job in (SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB):
|
||||
assert job.lease().total_seconds() > 600
|
||||
|
||||
|
||||
def test_the_size_check_job_follows_the_size_check_interval(monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_size_check_interval": "6h"})
|
||||
schedule = SPEND_LOG_SIZE_CHECK_JOB.schedule()
|
||||
assert not isinstance(schedule, InvalidSchedule)
|
||||
assert schedule.spec == "interval:21600s"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue