feat(BA-6581): wire IdleCheckerRepository into idle-check reconciler source

changes/12398.feature.md

@@ -0,0 +1 @@
+Wire IdleCheckerRepository into the idle-check reconciler source to fetch a normalized session/scope/checker snapshot

src/ai/backend/common/data/idle_checker/types.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import enum
+from typing import Self
+
+from pydantic import Field, model_validator
+
+from ai.backend.common.types import BackendAISchema
+
+
+class CheckerType(enum.StrEnum):
+    """Discriminator for the kind of idle checker; selects the concrete spec."""
+
+    SESSION_LIFETIME = "session_lifetime"
+    NETWORK_TIMEOUT = "network_timeout"
+    UTILIZATION = "utilization"
+
+
+class SessionLifetimeSpec(BackendAISchema):
+    """Config for ``CheckerType.SESSION_LIFETIME``.
+
+    Concrete fields (e.g. max lifetime) land with the checker-logic stories.
+    """
+
+
+class NetworkTimeoutSpec(BackendAISchema):
+    """Config for ``CheckerType.NETWORK_TIMEOUT``.
+
+    Concrete fields land with the checker-logic stories.
+    """
+
+
+class UtilizationSpec(BackendAISchema):
+    """Config for ``CheckerType.UTILIZATION``.
+
+    Concrete fields land with the checker-logic stories.
+    """
+
+
+class IdleCheckerSpec(BackendAISchema):
+    """Config payload stored in ``idle_checkers.spec`` as a single JSONB document.
+
+    ``type`` selects which sub-config is valid; exactly one sub-config field must
+    be present for the chosen type (validated below). Stored via ``PydanticColumn``.
+    """
+
+    type: CheckerType = Field(description="Idle checker kind; selects the sub-config.")
+    session_lifetime: SessionLifetimeSpec | None = Field(
+        default=None, description="session_lifetime config."
+    )
+    network: NetworkTimeoutSpec | None = Field(default=None, description="network_timeout config.")
+    utilization: UtilizationSpec | None = Field(default=None, description="utilization config.")
+
+    @model_validator(mode="after")
+    def validate_spec_matches_type(self) -> Self:
+        match self.type:
+            case CheckerType.SESSION_LIFETIME:
+                if self.session_lifetime is None:
+                    raise ValueError("session_lifetime is required for type=session_lifetime.")
+                if self.network or self.utilization:
+                    raise ValueError("Only session_lifetime is allowed for type=session_lifetime.")
+            case CheckerType.NETWORK_TIMEOUT:
+                if self.network is None:
+                    raise ValueError("network is required for type=network_timeout.")
+                if self.session_lifetime or self.utilization:
+                    raise ValueError("Only network is allowed for type=network_timeout.")
+            case CheckerType.UTILIZATION:
+                if self.utilization is None:
+                    raise ValueError("utilization is required for type=utilization.")
+                if self.session_lifetime or self.network:
+                    raise ValueError("Only utilization is allowed for type=utilization.")
+        return self

src/ai/backend/manager/data/idle_checker/types.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
-import enum
+from dataclasses import dataclass
+from datetime import datetime
 
+from ai.backend.common.types import SessionId
 
-class CheckerType(enum.StrEnum):
-    """Discriminator for the kind of idle checker; selects the concrete spec."""
 
-    SESSION_LIFETIME = "session_lifetime"
-    NETWORK_TIMEOUT = "network_timeout"
-    UTILIZATION = "utilization"
+@dataclass(frozen=True)
+class IdleCheckSession:
+    """Session fields needed to evaluate idle checkers."""
+
+    session_id: SessionId
+    created_at: datetime
+    starts_at: datetime | None

src/ai/backend/manager/dependencies/composer.py

