Merge 3710fa94e0 into b4a8bc715f

.env.example

@@ -30,3 +30,5 @@ INFINITY_API_KEY = ""
 LITELLM_MASTER_KEY = ""
 DATABASE_URL = "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
 STORE_MODEL_IN_DB = "True"
+
+PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS = "60"

litellm-proxy-extras/litellm_proxy_extras/migrations/20261009120000_proxy_worker_expiry/migration.sql

@@ -0,0 +1,2 @@
+SET LOCAL lock_timeout = '10s';
+ALTER TABLE "LiteLLM_ProxyWorkerHeartbeat" ADD COLUMN IF NOT EXISTS "expires_at" TIMESTAMP(3);

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1130,6 +1130,7 @@ model LiteLLM_ProxyWorkerHeartbeat {
   hostname          String
   started_at        DateTime @default(now())
   last_heartbeat_at DateTime @default(now())
+  expires_at        DateTime?
 }
 
 // Track the status of cron jobs running. Only allow one pod to run the job at a time

litellm/proxy/common_utils/scheduled_job_stagger.py

@@ -207,13 +207,14 @@ def _window_for(*, job_id: str, period_seconds: int | None, settings: ScheduledJ
     return min(limit for limit in limits if limit is not None)
 
 
-def _clamped_override(*, job_id: str, requested: int) -> int:
-    horizon: Final = DEFAULT_CRON_DEDUPE_SECONDS.get(job_id)
+def _clamped_override(*, job_id: str, requested: int, period_seconds: int | None) -> int:
+    horizon: Final = (
+        period_seconds if job_id == "proxy_worker_heartbeat_job" else DEFAULT_CRON_DEDUPE_SECONDS.get(job_id)
+    )
     if horizon is None or requested < horizon:
         return requested
     verbose_proxy_logger.warning(
-        "general_settings.%s.offsets[%s]=%ss would place replicas more than %ss apart, "
-        "which is long enough for a second replica to redo the run; using %ss instead",
+        "general_settings.%s.offsets[%s]=%ss exceeds the job's %ss scheduling window; using %ss instead",
         GENERAL_SETTINGS_KEY,
         job_id,
         requested,
@@ -233,7 +234,7 @@ def _offset_for(
 ) -> int:
     override: Final = settings.offsets.get(job_id)
     if override is not None:
-        return _clamped_override(job_id=job_id, requested=max(0, override))
+        return _clamped_override(job_id=job_id, requested=max(0, override), period_seconds=period_seconds)
     if not staggerable:
         return 0
     return offset_seconds(

litellm/proxy/db/proxy_worker_heartbeat.py

@@ -1,7 +1,7 @@
 """
 Live proxy worker census, one row per worker process.
 
-Every uvicorn worker upserts its own row on a fixed heartbeat, so counting
+Every uvicorn worker upserts its own row on a configurable heartbeat, so counting
 rows with a recent heartbeat answers "how many workers share this database?"
 without any coordination. The Admin UI's "no Redis" banner uses that count to
 hide itself for deployments that are provably a single worker, where per-worker
@@ -13,6 +13,7 @@ still agree.
 from __future__ import annotations
 
 import socket
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Final, Protocol
 
 from pydantic import TypeAdapter
@@ -20,6 +21,8 @@ from typing_extensions import LiteralString, ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
+from litellm.constants import APSCHEDULER_MISFIRE_GRACE_TIME
+from litellm.litellm_core_utils.env_utils import get_env_int
 from litellm.proxy.db.db_span import db_span
 from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
 
@@ -31,19 +34,20 @@ PROXY_WORKER_LIVENESS_WINDOW_SECONDS: Final = 3 * PROXY_WORKER_HEARTBEAT_INTERVA
 STALE_ROW_RETENTION_SECONDS: Final = 3600
 
 BEAT_SQL: Final = """
-INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at)
-VALUES ($1, $2, NOW())
-ON CONFLICT (worker_id) DO UPDATE SET last_heartbeat_at = NOW()
+INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at, expires_at)
+VALUES ($1, $2, NOW(), NOW() + make_interval(secs => $3))
+ON CONFLICT (worker_id) DO UPDATE SET last_heartbeat_at = NOW(), expires_at = EXCLUDED.expires_at
 """
 
 PRUNE_SQL: Final = """
 DELETE FROM "LiteLLM_ProxyWorkerHeartbeat"
 WHERE last_heartbeat_at < NOW() - make_interval(secs => $1)
+AND COALESCE(expires_at, last_heartbeat_at + make_interval(secs => $2)) <= NOW()
 """
 
 COUNT_SQL: Final = """
 SELECT COUNT(*)::int AS live_workers FROM "LiteLLM_ProxyWorkerHeartbeat"
-WHERE last_heartbeat_at > NOW() - make_interval(secs => $1)
+WHERE COALESCE(expires_at, last_heartbeat_at + make_interval(secs => $1)) > NOW()
 """
 
 DEREGISTER_SQL: Final = """
@@ -67,22 +71,65 @@ def _heartbeat_execute(prisma_client: PrismaClient) -> _HeartbeatExecute:
     return prisma_client.db.execute_raw
 
 
+class HeartbeatScheduler(Protocol):
+    def add_job(
+        self,
+        func: Callable[[], Awaitable[None]],
+        trigger: str,
+        *,
+        seconds: int,
+        id: str,
+        replace_existing: bool,
+        misfire_grace_time: int,
+    ) -> object: ...
+
+
+def _heartbeat_interval_seconds() -> int:
+    interval: Final = get_env_int("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS)
+    if interval >= 0:
+        return interval
+    verbose_proxy_logger.warning("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS must be non-negative; using 60 seconds")
+    return PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS
+
+
 class ProxyWorkerHeartbeat:
     def __init__(self, prisma_client: PrismaClient, worker_id: str | None = None) -> None:
         self.prisma_client: Final = prisma_client
         self.worker_id: Final[str] = worker_id or str(uuid.uuid4())
         self.hostname: Final = socket.gethostname()
+        self.interval_seconds: Final = _heartbeat_interval_seconds()
+
+    async def start(self, scheduler: HeartbeatScheduler) -> None:
+        if self.interval_seconds == 0:
+            return
+        await self.beat()
+        scheduler.add_job(
+            self.beat,
+            "interval",
+            seconds=self.interval_seconds,
+            id="proxy_worker_heartbeat_job",
+            replace_existing=True,
+            misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
+        )
 
     async def beat(self) -> None:
+        if self.interval_seconds == 0:
+            return
         try:
             async with db_span("proxy_worker_heartbeat", "LiteLLM_ProxyWorkerHeartbeat"):
-                _ = await _heartbeat_execute(self.prisma_client)(BEAT_SQL, self.worker_id, self.hostname)
+                _ = await _heartbeat_execute(self.prisma_client)(
+                    BEAT_SQL, self.worker_id, self.hostname, 3 * self.interval_seconds
+                )
             async with db_span("prune_proxy_worker_heartbeats", "LiteLLM_ProxyWorkerHeartbeat"):
-                _ = await _heartbeat_execute(self.prisma_client)(PRUNE_SQL, STALE_ROW_RETENTION_SECONDS)
+                _ = await _heartbeat_execute(self.prisma_client)(
+                    PRUNE_SQL, STALE_ROW_RETENTION_SECONDS, PROXY_WORKER_LIVENESS_WINDOW_SECONDS
+                )
         except Exception as beat_err:  # noqa: BLE001  # a missed heartbeat must never take down the worker
             verbose_proxy_logger.debug("Proxy worker heartbeat write failed: %s", beat_err)
 
     async def deregister(self) -> None:
+        if self.interval_seconds == 0:
+            return
         try:
             async with db_span("deregister_proxy_worker", "LiteLLM_ProxyWorkerHeartbeat"):
                 _ = await _heartbeat_execute(self.prisma_client)(DEREGISTER_SQL, self.worker_id)
@@ -96,6 +143,9 @@ async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None:
     cannot answer. Callers must treat None as "unknown", not as zero. Always
     counts on the primary: a lagging read replica must never undercount.
     """
+    interval_seconds: Final = _heartbeat_interval_seconds()
+    if interval_seconds == 0:
+        return None
     try:
         db: Final = prisma_client.db
         primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db

litellm/proxy/proxy_server.py

@@ -674,7 +674,7 @@ from litellm.proxy.db.master_key_migration import (
     migrate_if_requested,
 )
 from litellm.proxy.db.proxy_worker_heartbeat import (
-    PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
+    HeartbeatScheduler,
     ProxyWorkerHeartbeat,
 )
 from litellm.proxy.db.request_error_tracking import (
@@ -11635,14 +11635,8 @@ class ProxyStartupEvent:
 
         ### PROXY WORKER HEARTBEAT ###
         worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client)
-        await worker_heartbeat.beat()
-        scheduler.add_job(
-            worker_heartbeat.beat,
-            "interval",
-            seconds=PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
-            id="proxy_worker_heartbeat_job",
-            replace_existing=True,
-            misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
+        await worker_heartbeat.start(
+            cast(HeartbeatScheduler, scheduler)  # cast-ok: APScheduler has no stubs for its add_job API
         )
 
         ### RESET BUDGET ###

litellm/proxy/schema.prisma

@@ -1130,6 +1130,7 @@ model LiteLLM_ProxyWorkerHeartbeat {
   hostname          String
   started_at        DateTime @default(now())
   last_heartbeat_at DateTime @default(now())
+  expires_at        DateTime?
 }
 
 // Track the status of cron jobs running. Only allow one pod to run the job at a time

schema.prisma

@@ -1130,6 +1130,7 @@ model LiteLLM_ProxyWorkerHeartbeat {
   hostname          String
   started_at        DateTime @default(now())
   last_heartbeat_at DateTime @default(now())
+  expires_at        DateTime?
 }
 
 // Track the status of cron jobs running. Only allow one pod to run the job at a time

tests/integration/database/test_proxy_worker_heartbeat.py

@@ -0,0 +1,100 @@
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Final
+
+import psycopg
+import pytest
+import pytest_asyncio
+
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.db.proxy_worker_heartbeat import ProxyWorkerHeartbeat, count_live_proxy_workers
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+from tests.integration._support.database import read_rows, scratch_database, write_rows
+
+
+@pytest_asyncio.fixture(loop_scope="function")
+async def heartbeat_database(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[tuple[PrismaClient, str]]:
+    with scratch_database() as url:
+        write_rows(
+            'CREATE TABLE "LiteLLM_ProxyWorkerHeartbeat" (worker_id TEXT PRIMARY KEY, hostname TEXT NOT NULL, '
+            "started_at TIMESTAMP NOT NULL DEFAULT NOW(), last_heartbeat_at TIMESTAMP NOT NULL DEFAULT NOW())",
+            (),
+            database_url=url,
+        )
+        monkeypatch.setenv("DATABASE_URL", url)
+        monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
+        client: Final = PrismaClient(url, ProxyLogging(UserApiKeyCache()))
+        await client.connect()
+        try:
+            migration: Final = (
+                Path(__file__).resolve().parents[3]
+                / "litellm-proxy-extras/litellm_proxy_extras/migrations/20261009120000_proxy_worker_expiry/migration.sql"
+            )
+            with psycopg.connect(url) as connection:
+                _ = connection.execute(migration.read_text())
+                _ = connection.execute(migration.read_text())
+            yield client, url
+        finally:
+            await client.disconnect()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("slow_interval", "elapsed"), ((1800, 600), (7200, 3700)))
+async def test_fast_worker_neither_hides_nor_prunes_a_live_slow_worker(
+    heartbeat_database: tuple[PrismaClient, str], monkeypatch: pytest.MonkeyPatch, slow_interval: int, elapsed: int
+) -> None:
+    client, url = heartbeat_database
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", str(slow_interval))
+    slow: Final = ProxyWorkerHeartbeat(client, worker_id="slow")
+    await slow.beat()
+    write_rows(
+        'UPDATE "LiteLLM_ProxyWorkerHeartbeat" '
+        "SET last_heartbeat_at = last_heartbeat_at - make_interval(secs => %s::int), "
+        "expires_at = expires_at - make_interval(secs => %s::int)",
+        (str(elapsed), str(elapsed)),
+        database_url=url,
+    )
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", "60")
+    await ProxyWorkerHeartbeat(client, worker_id="fast").beat()
+
+    assert await count_live_proxy_workers(client) == 2
+    assert read_rows(
+        'SELECT worker_id FROM "LiteLLM_ProxyWorkerHeartbeat" ORDER BY worker_id', (), database_url=url
+    ) == [{"worker_id": "fast"}, {"worker_id": "slow"}]
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", str(slow_interval))
+    assert await count_live_proxy_workers(client) == 2
+
+
+@pytest.mark.asyncio
+async def test_expired_workers_are_excluded_and_legacy_rows_keep_the_default_window(
+    heartbeat_database: tuple[PrismaClient, str], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    client, url = heartbeat_database
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", "15")
+    await ProxyWorkerHeartbeat(client, worker_id="expired").beat()
+    write_rows(
+        'UPDATE "LiteLLM_ProxyWorkerHeartbeat" '
+        "SET last_heartbeat_at = NOW() - INTERVAL '50 seconds', expires_at = NOW() - INTERVAL '5 seconds'",
+        (),
+        database_url=url,
+    )
+    write_rows(
+        'INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at) '
+        "VALUES ('legacy-live', 'old-worker', NOW() - INTERVAL '120 seconds'), "
+        "('legacy-expired', 'old-worker', NOW() - INTERVAL '181 seconds'), "
+        "('legacy-stale', 'old-worker', NOW() - INTERVAL '2 hours')",
+        (),
+        database_url=url,
+    )
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", "1800")
+    await ProxyWorkerHeartbeat(client, worker_id="current").beat()
+
+    assert await count_live_proxy_workers(client) == 2
+    assert read_rows(
+        'SELECT worker_id FROM "LiteLLM_ProxyWorkerHeartbeat" ORDER BY worker_id', (), database_url=url
+    ) == [
+        {"worker_id": "current"},
+        {"worker_id": "expired"},
+        {"worker_id": "legacy-expired"},
+        {"worker_id": "legacy-live"},
+    ]

tests/unit/proxy/db/test_proxy_worker_heartbeat.py

@@ -1,7 +1,14 @@
 from collections.abc import Awaitable, Callable
+from datetime import timedelta
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.triggers.interval import IntervalTrigger
+
+from litellm.proxy._types import ScheduledJobStaggerSettings
+from litellm.proxy.common_utils.scheduled_job_stagger import apply_scheduled_job_stagger
 
 from litellm.proxy.db.proxy_worker_heartbeat import (
     BEAT_SQL,
@@ -17,6 +24,27 @@ from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
 from tests.unit.proxy.db.fake_prisma_engine import engine_call
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("offset", (0, 14, 15, 120))
+async def test_startup_heartbeat_survives_an_explicit_stagger_offset(
+    monkeypatch: pytest.MonkeyPatch, offset: int
+) -> None:
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", "15")
+    scheduler: Final = AsyncIOScheduler()
+    await ProxyWorkerHeartbeat(_prisma(), worker_id="staggered").start(scheduler)
+    job: Final = scheduler.get_job("proxy_worker_heartbeat_job")
+    assert job is not None
+    started_at: Final = job.trigger.start_date - timedelta(seconds=15)
+    apply_scheduled_job_stagger(
+        scheduler=scheduler,
+        settings=ScheduledJobStaggerSettings(offsets={"proxy_worker_heartbeat_job": offset}),
+        identity="test-worker",
+    )
+    next_fire: Final = job.trigger.get_next_fire_time(None, started_at)
+    assert next_fire is not None
+    assert next_fire - started_at < timedelta(seconds=30)
+
+
 def _prisma():
     prisma = MagicMock()
     prisma.db.execute_raw = AsyncMock()
@@ -24,14 +52,91 @@ def _prisma():
     return prisma
 
 
+@pytest.mark.asyncio
+async def test_disabled_heartbeat_does_not_access_the_database(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", "0")
+    prisma: Final = _prisma()
+    heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="disabled-worker")
+
+    await heartbeat.beat()
+    await heartbeat.deregister()
+
+    assert await count_live_proxy_workers(prisma) is None
+    prisma.db.execute_raw.assert_not_awaited()
+    prisma.db.query_raw.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("configured_interval", "expected_interval"),
+    ((None, 60), ("15", 15), ("1800", 1800), ("-1", 60), ("invalid", 60), ("", 60), ("1.5", 60)),
+)
+async def test_heartbeat_uses_configured_liveness_and_retention(
+    monkeypatch: pytest.MonkeyPatch, configured_interval: str | None, expected_interval: int
+) -> None:
+    if configured_interval is None:
+        monkeypatch.delenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", raising=False)
+    else:
+        monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", configured_interval)
+    prisma: Final = _prisma()
+    prisma.db.query_raw.return_value = [{"live_workers": 2}]
+    heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="configured-worker")
+
+    await heartbeat.beat()
+
+    assert prisma.db.execute_raw.call_args_list[0].args == (
+        BEAT_SQL,
+        "configured-worker",
+        heartbeat.hostname,
+        3 * expected_interval,
+    )
+    assert prisma.db.execute_raw.call_args_list[1].args == (
+        PRUNE_SQL,
+        STALE_ROW_RETENTION_SECONDS,
+        PROXY_WORKER_LIVENESS_WINDOW_SECONDS,
+    )
+    assert await count_live_proxy_workers(prisma) == 2
+    prisma.db.query_raw.assert_awaited_once_with(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("interval_seconds", (0, 15, 60, 1800))
+async def test_start_registers_only_enabled_heartbeats(monkeypatch: pytest.MonkeyPatch, interval_seconds: int) -> None:
+    monkeypatch.setenv("PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS", str(interval_seconds))
+    prisma: Final = _prisma()
+    heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="scheduled-worker")
+    scheduler: Final = AsyncIOScheduler()
+
+    await heartbeat.start(scheduler)
+
+    job: Final = scheduler.get_job("proxy_worker_heartbeat_job")
+    if interval_seconds == 0:
+        assert job is None
+        prisma.db.execute_raw.assert_not_awaited()
+        return
+
+    assert job is not None
+    assert isinstance(job.trigger, IntervalTrigger)
+    assert job.trigger.interval.total_seconds() == interval_seconds
+    assert prisma.db.execute_raw.await_count == 2
+    await job.func()
+    assert prisma.db.execute_raw.await_count == 4
+    assert prisma.db.execute_raw.call_args_list[2].args == (
+        BEAT_SQL,
+        "scheduled-worker",
+        heartbeat.hostname,
+        3 * interval_seconds,
+    )
+
+
 @pytest.mark.asyncio
 async def test_beat_upserts_own_row_then_prunes_stale_rows():
     prisma = _prisma()
     heartbeat = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1")
     await heartbeat.beat()
     calls = prisma.db.execute_raw.call_args_list
-    assert calls[0].args == (BEAT_SQL, "worker-1", heartbeat.hostname)
-    assert calls[1].args == (PRUNE_SQL, STALE_ROW_RETENTION_SECONDS)
+    assert calls[0].args == (BEAT_SQL, "worker-1", heartbeat.hostname, PROXY_WORKER_LIVENESS_WINDOW_SECONDS)
+    assert calls[1].args == (PRUNE_SQL, STALE_ROW_RETENTION_SECONDS, PROXY_WORKER_LIVENESS_WINDOW_SECONDS)
 
 
 @pytest.mark.asyncio