diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 5a32a1bf41..1257574958 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -35,6 +35,7 @@ from nemo_agents_plugin.fabric.session_manager import ( DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS, DEFAULT_MAX_CONCURRENT_INVOCATIONS, + DEFAULT_MAX_LIVE_SESSIONS, DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS, FabricSessionManager, FabricSessionStartError, @@ -57,12 +58,15 @@ class FabricServingSettings: """Operational settings for the Platform-owned Fabric server.""" max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS + max_live_sessions: int = DEFAULT_MAX_LIVE_SESSIONS idle_session_timeout_seconds: float = DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS session_cleanup_interval_seconds: float = DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS def __post_init__(self) -> None: if self.max_concurrent_invocations < 0: raise ValueError("max_concurrent_invocations must be greater than or equal to zero.") + if self.max_live_sessions < 0: + raise ValueError("max_live_sessions must be greater than or equal to zero.") if self.idle_session_timeout_seconds <= 0: raise ValueError("idle_session_timeout_seconds must be greater than zero.") if self.session_cleanup_interval_seconds <= 0: @@ -257,6 +261,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: base_dir=config_path.parent, session_registry=session_registry, max_concurrent_invocations=settings.max_concurrent_invocations, + max_live_sessions=settings.max_live_sessions, ) app.state.session_manager = session_manager cleanup_shutdown = asyncio.Event() @@ -390,6 +395,12 @@ def main(argv: list[str] | None = None) -> int: default=DEFAULT_MAX_CONCURRENT_INVOCATIONS, help="Maximum concurrent Fabric invocations; use 0 for unlimited.", ) + parser.add_argument( + "--max-live-sessions", + type=int, + default=DEFAULT_MAX_LIVE_SESSIONS, + help="Maximum concurrently live sessions; the least recently used idle one is evicted above it. Use 0 for unlimited.", + ) parser.add_argument( "--idle-session-timeout-seconds", type=float, @@ -412,6 +423,7 @@ def main(argv: list[str] | None = None) -> int: args.agent_config, settings=FabricServingSettings( max_concurrent_invocations=args.max_concurrent_invocations, + max_live_sessions=args.max_live_sessions, idle_session_timeout_seconds=args.idle_session_timeout_seconds, session_cleanup_interval_seconds=args.session_cleanup_interval_seconds, ), diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 23d1a27320..a279b8380d 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -35,6 +35,8 @@ DEFAULT_MAX_CONCURRENT_INVOCATIONS = 8 DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS = 30 * 60 DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS = 5 * 60 +# A caller that sends no session header cannot close what it opens, so the cap does it instead. +DEFAULT_MAX_LIVE_SESSIONS = DEFAULT_MAX_CONCURRENT_INVOCATIONS class FabricSessionStartError(RuntimeError): @@ -56,14 +58,18 @@ def __init__( session_registry: FabricSessionRegistry, fabric: Fabric | None = None, max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS, + max_live_sessions: int = DEFAULT_MAX_LIVE_SESSIONS, ) -> None: if max_concurrent_invocations < 0: raise ValueError("max_concurrent_invocations must be greater than or equal to zero.") + if max_live_sessions < 0: + raise ValueError("max_live_sessions must be greater than or equal to zero.") self._agent_config = agent_config self._base_dir = base_dir self._session_registry = session_registry self._fabric = fabric + self._max_live_sessions = max_live_sessions self._invocation_semaphore = ( asyncio.Semaphore(max_concurrent_invocations) if max_concurrent_invocations > 0 else None ) @@ -87,7 +93,7 @@ async def open_session(self) -> FabricRuntimeSession: raise FabricSessionStartError(f"Fabric runtime startup failed: {error}") from error try: - return await self._session_registry.register(runtime) + session = await self._session_registry.register(runtime) except BaseException: # A started runtime must not leak if registration fails or is cancelled. try: @@ -96,6 +102,22 @@ async def open_session(self) -> FabricRuntimeSession: logger.exception("Failed to stop Fabric runtime after session registration failed.") raise + await self._evict_over_capacity(keep=session.session_id) + return session + + async def _evict_over_capacity(self, *, keep: str) -> None: + """Stop the least recently used idle sessions above the live-session cap.""" + evicted = await self._session_registry.evict_over_capacity( + max_sessions=self._max_live_sessions, + keep=keep, + ) + for session in evicted: + logger.info("Evicting idle Fabric session %s to stay within the live-session cap.", session.session_id) + try: + await self._stop_session(session) + except FabricSessionStopError: + logger.exception("Failed to stop evicted Fabric session %s.", session.session_id) + async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession: """Open a new session or resolve an existing session by its opaque ID.""" if session_id is None: diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py index 1c7ed20230..467dd07b55 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py @@ -95,6 +95,27 @@ async def remove(self, session_id: str) -> FabricRuntimeSession | None: session.closing = True return session + async def evict_over_capacity(self, *, max_sessions: int, keep: str) -> list[FabricRuntimeSession]: + """Remove the least recently used idle sessions above *max_sessions*, never *keep*.""" + if max_sessions <= 0: + return [] + + evicted: list[FabricRuntimeSession] = [] + async with self._lock: + while len(self._sessions) > max_sessions: + candidates = [ + session + for session in self._sessions.values() + if session.session_id != keep and not session.invocation_lock.locked() + ] + if not candidates: + break + victim = min(candidates, key=lambda session: session.last_accessed_at) + victim.closing = True + evicted.append(victim) + del self._sessions[victim.session_id] + return evicted + async def remove_expired(self, *, idle_timeout_seconds: float) -> list[FabricRuntimeSession]: """Remove inactive sessions that are not currently invoking.""" cutoff = time.monotonic() - idle_timeout_seconds diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py index ddf5a791d8..9a54d5b490 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_registry.py @@ -117,6 +117,69 @@ async def test_remove_expired_removes_only_idle_sessions(monkeypatch: pytest.Mon assert await registry.count() == 2 +@pytest.mark.asyncio +async def test_evict_over_capacity_removes_least_recently_used_first() -> None: + registry = FabricSessionRegistry() + oldest = await registry.register(cast(Any, object()), session_id="oldest") + middle = await registry.register(cast(Any, object()), session_id="middle") + newest = await registry.register(cast(Any, object()), session_id="newest") + oldest.last_accessed_at = 10.0 + middle.last_accessed_at = 20.0 + newest.last_accessed_at = 30.0 + + evicted = await registry.evict_over_capacity(max_sessions=1, keep="newest") + + assert evicted == [oldest, middle] + assert oldest.closing is True + assert middle.closing is True + assert newest.closing is False + assert await registry.count() == 1 + + +@pytest.mark.asyncio +async def test_evict_over_capacity_keeps_the_session_just_opened() -> None: + registry = FabricSessionRegistry() + established = await registry.register(cast(Any, object()), session_id="established") + opened = await registry.register(cast(Any, object()), session_id="opened") + established.last_accessed_at = 90.0 + opened.last_accessed_at = 10.0 + + evicted = await registry.evict_over_capacity(max_sessions=1, keep="opened") + + assert evicted == [established] + assert await registry.count() == 1 + + +@pytest.mark.asyncio +async def test_evict_over_capacity_never_evicts_a_session_mid_invocation() -> None: + registry = FabricSessionRegistry() + busy = await registry.register(cast(Any, object()), session_id="busy") + opened = await registry.register(cast(Any, object()), session_id="opened") + busy.last_accessed_at = 10.0 + opened.last_accessed_at = 20.0 + + await busy.invocation_lock.acquire() + try: + evicted = await registry.evict_over_capacity(max_sessions=1, keep="opened") + finally: + busy.invocation_lock.release() + + assert evicted == [] + assert busy.closing is False + assert await registry.count() == 2 + + +@pytest.mark.asyncio +async def test_evict_over_capacity_is_a_no_op_under_the_cap_or_when_unlimited() -> None: + registry = FabricSessionRegistry() + await registry.register(cast(Any, object()), session_id="session-1") + await registry.register(cast(Any, object()), session_id="session-2") + + assert await registry.evict_over_capacity(max_sessions=4, keep="session-2") == [] + assert await registry.evict_over_capacity(max_sessions=0, keep="session-2") == [] + assert await registry.count() == 2 + + @pytest.mark.asyncio async def test_drain_removes_sessions_and_rejects_new_registrations() -> None: registry = FabricSessionRegistry()