Skip to content

Commit f496a0a

Browse files
jopemachineclaude
andcommitted
feat(BA-6882): expose kernel scheduling history via repository/service/adapter/REST v2/SDK/CLI
WIP — local checkpoint, not yet reviewed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 5c2b415 commit f496a0a

26 files changed

Lines changed: 1525 additions & 17 deletions

File tree

.claude/skills/bai-cli/SKILL.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ Check options with `--help`.
5555
- **service-catalog**: user(empty group) · admin(search)
5656
- **runtime-variant**: user(get, search) · admin(get, search, create, update, delete, bulk-delete)
5757
- **runtime-variant-preset**: user(get, search) · admin(get, search, create, update, delete)
58-
- **scheduling-history**: sub session / deployment / route — each (search, search-scoped)
58+
- **scheduling-history**: sub session / kernel / deployment / route — each (search, search-scoped). `kernel search-scoped` takes `--session-id` / `--kernel-id` options (at least one) instead of a positional scope ID.
5959
- **scheduling-handler**: admin(list)
6060

6161
### Storage

src/ai/backend/client/cli/v2/scheduling_history/commands.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
"""CLI commands for scheduling history management.
22
3-
Commands are organized into sub-groups: ``session``, ``deployment``,
3+
Commands are organized into sub-groups: ``session``, ``kernel``, ``deployment``,
44
and ``route``.
55
"""
66

@@ -9,6 +9,7 @@
99
import click
1010

1111
from .deployment import deployment
12+
from .kernel import kernel
1213
from .route import route
1314
from .session import session
1415

@@ -20,5 +21,6 @@ def scheduling_history() -> None:
2021

