diff --git a/.claude/skills/bai-cli/SKILL.md b/.claude/skills/bai-cli/SKILL.md index a9959b96a9a..810462e2c55 100644 --- a/.claude/skills/bai-cli/SKILL.md +++ b/.claude/skills/bai-cli/SKILL.md @@ -55,7 +55,7 @@ Check options with `--help`. - **service-catalog**: user(empty group) · admin(search) - **runtime-variant**: user(get, search) · admin(get, search, create, update, delete, bulk-delete) - **runtime-variant-preset**: user(get, search) · admin(get, search, create, update, delete) -- **scheduling-history**: sub session / deployment / route — each (search, search-scoped) +- **scheduling-history**: sub session / kernel / deployment / route — each (search, search-scoped) - **scheduling-handler**: admin(list) ### Storage diff --git a/changes/12866.feature.md b/changes/12866.feature.md new file mode 100644 index 00000000000..2c5b726ebdb --- /dev/null +++ b/changes/12866.feature.md @@ -0,0 +1 @@ +Add kernel scheduling-history queries to REST v2, the SDK, and the CLI (`bai scheduling-history kernel search / search-scoped`). diff --git a/src/ai/backend/client/cli/v2/scheduling_history/commands.py b/src/ai/backend/client/cli/v2/scheduling_history/commands.py index 04ef3968cfd..2fe971fbd03 100644 --- a/src/ai/backend/client/cli/v2/scheduling_history/commands.py +++ b/src/ai/backend/client/cli/v2/scheduling_history/commands.py @@ -1,6 +1,6 @@ """CLI commands for scheduling history management. -Commands are organized into sub-groups: ``session``, ``deployment``, +Commands are organized into sub-groups: ``session``, ``kernel``, ``deployment``, and ``route``. """ @@ -9,6 +9,7 @@ import click from .deployment import deployment +from .kernel import kernel from .route import route from .session import session @@ -20,5 +21,6 @@ def scheduling_history() -> None: # Register sub-groups scheduling_history.add_command(session) +scheduling_history.add_command(kernel) scheduling_history.add_command(deployment) scheduling_history.add_command(route) diff --git a/src/ai/backend/client/cli/v2/scheduling_history/kernel.py b/src/ai/backend/client/cli/v2/scheduling_history/kernel.py new file mode 100644 index 00000000000..5f5d3abf4f6 --- /dev/null +++ b/src/ai/backend/client/cli/v2/scheduling_history/kernel.py @@ -0,0 +1,250 @@ +"""CLI commands for kernel scheduling history.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING +from uuid import UUID + +import click + +from ai.backend.client.cli.v2.helpers import ( + create_v2_registry, + load_v2_config, + parse_order_options, + print_result, +) + +if TYPE_CHECKING: + from ai.backend.common.dto.manager.v2.scheduling_history.request import KernelHistoryFilter + +# Shared result choices for scheduling history filters +_RESULT_CHOICES = click.Choice( + ["SUCCESS", "FAILURE", "STALE", "NEED_RETRY", "EXPIRED", "GIVE_UP", "SKIPPED"], + case_sensitive=False, +) + + +def _build_kernel_history_filter( + *, + phase: str | None, + from_status: tuple[str, ...], + to_status: tuple[str, ...], + result: str | None, + error_code: str | None, + message: str | None, + kernel_id: str | None = None, + session_id: str | None = None, +) -> KernelHistoryFilter | None: + """Build a KernelHistoryFilter from explicit CLI options. + + ``kernel_id`` and ``session_id`` are omitted by callers whose scope already + narrows by that id. Returns None if no filter options were provided. + """ + from ai.backend.common.dto.manager.query import StringFilter, UUIDFilter + from ai.backend.common.dto.manager.v2.scheduling_history.request import ( + KernelHistoryFilter, + SchedulingResultFilter, + ) + from ai.backend.common.dto.manager.v2.scheduling_history.types import SchedulingResultType + + has_any = any( + opt is not None for opt in (kernel_id, session_id, phase, result, error_code, message) + ) + if not has_any and not from_status and not to_status: + return None + + return KernelHistoryFilter( + kernel_id=UUIDFilter(equals=UUID(kernel_id)) if kernel_id is not None else None, + session_id=UUIDFilter(equals=UUID(session_id)) if session_id is not None else None, + phase=StringFilter(contains=phase) if phase is not None else None, + from_status=list(from_status) if from_status else None, + to_status=list(to_status) if to_status else None, + result=( + SchedulingResultFilter(equals=SchedulingResultType(result)) + if result is not None + else None + ), + error_code=StringFilter(contains=error_code) if error_code is not None else None, + message=StringFilter(contains=message) if message is not None else None, + ) + + +@click.group() +def kernel() -> None: + """Kernel scheduling history commands.""" + + +@kernel.command() +@click.option("--limit", type=int, default=None, help="Maximum items to return.") +@click.option("--offset", type=int, default=None, help="Number of items to skip.") +@click.option("--kernel-id", type=str, default=None, help="Filter by kernel ID (UUID).") +@click.option("--session-id", type=str, default=None, help="Filter by session ID (UUID).") +@click.option("--phase", type=str, default=None, help="Filter by scheduling phase (contains).") +@click.option( + "--from-status", + type=str, + multiple=True, + help="Filter by from_status values (repeatable).", +) +@click.option( + "--to-status", + type=str, + multiple=True, + help="Filter by to_status values (repeatable).", +) +@click.option("--result", type=_RESULT_CHOICES, default=None, help="Filter by scheduling result.") +@click.option("--error-code", type=str, default=None, help="Filter by error code (contains).") +@click.option("--message", type=str, default=None, help="Filter by message (contains).") +@click.option( + "--order-by", + multiple=True, + help=( + "Order by field:direction (e.g., created_at:desc). Fields: created_at, " + "updated_at, phase, from_status, to_status, result, attempts." + ), +) +def search( + limit: int | None, + offset: int | None, + kernel_id: str | None, + session_id: str | None, + phase: str | None, + from_status: tuple[str, ...], + to_status: tuple[str, ...], + result: str | None, + error_code: str | None, + message: str | None, + order_by: tuple[str, ...], +) -> None: + """Search kernel scheduling histories (superadmin only).""" + from ai.backend.common.dto.manager.v2.scheduling_history.request import ( + AdminSearchKernelHistoriesInput, + KernelHistoryOrder, + ) + from ai.backend.common.dto.manager.v2.scheduling_history.types import KernelHistoryOrderField + + history_filter = _build_kernel_history_filter( + kernel_id=kernel_id, + session_id=session_id, + phase=phase, + from_status=from_status, + to_status=to_status, + result=result, + error_code=error_code, + message=message, + ) + + orders = ( + parse_order_options(order_by, KernelHistoryOrderField, KernelHistoryOrder) + if order_by + else None + ) + + async def _run() -> None: + registry = await create_v2_registry(load_v2_config()) + try: + result_data = await registry.scheduling_history.search_kernel_history( + AdminSearchKernelHistoriesInput( + filter=history_filter, + order=orders, + limit=limit, + offset=offset, + ), + ) + print_result(result_data) + finally: + await registry.close() + + asyncio.run(_run()) + + +@kernel.command(name="search-scoped") +@click.argument("kernel_id", type=str) +@click.option("--limit", type=int, default=None, help="Maximum items to return.") +@click.option("--offset", type=int, default=None, help="Number of items to skip.") +@click.option("--session-id", type=str, default=None, help="Filter by session ID (UUID).") +@click.option("--phase", type=str, default=None, help="Filter by scheduling phase (contains).") +@click.option( + "--from-status", + type=str, + multiple=True, + help="Filter by from_status values (repeatable).", +) +@click.option( + "--to-status", + type=str, + multiple=True, + help="Filter by to_status values (repeatable).", +) +@click.option("--result", type=_RESULT_CHOICES, default=None, help="Filter by scheduling result.") +@click.option("--error-code", type=str, default=None, help="Filter by error code (contains).") +@click.option("--message", type=str, default=None, help="Filter by message (contains).") +@click.option( + "--order-by", + multiple=True, + help=( + "Order by field:direction (e.g., created_at:desc). Fields: created_at, " + "updated_at, phase, from_status, to_status, result, attempts." + ), +) +def search_scoped( + kernel_id: str, + session_id: str | None, + limit: int | None, + offset: int | None, + phase: str | None, + from_status: tuple[str, ...], + to_status: tuple[str, ...], + result: str | None, + error_code: str | None, + message: str | None, + order_by: tuple[str, ...], +) -> None: + """Search kernel scheduling history scoped to KERNEL_ID.""" + from ai.backend.common.dto.manager.v2.rbac.types import UUIDScope + from ai.backend.common.dto.manager.v2.scheduling_history.request import ( + KernelHistoryOrder, + ScopedSearchKernelHistoriesInput, + ) + from ai.backend.common.dto.manager.v2.scheduling_history.types import ( + KernelHistoryOrderField, + KernelHistoryScopeDTO, + ) + + scope = KernelHistoryScopeDTO(kernel=[UUIDScope(value=UUID(kernel_id))]) + + # The scope already narrows by kernel ID, so it is not repeated in the filter. + history_filter = _build_kernel_history_filter( + session_id=session_id, + phase=phase, + from_status=from_status, + to_status=to_status, + result=result, + error_code=error_code, + message=message, + ) + + orders = ( + parse_order_options(order_by, KernelHistoryOrderField, KernelHistoryOrder) + if order_by + else None + ) + + async def _run() -> None: + registry = await create_v2_registry(load_v2_config()) + try: + result_data = await registry.scheduling_history.kernel_scoped_search( + ScopedSearchKernelHistoriesInput( + scope=scope, + filter=history_filter, + order=orders, + limit=limit, + offset=offset, + ), + ) + print_result(result_data) + finally: + await registry.close() + + asyncio.run(_run()) diff --git a/src/ai/backend/client/v2/domains_v2/scheduling_history.py b/src/ai/backend/client/v2/domains_v2/scheduling_history.py index 62f8367c242..005507a1974 100644 --- a/src/ai/backend/client/v2/domains_v2/scheduling_history.py +++ b/src/ai/backend/client/v2/domains_v2/scheduling_history.py @@ -7,13 +7,16 @@ from ai.backend.client.v2.base_domain import BaseDomainClient from ai.backend.common.dto.manager.v2.scheduling_history.request import ( AdminSearchDeploymentHistoriesInput, + AdminSearchKernelHistoriesInput, AdminSearchRouteHistoriesInput, AdminSearchSessionHistoriesInput, + ScopedSearchKernelHistoriesInput, ) from ai.backend.common.dto.manager.v2.scheduling_history.response import ( AdminSearchDeploymentHistoriesPayload, AdminSearchRouteHistoriesPayload, AdminSearchSessionHistoriesPayload, + SearchKernelHistoriesPayload, ) _PATH = "/v2/scheduling-history" @@ -46,6 +49,30 @@ async def session_scoped_search( response_model=AdminSearchSessionHistoriesPayload, ) + # ========== Kernel History ========== + + async def search_kernel_history( + self, request: AdminSearchKernelHistoriesInput + ) -> SearchKernelHistoriesPayload: + """Search kernel scheduling histories with admin scope.""" + return await self._client.typed_request( + "POST", + f"{_PATH}/kernels/admin/search", + request=request, + response_model=SearchKernelHistoriesPayload, + ) + + async def kernel_scoped_search( + self, request: ScopedSearchKernelHistoriesInput + ) -> SearchKernelHistoriesPayload: + """Search kernel scheduling histories within a session and/or kernel scope.""" + return await self._client.typed_request( + "POST", + f"{_PATH}/kernels/scoped/search", + request=request, + response_model=SearchKernelHistoriesPayload, + ) + # ========== Deployment History ========== async def search_deployment_history(