Skip to content

Commit 18ee026

Browse files
committed
refactor(BA-5650-G): resolve owner_id via current_user() in services
- Drop ``owner_id`` from 21 read/control session action dataclasses; keep it only on ``create_from_template`` / ``create_from_params`` / ``create_cluster`` (the three delegation-capable creation paths). - ``services/session/service.py``: add ``_requester_user_id()`` that pulls the caller's UUID from the ``current_user()`` context var, replacing 21 ``owner_id = action.owner_id`` sites. - ``services/session/lifecycle.py``: inject ``UserRepository`` via the constructor instead of instantiating it internally. - ``registry.py``: accept ``UserRepository`` and forward it into ``SessionLifecycleManager``. - ``dependencies/agents/{registry,composer}.py``: wire the repository through the dependency graph. - ``repositories/user/{repository,db_source/db_source}.py``: add ``get_main_access_key_by_id`` (renamed from ``_by_uuid``). - Minor adapter/service touch-ups in ``services/export`` and ``repositories/model_serving``.
1 parent 14c79d2 commit 18ee026

32 files changed

Lines changed: 116 additions & 162 deletions

changes/BA-5650-G.misc.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Drop `owner_id` from read/control session action dataclasses; keep it only on creation actions (`create_from_template`, `create_from_params`, `create_cluster`). Service layer resolves the requester's UUID via `current_user()` context var. `SessionLifecycleManager` receives `UserRepository` via constructor injection instead of constructing it internally.

