fix(gateway): bound the agent cache by memory, not just count and age
The per-session agent cache is capped at 128 entries with a 1h idle TTL, and neither bound knows how many bytes it holds. Each cached agent pins _session_messages -- the full transcript including tool output, tens of MB on a session with 100+ tool calls -- so a gateway serving many chats keeps every warm transcript resident: agents that took a turn inside the TTL are never idle-swept, and the idle sweep additionally defers finalizable sessions until they expire. RSS climbs until the cgroup throttles and SIGTERM can no longer flush inside systemd's stop timeout. Add the missing bound. Each session-expiry watcher tick compares the process's anonymous RSS against a budget and, when over, sheds LRU agents through the same soft-eviction path the cap enforcer uses, then runs malloc_trim so the freed arenas actually return to the OS. Evicted sessions rebuild their transcript from the persisted session on the next turn. Three classes of session are never shed: agents mid-turn, the most recently used ones, and any session whose transcript has not finished reaching disk (_last_flushed_db_idx vs len(_session_messages) -- the same divergence the FTS write-corruption guard reacts to when it preserves live history). memory_high_mb defaults to "auto", deriving the budget from the cgroup limit the gateway runs under, so a MemoryHigh/MemoryMax on the unit is respected without a second number to keep in sync. The two existing bounds become configurable alongside it under agent.agent_cache. protect_recent is clamped to half the cache: a couple of sessions can exhaust the budget on their own, and a fixed MRU guard would then protect everything and leave the gateway climbing with nothing it would shed. Fixes #80764
This commit is contained in:
HexLab98
kshitij
parent
f15a38ee73
commit
6bbe55dd09
|
|
@ -0,0 +1,281 @@
|
|||
"""Memory-pressure bounds for the gateway's per-session AIAgent cache.
|
||||
|
||||
The gateway caches one ``AIAgent`` per session so a long-lived conversation
|
||||
reuses its prompt prefix instead of rebuilding the system prompt every turn.
|
||||
Each cached agent also pins ``_session_messages`` — the full live transcript,
|
||||
tool outputs included, which is tens of MB on a tool-heavy session.
|
||||
|
||||
``gateway/run.py`` bounds that cache two ways, and both are blind to how much
|
||||
memory it actually holds:
|
||||
|
||||
* the LRU cap counts *entries*, not bytes, and 128 warm transcripts is
|
||||
several GB;
|
||||
* the idle TTL only sheds agents that went quiet for an hour, and it
|
||||
deliberately defers eviction for a finalizable session that has not expired
|
||||
yet, so a busy gateway hoards every transcript all day.
|
||||
|
||||
This module supplies the missing signal: the process's own anonymous RSS,
|
||||
compared against a budget derived from the cgroup limit the gateway actually
|
||||
runs under. ``GatewayRunner._sweep_agent_cache_under_pressure`` uses it to
|
||||
shed LRU transcripts through the existing soft-eviction path, which rebuilds
|
||||
from the persisted session on the next turn (#80764).
|
||||
|
||||
Everything here is pure or read-only so it can be tested without a gateway.
|
||||
Config lives under ``agent.agent_cache`` in ``config.yaml``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, List, Optional, Tuple
|
||||
|
||||
# Fraction of the resolved memory limit at which we start shedding
|
||||
# transcripts. Deliberately well under the limit: on the reported incident
|
||||
# the gateway hit cgroup ``memory.high`` throttling with swap full, and a
|
||||
# SIGTERM flush from there could not finish inside systemd's stop timeout.
|
||||
# Eviction has to happen while the process still has room to breathe.
|
||||
_AUTO_BUDGET_FRACTION = 0.65
|
||||
# Below this a "budget" is noise — small containers would evict on every pass
|
||||
# and never keep a warm prefix.
|
||||
_AUTO_BUDGET_FLOOR_MB = 512
|
||||
|
||||
_DEFAULT_MAX_EVICTIONS_PER_PASS = 16
|
||||
# Never let a pressure pass touch the hottest sessions: they are the ones
|
||||
# whose prompt cache is worth the most, and shedding them just moves the cost
|
||||
# to the next turn instead of removing it.
|
||||
_DEFAULT_PROTECT_RECENT = 8
|
||||
|
||||
_BYTES_PER_MB = 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AgentCacheBounds:
|
||||
"""Operator-facing bounds for the per-session agent cache.
|
||||
|
||||
``max_size`` and ``idle_ttl_secs`` are ``None`` when the operator did not
|
||||
set them, so ``gateway/run.py`` keeps using its module-level defaults.
|
||||
``memory_high_mb`` is ``None`` when pressure eviction is switched off.
|
||||
"""
|
||||
|
||||
max_size: Optional[int] = None
|
||||
idle_ttl_secs: Optional[float] = None
|
||||
memory_high_mb: Optional[int] = None
|
||||
max_evictions_per_pass: int = _DEFAULT_MAX_EVICTIONS_PER_PASS
|
||||
protect_recent: int = _DEFAULT_PROTECT_RECENT
|
||||
|
||||
|
||||
def _positive_int(value: Any) -> Optional[int]:
|
||||
if isinstance(value, bool) or value is None:
|
||||
return None
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return parsed if parsed > 0 else None
|
||||
|
||||
|
||||
def _positive_float(value: Any) -> Optional[float]:
|
||||
if isinstance(value, bool) or value is None:
|
||||
return None
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return parsed if parsed > 0 else None
|
||||
|
||||
|
||||
def _cgroup_limit_bytes() -> Optional[int]:
|
||||
"""Return the memory limit this process runs under, if it is cgroup-capped.
|
||||
|
||||
Prefers cgroup v2 ``memory.high`` (the throttling point — passing it is
|
||||
what stalled the reported shutdown) over ``memory.max``, and falls back to
|
||||
cgroup v1. ``max`` / absurd sentinel values mean "unlimited".
|
||||
"""
|
||||
if sys.platform != "linux":
|
||||
return None
|
||||
candidates = (
|
||||
"/sys/fs/cgroup/memory.high",
|
||||
"/sys/fs/cgroup/memory.max",
|
||||
"/sys/fs/cgroup/memory/memory.limit_in_bytes",
|
||||
)
|
||||
for candidate in candidates:
|
||||
try:
|
||||
raw = Path(candidate).read_text(encoding="utf-8").strip()
|
||||
except OSError:
|
||||
continue
|
||||
if not raw or raw == "max":
|
||||
continue
|
||||
try:
|
||||
limit = int(raw)
|
||||
except ValueError:
|
||||
continue
|
||||
# cgroup v1 reports "unlimited" as a near-2^63 sentinel.
|
||||
if limit <= 0 or limit >= (1 << 62):
|
||||
continue
|
||||
return limit
|
||||
return None
|
||||
|
||||
|
||||
def _total_memory_bytes() -> Optional[int]:
|
||||
try:
|
||||
return int(os.sysconf("SC_PAGE_SIZE")) * int(os.sysconf("SC_PHYS_PAGES"))
|
||||
except (OSError, ValueError, AttributeError):
|
||||
pass
|
||||
try:
|
||||
import psutil # type: ignore
|
||||
|
||||
return int(psutil.virtual_memory().total)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def resolve_memory_high_mb(setting: Any) -> Optional[int]:
|
||||
"""Resolve the ``memory_high_mb`` setting into an absolute MB budget.
|
||||
|
||||
``"auto"`` derives a budget from the cgroup limit the gateway runs under
|
||||
(or total RAM when uncapped), which is what makes this fix work out of the
|
||||
box on the containerised/systemd deployments where the leak bites. A
|
||||
positive number is taken literally; anything falsy disables the pass.
|
||||
"""
|
||||
if isinstance(setting, str):
|
||||
normalized = setting.strip().lower()
|
||||
if normalized != "auto":
|
||||
return (
|
||||
None
|
||||
if normalized in ("", "off", "none", "false", "disabled")
|
||||
else _positive_int(normalized)
|
||||
)
|
||||
elif isinstance(setting, bool):
|
||||
if not setting:
|
||||
return None
|
||||
else:
|
||||
return _positive_int(setting)
|
||||
|
||||
limit = _cgroup_limit_bytes() or _total_memory_bytes()
|
||||
if not limit:
|
||||
return None
|
||||
budget = int(limit * _AUTO_BUDGET_FRACTION / _BYTES_PER_MB)
|
||||
return budget if budget >= _AUTO_BUDGET_FLOOR_MB else None
|
||||
|
||||
|
||||
def resolve_agent_cache_bounds(config: Any) -> AgentCacheBounds:
|
||||
"""Read ``agent.agent_cache`` out of a raw config mapping.
|
||||
|
||||
Reads the *raw* user config (the gateway's loader does not deep-merge
|
||||
``DEFAULT_CONFIG``), so an absent key stays absent and the caller can tell
|
||||
"operator chose 128" from "operator said nothing".
|
||||
"""
|
||||
section: Any = None
|
||||
if isinstance(config, dict):
|
||||
agent_cfg = config.get("agent")
|
||||
if isinstance(agent_cfg, dict):
|
||||
section = agent_cfg.get("agent_cache")
|
||||
if not isinstance(section, dict):
|
||||
section = {}
|
||||
|
||||
max_evictions = _positive_int(section.get("max_evictions_per_pass"))
|
||||
protect_recent = section.get("protect_recent")
|
||||
protect_parsed = _positive_int(protect_recent)
|
||||
if protect_parsed is None and protect_recent == 0:
|
||||
protect_parsed = 0
|
||||
|
||||
return AgentCacheBounds(
|
||||
max_size=_positive_int(section.get("max_size")),
|
||||
idle_ttl_secs=_positive_float(section.get("idle_ttl_secs")),
|
||||
memory_high_mb=resolve_memory_high_mb(section.get("memory_high_mb", "auto")),
|
||||
max_evictions_per_pass=(
|
||||
max_evictions if max_evictions is not None else _DEFAULT_MAX_EVICTIONS_PER_PASS
|
||||
),
|
||||
protect_recent=(
|
||||
protect_parsed if protect_parsed is not None else _DEFAULT_PROTECT_RECENT
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def read_anon_rss_mb() -> Optional[int]:
|
||||
"""Return the process's anonymous resident memory in MB, or None.
|
||||
|
||||
Anonymous pages are the ones cached transcripts live in — the reported
|
||||
incident measured 11.0 GB of anon out of 11.0 GB total, so file-backed
|
||||
pages are noise here. ``collect_memory_snapshot`` already reads
|
||||
``/proc/self/status`` without a dependency; psutil covers everything else,
|
||||
where only total RSS is available.
|
||||
"""
|
||||
try:
|
||||
from hermes_cli.mem_trim import collect_memory_snapshot
|
||||
|
||||
snapshot = collect_memory_snapshot()
|
||||
anon_kib = snapshot.get("rss_anon_kib")
|
||||
if isinstance(anon_kib, int) and anon_kib > 0:
|
||||
return anon_kib // 1024
|
||||
rss_kib = snapshot.get("rss_kib")
|
||||
if isinstance(rss_kib, int) and rss_kib > 0:
|
||||
return rss_kib // 1024
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
import psutil # type: ignore
|
||||
|
||||
return int(psutil.Process(os.getpid()).memory_info().rss / _BYTES_PER_MB)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def transcript_persistence_caught_up(agent: Any) -> bool:
|
||||
"""True when the agent's live transcript is fully on disk.
|
||||
|
||||
Soft eviction drops ``_session_messages`` and rebuilds it from the
|
||||
persisted session next turn, so it is only safe once persistence has
|
||||
caught up. ``_last_flushed_db_idx`` is advanced to ``len(messages)`` by
|
||||
``AIAgent._flush_messages_to_session_db`` and only on a fully successful
|
||||
write — the same divergence the FTS write-corruption guard reacts to when
|
||||
it preserves live history over a lagging transcript. Unknown shapes are
|
||||
treated as *not* caught up: a skipped eviction costs memory, a wrong one
|
||||
costs the user their conversation.
|
||||
"""
|
||||
messages = getattr(agent, "_session_messages", None)
|
||||
if not isinstance(messages, list):
|
||||
return False
|
||||
flushed = getattr(agent, "_last_flushed_db_idx", None)
|
||||
if not isinstance(flushed, int) or isinstance(flushed, bool):
|
||||
return False
|
||||
return flushed >= len(messages)
|
||||
|
||||
|
||||
def plan_pressure_evictions(
|
||||
ordered_entries: Iterable[Tuple[str, Any]],
|
||||
*,
|
||||
is_evictable: Callable[[str, Any], bool],
|
||||
max_evictions: int,
|
||||
protect_recent: int = 0,
|
||||
) -> List[Tuple[str, Any]]:
|
||||
"""Choose which cached sessions to shed, least-recently-used first.
|
||||
|
||||
``ordered_entries`` must be in LRU→MRU order (the cache is an
|
||||
``OrderedDict`` kept in that order by ``move_to_end`` on every hit). The
|
||||
batch is capped so one pass cannot stall the gateway tearing down clients.
|
||||
|
||||
``protect_recent`` is an upper bound, clamped to half the cache: a handful
|
||||
of sessions can be big enough to exhaust the budget on their own (a single
|
||||
tool-heavy transcript runs to hundreds of MB), and a fixed guard would
|
||||
then protect the entire cache and leave the gateway climbing toward the
|
||||
OOM killer with nothing it is willing to shed.
|
||||
"""
|
||||
entries = list(ordered_entries)
|
||||
if max_evictions <= 0 or not entries:
|
||||
return []
|
||||
protect = min(max(protect_recent, 0), len(entries) // 2)
|
||||
if protect:
|
||||
entries = entries[:-protect]
|
||||
|
||||
plan: List[Tuple[str, Any]] = []
|
||||
for key, agent in entries:
|
||||
if len(plan) >= max_evictions:
|
||||
break
|
||||
if is_evictable(key, agent):
|
||||
plan.append((key, agent))
|
||||
return plan
|
||||
179
gateway/run.py
179
gateway/run.py
|
|
@ -71,6 +71,12 @@ from hermes_cli.fallback_config import get_fallback_chain
|
|||
# long-lived gateways (each AIAgent holds LLM clients, tool schemas,
|
||||
# memory providers, etc.). LRU order + idle TTL eviction are enforced
|
||||
# from _enforce_agent_cache_cap() and _session_expiry_watcher() below.
|
||||
#
|
||||
# These are the defaults; `agent.agent_cache.max_size` /
|
||||
# `agent.agent_cache.idle_ttl_secs` in config.yaml override them per
|
||||
# deployment. Neither bound knows how many BYTES a cached agent holds, so
|
||||
# _sweep_agent_cache_under_pressure() adds the missing memory-pressure valve
|
||||
# (see gateway/agent_cache_pressure.py).
|
||||
_AGENT_CACHE_MAX_SIZE = 128
|
||||
_AGENT_CACHE_IDLE_TTL_SECS = 3600.0 # evict agents idle for >1h
|
||||
_PLATFORM_CONNECT_TIMEOUT_SECS_DEFAULT = 30.0
|
||||
|
|
@ -12180,6 +12186,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
except Exception as _e:
|
||||
logger.debug("Idle agent sweep failed: %s", _e)
|
||||
|
||||
# Neither the LRU cap nor the idle TTL is aware of how much
|
||||
# memory a cached transcript costs, so a busy gateway keeps
|
||||
# every warm session's tool output resident until RSS hits the
|
||||
# cgroup limit (#80764). Shed LRU transcripts once the heap is
|
||||
# over budget; they reload from the persisted session on the
|
||||
# next turn.
|
||||
try:
|
||||
self._sweep_agent_cache_under_pressure()
|
||||
except Exception as _e:
|
||||
logger.debug("Agent cache pressure sweep failed: %s", _e)
|
||||
|
||||
# Periodically prune stale SessionStore entries. The
|
||||
# in-memory dict (and sessions.json) would otherwise grow
|
||||
# unbounded in gateways serving many rotating chats /
|
||||
|
|
@ -23731,8 +23748,156 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
if hasattr(agent, "_session_messages"):
|
||||
agent._session_messages = []
|
||||
|
||||
def _agent_cache_bounds(self):
|
||||
"""Operator-configured agent-cache bounds, resolved once per process.
|
||||
|
||||
Resolved lazily rather than in ``__init__`` so it also works for the
|
||||
``__new__``-constructed runners used by tests and by the slash-command
|
||||
mixin.
|
||||
"""
|
||||
bounds = getattr(self, "_agent_cache_bounds_cache", None)
|
||||
if bounds is None:
|
||||
from gateway.agent_cache_pressure import resolve_agent_cache_bounds
|
||||
|
||||
try:
|
||||
bounds = resolve_agent_cache_bounds(_load_gateway_config())
|
||||
except Exception as _e:
|
||||
from gateway.agent_cache_pressure import AgentCacheBounds
|
||||
|
||||
logger.debug("Agent cache bounds config read failed: %s", _e)
|
||||
bounds = AgentCacheBounds()
|
||||
self._agent_cache_bounds_cache = bounds
|
||||
return bounds
|
||||
|
||||
def _agent_cache_cap(self) -> int:
|
||||
"""Effective LRU cap — the configured override, else the default."""
|
||||
configured = self._agent_cache_bounds().max_size
|
||||
return configured if configured else _AGENT_CACHE_MAX_SIZE
|
||||
|
||||
def _agent_cache_idle_ttl(self) -> float:
|
||||
"""Effective idle TTL in seconds — configured override, else default."""
|
||||
configured = self._agent_cache_bounds().idle_ttl_secs
|
||||
return configured if configured else _AGENT_CACHE_IDLE_TTL_SECS
|
||||
|
||||
def _sweep_agent_cache_under_pressure(self) -> int:
|
||||
"""Shed cached transcripts once the gateway's own heap nears its budget.
|
||||
|
||||
The LRU cap counts entries and the idle sweep counts seconds; neither
|
||||
knows that one cached agent pins a full ``_session_messages``
|
||||
transcript — tens of MB on a session with 100+ tool calls. A gateway
|
||||
serving many chats therefore holds every warm transcript indefinitely:
|
||||
agents that took a turn within the TTL are never idle-swept, and the
|
||||
sweep additionally defers finalizable sessions until they expire. RSS
|
||||
climbs until the cgroup throttles and SIGTERM can no longer flush
|
||||
inside systemd's stop timeout (#80764).
|
||||
|
||||
This is the missing valve. Above the configured anonymous-RSS budget
|
||||
it evicts LRU agents through the same soft path the cap enforcer uses,
|
||||
so the transcript is dropped and rebuilt from the persisted session on
|
||||
the next turn. Three things are never touched: agents mid-turn (their
|
||||
clients and sandboxes are in use), the most recently used sessions
|
||||
(whose prompt cache is worth the most), and any session whose live
|
||||
transcript has not finished reaching disk.
|
||||
|
||||
Returns the number of entries evicted (0 when memory is fine).
|
||||
"""
|
||||
from gateway.agent_cache_pressure import (
|
||||
plan_pressure_evictions,
|
||||
read_anon_rss_mb,
|
||||
transcript_persistence_caught_up,
|
||||
)
|
||||
|
||||
bounds = self._agent_cache_bounds()
|
||||
if not bounds.memory_high_mb:
|
||||
return 0
|
||||
_cache = getattr(self, "_agent_cache", None)
|
||||
_lock = getattr(self, "_agent_cache_lock", None)
|
||||
if not _cache or _lock is None:
|
||||
# Nothing cached — whatever is using the heap, it isn't us, and
|
||||
# warning about it every tick would point at the wrong subsystem.
|
||||
return 0
|
||||
|
||||
rss_mb = read_anon_rss_mb()
|
||||
if rss_mb is None or rss_mb < bounds.memory_high_mb:
|
||||
return 0
|
||||
|
||||
running_ids = {
|
||||
id(a)
|
||||
for _, a in self._running_agent_items()
|
||||
if a is not None and a is not _AGENT_PENDING_SENTINEL
|
||||
}
|
||||
|
||||
def _is_evictable(key: str, agent: Any) -> bool:
|
||||
if agent is None or agent is _AGENT_PENDING_SENTINEL:
|
||||
return False
|
||||
if id(agent) in running_ids:
|
||||
return False
|
||||
return transcript_persistence_caught_up(agent)
|
||||
|
||||
with _lock:
|
||||
ordered = [
|
||||
(key, entry[0] if isinstance(entry, tuple) and entry else entry)
|
||||
for key, entry in _cache.items()
|
||||
]
|
||||
plan = plan_pressure_evictions(
|
||||
ordered,
|
||||
is_evictable=_is_evictable,
|
||||
max_evictions=bounds.max_evictions_per_pass,
|
||||
protect_recent=bounds.protect_recent,
|
||||
)
|
||||
for key, _ in plan:
|
||||
_cache.pop(key, None)
|
||||
|
||||
if not plan:
|
||||
_mid_turn = sum(1 for _, a in ordered if a is not None and id(a) in running_ids)
|
||||
logger.warning(
|
||||
"Agent cache pressure: anon RSS %dMB over budget %dMB but no "
|
||||
"evictable session (%d cached, %d mid-turn) — memory will keep "
|
||||
"climbing until those turns finish.",
|
||||
rss_mb, bounds.memory_high_mb, len(ordered), _mid_turn,
|
||||
)
|
||||
return 0
|
||||
|
||||
logger.warning(
|
||||
"Agent cache pressure: anon RSS %dMB over budget %dMB — evicting "
|
||||
"%d LRU session(s): %s",
|
||||
rss_mb, bounds.memory_high_mb, len(plan),
|
||||
", ".join(key for key, _ in plan),
|
||||
)
|
||||
try:
|
||||
threading.Thread(
|
||||
target=self._release_pressure_batch,
|
||||
args=(plan,),
|
||||
daemon=True,
|
||||
name="agent-cache-pressure",
|
||||
).start()
|
||||
except Exception:
|
||||
self._release_pressure_batch(plan)
|
||||
return len(plan)
|
||||
|
||||
def _release_pressure_batch(self, plan: List[tuple]) -> None:
|
||||
"""Release a pressure-evicted batch, then return the heap to the OS.
|
||||
|
||||
Sequential on one daemon thread rather than a thread per agent: the
|
||||
batch is already capped, and the point of the pass is to reclaim
|
||||
memory, not to race N teardowns. The trailing ``malloc_trim`` is what
|
||||
turns "Python dropped the transcript" into "RSS actually fell" —
|
||||
without it glibc keeps the freed arenas and the cgroup never notices.
|
||||
"""
|
||||
for key, agent in plan:
|
||||
try:
|
||||
self._commit_then_release_soft(agent, key)
|
||||
except Exception as _e:
|
||||
logger.debug("Pressure release failed for %s: %s", key, _e)
|
||||
try:
|
||||
from hermes_cli.mem_trim import trim_memory
|
||||
|
||||
trim_memory(force=True, reason="agent_cache_pressure")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _enforce_agent_cache_cap(self) -> None:
|
||||
"""Evict oldest cached agents when cache exceeds _AGENT_CACHE_MAX_SIZE.
|
||||
"""Evict oldest cached agents when cache exceeds the LRU cap.
|
||||
|
||||
Must be called with _agent_cache_lock held. Resource cleanup
|
||||
(memory provider shutdown, tool resource close) is scheduled
|
||||
|
|
@ -23772,7 +23937,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
# already-cached long-running one. The cache may therefore stay
|
||||
# temporarily over cap; it will re-check on the next insert,
|
||||
# after active turns have finished.
|
||||
excess = max(0, len(_cache) - _AGENT_CACHE_MAX_SIZE)
|
||||
cap = self._agent_cache_cap()
|
||||
excess = max(0, len(_cache) - cap)
|
||||
evict_plan: List[tuple] = [] # [(key, agent), ...]
|
||||
if excess > 0:
|
||||
ordered_keys = list(_cache.keys())
|
||||
|
|
@ -23786,12 +23952,12 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
for key, _ in evict_plan:
|
||||
_cache.pop(key, None)
|
||||
|
||||
remaining_over_cap = len(_cache) - _AGENT_CACHE_MAX_SIZE
|
||||
remaining_over_cap = len(_cache) - cap
|
||||
if remaining_over_cap > 0:
|
||||
logger.warning(
|
||||
"Agent cache over cap (%d > %d); %d excess slot(s) held by "
|
||||
"mid-turn agents — will re-check on next insert.",
|
||||
len(_cache), _AGENT_CACHE_MAX_SIZE, remaining_over_cap,
|
||||
len(_cache), cap, remaining_over_cap,
|
||||
)
|
||||
|
||||
for key, agent in evict_plan:
|
||||
|
|
@ -23814,7 +23980,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
).start()
|
||||
|
||||
def _sweep_idle_cached_agents(self) -> int:
|
||||
"""Evict cached agents whose AIAgent has been idle > _AGENT_CACHE_IDLE_TTL_SECS.
|
||||
"""Evict cached agents whose AIAgent has been idle past the idle TTL.
|
||||
|
||||
Safe to call from the session expiry watcher without holding the
|
||||
cache lock — acquires it internally. Returns the number of entries
|
||||
|
|
@ -23829,6 +23995,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
if _cache is None or _lock is None:
|
||||
return 0
|
||||
now = time.time()
|
||||
idle_ttl = self._agent_cache_idle_ttl()
|
||||
to_evict: List[tuple] = []
|
||||
running_ids = {
|
||||
id(a)
|
||||
|
|
@ -23845,7 +24012,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
|||
last_activity = getattr(agent, "_last_activity_ts", None)
|
||||
if last_activity is None:
|
||||
continue
|
||||
if (now - last_activity) > _AGENT_CACHE_IDLE_TTL_SECS:
|
||||
if (now - last_activity) > idle_ttl:
|
||||
# Check whether the session has actually expired in the
|
||||
# session store. If it hasn't (e.g. daily-reset mode
|
||||
# where the reset fires hours after the user's last
|
||||
|
|
|
|||
|
|
@ -40,6 +40,29 @@ DEFAULT_CONFIG = {
|
|||
# rejected with a resend notice rather than run without serialization.
|
||||
# Non-positive values fall back to 1800 seconds.
|
||||
"gateway_turn_lease_timeout": 1800,
|
||||
# Per-session AIAgent cache in the gateway. Each cached agent keeps a
|
||||
# warm prompt prefix AND the session's full transcript, so the cache
|
||||
# trades memory for cost: too small and every turn re-pays an uncached
|
||||
# prompt, too large and tool-heavy transcripts fill the heap.
|
||||
"agent_cache": {
|
||||
# LRU entry cap.
|
||||
"max_size": 128,
|
||||
# Evict an agent that has been idle this long (seconds).
|
||||
"idle_ttl_secs": 3600,
|
||||
# Anonymous-RSS budget (MB) above which the gateway starts shedding
|
||||
# least-recently-used transcripts, which reload from the persisted
|
||||
# session on the next turn. "auto" derives the budget from the
|
||||
# cgroup memory limit the gateway runs under (or total RAM when
|
||||
# uncapped); a number sets it explicitly; 0/off disables the pass
|
||||
# and lets memory grow to whatever the two bounds above allow.
|
||||
"memory_high_mb": "auto",
|
||||
# Upper bound on how many sessions one pressure pass sheds, so a
|
||||
# burst of teardowns cannot stall the gateway.
|
||||
"max_evictions_per_pass": 16,
|
||||
# Most-recently-used sessions the pressure pass never touches —
|
||||
# they are the ones actively paying for a warm prompt cache.
|
||||
"protect_recent": 8,
|
||||
},
|
||||
# Force-interrupt budget once gateway stop()/drain has begun
|
||||
# (seconds). Applies to SIGTERM/external stop and to the final
|
||||
# phase of in-band restart after any after-turn wait. 0 = interrupt
|
||||
|
|
|
|||
|
|
@ -0,0 +1,395 @@
|
|||
"""Memory-pressure eviction for the gateway agent cache (#80764).
|
||||
|
||||
The LRU cap counts entries and the idle sweep counts seconds, so a gateway
|
||||
serving many warm sessions holds every full transcript resident until the
|
||||
cgroup kills it. These tests pin the pressure valve that sheds them, and the
|
||||
three things it must never shed: a mid-turn agent, the most-recently-used
|
||||
sessions, and a session whose transcript has not finished reaching disk.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from gateway.agent_cache_pressure import (
|
||||
AgentCacheBounds,
|
||||
plan_pressure_evictions,
|
||||
resolve_agent_cache_bounds,
|
||||
resolve_memory_high_mb,
|
||||
transcript_persistence_caught_up,
|
||||
)
|
||||
|
||||
|
||||
class TestBoundsResolution:
|
||||
"""Absent config must stay absent so gateway/run.py keeps its defaults."""
|
||||
|
||||
def test_absent_section_leaves_lru_bounds_unset(self):
|
||||
bounds = resolve_agent_cache_bounds({})
|
||||
assert bounds.max_size is None
|
||||
assert bounds.idle_ttl_secs is None
|
||||
|
||||
def test_configured_values_are_honoured(self):
|
||||
bounds = resolve_agent_cache_bounds(
|
||||
{
|
||||
"agent": {
|
||||
"agent_cache": {
|
||||
"max_size": 32,
|
||||
"idle_ttl_secs": 600,
|
||||
"memory_high_mb": 2048,
|
||||
"max_evictions_per_pass": 4,
|
||||
"protect_recent": 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
assert bounds.max_size == 32
|
||||
assert bounds.idle_ttl_secs == 600.0
|
||||
assert bounds.memory_high_mb == 2048
|
||||
assert bounds.max_evictions_per_pass == 4
|
||||
assert bounds.protect_recent == 2
|
||||
|
||||
def test_garbage_values_fall_back_to_defaults(self):
|
||||
"""A typo in config.yaml must not disable the cache or crash startup."""
|
||||
bounds = resolve_agent_cache_bounds(
|
||||
{"agent": {"agent_cache": {"max_size": "lots", "idle_ttl_secs": -5}}}
|
||||
)
|
||||
assert bounds.max_size is None
|
||||
assert bounds.idle_ttl_secs is None
|
||||
assert bounds.max_evictions_per_pass > 0
|
||||
|
||||
def test_protect_recent_zero_is_respected(self):
|
||||
"""0 means "shed anything", which is distinct from "unset"."""
|
||||
bounds = resolve_agent_cache_bounds(
|
||||
{"agent": {"agent_cache": {"protect_recent": 0}}}
|
||||
)
|
||||
assert bounds.protect_recent == 0
|
||||
|
||||
|
||||
class TestMemoryBudgetResolution:
|
||||
@pytest.mark.parametrize("setting", [0, False, None, "off", "none", ""])
|
||||
def test_falsy_settings_disable_the_pass(self, setting):
|
||||
assert resolve_memory_high_mb(setting) is None
|
||||
|
||||
@pytest.mark.parametrize("setting", [4096, "4096", 4096.0])
|
||||
def test_explicit_budget_is_taken_literally(self, setting):
|
||||
assert resolve_memory_high_mb(setting) == 4096
|
||||
|
||||
def test_auto_derives_a_budget_below_the_cgroup_limit(self, monkeypatch):
|
||||
"""The budget must leave headroom: hitting memory.high is what makes
|
||||
the shutdown flush time out in the first place."""
|
||||
import gateway.agent_cache_pressure as acp
|
||||
|
||||
limit_mb = 10 * 1024
|
||||
monkeypatch.setattr(acp, "_cgroup_limit_bytes", lambda: limit_mb * 1024 * 1024)
|
||||
|
||||
budget = resolve_memory_high_mb("auto")
|
||||
|
||||
assert budget is not None
|
||||
assert 0 < budget < limit_mb
|
||||
|
||||
def test_auto_is_disabled_when_no_limit_is_discoverable(self, monkeypatch):
|
||||
import gateway.agent_cache_pressure as acp
|
||||
|
||||
monkeypatch.setattr(acp, "_cgroup_limit_bytes", lambda: None)
|
||||
monkeypatch.setattr(acp, "_total_memory_bytes", lambda: None)
|
||||
|
||||
assert resolve_memory_high_mb("auto") is None
|
||||
|
||||
|
||||
class TestPersistenceGuard:
|
||||
"""Soft eviction drops the transcript, so it may only run once the
|
||||
transcript is durable. Exercised against the real AIAgent flush."""
|
||||
|
||||
def _agent(self, tmp_path, session_id):
|
||||
from hermes_state import SessionDB
|
||||
from run_agent import AIAgent
|
||||
|
||||
db = SessionDB(db_path=tmp_path / "sessions.db")
|
||||
agent = AIAgent(
|
||||
model="anthropic/claude-sonnet-4",
|
||||
api_key="test",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
provider="openrouter",
|
||||
max_iterations=5,
|
||||
quiet_mode=True,
|
||||
skip_context_files=True,
|
||||
skip_memory=True,
|
||||
session_id=session_id,
|
||||
session_db=db,
|
||||
)
|
||||
db.create_session(session_id, source="telegram")
|
||||
agent._session_db_created = True
|
||||
return agent
|
||||
|
||||
def test_fresh_agent_holds_nothing_to_lose(self, tmp_path):
|
||||
agent = self._agent(tmp_path, "fresh")
|
||||
try:
|
||||
assert transcript_persistence_caught_up(agent) is True
|
||||
finally:
|
||||
agent.close()
|
||||
|
||||
def test_unflushed_turn_blocks_eviction_then_flush_unblocks_it(self, tmp_path):
|
||||
agent = self._agent(tmp_path, "lagging")
|
||||
try:
|
||||
messages = [
|
||||
{"role": "user", "content": "read the logs"},
|
||||
{"role": "assistant", "content": "done"},
|
||||
]
|
||||
agent._session_messages = messages
|
||||
|
||||
assert transcript_persistence_caught_up(agent) is False, (
|
||||
"a transcript that never reached disk must not be dropped — "
|
||||
"the session would come back with amnesia"
|
||||
)
|
||||
|
||||
assert agent._flush_messages_to_session_db(messages) is True
|
||||
assert transcript_persistence_caught_up(agent) is True
|
||||
finally:
|
||||
agent.close()
|
||||
|
||||
def test_unknown_shapes_are_treated_as_unsafe(self):
|
||||
assert transcript_persistence_caught_up(object()) is False
|
||||
assert transcript_persistence_caught_up(None) is False
|
||||
|
||||
|
||||
class TestEvictionPlanner:
|
||||
def _entries(self, n):
|
||||
return [(f"s{i}", MagicMock()) for i in range(n)]
|
||||
|
||||
def test_evicts_least_recently_used_first(self):
|
||||
entries = self._entries(6)
|
||||
plan = plan_pressure_evictions(
|
||||
entries, is_evictable=lambda k, a: True, max_evictions=2, protect_recent=0
|
||||
)
|
||||
assert [key for key, _ in plan] == ["s0", "s1"]
|
||||
|
||||
def test_never_touches_the_protected_tail(self):
|
||||
entries = self._entries(10)
|
||||
plan = plan_pressure_evictions(
|
||||
entries, is_evictable=lambda k, a: True, max_evictions=10, protect_recent=3
|
||||
)
|
||||
assert [key for key, _ in plan] == ["s0", "s1", "s2", "s3", "s4", "s5", "s6"]
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 2, 3, 5])
|
||||
def test_a_small_cache_of_large_transcripts_is_still_shedable(self, size):
|
||||
"""A fixed MRU guard would protect the whole cache when a couple of
|
||||
sessions are big enough to blow the budget on their own — the gateway
|
||||
would then climb toward the OOM killer with nothing it would shed."""
|
||||
plan = plan_pressure_evictions(
|
||||
self._entries(size),
|
||||
is_evictable=lambda k, a: True,
|
||||
max_evictions=10,
|
||||
protect_recent=8,
|
||||
)
|
||||
assert plan, f"nothing evictable with {size} cached session(s)"
|
||||
assert len(plan) <= size
|
||||
|
||||
def test_protection_still_keeps_the_hottest_session(self):
|
||||
plan = plan_pressure_evictions(
|
||||
self._entries(4),
|
||||
is_evictable=lambda k, a: True,
|
||||
max_evictions=10,
|
||||
protect_recent=8,
|
||||
)
|
||||
assert "s3" not in [key for key, _ in plan]
|
||||
|
||||
def test_skipped_candidates_do_not_consume_the_batch(self):
|
||||
"""Skipping a protected session must not shrink the batch — otherwise
|
||||
one wedged session throttles the whole pass."""
|
||||
entries = self._entries(6)
|
||||
plan = plan_pressure_evictions(
|
||||
entries,
|
||||
is_evictable=lambda k, a: k != "s0",
|
||||
max_evictions=2,
|
||||
protect_recent=0,
|
||||
)
|
||||
assert [key for key, _ in plan] == ["s1", "s2"]
|
||||
|
||||
|
||||
class TestGatewayPressureSweep:
|
||||
"""End-to-end against the real GatewayRunner method."""
|
||||
|
||||
def _runner(self, bounds=None):
|
||||
from gateway.run import GatewayRunner
|
||||
|
||||
runner = GatewayRunner.__new__(GatewayRunner)
|
||||
runner._agent_cache = OrderedDict()
|
||||
runner._agent_cache_lock = threading.Lock()
|
||||
runner._running_agents = {}
|
||||
runner._agent_cache_bounds_cache = bounds or AgentCacheBounds(
|
||||
memory_high_mb=1000, max_evictions_per_pass=8, protect_recent=1
|
||||
)
|
||||
return runner
|
||||
|
||||
def _cached_agent(self, *, persisted=True, messages=2):
|
||||
agent = MagicMock()
|
||||
agent._session_messages = [{"role": "user", "content": "x"}] * messages
|
||||
agent._last_flushed_db_idx = messages if persisted else 0
|
||||
return agent
|
||||
|
||||
def _at_rss(self, monkeypatch, mb):
|
||||
import gateway.agent_cache_pressure as acp
|
||||
|
||||
monkeypatch.setattr(acp, "read_anon_rss_mb", lambda: mb)
|
||||
|
||||
def test_no_eviction_below_budget(self, monkeypatch):
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 400)
|
||||
for i in range(5):
|
||||
runner._agent_cache[f"s{i}"] = (self._cached_agent(), "sig")
|
||||
|
||||
assert runner._sweep_agent_cache_under_pressure() == 0
|
||||
assert len(runner._agent_cache) == 5
|
||||
|
||||
def test_over_budget_sheds_lru_and_frees_the_transcript(self, monkeypatch):
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 4000)
|
||||
released: list = []
|
||||
runner._commit_then_release_soft = lambda agent, key: (
|
||||
released.append(key),
|
||||
setattr(agent, "_session_messages", []),
|
||||
)
|
||||
|
||||
for i in range(4):
|
||||
runner._agent_cache[f"s{i}"] = (self._cached_agent(), "sig")
|
||||
oldest = runner._agent_cache["s0"][0]
|
||||
|
||||
evicted = runner._sweep_agent_cache_under_pressure()
|
||||
|
||||
assert evicted == 3 # protect_recent=1 keeps the newest
|
||||
assert "s0" not in runner._agent_cache
|
||||
assert "s3" in runner._agent_cache
|
||||
_wait_for(lambda: released == ["s0", "s1", "s2"])
|
||||
assert oldest._session_messages == []
|
||||
|
||||
def test_mid_turn_session_is_never_evicted(self, monkeypatch):
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 4000)
|
||||
runner._commit_then_release_soft = lambda agent, key: None
|
||||
|
||||
active = self._cached_agent()
|
||||
runner._agent_cache["s-active"] = (active, "sig")
|
||||
runner._agent_cache["s-idle"] = (self._cached_agent(), "sig")
|
||||
runner._agent_cache["s-new"] = (self._cached_agent(), "sig")
|
||||
runner._running_agents["s-active"] = active
|
||||
|
||||
runner._sweep_agent_cache_under_pressure()
|
||||
|
||||
assert "s-active" in runner._agent_cache, (
|
||||
"evicting a mid-turn agent tears down the clients and sandbox the "
|
||||
"running request is using"
|
||||
)
|
||||
assert "s-idle" not in runner._agent_cache
|
||||
|
||||
def test_lagging_persistence_blocks_eviction(self, monkeypatch):
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 4000)
|
||||
runner._commit_then_release_soft = lambda agent, key: None
|
||||
|
||||
runner._agent_cache["s-lagging"] = (
|
||||
self._cached_agent(persisted=False), "sig",
|
||||
)
|
||||
runner._agent_cache["s-durable"] = (self._cached_agent(), "sig")
|
||||
runner._agent_cache["s-new"] = (self._cached_agent(), "sig")
|
||||
|
||||
runner._sweep_agent_cache_under_pressure()
|
||||
|
||||
assert "s-lagging" in runner._agent_cache, (
|
||||
"dropping a transcript that never reached disk loses the "
|
||||
"conversation the FTS guard exists to protect"
|
||||
)
|
||||
assert "s-durable" not in runner._agent_cache
|
||||
|
||||
def test_empty_cache_is_a_no_op(self, monkeypatch):
|
||||
"""Heap pressure with nothing cached is somebody else's problem."""
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 999_999)
|
||||
|
||||
assert runner._sweep_agent_cache_under_pressure() == 0
|
||||
|
||||
def test_all_candidates_skipped_reports_zero_without_raising(self, monkeypatch):
|
||||
runner = self._runner()
|
||||
self._at_rss(monkeypatch, 4000)
|
||||
runner._commit_then_release_soft = lambda agent, key: None
|
||||
for i in range(3):
|
||||
runner._agent_cache[f"s{i}"] = (
|
||||
self._cached_agent(persisted=False), "sig",
|
||||
)
|
||||
|
||||
assert runner._sweep_agent_cache_under_pressure() == 0
|
||||
assert len(runner._agent_cache) == 3
|
||||
|
||||
def test_disabled_budget_is_a_no_op(self, monkeypatch):
|
||||
runner = self._runner(bounds=AgentCacheBounds(memory_high_mb=None))
|
||||
self._at_rss(monkeypatch, 999_999)
|
||||
runner._agent_cache["s0"] = (self._cached_agent(), "sig")
|
||||
|
||||
assert runner._sweep_agent_cache_under_pressure() == 0
|
||||
assert "s0" in runner._agent_cache
|
||||
|
||||
|
||||
class TestConfiguredBoundsReachTheCache:
|
||||
"""The two existing bounds must be operator-tunable, and must keep their
|
||||
built-in values when config.yaml says nothing."""
|
||||
|
||||
def _runner(self, bounds):
|
||||
from gateway.run import GatewayRunner
|
||||
|
||||
runner = GatewayRunner.__new__(GatewayRunner)
|
||||
runner._agent_cache_bounds_cache = bounds
|
||||
return runner
|
||||
|
||||
def test_unset_config_keeps_the_built_in_defaults(self):
|
||||
from gateway import run as gw_run
|
||||
|
||||
runner = self._runner(AgentCacheBounds())
|
||||
assert runner._agent_cache_cap() == gw_run._AGENT_CACHE_MAX_SIZE
|
||||
assert runner._agent_cache_idle_ttl() == gw_run._AGENT_CACHE_IDLE_TTL_SECS
|
||||
|
||||
def test_configured_cap_bounds_the_real_enforcer(self):
|
||||
"""A configured cap must actually shrink the cache, not just report."""
|
||||
runner = self._runner(AgentCacheBounds(max_size=2))
|
||||
runner._agent_cache = OrderedDict()
|
||||
runner._agent_cache_lock = threading.Lock()
|
||||
runner._running_agents = {}
|
||||
runner._release_evicted_agent_soft = lambda agent: None
|
||||
runner._commit_then_release_soft = lambda agent, key: None
|
||||
|
||||
with runner._agent_cache_lock:
|
||||
for i in range(5):
|
||||
runner._agent_cache[f"s{i}"] = (MagicMock(), "sig")
|
||||
runner._enforce_agent_cache_cap()
|
||||
|
||||
assert len(runner._agent_cache) == 2
|
||||
assert list(runner._agent_cache) == ["s3", "s4"]
|
||||
|
||||
def test_configured_idle_ttl_drives_the_real_sweep(self):
|
||||
import time as _t
|
||||
|
||||
runner = self._runner(AgentCacheBounds(idle_ttl_secs=0.01))
|
||||
runner._agent_cache = OrderedDict()
|
||||
runner._agent_cache_lock = threading.Lock()
|
||||
runner._running_agents = {}
|
||||
runner._release_evicted_agent_soft = lambda agent: None
|
||||
runner.session_store = None
|
||||
|
||||
stale = MagicMock()
|
||||
stale._last_activity_ts = _t.time() - 5.0
|
||||
runner._agent_cache["s-stale"] = (stale, "sig")
|
||||
|
||||
assert runner._sweep_idle_cached_agents() == 1
|
||||
assert "s-stale" not in runner._agent_cache
|
||||
|
||||
|
||||
def _wait_for(predicate, timeout: float = 3.0) -> None:
|
||||
"""Wait for a background release thread to finish its work."""
|
||||
import time as _t
|
||||
|
||||
deadline = _t.time() + timeout
|
||||
while _t.time() < deadline:
|
||||
if predicate():
|
||||
return
|
||||
_t.sleep(0.02)
|
||||
assert predicate(), "background release did not complete in time"
|
||||
Loading…
Reference in New Issue