Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .claude/skills/bai-cli/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions changes/12866.feature.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add kernel scheduling-history queries to REST v2, the SDK, and the CLI (`bai scheduling-history kernel search / search-scoped`).
Original file line number Diff line number Diff line change
@@ -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``.
"""

Expand All @@ -9,6 +9,7 @@
import click

from .deployment import deployment
from .kernel import kernel
from .route import route
from .session import session

Expand All @@ -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)
250 changes: 250 additions & 0 deletions src/ai/backend/client/cli/v2/scheduling_history/kernel.py
Original file line number Diff line number Diff line change
@@ -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())
27 changes: 27 additions & 0 deletions src/ai/backend/client/v2/domains_v2/scheduling_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -46,6 +49,30 @@ async def session_scoped_search(
response_model=AdminSearchSessionHistoriesPayload,
)

# ========== Kernel History ==========

async def search_kernel_history(
self, request: AdminSearchKernelHistoriesInput
) -> SearchKernelHistoriesPayload:
Comment thread
jopemachine marked this conversation as resolved.
"""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:
Comment thread
jopemachine marked this conversation as resolved.
"""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(
Expand Down
Loading