2122
# Register sub-groups
2223
scheduling_history.add_command(session)
24+
scheduling_history.add_command(kernel)
2325
scheduling_history.add_command(deployment)
2426
scheduling_history.add_command(route)
Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
1+
"""CLI commands for kernel scheduling history."""
2+
3+
from __future__ import annotations
4+
5+
import asyncio
6+
from typing import TYPE_CHECKING
7+
from uuid import UUID
8+
9+
import click
10+
11+
from ai.backend.client.cli.v2.helpers import (
12+
create_v2_registry,
13+
load_v2_config,
14+
parse_order_options,
15+
print_result,
16+
)
17+
18+
if TYPE_CHECKING:
19+
from ai.backend.common.dto.manager.v2.scheduling_history.request import KernelHistoryFilter
20+
21+
# Shared result choices for scheduling history filters
22+
_RESULT_CHOICES = click.Choice(
23+
["SUCCESS", "FAILURE", "STALE", "NEED_RETRY", "EXPIRED", "GIVE_UP", "SKIPPED"],
24+
case_sensitive=False,
25+
)
26+
27+
28+
def _build_kernel_history_filter(
29+
kernel_id: str | None,
30+
session_id: str | None,
31+
phase: str | None,
32+
from_status: tuple[str, ...],
33+
to_status: tuple[str, ...],
34+
result: str | None,
35+
error_code: str | None,
36+
message: str | None,
37+
) -> KernelHistoryFilter | None:
38+
"""Build a KernelHistoryFilter from explicit CLI options.
39+
40+
Returns None if no filter options were provided.
41+
"""
42+
from ai.backend.common.dto.manager.query import StringFilter, UUIDFilter
43+
from ai.backend.common.dto.manager.v2.scheduling_history.request import (
44+
KernelHistoryFilter,
45+
SchedulingResultFilter,
46+
)
47+
from ai.backend.common.dto.manager.v2.scheduling_history.types import SchedulingResultType
48+
49+
has_any = any(
50+
opt is not None for opt in (kernel_id, session_id, phase, result, error_code, message)
51+
)
52+
if not has_any and not from_status and not to_status:
53+
return None
54+
55+
return KernelHistoryFilter(
56+
kernel_id=UUIDFilter(equals=UUID(kernel_id)) if kernel_id is not None else None,
57+
session_id=UUIDFilter(equals=UUID(session_id)) if session_id is not None else None,
58+
phase=StringFilter(contains=phase) if phase is not None else None,
59+
from_status=list(from_status) if from_status else None,
60+
to_status=list(to_status) if to_status else None,
61+
result=(
62+
SchedulingResultFilter(equals=SchedulingResultType(result))
63+
if result is not None
64+
else None
65+
),
66+
error_code=StringFilter(contains=error_code) if error_code is not None else None,
67+
message=StringFilter(contains=message) if message is not None else None,
68+
)
69+
70+
71+
@click.group()
72+
def kernel() -> None:
73+
"""Kernel scheduling history commands."""
74+
75+
76+
@kernel.command()
77+
@click.option("--limit", type=int, default=None, help="Maximum items to return.")
78+
@click.option("--offset", type=int, default=None, help="Number of items to skip.")
79+
@click.option("--kernel-id", type=str, default=None, help="Filter by kernel ID (UUID).")
80+
@click.option("--session-id", type=str, default=None, help="Filter by session ID (UUID).")
81+
@click.option("--phase", type=str, default=None, help="Filter by scheduling phase (contains).")
82+
@click.option(
83+
"--from-status",
84+
type=str,
85+
multiple=True,
86+
help="Filter by from_status values (repeatable).",
87+
)
88+
@click.option(
89+
"--to-status",
90+
type=str,
91+
multiple=True,
92+
help="Filter by to_status values (repeatable).",
93+
)
94+
@click.option("--result", type=_RESULT_CHOICES, default=None, help="Filter by scheduling result.")
95+
@click.option("--error-code", type=str, default=None, help="Filter by error code (contains).")
96+
@click.option("--message", type=str, default=None, help="Filter by message (contains).")
97+
@click.option(
98+
"--order-by",
99+
multiple=True,
100+
help="Order by field:direction (e.g., created_at:desc). Fields: created_at, updated_at.",
101+
)
102+
def search(
103+
limit: int | None,
104+
offset: int | None,
105+
kernel_id: str | None,
106+
session_id: str | None,
107+
phase: str | None,
108+
from_status: tuple[str, ...],
109+
to_status: tuple[str, ...],
110+
result: str | None,
111+
error_code: str | None,
112+
message: str | None,
113+
order_by: tuple[str, ...],
114+
) -> None:
115+
"""Search kernel scheduling histories (superadmin only)."""
116+
from ai.backend.common.dto.manager.v2.scheduling_history.request import (
117+
AdminSearchKernelHistoriesInput,
118+
KernelHistoryOrder,
119+
)
120+
from ai.backend.common.dto.manager.v2.scheduling_history.types import KernelHistoryOrderField
121+
122+
history_filter = _build_kernel_history_filter(
123+
kernel_id, session_id, phase, from_status, to_status, result, error_code, message
124+
)
125+
126+
orders = (
127+
parse_order_options(order_by, KernelHistoryOrderField, KernelHistoryOrder)
128+
if order_by
129+
else None
130+
)
131+
132+
async def _run() -> None:
133+
registry = await create_v2_registry(load_v2_config())
134+
try:
135+
result_data = await registry.scheduling_history.search_kernel_history(
136+
AdminSearchKernelHistoriesInput(
137+
filter=history_filter,
138+
order=orders,
139+
limit=limit,
140+
offset=offset,
141+
),
142+
)
143+
print_result(result_data)
144+
finally:
145+
await registry.close()
146+
147+
asyncio.run(_run())
148+
149+
150+
@kernel.command(name="search-scoped")
151+
@click.option("--session-id", type=str, default=None, help="Scope to the kernels of this session.")
152+
@click.option("--kernel-id", type=str, default=None, help="Scope to this kernel.")
153+
@click.option("--limit", type=int, default=None, help="Maximum items to return.")
154+
@click.option("--offset", type=int, default=None, help="Number of items to skip.")
155+
@click.option("--phase", type=str, default=None, help="Filter by scheduling phase (contains).")
156+
@click.option(
157+
"--from-status",
158+
type=str,
159+
multiple=True,
160+
help="Filter by from_status values (repeatable).",
161+
)
162+
@click.option(
163+
"--to-status",
164+
type=str,
165+
multiple=True,
166+
help="Filter by to_status values (repeatable).",
167+
)
168+
@click.option("--result", type=_RESULT_CHOICES, default=None, help="Filter by scheduling result.")
169+
@click.option("--error-code", type=str, default=None, help="Filter by error code (contains).")
170+
@click.option("--message", type=str, default=None, help="Filter by message (contains).")
171+
@click.option(
172+
"--order-by",
173+
multiple=True,
174+
help="Order by field:direction (e.g., created_at:desc). Fields: created_at, updated_at.",
175+
)
176+
def search_scoped(
177+
session_id: str | None,
178+
kernel_id: str | None,
179+
limit: int | None,
180+
offset: int | None,
181+
phase: str | None,
182+
from_status: tuple[str, ...],
183+
to_status: tuple[str, ...],
184+
result: str | None,
185+
error_code: str | None,
186+
message: str | None,
187+
order_by: tuple[str, ...],
188+
) -> None:
189+
"""Search kernel scheduling history within a session and/or kernel scope.
190+
191+
The scope takes two axes rather than one positional ID, so it is given as options.
192+
At least one of --session-id / --kernel-id is required.
193+
"""
194+
from ai.backend.common.dto.manager.v2.scheduling_history.request import (
195+
KernelHistoryOrder,
196+
ScopedSearchKernelHistoriesInput,
197+
)
198+
from ai.backend.common.dto.manager.v2.scheduling_history.types import (
199+
KernelHistoryOrderField,
200+
KernelHistoryScopeDTO,
201+
)
202+
203+
if session_id is None and kernel_id is None:
204+
raise click.UsageError("At least one of --session-id or --kernel-id is required.")
205+
206+
scope = KernelHistoryScopeDTO(
207+
session_id=UUID(session_id) if session_id is not None else None,
208+
kernel_id=UUID(kernel_id) if kernel_id is not None else None,
209+
)
210+
211+
# The scope already narrows by id, so ids are not repeated in the filter.
212+
history_filter = _build_kernel_history_filter(
213+
None, None, phase, from_status, to_status, result, error_code, message
214+
)
215+
216+
orders = (
217+
parse_order_options(order_by, KernelHistoryOrderField, KernelHistoryOrder)
218+
if order_by
219+
else None
220+
)
221+
222+
async def _run() -> None:
223+
registry = await create_v2_registry(load_v2_config())
224+
try:
225+
result_data = await registry.scheduling_history.kernel_scoped_search(
226+
ScopedSearchKernelHistoriesInput(
227+
scope=scope,
228+
filter=history_filter,
229+
order=orders,
230+
limit=limit,
231+
offset=offset,
232+
),
233+
)
234+
print_result(result_data)
235+
finally:
236+
await registry.close()
237+
238+
asyncio.run(_run())

