refactor(BA-1004): Sweeper

Author: rapsealk

src/ai/backend/manager/cleanup/kernel.py

@@ -7,7 +7,7 @@
 from collections.abc import Iterable, Mapping, Sequence
 from contextlib import asynccontextmanager as actxmgr
 from dataclasses import dataclass, field
-from datetime import datetime, timedelta
+from datetime import datetime
 from decimal import Decimal
 from typing import (
     TYPE_CHECKING,
@@ -1312,8 +1312,10 @@ async def get_network_ref(self, db_sess: SASession) -> str | None:
                 return None
 
     @classmethod
-    def get_status_elapsed_time(cls, status: SessionStatus, now: datetime) -> timedelta:
-        return now - cls.status_history[status.name].astext.cast(sa.types.DateTime(timezone=True))
+    def get_status_elapsed_time(
+        cls, status: SessionStatus, until: datetime
+    ) -> sa.sql.elements.BinaryExpression:
+        return until - cls.status_history[status.name].astext.cast(sa.types.DateTime(timezone=True))
 
 
 class SessionLifecycleManager:

src/ai/backend/manager/cleanup/session.py

@@ -92,10 +92,10 @@
     CleanupContext,
     WebRequestHandler,
 )
-from .cleanup import stale_kernel_collection_ctx, stale_session_collection_ctx
 from .config import LocalConfig, SharedConfig, volume_config_iv
 from .config import load as load_config
 from .exceptions import InvalidArgument
+from .sweeper import stale_kernel_collection_ctx, stale_session_collection_ctx
 from .types import DistributedLockFactory
 
 VALID_VERSIONS: Final = frozenset([

src/ai/backend/manager/models/session.py

@@ -0,0 +1,7 @@
+import abc
+
+
+class AbstractSweeper(abc.ABC):
+    @abc.abstractmethod
+    async def sweep(self, *args) -> None:
+        raise NotImplementedError

src/ai/backend/manager/server.py

@@ -0,0 +1,103 @@
+import asyncio
+import logging
+from collections import defaultdict
+from contextlib import asynccontextmanager as actxmgr
+from contextlib import suppress
+from typing import AsyncIterator, override
+
+import aiotools
+import sqlalchemy as sa
+from sqlalchemy.orm import load_only, noload
+
+from ai.backend.common.events import KernelLifecycleEventReason
+from ai.backend.common.validators import TimeDelta
+from ai.backend.logging import BraceStyleAdapter
+
+from ..api.context import RootContext
+from ..config import session_hang_tolerance_iv
+from ..models import DEAD_KERNEL_STATUSES, DEAD_SESSION_STATUSES, KernelRow, SessionRow
+from .base import AbstractSweeper
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+
+class KernelSweeper(AbstractSweeper):
+    _root_ctx: RootContext
+
+    def __init__(self, root_ctx: RootContext) -> None:
+        self._root_ctx = root_ctx
+
+    @override
+    async def sweep(self, *args) -> None:
+        query = (
+            sa.select(KernelRow)
+            .join(SessionRow, KernelRow.session_id == SessionRow.id)
+            .where(KernelRow.status.not_in(DEAD_KERNEL_STATUSES))
+            .where(SessionRow.status.in_(DEAD_SESSION_STATUSES))
+            .options(
+                noload("*"),
+                load_only(
+                    KernelRow.id,
+                    KernelRow.session_id,
+                    KernelRow.agent,
+                    KernelRow.agent_addr,
+                    KernelRow.container_id,
+                ),
+            )
+        )
+
+        async with self._root_ctx.db.begin_readonly() as conn:
+            result = await conn.execute(query)
+            kernels = result.fetchall()
+
+        kernels_per_session = defaultdict(list)
+        for kernel in kernels:
+            kernels_per_session[kernel.session_id].append(kernel)
+
+        await asyncio.gather(
+            *[
+                asyncio.create_task(
+                    self._root_ctx.registry.destroy_session_lowlevel(
+                        session_id,
+                        [
+                            {
+                                "id": kernel.id,
+                                "session_id": kernel.session_id,
+                                "agent": kernel.agent,
+                                "agent_addr": kernel.agent_addr,
+                                "container_id": kernel.container_id,
+                            }
+                            for kernel in kernels_
+                        ],
+                        reason=KernelLifecycleEventReason.HANG_TIMEOUT,
+                    )
+                )
+                for session_id, kernels_ in kernels_per_session.items()
+            ],
+            return_exceptions=False,
+        )
+
+
+@actxmgr
+async def stale_kernel_collection_ctx(root_ctx: RootContext) -> AsyncIterator[None]:
+    session_hang_tolerance = session_hang_tolerance_iv.check(
+        await root_ctx.shared_config.etcd.get_prefix_dict("config/session/hang-tolerance")
+    )
+    default_interval_sec = 60.0
+    interval_sec = float("inf")
+    threshold: TimeDelta
+    for threshold in session_hang_tolerance["threshold"].values():
+        interval_sec = min(interval_sec, threshold.seconds)
+    if interval_sec == float("inf"):
+        interval_sec = default_interval_sec
+    task = aiotools.create_timer(
+        KernelSweeper(root_ctx).sweep,
+        interval=interval_sec,
+    )
+
+    yield
+
+    if not task.done():
+        task.cancel()
+        with suppress(asyncio.CancelledError):
+            await task