feat(BA-5985): scan routes via BatchQuerier with traffic-status filter (#11534)

Author: fregataa

changes/11534.feature.md

@@ -0,0 +1 @@
+Route coordinator now scans lifecycle routes via `BatchQuerier`, and `RouteTargetStatuses` gains an explicit traffic-status filter axis so handlers can target only routes whose `traffic_status` is in a given list.

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

@@ -251,10 +251,15 @@ class DeploymentTargetStatuses:
 
 @dataclass(frozen=True)
 class RouteTargetStatuses:
-    """Target statuses for route handler filtering (lifecycle x health)."""
+    """Target statuses for route handler filtering (lifecycle x health x traffic).
+
+    ``traffic=None`` skips the traffic-status predicate. Pass a non-empty list
+    to restrict to specific ``RouteTrafficStatus`` values.
+    """
 
     lifecycle: list[RouteStatus]
     health: list[RouteHealthStatus]
+    traffic: list[RouteTrafficStatus] | None = None
 
 
 @dataclass(frozen=True)

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

@@ -1584,23 +1584,24 @@ async def scale_routes(
 
     # Route operations
 
-    async def get_routes_by_statuses(
+    async def search_route_datas(
         self,
-        statuses: list[RouteStatus],
-        health_statuses: list[RouteHealthStatus],
+        *,
+        querier: BatchQuerier,
     ) -> list[RouteData]:
-        """Get routes by lifecycle and health statuses."""
-        async with self._begin_readonly_session_read_committed() as db_sess:
-            query = sa.select(RoutingRow).where(
-                RoutingRow.status.in_(statuses),
-                RoutingRow.health_status.in_(health_statuses),
-            )
-            result = await db_sess.execute(query)
-            rows: Sequence[RoutingRow] = result.scalars().all()
+        """Search routes via :class:`BatchQuerier`.
 
-            route_data_list: list[RouteData] = []
-            for row in rows:
-                route_data = RouteData(
+        The caller composes ``querier`` with every filter that applies
+        (lifecycle / health / traffic_status / endpoint id set, etc.).
+        Pagination is part of the querier — pass ``NoPagination`` for
+        unbounded scans.
+        """
+        async with self._begin_readonly_session_read_committed() as db_sess:
+            query = sa.select(RoutingRow)
+            query_result = await execute_batch_querier(db_sess, query, querier)
+            route_rows: list[RoutingRow] = [row.RoutingRow for row in query_result.rows]
+            return [
+                RouteData(
                     route_id=row.id,
                     deployment_id=row.endpoint,
                     session_id=SessionId(row.session) if row.session else None,
@@ -1613,9 +1614,8 @@ async def get_routes_by_statuses(
                     replica_port=row.replica_port,
                     error_data=row.error_data or {},
                 )
-                route_data_list.append(route_data)
-
-            return route_data_list
+                for row in route_rows
+            ]
 
     async def update_route_status_bulk(
         self,

src/ai/backend/manager/repositories/deployment/repository.py

@@ -67,7 +67,6 @@
     ModelDeploymentAutoScalingRuleData,
     ModelRevisionData,
     RevisionSearchResult,
-    RouteHealthStatus,
     RouteInfo,
     RouteSearchResult,
     RouteStatus,
@@ -629,13 +628,18 @@ async def scale_routes(
     # Route operations
 
     @deployment_repository_resilience.apply()
-    async def get_routes_by_statuses(
+    async def search_route_datas(
         self,
-        statuses: list[RouteStatus],
-        health_statuses: list[RouteHealthStatus],
+        *,
+        querier: BatchQuerier,
     ) -> list[RouteData]:
-        """Get routes by lifecycle and health statuses."""
-        return await self._db_source.get_routes_by_statuses(statuses, health_statuses)
+        """Search routes via :class:`BatchQuerier`.
+
+        The caller composes ``querier`` with every filter that applies;
+        pagination is part of the querier (use ``NoPagination`` for
+        unbounded scans).
+        """
+        return await self._db_source.search_route_datas(querier=querier)
 
     @deployment_repository_resilience.apply()
     async def update_route_status_bulk(

src/ai/backend/manager/sokovan/deployment/route/coordinator.py

@@ -24,6 +24,7 @@
 from ai.backend.manager.data.session.types import SchedulingResult
 from ai.backend.manager.models.routing import RoutingRow
 from ai.backend.manager.models.routing.conditions import RouteConditions
+from ai.backend.manager.repositories.base import BatchQuerier, NoPagination, QueryCondition
 from ai.backend.manager.repositories.base.creator import BulkCreator
 from ai.backend.manager.repositories.base.updater import BatchUpdater
 from ai.backend.manager.repositories.deployment import DeploymentRepository
@@ -197,11 +198,16 @@ async def process_route_lifecycle(
                 lock_lifetime = self._config_provider.config.manager.session_schedule_lock_lifetime
                 await stack.enter_async_context(self._lock_factory(handler.lock_id, lock_lifetime))
 
-            # Get routes by target lifecycle + health statuses
+            # Get routes by target lifecycle + health (+ optional traffic) statuses
             target = handler.target_statuses()
-            routes = await self._deployment_repository.get_routes_by_statuses(
-                target.lifecycle,
-                target.health,
+            conditions: list[QueryCondition] = [
+                RouteConditions.by_lifecycle_statuses(target.lifecycle),
+                RouteConditions.by_health_statuses(target.health),
+            ]
+            if target.traffic is not None:
+                conditions.append(RouteConditions.by_traffic_statuses(target.traffic))
+            routes = await self._deployment_repository.search_route_datas(
+                querier=BatchQuerier(pagination=NoPagination(), conditions=conditions),
             )
             if not routes:
                 log.trace("No routes to process for handler: {}", handler.name())
@@ -230,9 +236,14 @@ async def _process_observer(self, observer: RouteObserver) -> None:
         changing route status in DB.
         """
         try:
-            routes = await self._deployment_repository.get_routes_by_statuses(
-                [RouteStatus.RUNNING],
-                list(RouteHealthStatus),
+            routes = await self._deployment_repository.search_route_datas(
+                querier=BatchQuerier(
+                    pagination=NoPagination(),
+                    conditions=[
+                        RouteConditions.by_lifecycle_statuses([RouteStatus.RUNNING]),
+                        RouteConditions.by_health_statuses(list(RouteHealthStatus)),
+                    ],
+                ),
             )
             if not routes:
                 return

tests/unit/manager/sokovan/deployment/route/test_coordinator_history.py

@@ -75,7 +75,7 @@ def sample_route_execution_error(
 def mock_deployment_repository() -> AsyncMock:
     """Mock DeploymentRepository with minimal implementation."""
     mock = AsyncMock(spec=DeploymentRepository)
-    mock.get_routes_by_statuses = AsyncMock(return_value=[])
+    mock.search_route_datas = AsyncMock(return_value=[])
     mock.update_route_status_bulk_with_history = AsyncMock(return_value=0)
     return mock
 
@@ -285,7 +285,7 @@ def coordinator_with_provisioning_routes(
     sample_route_data: RouteData,
 ) -> Generator[RouteCoordinator, None, None]:
     """Coordinator with PROVISIONING routes available."""
-    mock_deployment_repository.get_routes_by_statuses = AsyncMock(return_value=[sample_route_data])
+    mock_deployment_repository.search_route_datas = AsyncMock(return_value=[sample_route_data])
 
     coordinator = RouteCoordinator(
         valkey_schedule=mock_valkey_schedule,
@@ -313,7 +313,7 @@ def coordinator_without_routes(
     mock_service_discovery: MagicMock,
 ) -> Generator[RouteCoordinator, None, None]:
     """Coordinator with no routes available."""
-    mock_deployment_repository.get_routes_by_statuses = AsyncMock(return_value=[])
+    mock_deployment_repository.search_route_datas = AsyncMock(return_value=[])
 
     coordinator = RouteCoordinator(
         valkey_schedule=mock_valkey_schedule,
@@ -382,7 +382,7 @@ async def test_records_history_on_stale(
             RouteLifecycleType.HEALTH_CHECK: mock_handler_with_stale
         }
         # Update repository mock to return routes for HEALTHY status
-        mock_deployment_repository.get_routes_by_statuses = AsyncMock(
+        mock_deployment_repository.search_route_datas = AsyncMock(
             return_value=[
                 RouteData(
                     route_id=uuid4(),