feat(BA-6894): implement network timeout idle checker

changes/12878.feature.md

@@ -0,0 +1 @@
+Add a reconciler-based network timeout idle checker using last-access and active-connection signals.

src/ai/backend/common/clients/valkey_client/valkey_live/client.py

@@ -30,7 +30,7 @@
     RetryArgs,
     RetryPolicy,
 )
-from ai.backend.common.types import ValkeyTarget
+from ai.backend.common.types import SessionId, ValkeyTarget
 from ai.backend.logging.utils import BraceStyleAdapter
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
@@ -249,6 +249,24 @@ async def count_active_connections(self, session_id: str) -> int:
                 InfBound.POS_INF,
             )
 
+    @valkey_live_resilience.apply()
+    async def count_active_connections_batch(
+        self,
+        session_ids: Sequence[SessionId],
+    ) -> dict[SessionId, int]:
+        """Count active connections for multiple sessions in one batch."""
+        if not session_ids:
+            return {}
+        batch = self._create_batch()
+        for session_id in session_ids:
+            batch.zcount(
+                self._active_app_connection_key(str(session_id)),
+                InfBound.NEG_INF,
+                InfBound.POS_INF,
+            )
+        results = cast(list[int], await self._execute_batch(batch))
+        return dict(zip(session_ids, results, strict=True))
+
     @valkey_live_resilience.apply()
     async def add_scheduler_metadata(
         self,

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

@@ -30,10 +30,16 @@ class SessionLifetimeSpec(BackendAISchema):
 
 
 class NetworkTimeoutSpec(BackendAISchema):
-    """Config for ``CheckerType.NETWORK_TIMEOUT``.
+    """Config for ``CheckerType.NETWORK_TIMEOUT``."""
 
-    Concrete fields land with the checker-logic stories.
-    """
+    idle_timeout_seconds: int = Field(
+        ge=0,
+        description=(
+            "Maximum time in seconds that an interactive session may have neither "
+            "access nor active connections. Zero disables this checker definition. "
+            "This is the sole network idle timeout used by the reconciler idle checker."
+        ),
+    )
 
 
 class UtilizationSpec(BackendAISchema):

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

@@ -274,6 +274,7 @@ async def compose(
                 distributed_lock_factory=domain.distributed_lock_factory,
                 valkey_profile_target=config.redis.to_valkey_profile_target(),
                 valkey_schedule=infrastructure.valkey.schedule,
+                valkey_live=infrastructure.valkey.live,
                 valkey_stat=infrastructure.valkey.stat,
                 pidx=setup_input.pidx,
                 scheduler_repository=domain.repositories.scheduler.repository,

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

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import override
 
+from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
 from ai.backend.common.clients.valkey_client.valkey_schedule import ValkeyScheduleClient
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.dependencies import DependencyComposer, DependencyStack
@@ -58,6 +59,7 @@ class OrchestrationInput:
     distributed_lock_factory: DistributedLockFactory
     valkey_profile_target: ValkeyProfileTarget
     valkey_schedule: ValkeyScheduleClient
+    valkey_live: ValkeyLiveClient
     valkey_stat: ValkeyStatClient
     pidx: int
     # Sokovan-specific
@@ -153,6 +155,7 @@ async def compose(
             network_plugin_ctx=setup_input.network_plugin_ctx,
             event_producer=setup_input.event_producer,
             valkey_schedule=setup_input.valkey_schedule,
+            valkey_live=setup_input.valkey_live,
             valkey_stat=setup_input.valkey_stat,
             agent_selector=setup_input.agent_selector,
             scheduling_controller=setup_input.scheduling_controller,

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

@@ -10,6 +10,7 @@
     ClientPool,
     tcp_client_session_factory,
 )
+from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
 from ai.backend.common.clients.valkey_client.valkey_schedule import ValkeyScheduleClient
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.dependencies import NonMonitorableDependencyProvider
@@ -73,6 +74,7 @@ class SokovanOrchestratorInput:
     network_plugin_ctx: NetworkPluginContext
     event_producer: EventProducer
     valkey_schedule: ValkeyScheduleClient
+    valkey_live: ValkeyLiveClient
     valkey_stat: ValkeyStatClient
     agent_selector: AgentSelector
     # Controller dependencies
@@ -198,6 +200,7 @@ async def provide(
         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_live=setup_input.valkey_live,
             valkey_schedule=setup_input.valkey_schedule,
             lock_factory=setup_input.distributed_lock_factory,
             config_provider=setup_input.config_provider,

src/ai/backend/manager/sokovan/idle_check/checkers/network_timeout.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from decimal import Decimal
+from typing import override
+
+from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
+from ai.backend.common.types import SessionId
+from ai.backend.logging import BraceStyleAdapter
+from ai.backend.manager.data.idle_checker.types import IdleCheckSession
+from ai.backend.manager.sokovan.idle_check.checkers.base import (
+    CheckerAssignment,
+    IdleChecker,
+    IdleCheckerContext,
+    IdleJudgment,
+)
+
+log = BraceStyleAdapter(logging.getLogger(__name__))
+
+
+@dataclass(frozen=True)
+class _NetworkIdleState:
+    last_access: Decimal | None
+    active_connections: int
+
+
+class NetworkTimeoutChecker(IdleChecker):
+    """Judge interactive sessions from their shared network-activity markers."""
+
+    _valkey_live: ValkeyLiveClient
+
+    def __init__(self, valkey_live: ValkeyLiveClient) -> None:
+        self._valkey_live = valkey_live
+
+    @override
+    async def judge(
+        self,
+        assignments: Sequence[CheckerAssignment],
+        *,
+        context: IdleCheckerContext,
+    ) -> Sequence[IdleJudgment]:
+        # Fetch session states in one batch to avoid repeated I/O per assignment.
+        sessions: list[IdleCheckSession] = []
+        for assignment in assignments:
+            sessions.extend(assignment.sessions)
+        states = await self._prepare_states(sessions)
+
+        # Judge each assignment in one pass, using the pre-fetched states.
+        judgments: list[IdleJudgment] = []
+        for assignment in assignments:
+            network_spec = assignment.definition.spec.network
+            if network_spec is None:
+                log.error(
+                    "Network timeout checker {} has mismatched spec type: {}",
+                    assignment.definition.checker_id,
+                    assignment.definition.spec.type,
+                )
+                continue
+            # Skip sessions with no timeout configured (0 means "never idle").
+            if network_spec.idle_timeout_seconds == 0:
+                continue
+            idle_timeout_seconds = Decimal(network_spec.idle_timeout_seconds)
+            for session in assignment.sessions:
+                state = states[session.session_id]
+                if state.last_access is None:
+                    continue
+                current_time = Decimal(str(context.current_time.timestamp()))
+                idle_seconds = (current_time - state.last_access).normalize()
+                is_idle = (state.active_connections == 0) and (idle_seconds >= idle_timeout_seconds)
+                last_access_at = datetime.fromtimestamp(float(state.last_access), tz=UTC)
+                judgments.append(
+                    IdleJudgment(
+                        checker_id=assignment.definition.checker_id,
+                        session_id=session.session_id,
+                        is_idle=is_idle,
+                        message=(
+                            "Network timeout exceeded: "
+                            f"idle_timeout_seconds={idle_timeout_seconds:f}, "
+                            f"last_access_at={last_access_at:%Y-%m-%d %H:%M:%S} UTC, "
+                            f"idle_seconds={idle_seconds:f} "
+                        ),
+                    )
+                )
+        return judgments
+
+    async def _prepare_states(
+        self,
+        sessions: Sequence[IdleCheckSession],
+    ) -> dict[SessionId, _NetworkIdleState]:
+        session_ids = list(dict.fromkeys(session.session_id for session in sessions))
+        last_access_values, active_connection_counts = await asyncio.gather(
+            self._valkey_live.get_multiple_live_data([
+                f"session.{session_id}.last_access" for session_id in session_ids
+            ]),
+            self._valkey_live.count_active_connections_batch(session_ids),
+        )
+        states: dict[SessionId, _NetworkIdleState] = {}
+        for session_id, raw_last_access in zip(
+            session_ids,
+            last_access_values,
+            strict=True,
+        ):
+            if raw_last_access is None or raw_last_access == b"0":
+                last_access = None
+            else:
+                last_access = Decimal(raw_last_access.decode("utf-8"))
+            states[session_id] = _NetworkIdleState(
+                last_access=last_access,
+                active_connections=active_connection_counts[session_id],
+            )
+        return states

src/ai/backend/manager/sokovan/stages/factory.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
 from ai.backend.common.clients.valkey_client.valkey_schedule import ValkeyScheduleClient
 from ai.backend.manager.config.provider import ManagerConfigProvider
 from ai.backend.manager.repositories.idle_checker.repository import IdleCheckerRepository
@@ -21,6 +22,7 @@ def build_reconciler_coordinator(
     *,
     replica_group_repository: ReplicaGroupRepository,
     idle_checker_repository: IdleCheckerRepository,
+    valkey_live: ValkeyLiveClient,
     valkey_schedule: ValkeyScheduleClient,
     lock_factory: DistributedLockFactory,
     config_provider: ManagerConfigProvider,
@@ -30,7 +32,7 @@ def build_reconciler_coordinator(
         build_group_rolling_stage(replica_group_repository),
         build_group_draining_stage(replica_group_repository),
         build_group_autoscale_stage(replica_group_repository),
-        build_idle_check_stage(idle_checker_repository),
+        build_idle_check_stage(idle_checker_repository, valkey_live),
     ]
     stages: dict[str, ReconcilerStageRunner] = {}
     task_specs: list[ReconcilerTaskSpec] = []

src/ai/backend/manager/sokovan/stages/idle_check.py

@@ -4,6 +4,7 @@
 
 from collections.abc import Mapping
 
+from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
 from ai.backend.common.data.idle_checker.types import CheckerType
 from ai.backend.common.events.event_types.schedule.anycast import (
     DoReconcileProcessEvent,
@@ -14,6 +15,7 @@
 from ai.backend.manager.repositories.idle_checker.repository import IdleCheckerRepository
 from ai.backend.manager.sokovan.idle_check.applier import IdleCheckApplier
 from ai.backend.manager.sokovan.idle_check.checkers.base import IdleChecker
+from ai.backend.manager.sokovan.idle_check.checkers.network_timeout import NetworkTimeoutChecker
 from ai.backend.manager.sokovan.idle_check.checkers.session_lifetime import (
     SessionLifetimeChecker,
 )
@@ -34,6 +36,7 @@
 
 def build_idle_check_stage(
     idle_checker_repository: IdleCheckerRepository,
+    valkey_live: ValkeyLiveClient,
 ) -> ReconcilerStageRegistration:
     reconcile_type = "idle_check"
     # Termination runs through the scheduler lifecycle (mark_sessions_for_termination in
@@ -53,6 +56,7 @@ def build_idle_check_stage(
     )
     checkers: Mapping[CheckerType, IdleChecker] = {
         CheckerType.SESSION_LIFETIME: SessionLifetimeChecker(),
+        CheckerType.NETWORK_TIMEOUT: NetworkTimeoutChecker(valkey_live),
     }
     stage = ReconcilerStage(
         handler=IdleCheckReconcileHandler(checkers),

tests/unit/common/clients/valkey_client/test_valkey_live_client.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
 import random
+from uuid import uuid4
+
+import pytest
 
 from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
+from ai.backend.common.types import SessionId
 
 
 async def test_valkey_live_hset_operations(test_valkey_live: ValkeyLiveClient) -> None:
@@ -44,6 +48,42 @@ async def test_valkey_live_multiple_data_operations(test_valkey_live: ValkeyLive
     assert all(result is not None for result in results)
 
 
+class TestCountActiveConnectionsBatch:
+    @pytest.fixture()
+    async def session_ids_with_active_connections(
+        self,
+        test_valkey_live: ValkeyLiveClient,
+    ) -> list[SessionId]:
+        session_ids = [SessionId(uuid4()) for _ in range(3)]
+        await test_valkey_live.update_connection_tracker(str(session_ids[0]), "ssh", "stream-1")
+        await test_valkey_live.update_connection_tracker(str(session_ids[0]), "ssh", "stream-2")
+        await test_valkey_live.update_connection_tracker(str(session_ids[1]), "jupyter", "stream-1")
+        return session_ids
+
+    async def test_returns_connection_counts_by_session(
+        self,
+        test_valkey_live: ValkeyLiveClient,
+        session_ids_with_active_connections: list[SessionId],
+    ) -> None:
+        result = await test_valkey_live.count_active_connections_batch(
+            session_ids_with_active_connections
+        )
+
+        assert result == {
+            session_ids_with_active_connections[0]: 2,
+            session_ids_with_active_connections[1]: 1,
+            session_ids_with_active_connections[2]: 0,
+        }
+
+    async def test_returns_empty_mapping_for_empty_input(
+        self,
+        test_valkey_live: ValkeyLiveClient,
+    ) -> None:
+        result = await test_valkey_live.count_active_connections_batch([])
+
+        assert result == {}
+
+
 async def test_valkey_live_client_lifecycle(test_valkey_live: ValkeyLiveClient) -> None:
     """Test client lifecycle management"""
     # The client should be created and working

tests/unit/manager/dependencies/orchestration/test_composer.py

@@ -41,6 +41,7 @@ async def test_compose_creates_all_resources(
             distributed_lock_factory=MagicMock(),
             valkey_profile_target=MagicMock(),
             valkey_schedule=MagicMock(),
+            valkey_live=MagicMock(),
             valkey_stat=MagicMock(),
             pidx=0,
             scheduler_repository=MagicMock(),
@@ -121,6 +122,7 @@ async def _provide(setup_input: object) -> AsyncIterator[MagicMock]:
             distributed_lock_factory=MagicMock(),
             valkey_profile_target=MagicMock(),
             valkey_schedule=MagicMock(),
+            valkey_live=MagicMock(),
             valkey_stat=MagicMock(),
             pidx=0,
             scheduler_repository=MagicMock(),

tests/unit/manager/dependencies/orchestration/test_sokovan.py

@@ -59,6 +59,7 @@ async def test_provide_creates_orchestrator(
             network_plugin_ctx=MagicMock(),
             event_producer=MagicMock(),
             valkey_schedule=MagicMock(),
+            valkey_live=MagicMock(),
             valkey_stat=MagicMock(),
             agent_selector=MagicMock(),
             scheduling_controller=MagicMock(),
@@ -127,6 +128,7 @@ async def test_provide_passes_correct_dependencies(
             network_plugin_ctx=network_plugin_ctx,
             event_producer=event_producer,
             valkey_schedule=valkey_schedule,
+            valkey_live=MagicMock(),
             valkey_stat=MagicMock(),
             agent_selector=agent_selector,
             scheduling_controller=MagicMock(),

tests/unit/manager/repositories/idle_checker/test_repository.py

@@ -384,7 +384,7 @@ async def seeded_idle_check_data(
                 CheckerType.NETWORK_TIMEOUT,
                 IdleCheckerSpec(
                     type=CheckerType.NETWORK_TIMEOUT,
-                    network=NetworkTimeoutSpec(),
+                    network=NetworkTimeoutSpec(idle_timeout_seconds=3600),
                 ),
                 ScopeType.PROJECT,
                 project_scope.project_id,

tests/unit/manager/sokovan/idle_check/test_handler.py

@@ -47,7 +47,8 @@
         session_lifetime=SessionLifetimeSpec(max_lifetime_seconds=3600),
     ),
     CheckerType.NETWORK_TIMEOUT: IdleCheckerSpec(
-        type=CheckerType.NETWORK_TIMEOUT, network=NetworkTimeoutSpec()
+        type=CheckerType.NETWORK_TIMEOUT,
+        network=NetworkTimeoutSpec(idle_timeout_seconds=3600),
     ),
 }