@@ -270,6 +270,7 @@ async def compose(
                 scheduler_repository=domain.repositories.scheduler.repository,
                 deployment_repository=domain.repositories.deployment.repository,
                 replica_group_repository=domain.repositories.replica_group.repository,
+                idle_checker_repository=domain.repositories.idle_checker.repository,
                 fair_share_repository=domain.repositories.fair_share.repository,
                 resource_usage_repository=domain.repositories.resource_usage_history.repository,
                 agent_client_pool=agents.agent_client_pool,

src/ai/backend/manager/dependencies/orchestration/composer.py

@@ -21,6 +21,7 @@
 from ai.backend.manager.plugin.network import NetworkPluginContext
 from ai.backend.manager.repositories.deployment.repository import DeploymentRepository
 from ai.backend.manager.repositories.fair_share import FairShareRepository
+from ai.backend.manager.repositories.idle_checker.repository import IdleCheckerRepository
 from ai.backend.manager.repositories.prometheus_query_preset.repository import (
     PrometheusQueryPresetRepository,
 )
@@ -62,6 +63,7 @@ class OrchestrationInput:
     scheduler_repository: SchedulerRepository
     deployment_repository: DeploymentRepository
     replica_group_repository: ReplicaGroupRepository
+    idle_checker_repository: IdleCheckerRepository
     fair_share_repository: FairShareRepository
     resource_usage_repository: ResourceUsageHistoryRepository
     agent_client_pool: AgentClientPool
@@ -140,6 +142,7 @@ async def compose(
             scheduler_repository=setup_input.scheduler_repository,
             deployment_repository=setup_input.deployment_repository,
             replica_group_repository=setup_input.replica_group_repository,
+            idle_checker_repository=setup_input.idle_checker_repository,
             fair_share_repository=setup_input.fair_share_repository,
             resource_usage_repository=setup_input.resource_usage_repository,
             config_provider=setup_input.config_provider,

src/ai/backend/manager/dependencies/orchestration/sokovan.py

@@ -23,6 +23,7 @@
 from ai.backend.manager.plugin.network import NetworkPluginContext
 from ai.backend.manager.repositories.deployment.repository import DeploymentRepository
 from ai.backend.manager.repositories.fair_share import FairShareRepository
+from ai.backend.manager.repositories.idle_checker.repository import IdleCheckerRepository
 from ai.backend.manager.repositories.prometheus_query_preset.repository import (
     PrometheusQueryPresetRepository,
 )
@@ -62,6 +63,7 @@ class SokovanOrchestratorInput:
     scheduler_repository: SchedulerRepository
     deployment_repository: DeploymentRepository
     replica_group_repository: ReplicaGroupRepository
+    idle_checker_repository: IdleCheckerRepository
     fair_share_repository: FairShareRepository
     resource_usage_repository: ResourceUsageHistoryRepository
     config_provider: ManagerConfigProvider
@@ -192,6 +194,7 @@ async def provide(
         # Reconciler coordinator: sokovan owns its stage assembly (DI just passes deps).
         reconciler_coordinator, reconciler_task_specs = build_reconciler_coordinator(
             replica_group_repository=setup_input.replica_group_repository,
+            idle_checker_repository=setup_input.idle_checker_repository,
             valkey_schedule=setup_input.valkey_schedule,
             lock_factory=setup_input.distributed_lock_factory,
             config_provider=setup_input.config_provider,

src/ai/backend/manager/models/alembic/versions/d3f8a1c45e9b_create_idle_checker_tables.py

@@ -2,10 +2,10 @@
 
 Introduce the persistence foundation for the first-class idle checker
 (BEP-1054). ``idle_checkers`` holds a reusable, scope-free checker spec
-(``checker_type`` + JSONB ``spec``); ``idle_checker_bindings`` is the
-scope <-> checker association carrying its own ``enabled`` flag. Scope is a
-polymorphic ``(scope_type, scope_id)`` string pair validated on the write
-path rather than by a DB foreign key.
+(``checker_type`` + target session types + JSONB ``spec``);
+``idle_checker_bindings`` is the scope <-> checker association carrying its
+own ``enabled`` flag. Scope is a polymorphic ``(scope_type, scope_id)`` string
+pair validated on the write path rather than by a DB foreign key.
 
 Revision ID: d3f8a1c45e9b
 Revises: a8e06485829f
@@ -34,6 +34,11 @@ def upgrade() -> None:
         sa.Column("name", sa.String(length=128), nullable=False),
         sa.Column("description", sa.Text(), nullable=True),
         sa.Column("checker_type", sa.String(length=64), nullable=False),
+        sa.Column(
+            "target_session_types",
+            postgresql.ARRAY(sa.String(length=64)),
+            nullable=False,
+        ),
         sa.Column("spec", postgresql.JSONB(), nullable=False),
         sa.Column(
             "created_at",

src/ai/backend/manager/models/idle_checker/conditions.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import sqlalchemy as sa
+
+from ai.backend.manager.models.clauses import QueryCondition
+
+from .row import IdleCheckerBindingRow
+
+
+class IdleCheckerBindingConditions:
+    @staticmethod
+    def enabled() -> QueryCondition:
+        def inner() -> sa.sql.expression.ColumnElement[bool]:
+            return IdleCheckerBindingRow.enabled == sa.true()
+
+        return inner

src/ai/backend/manager/models/idle_checker/row.py

@@ -6,10 +6,10 @@
 import sqlalchemy as sa
 from sqlalchemy.orm import Mapped, mapped_column
 
+from ai.backend.common.data.idle_checker.types import CheckerType, IdleCheckerSpec
 from ai.backend.common.identifier.idle_checker import IdleCheckerID
-from ai.backend.manager.data.idle_checker.types import CheckerType
-from ai.backend.manager.models.base import GUID, ABCColumn, Base, StrEnumType
-from ai.backend.manager.models.idle_checker.spec import IdleCheckerSpecABC
+from ai.backend.common.types import SessionTypes
+from ai.backend.manager.models.base import GUID, Base, PydanticColumn, StrEnumType
 
 
 class IdleCheckerRow(Base):  # type: ignore[misc]
@@ -23,8 +23,13 @@ class IdleCheckerRow(Base):  # type: ignore[misc]
     checker_type: Mapped[CheckerType] = mapped_column(
         "checker_type", StrEnumType(CheckerType), nullable=False
     )
-    spec: Mapped[IdleCheckerSpecABC] = mapped_column(
-        "spec", ABCColumn(IdleCheckerSpecABC), nullable=False
+    target_session_types: Mapped[list[SessionTypes]] = mapped_column(
+        "target_session_types",
+        sa.ARRAY(StrEnumType(SessionTypes, use_name=True)),
+        nullable=False,
+    )
+    spec: Mapped[IdleCheckerSpec] = mapped_column(
+        "spec", PydanticColumn(IdleCheckerSpec), nullable=False
     )
     created_at: Mapped[datetime] = mapped_column(
         "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()

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

@@ -15,8 +15,10 @@
         UUIDInMatchSpec,
     )
 
-from ai.backend.common.types import SessionId
+from ai.backend.common.data.permission.types import ScopeType
+from ai.backend.common.types import SessionId, SessionTypes
 from ai.backend.manager.data.kernel.types import KernelStatus
+from ai.backend.manager.data.permission.id import ScopeId
 from ai.backend.manager.data.session.types import KernelMatchType, SessionStatus
 from ai.backend.manager.models.clauses import QueryCondition
 from ai.backend.manager.models.condition_utils import make_string_in_factory
@@ -42,6 +44,30 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:
 
         return inner
 
+    @staticmethod
+    def by_idle_check_candidates(
+        candidates: Collection[tuple[ScopeId, Collection[SessionTypes]]],
+    ) -> QueryCondition:
+        scope_columns = {
+            ScopeType.RESOURCE_GROUP: SessionRow.resource_group_id,
+            ScopeType.PROJECT: SessionRow.group_id,
+            ScopeType.DOMAIN: SessionRow.domain_id,
+        }
+        candidate_conditions = []
+        for scope, target_session_types in candidates:
+            scope_column = scope_columns[scope.scope_type]
+            candidate_conditions.append(
+                sa.and_(
+                    scope_column == scope.scope_id,
+                    SessionRow.session_type.in_(target_session_types),
+                )
+            )
+
+        def inner() -> sa.sql.expression.ColumnElement[bool]:
+            return sa.or_(*candidate_conditions) if candidate_conditions else sa.false()
+
+        return inner
+
     @staticmethod
     def by_scaling_group(scaling_group: str) -> QueryCondition:
         def inner() -> sa.sql.expression.ColumnElement[bool]:

src/ai/backend/manager/repositories/idle_checker/db_source/db_source.py

@@ -0,0 +1,102 @@
+"""DB reads backing the idle-check Source: the per-session checker batch."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import cast
+
+import sqlalchemy as sa
+
+from ai.backend.common.data.idle_checker.types import CheckerType, IdleCheckerSpec
+from ai.backend.common.data.permission.types import ScopeType
+from ai.backend.common.identifier.idle_checker import IdleCheckerID
+from ai.backend.common.types import SessionId, SessionTypes
+from ai.backend.manager.data.idle_checker.types import IdleCheckSession
+from ai.backend.manager.data.permission.id import ScopeId
+from ai.backend.manager.models.idle_checker.row import IdleCheckerBindingRow, IdleCheckerRow
+from ai.backend.manager.models.session.row import SessionRow
+from ai.backend.manager.repositories.base import BatchQuerier
+from ai.backend.manager.repositories.idle_checker.types import (
+    BoundCheckerData,
+    IdleCheckerDefinitionData,
+    IdleCheckSessionData,
+)
+from ai.backend.manager.repositories.ops import DBOpsProvider
+
+
+class IdleCheckerDBSource:
+    _ops: DBOpsProvider
+
+    def __init__(self, ops_provider: DBOpsProvider) -> None:
+        self._ops = ops_provider
+
+    async def fetch_bound_checkers(
+        self,
+        querier: BatchQuerier,
+    ) -> Sequence[BoundCheckerData]:
+        binding_query = sa.select(
+            IdleCheckerBindingRow.scope_type,
+            IdleCheckerBindingRow.scope_id,
+            IdleCheckerBindingRow.created_at,
+            IdleCheckerBindingRow.idle_checker_id,
+            IdleCheckerRow.checker_type,
+            IdleCheckerRow.target_session_types,
+            IdleCheckerRow.spec,
+        ).join(IdleCheckerRow, IdleCheckerBindingRow.idle_checker_id == IdleCheckerRow.id)
+
+        async with self._ops.read_ops() as r:
+            binding_rows = (await r.batch_query_in_global(binding_query, querier)).rows
+        bound_checkers: list[BoundCheckerData] = []
+        for row in binding_rows:
+            scope = ScopeId(
+                scope_type=ScopeType(row.scope_type),
+                scope_id=str(row.scope_id),
+            )
+            checker = IdleCheckerDefinitionData(
+                checker_id=cast(IdleCheckerID, row.idle_checker_id),
+                checker_type=cast(CheckerType, row.checker_type),
+                target_session_types=frozenset(
+                    cast(Sequence[SessionTypes], row.target_session_types)
+                ),
+                spec=cast(IdleCheckerSpec, row.spec),
+            )
+            bound_checkers.append(
+                BoundCheckerData(
+                    scope=scope,
+                    binding_created_at=row.created_at,
+                    checker=checker,
+                )
+            )
+        return bound_checkers
+
+    async def fetch_candidate_sessions(
+        self,
+        querier: BatchQuerier,
+    ) -> Sequence[IdleCheckSessionData]:
+        session_query = sa.select(
+            SessionRow.id,
+            SessionRow.created_at,
+            SessionRow.starts_at,
+            SessionRow.session_type,
+            SessionRow.resource_group_id,
+            SessionRow.group_id,
+            SessionRow.domain_id,
+        )
+        async with self._ops.read_ops() as r:
+            session_rows = (await r.batch_query_in_global(session_query, querier)).rows
+        return tuple(
+            IdleCheckSessionData(
+                session=IdleCheckSession(
+                    session_id=SessionId(row.id),
+                    created_at=row.created_at,
+                    starts_at=row.starts_at,
+                ),
+                session_type=row.session_type,
+                scopes=(
+                    ScopeId(ScopeType.RESOURCE_GROUP, str(row.resource_group_id)),
+                    ScopeId(ScopeType.PROJECT, str(row.group_id)),
+                    ScopeId(ScopeType.DOMAIN, str(row.domain_id)),
+                ),
+            )
+            for row in session_rows
+        )