refactor(BA-1004): Move elapsed time function

Author: rapsealk

src/ai/backend/manager/cleanup/BUILD

@@ -0,0 +1 @@
+python_sources(name="src")

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

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from contextlib import asynccontextmanager as actxmgr
 from contextlib import suppress
-from typing import AsyncIterator, Any, Mapping, Protocol, Sequence
+from typing import Any, AsyncIterator, Mapping, Protocol, Sequence
 
 import aiotools
 import sqlalchemy as sa
@@ -17,11 +17,9 @@
 
 from ..api.context import RootContext
 from ..config import session_hang_tolerance_iv
-from ..models import KernelRow, SessionRow, DEAD_SESSION_STATUSES, DEAD_KERNEL_STATUSES
+from ..models import DEAD_KERNEL_STATUSES, DEAD_SESSION_STATUSES, KernelRow, SessionRow
 from ..models.utils import ExtendedAsyncSAEngine
 
-__all__ = ("stale_kernel_collection_ctx",)
-
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
 
@@ -103,10 +101,10 @@ async def stale_kernel_collection_ctx(root_ctx: RootContext) -> AsyncIterator[No
         await root_ctx.shared_config.etcd.get_prefix_dict("config/session/hang-tolerance")
     )
     default_interval_sec = 60.0
-    interval_sec: float = float("inf")
+    interval_sec = float("inf")
     threshold: TimeDelta
     for threshold in session_hang_tolerance["threshold"].values():
-        interval_sec = min(interval_sec, threshold.total_seconds())
+        interval_sec = min(interval_sec, threshold.seconds)
     if interval_sec == float("inf"):
         interval_sec = default_interval_sec
     task = aiotools.create_timer(
@@ -120,7 +118,7 @@ async def stale_kernel_collection_ctx(root_ctx: RootContext) -> AsyncIterator[No
 
     yield
 
-    if not task.done:
+    if not task.done():
         task.cancel()
         with suppress(asyncio.CancelledError):
             await task

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

@@ -4,7 +4,7 @@
 from contextlib import asynccontextmanager as actxmgr
 from contextlib import suppress
 from datetime import datetime, timedelta
-from typing import Any, AsyncIterator, Mapping, Optional, Protocol, Sequence
+from typing import Any, AsyncIterator, Mapping, Optional, Protocol
 
 import aiotools
 import sqlalchemy as sa
@@ -22,8 +22,6 @@
 from ..models.session import SessionStatus
 from ..models.utils import ExtendedAsyncSAEngine
 
-__all__ = ("stale_session_collection_ctx",)
-
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
 
@@ -38,7 +36,7 @@ async def destroy_session(
     ) -> Mapping[str, Any]: ...
 
 
-def get_interval(
+def _get_interval(
     threshold: TimeDelta,
     *,
     max_interval: float = timedelta(hours=1).total_seconds(),
@@ -60,11 +58,7 @@ async def handle_stale_sessions(
     query = (
         sa.select(SessionRow)
         .where(SessionRow.status == status)
-        .where(
-            now
-            - SessionRow.status_history[status.name].astext.cast(sa.types.DateTime(timezone=True))
-            > threshold
-        )
+        .where(SessionRow.get_status_elapsed_time(status, now).total_seconds() > threshold.seconds)
         .options(
             noload("*"),
             load_only(SessionRow.id, SessionRow.name, SessionRow.access_key),
@@ -109,7 +103,7 @@ async def stale_session_collection_ctx(root_ctx: RootContext) -> AsyncIterator[N
             log.warning(f"Invalid session status for hang-threshold: '{raw_status}'")
             continue
 
-        interval = get_interval(threshold)
+        interval = _get_interval(threshold)
         tasks.append(
             aiotools.create_timer(
                 functools.partial(
@@ -126,7 +120,7 @@ async def stale_session_collection_ctx(root_ctx: RootContext) -> AsyncIterator[N
     yield
 
     for task in tasks:
-        if not task.done:
+        if not task.done():
             task.cancel()
             with suppress(asyncio.CancelledError):
                 await task

src/ai/backend/manager/models/session.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
+from datetime import datetime, timedelta
 from decimal import Decimal
 from typing import (
     TYPE_CHECKING,
@@ -1311,6 +1311,10 @@ async def get_network_ref(self, db_sess: SASession) -> str | None:
             case _:
                 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))
+
 
 class SessionLifecycleManager:
     status_set_key = "session_status_update"