fix: roll up the open deflake fixes for the MCP logging queue, PTU rollup, license gate, and pricing test isolation

litellm/proxy/proxy_server.py

@@ -9240,6 +9240,7 @@ async def _scheduled_ptu_rollup() -> None:
                     prisma_client,
                     pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
                     alert=_alert_ptu_rollup_failure,
+                    router=llm_router,
                 )
 
             scheduler.add_job(

litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py

@@ -14,7 +14,6 @@
 
 import asyncio
 import json
-import sys
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from datetime import date, datetime, time, timedelta, timezone
@@ -326,16 +325,6 @@ class _LoadedDeployments:
     scanned_ids: frozenset[str]
 
 
-def _running_router() -> object | None:
-    """The proxy's router, or None outside a running proxy.
-
-    Read out of ``sys.modules`` rather than imported, so a rollup driven from a test or a
-    script does not pull the whole proxy server in behind it.
-    """
-    proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server")
-    return getattr(proxy_server, "llm_router", None) if proxy_server is not None else None
-
-
 def _config_deployments(router: object | None, *, owned_by_db: frozenset[str]) -> tuple[_PTUDeployment, ...]:
     """Deployments the router holds that no ``LiteLLM_ProxyModelTable`` row owns.
 
@@ -356,15 +345,17 @@ def _config_deployments(router: object | None, *, owned_by_db: frozenset[str]) -
     )
 
 
-async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments:
+async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | None) -> _LoadedDeployments:
     """Every deployment carrying valid manual PTU config, and every id the scan saw.
 
     Reserved capacity is billed by the provider whichever file declared it, so a
     deployment the proxy only knows from config.yaml accrues alongside the stored ones.
+    The router is handed in rather than read off the proxy module, so a run prices exactly
+    the deployments its caller declares and nothing a co-resident process left behind.
     """
     rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many()
     db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or "")))
-    config_records: Final = _config_deployments(_running_router(), owned_by_db=db_ids)
+    config_records: Final = _config_deployments(router, owned_by_db=db_ids)
     models: Final = tuple(
         parsed for parsed in (_parse_ptu_model(row) for row in (*rows, *config_records)) if parsed is not None
     )
@@ -380,6 +371,7 @@ async def run_ptu_flat_cost_rollup(
     prisma_client: "PrismaClient",
     target_date: date | None = None,
     may_prune: bool = True,
+    router: object | None = None,
 ) -> RollupResult:
     """Rollup one UTC day of flat PTU cost across all PTU-configured model deployments.
 
@@ -406,7 +398,7 @@ async def run_ptu_flat_cost_rollup(
     date_str: Final = day.isoformat()
     run_started: Final = datetime.now(timezone.utc)
 
-    loaded: Final = await _load_ptu_models(prisma_client)
+    loaded: Final = await _load_ptu_models(prisma_client, router=router)
     ptu_models: Final = loaded.models
     charges: Final = _aggregate_charges(ptu_models, day)
 
@@ -527,6 +519,7 @@ async def _existing_sentinel_keys(
 async def run_ptu_flat_cost_backfill(
     prisma_client: "PrismaClient",
     today: date | None = None,
+    router: object | None = None,
 ) -> BackfillResult:
     """Price the elapsed days of every PTU window that carry no sentinel row yet.
 
@@ -546,7 +539,7 @@ async def run_ptu_flat_cost_backfill(
         verbose_proxy_logger.warning("PTU backfill: prisma_client is None, skipping")
         return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0)
 
-    ptu_models: Final = (await _load_ptu_models(prisma_client)).models
+    ptu_models: Final = (await _load_ptu_models(prisma_client, router=router)).models
     days: Final = _backfill_window(ptu_models, end)
 
     if not days:
@@ -591,6 +584,7 @@ async def run_scheduled_ptu_rollup(
     pod_lock_manager: "PodLockManager | None" = None,
     target_date: date | None = None,
     alert: Callable[[str], Awaitable[None]] | None = None,
+    router: object | None = None,
 ) -> RollupResult | None:
     """Run the daily rollup under a cross-pod lock so only one proxy reconciles a day.
 
@@ -615,7 +609,7 @@ async def run_scheduled_ptu_rollup(
         return None
 
     if pod_lock_manager is None or pod_lock_manager.redis_cache is None:
-        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False)
+        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False, router=router)
 
     if not await pod_lock_manager.acquire_lock(cronjob_id=PTU_ROLLUP_JOB_ID, ttl=PTU_ROLLUP_LOCK_TTL_SECONDS):
         if await _lock_is_held(pod_lock_manager):
@@ -629,10 +623,10 @@ async def run_scheduled_ptu_rollup(
             "PTU rollup: could not take the rollup lock and no other pod holds it, "
             "running unguarded rather than skipping the day"
         )
-        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False)
+        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False, router=router)
 
     try:
