Skip to content

Commit 574c93c

Browse files
fregataaclaude
andauthored
refactor(BA-5069): Apply RBAC Creator pattern to Kernel (#10013)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 5619d3d commit 574c93c

3 files changed

Lines changed: 36 additions & 5 deletions

File tree

changes/10013.enhance.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Apply RBAC Creator pattern to auto sub-entity (BA-5069)

src/ai/backend/manager/repositories/scheduler/db_source/db_source.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,9 @@
8686
)
8787
from ai.backend.manager.repositories.base.creator import BulkCreator
8888
from ai.backend.manager.repositories.base.rbac.entity_creator import (
89+
RBACBulkEntityCreator,
8990
RBACEntityCreator,
91+
execute_rbac_bulk_entity_creator,
9092
execute_rbac_entity_creator,
9193
)
9294
from ai.backend.manager.repositories.base.updater import BatchUpdater, execute_batch_updater
@@ -130,6 +132,7 @@
130132
SessionSchedulingHistoryCreatorSpec,
131133
)
132134
from ai.backend.manager.repositories.session.creators import (
135+
KernelRowCreatorSpec,
133136
SessionRowCreatorSpec,
134137
)
135138
from ai.backend.manager.sokovan.data import (
@@ -1295,7 +1298,7 @@ async def enqueue_session(
12951298
)
12961299

12971300
# Create kernel rows
1298-
kernels = []
1301+
kernel_specs = []
12991302
for kernel in session_data.kernels:
13001303
kernel_row = KernelRow(
13011304
id=kernel.id,
@@ -1342,7 +1345,7 @@ async def enqueue_session(
13421345
main_gid=kernel.main_gid,
13431346
gids=kernel.gids,
13441347
)
1345-
kernels.append(kernel_row)
1348+
kernel_specs.append(KernelRowCreatorSpec(row=kernel_row))
13461349

13471350
# Use RBACEntityCreator to create session with RBAC scope association
13481351
rbac_creator = RBACEntityCreator(
@@ -1361,11 +1364,19 @@ async def enqueue_session(
13611364
)
13621365
await execute_rbac_entity_creator(db_sess, rbac_creator)
13631366

1364-
db_sess.add_all(kernels)
1365-
await db_sess.flush()
1367+
# Use RBACBulkEntityCreator to create kernels with RBAC scope association
1368+
kernel_rbac_creator = RBACBulkEntityCreator(
1369+
specs=kernel_specs,
1370+
element_type=RBACElementType.KERNEL,
1371+
scope_ref=RBACElementRef(
1372+
element_type=RBACElementType.SESSION,
1373+
element_id=str(session_data.id),
1374+
),
1375+
)
1376+
kernel_result = await execute_rbac_bulk_entity_creator(db_sess, kernel_rbac_creator)
13661377

13671378
# Record requested resources in normalized resource_allocations table
1368-
for kernel_row in kernels:
1379+
for kernel_row in kernel_result.rows:
13691380
quantities = resource_slot_to_quantities(kernel_row.requested_slots)
13701381
if quantities:
13711382
await db_sess.execute(

src/ai/backend/manager/repositories/session/creators.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from dataclasses import dataclass
66
from typing import override
77

8+
from ai.backend.manager.models.kernel import KernelRow
89
from ai.backend.manager.models.session import SessionRow
910
from ai.backend.manager.repositories.base.creator import CreatorSpec
1011

@@ -25,3 +26,21 @@ class SessionRowCreatorSpec(CreatorSpec[SessionRow]):
2526
@override
2627
def build_row(self) -> SessionRow:
2728
return self.row
29+
30+
31+
@dataclass
32+
class KernelRowCreatorSpec(CreatorSpec[KernelRow]):
33+
"""CreatorSpec that wraps a pre-built KernelRow.
34+
35+
This spec is designed for retrofitting existing code that already builds
36+
KernelRow instances. It simply returns the provided row in build_row().
37+
38+
Kernel is a sub-entity of Session. Use RBACElementType.SESSION as
39+
the scope_ref element_type when creating the RBAC association.
40+
"""
41+
42+
row: KernelRow
43+
44+
@override
45+
def build_row(self) -> KernelRow:
46+
return self.row

0 commit comments

Comments
 (0)