feat: Add KeyPairResourcePolicy.max_session_lifetime

changes/583.feature.md

@@ -0,0 +1 @@
+Add `KeyPairResourcePolicy.max_session_lifetime` to force-terminate sessions after a specified timeout

fixtures/example-keypairs.json

@@ -114,6 +114,7 @@
             "name": "default",
             "default_for_unspecified": "UNLIMITED",
             "total_resource_slots": {},
+            "max_session_lifetime": 0,
             "max_concurrent_sessions": 5,
             "max_containers_per_session": 1,
             "max_vfolder_count": 10,

src/ai/backend/manager/api/context.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Sequence, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import attr
 
@@ -13,7 +13,7 @@
 
     from ..models.storage import StorageSessionManager
     from ..models.utils import ExtendedAsyncSAEngine
-    from ..idle import BaseIdleChecker
+    from ..idle import IdleCheckerHost
     from ..plugin.webapp import WebappPluginContext
     from ..registry import AgentRegistry
     from ..config import LocalConfig, SharedConfig
@@ -41,7 +41,7 @@ class RootContext(BaseContext):
     cors_options: CORSOptions
 
     webapp_plugin_ctx: WebappPluginContext
-    idle_checkers: Sequence[BaseIdleChecker]
+    idle_checker_host: IdleCheckerHost
     storage_manager: StorageSessionManager
     hook_plugin_ctx: HookPluginContext
 

src/ai/backend/manager/api/stream.py

@@ -496,11 +496,10 @@ async def add_conn_track() -> None:
                 # aioredis' ZADD implementation flattens mapping in value-key order
                 lambda r: r.zadd(conn_tracker_key, {conn_tracker_val: now}),
             )
-            for idle_checker in root_ctx.idle_checkers:
-                await idle_checker.update_app_streaming_status(
-                    kernel_id,
-                    AppStreamingStatus.HAS_ACTIVE_CONNECTIONS,
-                )
+            await root_ctx.idle_checker_host.update_app_streaming_status(
+                kernel_id,
+                AppStreamingStatus.HAS_ACTIVE_CONNECTIONS,
+            )
 
     async def clear_conn_track() -> None:
         async with app_ctx.conn_tracker_lock:
@@ -516,11 +515,10 @@ async def clear_conn_track() -> None:
                 ),
             )
             if remaining_count == 0:
-                for idle_checker in root_ctx.idle_checkers:
-                    await idle_checker.update_app_streaming_status(
-                        kernel_id,
-                        AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
-                    )
+                await root_ctx.idle_checker_host.update_app_streaming_status(
+                    kernel_id,
+                    AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
+                )
 
     try:
         await asyncio.shield(database_ptask_group.create_task(
@@ -655,11 +653,10 @@ async def stream_conn_tracker_gc(root_ctx: RootContext, app_ctx: PrivateContext)
                     log.debug(f"conn_tracker: gc {session_id} "
                               f"removed/remaining = {removed_count}/{remaining_count}")
                     if prev_remaining_count > 0 and remaining_count == 0:
-                        for idle_checker in root_ctx.idle_checkers:
-                            await idle_checker.update_app_streaming_status(
-                                session_id,
-                                AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
-                            )
+                        await root_ctx.idle_checker_host.update_app_streaming_status(
+                            session_id,
+                            AppStreamingStatus.NO_ACTIVE_CONNECTIONS,
+                        )
             await asyncio.sleep(10)
     except asyncio.CancelledError:
         pass

src/ai/backend/manager/config.py

@@ -94,6 +94,8 @@
            - time-window: "12h"  # time window to average utilization
                                  # a session will not be terminated until this time
            - initial-grace-period: "5m" # time to allow to be idle for first
+         # "session_lifetime" does not have etcd config but it is configured via
+         # the keypair_resource_polices table.
      + resource_slots
        - {"cuda.device"}: {"count"}
        - {"cuda.mem"}: {"bytes"}