diff --git a/.env.example b/.env.example index dc1fd5a6ccb..eb659202d25 100644 --- a/.env.example +++ b/.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" diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261009120000_proxy_worker_expiry/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261009120000_proxy_worker_expiry/migration.sql new file mode 100644 index 00000000000..81d330a20bb --- /dev/null +++ b/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); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 78f16d9f1fd..b6e55800480 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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 diff --git a/litellm/proxy/common_utils/scheduled_job_stagger.py b/litellm/proxy/common_utils/scheduled_job_stagger.py index e48e9686f13..91c613cb215 100644 --- a/litellm/proxy/common_utils/scheduled_job_stagger.py +++ b/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( diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 3214c0974b4..522703d5b0f 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9a80761b477..7fe823ea872 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 ### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 78f16d9f1fd..b6e55800480 100644 --- a/litellm/proxy/schema.prisma +++ b/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 diff --git a/schema.prisma b/schema.prisma index 78f16d9f1fd..b6e55800480 100644 --- a/schema.prisma +++ b/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 diff --git a/tests/integration/database/test_proxy_worker_heartbeat.py b/tests/integration/database/test_proxy_worker_heartbeat.py new file mode 100644 index 00000000000..7f63af1fa12 --- /dev/null +++ b/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"}, + ] diff --git a/tests/unit/proxy/db/test_proxy_worker_heartbeat.py b/tests/unit/proxy/db/test_proxy_worker_heartbeat.py index 967e4d3471f..87952f0d608 100644 --- a/tests/unit/proxy/db/test_proxy_worker_heartbeat.py +++ b/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