src/ai/backend/client/v2/domains_v2/scheduling_history.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,14 @@
77
from ai.backend.client.v2.base_domain import BaseDomainClient
88
from ai.backend.common.dto.manager.v2.scheduling_history.request import (
99
AdminSearchDeploymentHistoriesInput,
10+
AdminSearchKernelHistoriesInput,
1011
AdminSearchRouteHistoriesInput,
1112
AdminSearchSessionHistoriesInput,
13+
ScopedSearchKernelHistoriesInput,
1214
)
1315
from ai.backend.common.dto.manager.v2.scheduling_history.response import (
1416
AdminSearchDeploymentHistoriesPayload,
17+
AdminSearchKernelHistoriesPayload,
1518
AdminSearchRouteHistoriesPayload,
1619
AdminSearchSessionHistoriesPayload,
1720
)
@@ -46,6 +49,30 @@ async def session_scoped_search(
4649
response_model=AdminSearchSessionHistoriesPayload,
4750
)
4851

52+
# ========== Kernel History ==========
53+
54+
async def search_kernel_history(
55+
self, request: AdminSearchKernelHistoriesInput
56+
) -> AdminSearchKernelHistoriesPayload:
57+
"""Search kernel scheduling histories with admin scope."""
58+
return await self._client.typed_request(
59+
"POST",
60+
f"{_PATH}/kernels/search",
61+
request=request,
62+
response_model=AdminSearchKernelHistoriesPayload,
63+
)
64+
65+
async def kernel_scoped_search(
66+
self, request: ScopedSearchKernelHistoriesInput
67+
) -> AdminSearchKernelHistoriesPayload:
68+
"""Search kernel scheduling histories within a session and/or kernel scope."""
69+
return await self._client.typed_request(
70+
"POST",
71+
f"{_PATH}/kernels/scoped/search",
72+
request=request,
73+
response_model=AdminSearchKernelHistoriesPayload,
74+
)
75+
4976
# ========== Deployment History ==========
5077

5178
async def search_deployment_history(

src/ai/backend/common/data/permission/types.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,9 @@ class EntityType(enum.StrEnum):
173173
SESSION_DIRECT_ACCESS = "session:direct_access"
174174
SESSION_HISTORY = "session:history"
175175
SESSION_SCOPED_HISTORY = "session:scoped_history"
176+
# Kernel sub
177+
KERNEL_HISTORY = "kernel:history"
178+
KERNEL_SCOPED_HISTORY = "kernel:scoped_history"
176179
# Deployment sub
177180
DEPLOYMENT_REPLICA = "deployment:replica"
178181
DEPLOYMENT_ROUTE = "deployment:route"

0 commit comments

Comments
 (0)