diff --git a/changes/11534.feature.md b/changes/11534.feature.md new file mode 100644 index 00000000000..e1b879496f7 --- /dev/null +++ b/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. diff --git a/src/ai/backend/manager/data/deployment/types.py b/src/ai/backend/manager/data/deployment/types.py index a4cc438218f..bf62f5edd18 100644 --- a/src/ai/backend/manager/data/deployment/types.py +++ b/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) diff --git a/src/ai/backend/manager/repositories/deployment/db_source/db_source.py b/src/ai/backend/manager/repositories/deployment/db_source/db_source.py index af2c9bdb29f..b02d8fc2500 100644 --- a/src/ai/backend/manager/repositories/deployment/db_source/db_source.py +++ b/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, diff --git a/src/ai/backend/manager/repositories/deployment/repository.py b/src/ai/backend/manager/repositories/deployment/repository.py index 81d53f635e9..4c65fcbeab4 100644 --- a/src/ai/backend/manager/repositories/deployment/repository.py +++ b/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( diff --git a/src/ai/backend/manager/sokovan/deployment/route/coordinator.py b/src/ai/backend/manager/sokovan/deployment/route/coordinator.py index 1e8942e7276..bdbe432136b 100644 --- a/src/ai/backend/manager/sokovan/deployment/route/coordinator.py +++ b/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 diff --git a/tests/unit/manager/sokovan/deployment/route/test_coordinator_history.py b/tests/unit/manager/sokovan/deployment/route/test_coordinator_history.py index 27b908f5654..9416425f963 100644 --- a/tests/unit/manager/sokovan/deployment/route/test_coordinator_history.py +++ b/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(),