src/ai/backend/manager/dependencies/agents/composer.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from ai.backend.manager.registry import AgentRegistry
2121
from ai.backend.manager.repositories.deployment.repository import DeploymentRepository
2222
from ai.backend.manager.repositories.scheduler.repository import SchedulerRepository
23+
from ai.backend.manager.repositories.user.repository import UserRepository
2324
from ai.backend.manager.sokovan.deployment.deployment_controller import DeploymentController
2425
from ai.backend.manager.sokovan.deployment.revision_generator.registry import (
2526
RevisionGeneratorRegistry,
@@ -188,6 +189,7 @@ async def compose(
188189
hook_plugin_ctx=setup_input.hook_plugin_ctx,
189190
network_plugin_ctx=setup_input.network_plugin_ctx,
190191
scheduling_controller=scheduling_controller,
192+
user_repository=UserRepository(setup_input.db),
191193
debug=setup_input.config_provider.config.debug.enabled,
192194
manager_public_key=setup_input.agent_cache.manager_public_key,
193195
manager_secret_key=setup_input.agent_cache.manager_secret_key,

src/ai/backend/manager/dependencies/agents/registry.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
2020
from ai.backend.manager.plugin.network import NetworkPluginContext
2121
from ai.backend.manager.registry import AgentRegistry
22+
from ai.backend.manager.repositories.user.repository import UserRepository
2223
from ai.backend.manager.sokovan.scheduling_controller.scheduling_controller import (
2324
SchedulingController,
2425
)
@@ -41,6 +42,7 @@ class AgentRegistryInput:
4142
hook_plugin_ctx: HookPluginContext
4243
network_plugin_ctx: NetworkPluginContext
4344
scheduling_controller: SchedulingController
45+
user_repository: UserRepository
4446
debug: bool
4547
manager_public_key: PublicKey
4648
manager_secret_key: SecretKey
@@ -82,6 +84,7 @@ async def provide(self, setup_input: AgentRegistryInput) -> AsyncIterator[AgentR
8284
setup_input.hook_plugin_ctx,
8385
setup_input.network_plugin_ctx,
8486
setup_input.scheduling_controller,
87+
setup_input.user_repository,
8588
debug=setup_input.debug,
8689
manager_public_key=setup_input.manager_public_key,
8790
manager_secret_key=setup_input.manager_secret_key,

src/ai/backend/manager/registry.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@
116116
from ai.backend.manager.plugin.network import NetworkPluginContext
117117
from ai.backend.manager.repositories.resource_slot import ResourceSlotRepository
118118
from ai.backend.manager.repositories.scheduler.types.session_creation import SessionCreationSpec
119+
from ai.backend.manager.repositories.user.repository import UserRepository
119120
from ai.backend.manager.sokovan.scheduling_controller import SchedulingController
120121

121122
from .agent_cache import AgentRPCCache
@@ -221,6 +222,7 @@ def __init__(
221222
hook_plugin_ctx: HookPluginContext,
222223
network_plugin_ctx: NetworkPluginContext,
223224
scheduling_controller: SchedulingController,
225+
user_repository: UserRepository,
224226
*,
225227
debug: bool = False,
226228
manager_public_key: PublicKey,
@@ -252,6 +254,7 @@ def __init__(
252254
event_producer,
253255
hook_plugin_ctx,
254256
self,
257+
user_repository,
255258
)
256259
self._client_pool = ClientPool(tcp_client_session_factory)
257260

@@ -460,7 +463,7 @@ async def create_session(
460463
sess = await SessionRow.get_session(
461464
db_session,
462465
session_name,
463-
owner_access_key,
466+
owner_id=user_scope.user_uuid,
464467
kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
465468
)
466469
if sess.main_kernel.image is None:
@@ -687,7 +690,7 @@ async def create_cluster(
687690
await SessionRow.get_session(
688691
db_sess,
689692
session_name,
690-
owner_access_key,
693+
owner_id=user_scope.user_uuid,
691694
)
692695
except SessionNotFound:
693696
pass

src/ai/backend/manager/repositories/model_serving/repository.py

Lines changed: 2 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
import sqlalchemy as sa
77
from pydantic import HttpUrl
8-
from ruamel.yaml import YAML
98
from sqlalchemy.exc import IntegrityError, NoResultFound, StatementError
109
from sqlalchemy.ext.asyncio import AsyncSession as SASession
1110
from sqlalchemy.orm import selectinload
@@ -40,7 +39,7 @@
4039
UserData,
4140
)
4241
from ai.backend.manager.data.permission.types import RBACElementRef
43-
from ai.backend.manager.data.vfolder.types import VFolderLocation, VFolderOwnershipType
42+
from ai.backend.manager.data.vfolder.types import VFolderOwnershipType
4443
from ai.backend.manager.errors.common import ObjectNotFound
4544
from ai.backend.manager.errors.resource import DatabaseConnectionUnavailable
4645
from ai.backend.manager.errors.service import EndpointNotFound
@@ -81,9 +80,6 @@
8180
execute_rbac_entity_creator,
8281
)
8382
from ai.backend.manager.repositories.deployment.creators import DeploymentPolicyCreatorSpec
84-
from ai.backend.manager.repositories.deployment.storage_source.storage_source import (
85-
DeploymentStorageSource,
86-
)
8783
from ai.backend.manager.repositories.model_serving.updaters import EndpointUpdaterSpec
8884
from ai.backend.manager.services.model_serving.actions.modify_endpoint import ModifyEndpointAction
8985
from ai.backend.manager.services.model_serving.exceptions import (
@@ -738,7 +734,7 @@ async def get_session_by_id(
738734
async with self._db.begin_readonly_session_read_committed() as session:
739735
try:
740736
return await SessionRow.get_session(
741-
session, session_id, None, kernel_loading_strategy=kernel_loading_strategy
737+
session, session_id, kernel_loading_strategy=kernel_loading_strategy
742738
)
743739
except NoResultFound:
744740
return None
@@ -832,15 +828,6 @@ async def _do_mutate() -> MutationResult:
832828
if current_rev is None:
833829
raise InvalidAPIParameters("Endpoint has no current revision")
834830

835-
# Re-read model definition from vfolder to pick up file changes
836-
refreshed_model_definition = await self._fetch_model_definition_from_vfolder(
837-
db_session,
838-
storage_manager,
839-
current_rev.model,
840-
spec.model_definition_path.optional_value()
841-
or current_rev.model_definition_path,
842-
)
843-
844831
# Resolve image if changed
845832
image_id = current_rev.image
846833
image_ref = spec.image.optional_value()
@@ -873,7 +860,6 @@ async def _do_mutate() -> MutationResult:
873860
if spec.model_definition_path.optional_value() is not None
874861
else current_rev.model_definition_path
875862
),
876-
model_definition=refreshed_model_definition or current_rev.model_definition,
877863
resource_group=endpoint_row.resource_group,
878864
resource_opts=(
879865
spec.resource_opts.optional_value()
@@ -955,51 +941,6 @@ async def _do_mutate() -> MutationResult:
955941
except Exception:
956942
raise
957943

958-
async def _fetch_model_definition_from_vfolder(
959-
self,
960-
db_session: SASession,
961-
storage_manager: StorageSessionManager,
962-
vfolder_id: uuid.UUID | None,
963-
model_definition_path: str | None,
964-
) -> dict[str, Any] | None:
965-
"""Re-read model definition file from the vfolder storage.
966-
967-
Returns the parsed YAML content, or None if the file cannot be read.
968-
"""
969-
if vfolder_id is None:
970-
return None
971-
try:
972-
vf_query = sa.select(
973-
VFolderRow.id,
974-
VFolderRow.host,
975-
VFolderRow.quota_scope_id,
976-
VFolderRow.ownership_type,
977-
VFolderRow.usage_mode,
978-
).where(VFolderRow.id == vfolder_id)
979-
vf_result = await db_session.execute(vf_query)
980-
vf_row = vf_result.one_or_none()
981-
if vf_row is None:
982-
return None
983-
984-
vfolder_location = VFolderLocation(
985-
id=vf_row.id,
986-
quota_scope_id=vf_row.quota_scope_id,
987-
host=vf_row.host,
988-
ownership_type=vf_row.ownership_type,
989-
usage_mode=vf_row.usage_mode,
990-
)
991-
candidates = (
992-
[model_definition_path]
993-
if model_definition_path
994-
else ["model-definition.yaml", "model-definition.yml"]
995-
)
996-
storage_source = DeploymentStorageSource(storage_manager)
997-
content = await storage_source.fetch_definition_file(vfolder_location, candidates)
998-
yaml = YAML()
999-
return cast(dict[str, Any], yaml.load(content))
1000-
except Exception:
1001-
return None
1002-
1003944
@model_serving_repository_resilience.apply()
1004945
async def search_auto_scaling_rules(
1005946
self,

src/ai/backend/manager/services/session/actions/commit_session.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
from dataclasses import dataclass
33
from typing import Any, override
44

5-
from ai.backend.common.types import AccessKey
65
from ai.backend.manager.actions.action import BaseActionResult
76
from ai.backend.manager.actions.types import ActionOperationType
87
from ai.backend.manager.data.session.types import SessionData
@@ -12,7 +11,6 @@
1211
@dataclass
1312
class CommitSessionAction(SessionCommitAction):
1413
session_name: str
15-
owner_access_key: AccessKey
1614
filename: str | None
1715

1816
@override

src/ai/backend/manager/services/session/actions/complete.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
from typing import Any, override
44

55
from ai.backend.common.dto.agent.response import CodeCompletionResp
6-
from ai.backend.common.types import AccessKey
76
from ai.backend.manager.actions.action import BaseActionResult
87
from ai.backend.manager.actions.types import ActionOperationType
98
from ai.backend.manager.data.session.types import SessionData
@@ -14,7 +13,6 @@
1413
@dataclass
1514
class CompleteAction(SessionAction):
1615
session_name: str
17-
owner_access_key: AccessKey
1816
code: str
1917
# TODO: Add type
2018
options: Mapping[str, Any] | None

src/ai/backend/manager/services/session/actions/convert_session_to_image.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
from typing import override
44

55
from ai.backend.common.data.session.types import CustomizedImageVisibilityScope
6-
from ai.backend.common.types import AccessKey
76
from ai.backend.manager.actions.action import BaseActionResult
87
from ai.backend.manager.actions.types import ActionOperationType
98
from ai.backend.manager.data.session.types import SessionData
@@ -13,7 +12,6 @@
1312
@dataclass
1413
class ConvertSessionToImageAction(SessionAction):
1514
session_name: str
16-
owner_access_key: AccessKey
1715
image_name: str
1816
image_visibility: CustomizedImageVisibilityScope
1917
image_owner_id: uuid.UUID

src/ai/backend/manager/services/session/actions/create_cluster.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class CreateClusterAction(SessionScopeAction):
3030
domain_name: str
3131
scaling_group_name: str
3232
requester_access_key: AccessKey
33-
owner_access_key: AccessKey
33+
owner_id: uuid.UUID
3434
tag: str
3535
enqueue_only: bool
3636
keypair_resource_policy: dict[str, Any] | None

src/ai/backend/manager/services/session/actions/create_from_params.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class CreateFromParamsActionParams:
3030
tag: str
3131
priority: int
3232
is_preemptible: bool
33-
owner_access_key: AccessKey
33+
owner_id: uuid.UUID
3434
enqueue_only: bool
3535
max_wait_seconds: int
3636
starts_at: str | None

0 commit comments

Comments
 (0)