diff --git a/changes/12868.feature.md b/changes/12868.feature.md new file mode 100644 index 00000000000..c151c6acc8f --- /dev/null +++ b/changes/12868.feature.md @@ -0,0 +1 @@ +Add the kernel scheduling-history repository search scope, which scopes a query by session, by kernel, or by both, and rejects an empty scope rather than degenerating into a system-wide search. diff --git a/src/ai/backend/manager/repositories/scheduling_history/__init__.py b/src/ai/backend/manager/repositories/scheduling_history/__init__.py index fe14354a5c3..a6063fea478 100644 --- a/src/ai/backend/manager/repositories/scheduling_history/__init__.py +++ b/src/ai/backend/manager/repositories/scheduling_history/__init__.py @@ -8,6 +8,7 @@ from .repository import SchedulingHistoryRepository from .types import ( DeploymentHistorySearchScope, + KernelSchedulingHistorySearchScope, RouteHistorySearchScope, SessionSchedulingHistorySearchScope, ) @@ -16,6 +17,7 @@ "DeploymentHistoryCreatorSpec", "DeploymentHistorySearchScope", "KernelSchedulingHistoryCreatorSpec", + "KernelSchedulingHistorySearchScope", "RouteHistoryCreatorSpec", "RouteHistorySearchScope", "SchedulingHistoryRepositories", diff --git a/src/ai/backend/manager/repositories/scheduling_history/db_source/db_source.py b/src/ai/backend/manager/repositories/scheduling_history/db_source/db_source.py index 8e2ed5aa254..a4224356250 100644 --- a/src/ai/backend/manager/repositories/scheduling_history/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/scheduling_history/db_source/db_source.py @@ -28,6 +28,7 @@ ) from ai.backend.manager.repositories.scheduling_history.types import ( DeploymentHistorySearchScope, + KernelSchedulingHistorySearchScope, RouteHistorySearchScope, SessionSchedulingHistorySearchScope, ) @@ -94,7 +95,7 @@ async def search_session_scoped_history( has_previous_page=result.has_previous_page, ) - # ========== Kernel History ========== + # ========== Kernel History (Admin) ========== async def search_kernel_history( self, @@ -119,6 +120,28 @@ async def search_kernel_history( has_previous_page=result.has_previous_page, ) + # ========== Kernel History (Scoped) ========== + + async def search_kernel_scoped_history( + self, + querier: BatchQuerier, + scope: KernelSchedulingHistorySearchScope, + ) -> KernelSchedulingHistoryListResult: + """Search kernel scheduling history within scope.""" + async with self._db.begin_readonly_session() as db_sess: + query = sa.select(KernelSchedulingHistoryRow) + + result = await execute_batch_querier(db_sess, query, querier, scopes=[scope]) + + items = [row.KernelSchedulingHistoryRow.to_data() for row in result.rows] + + return KernelSchedulingHistoryListResult( + items=items, + total_count=result.total_count, + has_next_page=result.has_next_page, + has_previous_page=result.has_previous_page, + ) + # ========== Deployment History (Admin) ========== async def search_deployment_history( diff --git a/src/ai/backend/manager/repositories/scheduling_history/repository.py b/src/ai/backend/manager/repositories/scheduling_history/repository.py index ed6308adea4..a97440c1482 100644 --- a/src/ai/backend/manager/repositories/scheduling_history/repository.py +++ b/src/ai/backend/manager/repositories/scheduling_history/repository.py @@ -26,6 +26,7 @@ from .db_source import SchedulingHistoryDBSource from .types import ( DeploymentHistorySearchScope, + KernelSchedulingHistorySearchScope, RouteHistorySearchScope, SessionSchedulingHistorySearchScope, ) @@ -81,7 +82,7 @@ async def search_session_scoped_history( """Search session scheduling history within scope.""" return await self._db_source.search_session_scoped_history(querier, scope) - # ========== Kernel History ========== + # ========== Kernel History (Admin) ========== @scheduling_history_repository_resilience.apply() async def search_kernel_history( @@ -91,6 +92,17 @@ async def search_kernel_history( """Search kernel scheduling history with pagination.""" return await self._db_source.search_kernel_history(querier) + # ========== Kernel History (Scoped) ========== + + @scheduling_history_repository_resilience.apply() + async def search_kernel_scoped_history( + self, + querier: BatchQuerier, + scope: KernelSchedulingHistorySearchScope, + ) -> KernelSchedulingHistoryListResult: + """Search kernel scheduling history within scope.""" + return await self._db_source.search_kernel_scoped_history(querier, scope) + # ========== Deployment History (Admin) ========== @scheduling_history_repository_resilience.apply() diff --git a/src/ai/backend/manager/repositories/scheduling_history/types.py b/src/ai/backend/manager/repositories/scheduling_history/types.py index 150834e4043..3ab0d53f058 100644 --- a/src/ai/backend/manager/repositories/scheduling_history/types.py +++ b/src/ai/backend/manager/repositories/scheduling_history/types.py @@ -6,16 +6,24 @@ from typing import Any, override from uuid import UUID +import sqlalchemy as sa + from ai.backend.common.data.filter_specs import UUIDEqualMatchSpec from ai.backend.common.identifier.replica import ReplicaID +from ai.backend.common.types import KernelId, SessionId from ai.backend.manager.errors.deployment import EndpointNotFound -from ai.backend.manager.errors.kernel import SessionNotFound +from ai.backend.manager.errors.kernel import ( + KernelNotFound, + SessionNotFound, +) from ai.backend.manager.errors.service import RouteNotFound from ai.backend.manager.models.clauses import QueryCondition from ai.backend.manager.models.endpoint import EndpointRow +from ai.backend.manager.models.kernel.row import KernelRow from ai.backend.manager.models.routing import RoutingRow from ai.backend.manager.models.scheduling_history.conditions import ( DeploymentHistoryConditions, + KernelSchedulingHistoryConditions, RouteHistoryConditions, SessionSchedulingHistoryConditions, ) @@ -24,6 +32,7 @@ __all__ = ( "SessionSchedulingHistorySearchScope", + "KernelSchedulingHistorySearchScope", "DeploymentHistorySearchScope", "RouteHistorySearchScope", ) @@ -62,6 +71,69 @@ def existence_checks(self) -> list[ExistenceCheck[Any]]: ] +# Kernel Scheduling History Scope + + +@dataclass(frozen=True) +class KernelSchedulingHistorySearchScope(SearchScope): + """Scope for kernel scheduling history search. + + Either axis may be given; when both are, they intersect. The request DTO + rejects an empty scope before one is built. + """ + + session_id: SessionId | None = None + """Restrict to the kernels of this session.""" + + kernel_id: KernelId | None = None + """Restrict to this kernel.""" + + @override + def to_condition(self) -> QueryCondition: + """Convert scope to a query condition for KernelSchedulingHistoryRow.""" + conditions: list[QueryCondition] = [] + if self.session_id is not None: + conditions.append( + KernelSchedulingHistoryConditions.by_session_id_filter( + UUIDEqualMatchSpec(value=self.session_id, negated=False) + ) + ) + if self.kernel_id is not None: + conditions.append( + KernelSchedulingHistoryConditions.by_kernel_id_filter( + UUIDEqualMatchSpec(value=self.kernel_id, negated=False) + ) + ) + + def inner() -> sa.sql.expression.ColumnElement[bool]: + return sa.and_(*(cond() for cond in conditions)) + + return inner + + @property + @override + def existence_checks(self) -> list[ExistenceCheck[Any]]: + """Check that each scoped entity exists.""" + checks: list[ExistenceCheck[Any]] = [] + if self.session_id is not None: + checks.append( + ExistenceCheck( + column=SessionRow.id, + value=self.session_id, + error=SessionNotFound(str(self.session_id)), + ) + ) + if self.kernel_id is not None: + checks.append( + ExistenceCheck( + column=KernelRow.id, + value=self.kernel_id, + error=KernelNotFound(str(self.kernel_id)), + ) + ) + return checks + + # Deployment History Scope