mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(warm-pool): async maintenance loop refills + reaps team pools
This commit is contained in:
Ishaan Jaffer
parent
92e8047805
commit
cf8f240386
1 changed files with 402 additions and 0 deletions
402
litellm/proxy/agent_session_endpoints/warm_pool/manager.py
Normal file
402
litellm/proxy/agent_session_endpoints/warm_pool/manager.py
Normal file
|
|
@ -0,0 +1,402 @@
|
|||
"""Warm pool maintenance loop (LIT-2890).
|
||||
|
||||
Per-team background task that:
|
||||
1. Reads `LiteLLM_AgentVMConfig.warm_pool_size` for every team where the
|
||||
feature is enabled
|
||||
2. Counts how many `LiteLLM_AgentVM` rows are currently `state='warm'`
|
||||
3. Spawns the deficit by calling `EC2Provider.provision(mode='warm')`
|
||||
4. Reaps `state='warm'` rows older than `max_idle_minutes` (validation #8)
|
||||
|
||||
Runs every `WARM_POOL_TICK_SECONDS` seconds. Designed to be started ONCE
|
||||
on proxy boot via `WarmPoolManager.start()` from `proxy_server.py`.
|
||||
|
||||
The hot path (`POST /v2/sessions`) calls `attach.attach_warm_vm` directly
|
||||
and never blocks on the maintenance loop — the loop only refills the pool
|
||||
back to `warm_pool_size` after an attach.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.agent_session_endpoints.vm_providers.base import (
|
||||
AwsCreds,
|
||||
Ec2Config,
|
||||
ProvisionContext,
|
||||
Repo,
|
||||
)
|
||||
from litellm.proxy.agent_session_endpoints.vm_providers.factory import (
|
||||
get_vm_provider as build_vm_provider,
|
||||
)
|
||||
from litellm.proxy.agent_session_endpoints.vm_providers.team_config import (
|
||||
get_team_vm_config,
|
||||
)
|
||||
|
||||
# Tick rate for the maintenance loop. Per LIT-2890 spec: every 30s.
|
||||
# Tunable via `WARM_POOL_TICK_SECONDS` env var for tests (set to 1 there).
|
||||
WARM_POOL_TICK_SECONDS = 30
|
||||
|
||||
# How many spawn calls we issue concurrently per team per tick. Cap so a
|
||||
# single team can't saturate the proxy's thread pool.
|
||||
MAX_CONCURRENT_SPAWNS_PER_TEAM = 5
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _new_pool_vm_id() -> str:
|
||||
"""Internal id used while the EC2 provision is in flight.
|
||||
|
||||
The row is rekeyed to the EC2 instance id when the provider returns;
|
||||
callers should use this id as a placeholder only.
|
||||
"""
|
||||
return f"vm_pending_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
class WarmPoolManager:
|
||||
"""Async background task that refills + reaps warm-pool VMs.
|
||||
|
||||
Lifecycle:
|
||||
- `start()` schedules the loop on the running event loop
|
||||
- `stop()` cancels the loop and awaits its completion
|
||||
- `tick()` runs ONE iteration (called by start; useful in tests)
|
||||
|
||||
Idempotent restart: calling `start()` while running is a no-op.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prisma_getter: Optional[Any] = None,
|
||||
provider_settings: Optional[Dict[str, Any]] = None,
|
||||
tick_seconds: float = WARM_POOL_TICK_SECONDS,
|
||||
) -> None:
|
||||
self._tick_seconds = tick_seconds
|
||||
self._provider_settings = provider_settings or {}
|
||||
# Allow tests to inject a prisma client without spinning up the proxy.
|
||||
self._prisma_getter = prisma_getter
|
||||
self._task: Optional[asyncio.Task[None]] = None
|
||||
self._stopping = asyncio.Event()
|
||||
|
||||
# ---------- public lifecycle ----------
|
||||
|
||||
def start(self) -> None:
|
||||
"""Schedule the loop on the running event loop. Idempotent."""
|
||||
if self._task is not None and not self._task.done():
|
||||
return
|
||||
self._stopping.clear()
|
||||
loop = asyncio.get_event_loop()
|
||||
self._task = loop.create_task(self._run_forever())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Cancel the loop and wait for it to exit."""
|
||||
self._stopping.set()
|
||||
if self._task is None:
|
||||
return
|
||||
self._task.cancel()
|
||||
try:
|
||||
await self._task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
self._task = None
|
||||
|
||||
# ---------- core loop ----------
|
||||
|
||||
async def _run_forever(self) -> None:
|
||||
verbose_proxy_logger.info(
|
||||
"warm_pool.manager: started (tick=%ss)", self._tick_seconds
|
||||
)
|
||||
while not self._stopping.is_set():
|
||||
try:
|
||||
await self.tick()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.exception(
|
||||
"warm_pool.manager: tick failed: %s", exc
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._stopping.wait(), timeout=self._tick_seconds
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
verbose_proxy_logger.info("warm_pool.manager: stopped")
|
||||
|
||||
async def tick(self) -> None:
|
||||
"""Run one maintenance iteration: refill + reap, per team."""
|
||||
prisma_client = self._get_prisma_client()
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
||||
teams = await self._enabled_teams(prisma_client)
|
||||
if not teams:
|
||||
return
|
||||
|
||||
# Per-team work: refill deficit, reap stale.
|
||||
await asyncio.gather(
|
||||
*(self._maintain_team(prisma_client, team) for team in teams),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
# ---------- per-team work ----------
|
||||
|
||||
async def _maintain_team(self, prisma_client: Any, vm_config: Any) -> None:
|
||||
team_id = vm_config.team_id
|
||||
desired = int(getattr(vm_config, "warm_pool_size", 0) or 0)
|
||||
max_idle_minutes = int(getattr(vm_config, "max_idle_minutes", 30) or 30)
|
||||
|
||||
# Reap first — frees slots before we refill.
|
||||
await self._reap_stale_warm(prisma_client, team_id, max_idle_minutes)
|
||||
|
||||
# Then count what's still warm + provisioning, spawn the deficit.
|
||||
existing = await prisma_client.db.litellm_agentvm.find_many(
|
||||
where={
|
||||
"team_id": team_id,
|
||||
"state": {"in": ["provisioning", "warm"]},
|
||||
}
|
||||
)
|
||||
deficit = max(0, desired - len(existing))
|
||||
if deficit == 0:
|
||||
# Cap: if we have more warm than desired (e.g. user shrunk pool),
|
||||
# terminate the oldest excess.
|
||||
await self._shrink_to_desired(prisma_client, team_id, existing, desired)
|
||||
return
|
||||
|
||||
spawn_n = min(deficit, MAX_CONCURRENT_SPAWNS_PER_TEAM)
|
||||
verbose_proxy_logger.info(
|
||||
"warm_pool.manager: team=%s desired=%s have=%s spawning=%s",
|
||||
team_id,
|
||||
desired,
|
||||
len(existing),
|
||||
spawn_n,
|
||||
)
|
||||
await asyncio.gather(
|
||||
*(self._spawn_one(prisma_client, vm_config) for _ in range(spawn_n)),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
async def _spawn_one(self, prisma_client: Any, vm_config: Any) -> None:
|
||||
"""Provision one warm VM end-to-end: row -> provider call -> rekey."""
|
||||
team_id = vm_config.team_id
|
||||
placeholder_id = _new_pool_vm_id()
|
||||
try:
|
||||
await prisma_client.db.litellm_agentvm.create(
|
||||
data={
|
||||
"id": placeholder_id,
|
||||
"provider": "ec2",
|
||||
"region": getattr(vm_config, "aws_region", None),
|
||||
"state": "provisioning",
|
||||
"team_id": team_id,
|
||||
"pool_id": team_id,
|
||||
"metadata": {"placeholder": True},
|
||||
}
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.exception(
|
||||
"warm_pool.manager: failed to insert placeholder team=%s: %s",
|
||||
team_id,
|
||||
exc,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
team_resolved = await get_team_vm_config(team_id, prisma_client)
|
||||
aws_creds = team_resolved.aws_creds
|
||||
ec2_config = team_resolved.ec2_config
|
||||
|
||||
provider = build_vm_provider(
|
||||
{"vm_provider": "ec2", "ec2": self._provider_settings}
|
||||
)
|
||||
ctx = ProvisionContext(
|
||||
session_id=placeholder_id,
|
||||
team_id=team_id,
|
||||
agent_id=None,
|
||||
repos=[],
|
||||
env_vars={},
|
||||
secrets={},
|
||||
aws_creds=aws_creds,
|
||||
ec2_config=ec2_config,
|
||||
daemon_jwt=None,
|
||||
daemon_base_url=None,
|
||||
mode="warm",
|
||||
)
|
||||
handle = await provider.provision(ctx)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.exception(
|
||||
"warm_pool.manager: provision failed team=%s placeholder=%s: %s",
|
||||
team_id,
|
||||
placeholder_id,
|
||||
exc,
|
||||
)
|
||||
try:
|
||||
await prisma_client.db.litellm_agentvm.update(
|
||||
where={"id": placeholder_id},
|
||||
data={
|
||||
"state": "terminated",
|
||||
"terminated_at": _now(),
|
||||
"metadata": {"error": type(exc).__name__},
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
# Rekey placeholder -> real EC2 id, mark warm. We delete + create
|
||||
# because the id is part of the primary key.
|
||||
try:
|
||||
await prisma_client.db.litellm_agentvm.delete(where={"id": placeholder_id})
|
||||
await prisma_client.db.litellm_agentvm.create(
|
||||
data={
|
||||
"id": handle.vm_id,
|
||||
"provider": handle.provider,
|
||||
"region": handle.region,
|
||||
"state": "warm",
|
||||
"team_id": team_id,
|
||||
"pool_id": team_id,
|
||||
"warmed_at": _now(),
|
||||
"metadata": handle.metadata or {},
|
||||
}
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"warm_pool.manager: warmed team=%s vm_id=%s", team_id, handle.vm_id
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.exception(
|
||||
"warm_pool.manager: rekey failed team=%s placeholder=%s real=%s: %s",
|
||||
team_id,
|
||||
placeholder_id,
|
||||
handle.vm_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
# ---------- reaping ----------
|
||||
|
||||
async def _reap_stale_warm(
|
||||
self, prisma_client: Any, team_id: str, max_idle_minutes: int
|
||||
) -> None:
|
||||
cutoff = _now() - timedelta(minutes=max_idle_minutes)
|
||||
stale = await prisma_client.db.litellm_agentvm.find_many(
|
||||
where={
|
||||
"team_id": team_id,
|
||||
"state": "warm",
|
||||
"warmed_at": {"lt": cutoff},
|
||||
}
|
||||
)
|
||||
if not stale:
|
||||
return
|
||||
verbose_proxy_logger.info(
|
||||
"warm_pool.manager: reaping team=%s count=%s cutoff=%s",
|
||||
team_id,
|
||||
len(stale),
|
||||
cutoff.isoformat(),
|
||||
)
|
||||
for row in stale:
|
||||
await self._terminate_vm(prisma_client, row)
|
||||
|
||||
async def _shrink_to_desired(
|
||||
self,
|
||||
prisma_client: Any,
|
||||
team_id: str,
|
||||
existing: List[Any],
|
||||
desired: int,
|
||||
) -> None:
|
||||
warm_only = [r for r in existing if r.state == "warm"]
|
||||
if len(warm_only) <= desired:
|
||||
return
|
||||
warm_only.sort(key=lambda r: r.warmed_at or r.created_at)
|
||||
excess = warm_only[: len(warm_only) - desired]
|
||||
for row in excess:
|
||||
await self._terminate_vm(prisma_client, row)
|
||||
|
||||
async def _terminate_vm(self, prisma_client: Any, row: Any) -> None:
|
||||
"""Best-effort terminate: mark `terminating`, fire provider, mark `terminated`."""
|
||||
try:
|
||||
await prisma_client.db.litellm_agentvm.update(
|
||||
where={"id": row.id},
|
||||
data={"state": "terminating"},
|
||||
)
|
||||
try:
|
||||
team_resolved = await get_team_vm_config(row.team_id, prisma_client)
|
||||
aws_creds = team_resolved.aws_creds
|
||||
except Exception:
|
||||
aws_creds = None
|
||||
if aws_creds is not None and row.provider == "ec2":
|
||||
provider = build_vm_provider(
|
||||
{"vm_provider": "ec2", "ec2": self._provider_settings}
|
||||
)
|
||||
from litellm.proxy.agent_session_endpoints.vm_providers.base import (
|
||||
VMHandle,
|
||||
)
|
||||
|
||||
handle = VMHandle(
|
||||
vm_id=row.id,
|
||||
provider="ec2",
|
||||
region=row.region,
|
||||
metadata=dict(row.metadata or {}),
|
||||
)
|
||||
await provider.terminate(handle, aws_creds=aws_creds)
|
||||
await prisma_client.db.litellm_agentvm.update(
|
||||
where={"id": row.id},
|
||||
data={"state": "terminated", "terminated_at": _now()},
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"warm_pool.manager: terminate failed vm=%s: %s", row.id, exc
|
||||
)
|
||||
|
||||
# ---------- helpers ----------
|
||||
|
||||
async def _enabled_teams(self, prisma_client: Any) -> List[Any]:
|
||||
"""Return the list of `LiteLLM_AgentVMConfig` rows with warm pool on."""
|
||||
return await prisma_client.db.litellm_agentvmconfig.find_many(
|
||||
where={
|
||||
"warm_pool_enabled": True,
|
||||
"warm_pool_size": {"gt": 0},
|
||||
}
|
||||
)
|
||||
|
||||
def _get_prisma_client(self) -> Any:
|
||||
if self._prisma_getter is not None:
|
||||
return self._prisma_getter()
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
return prisma_client
|
||||
|
||||
|
||||
# Module-level singleton. `proxy_server.py` calls `get_warm_pool_manager().start()`
|
||||
# on startup and `.stop()` on shutdown.
|
||||
_MANAGER_SINGLETON: Optional[WarmPoolManager] = None
|
||||
|
||||
|
||||
def get_warm_pool_manager() -> WarmPoolManager:
|
||||
"""Return (or create) the process-wide warm-pool manager."""
|
||||
global _MANAGER_SINGLETON
|
||||
if _MANAGER_SINGLETON is None:
|
||||
_MANAGER_SINGLETON = WarmPoolManager()
|
||||
return _MANAGER_SINGLETON
|
||||
|
||||
|
||||
def reset_warm_pool_manager() -> None:
|
||||
"""Test helper: clear the singleton between tests."""
|
||||
global _MANAGER_SINGLETON
|
||||
_MANAGER_SINGLETON = None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_CONCURRENT_SPAWNS_PER_TEAM",
|
||||
"WARM_POOL_TICK_SECONDS",
|
||||
"WarmPoolManager",
|
||||
"get_warm_pool_manager",
|
||||
"reset_warm_pool_manager",
|
||||
]
|
||||
|
||||
|
||||
# Re-export Repo / AwsCreds so tests can build fixtures without crawling deep
|
||||
# into the providers package.
|
||||
__all__ += ["AwsCreds", "Ec2Config", "Repo"]
|
||||
Loading…
Add table
Reference in a new issue