-        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=True)
+        return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=True, router=router)
     finally:
         await pod_lock_manager.release_lock(cronjob_id=PTU_ROLLUP_JOB_ID)
 
@@ -657,6 +651,7 @@ async def _run_and_alert(
     target_date: date | None,
     alert: "Callable[[str], Awaitable[None]] | None",
     may_prune: bool = True,
+    router: object | None = None,
 ) -> RollupResult:
     """Reconcile the day, catch up any days left unpriced, and alert on charges that did not land.
 
@@ -669,7 +664,9 @@ async def _run_and_alert(
     explicit date means reconcile exactly that day, so it stays a single-day operation.
     Its failure is contained: the day's own result is returned either way.
     """
-    result: Final = await run_ptu_flat_cost_rollup(prisma_client, target_date=target_date, may_prune=may_prune)
+    result: Final = await run_ptu_flat_cost_rollup(
+        prisma_client, target_date=target_date, may_prune=may_prune, router=router
+    )
     if result.rows_failed:
         await _deliver_alert(
             alert,
@@ -686,22 +683,23 @@ async def _run_and_alert(
             "by the provider with nothing attributing it here. Extend the window, or retire the deployment.",
         )
     if target_date is None:
-        await _backfill_and_alert(prisma_client, alert=alert)
+        await _backfill_and_alert(prisma_client, alert=alert, router=router)
     return result
 
 
 async def _backfill_and_alert(
     prisma_client: "PrismaClient",
     *,
     alert: "Callable[[str], Awaitable[None]] | None",
+    router: object | None = None,
 ) -> None:
     """Catch up unpriced PTU days, alerting on charges that did not land.
 
     Never raises: the day's own rollup has already run and its result must reach the
     caller whatever the catch-up pass does.
     """
     try:
-        backfill: Final = await run_ptu_flat_cost_backfill(prisma_client)
+        backfill: Final = await run_ptu_flat_cost_backfill(prisma_client, router=router)
     except Exception as exc:  # noqa: BLE001  # the catch-up pass must not fail the day's rollup
         verbose_proxy_logger.error("PTU backfill: catch-up pass failed, the day's rollup still stands: %s", exc)
         return

tests/code_coverage_tests/check_licenses.py

@@ -5,8 +5,9 @@
 from pathlib import Path
 import re
 import sys
+import time
 import tomllib
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Callable, Dict, Final, List, Optional, Protocol, Set, Tuple
 
 from packaging.requirements import Requirement
 import requests
@@ -37,6 +38,13 @@
 # of the identifier, not an operator.
 _SPDX_OPERATOR_SPLIT = re.compile(r"\s+(?:OR|AND)\s+")
 _SPDX_WITH_SUFFIX = re.compile(r"\s+WITH\s+.*", re.DOTALL)
+_PYPI_FETCH_ATTEMPTS: Final[int] = 3
+_PYPI_FETCH_BACKOFF_SECONDS: Final[float] = 0.5
+
+
+class _HttpGet(Protocol):
+    def __call__(self, url: str, *, timeout: float) -> requests.Response:
+        ...
 
 
 @dataclass
@@ -50,7 +58,10 @@ class PackageLicense:
 
 class LicenseChecker:
     def __init__(
-        self, config_file: Path = Path("./tests/code_coverage_tests/liccheck.ini")
+        self,
+        config_file: Path = Path("./tests/code_coverage_tests/liccheck.ini"),
+        http_get: Optional[_HttpGet] = None,
+        sleep: Optional[Callable[[float], None]] = None,
     ):
         if not config_file.exists():
             print(f"Error: Config file {config_file} not found")
@@ -79,6 +90,8 @@ def __init__(
 
         # Track package results
         self.package_results: List[PackageLicense] = []
+        self._http_get = http_get
+        self._sleep = sleep
 
     @staticmethod
     def _normalize_package_name(package_name: str) -> str:
@@ -123,21 +136,38 @@ def get_package_license_from_pypi(
         last resort derives the license from the ``License :: OSI Approved ::
         ...`` trove classifiers.
         """
-        try:
-            url = f"https://pypi.org/pypi/{package_name}/{version}/json"
-            response = requests.get(url, timeout=10)
-            response.raise_for_status()
-            info = response.json().get("info", {}) or {}
-            return (
-                info.get("license_expression")
-                or info.get("license")
-                or self._license_from_classifiers(info.get("classifiers") or [])
-            )
-        except Exception as e:
-            print(
-                f"Warning: Failed to fetch license for {package_name} {version}: {str(e)}"
-            )
-            return None
+        url = f"https://pypi.org/pypi/{package_name}/{version}/json"
+        http_get = self._http_get if self._http_get is not None else requests.get
+        sleep = self._sleep if self._sleep is not None else time.sleep
+
+        for attempt in range(_PYPI_FETCH_ATTEMPTS):
+            try:
+                response = http_get(url, timeout=10)
+                response.raise_for_status()
+                info = response.json().get("info", {}) or {}
+                return (
+                    info.get("license_expression")
+                    or info.get("license")
+                    or self._license_from_classifiers(info.get("classifiers") or [])
+                )
+            except Exception as error:
+                if self._is_retryable_pypi_error(error) and attempt < _PYPI_FETCH_ATTEMPTS - 1:
+                    sleep(_PYPI_FETCH_BACKOFF_SECONDS)
+                    continue
+                print(
+                    f"Warning: Failed to fetch license for {package_name} {version}: {str(error)}"
+                )
+                return None
+        return None
+
+    @staticmethod
+    def _is_retryable_pypi_error(error: Exception) -> bool:
+        if isinstance(error, (requests.ConnectionError, requests.Timeout)):
+            return True
+        if not isinstance(error, requests.HTTPError) or error.response is None:
+            return False
+        status_code = error.response.status_code
+        return status_code == 429 or status_code >= 500
 
     @staticmethod
     def _license_from_classifiers(classifiers: List[str]) -> Optional[str]:

tests/mcp_tests/conftest.py

@@ -45,6 +45,22 @@ def setup_and_teardown():
     asyncio.set_event_loop(None)  # Remove the reference to the loop
 
 
+@pytest.fixture(scope="function", autouse=True)
+async def drain_logging_worker():
+    """
+    The logging queue is bound to the running loop, so anything left queued when a test's loop
+    goes away is carried onto the next test's loop and fires against its callbacks.
+    """
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    yield
+
+    try:
+        await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10)
+    except asyncio.TimeoutError:
+        pass
+
+
 def pytest_collection_modifyitems(config, items):
     # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests
     custom_logger_tests = [

tests/test_litellm/conftest.py

@@ -375,6 +375,9 @@ def isolate_litellm_state():
         litellm.in_memory_llm_clients_cache.flush_cache()
     image_handling_module.in_memory_cache.flush_cache()
     _reset_module_level_aws_auth_caches()
+    # litellm.get_model_info() memoizes ModelInfo built from litellm.model_cost, so a
+    # test that rebinds the cost map leaves later tests pricing against the old map.
+    litellm_utils_module._invalidate_model_cost_lowercase_map()
 
     # Clear all callback lists to prevent cross-test contamination
     if hasattr(litellm, "callbacks"):
@@ -418,6 +421,7 @@ def isolate_litellm_state():
 
     litellm_utils_module._runtime_registered_model_cost.clear()
     litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
+    litellm_utils_module._invalidate_model_cost_lowercase_map()
 
     for _router in tuple(litellm_router_module._live_routers):
         litellm_router_module._live_routers.discard(_router)