diff --git a/changes/11043.enhance.md b/changes/11043.enhance.md new file mode 100644 index 00000000000..3ad892091f3 --- /dev/null +++ b/changes/11043.enhance.md @@ -0,0 +1 @@ +Rename `UserPermission.user_uuid` to `owner_id` and add `main_access_key` field. `KernelRow.to_kernel_info` populates `main_access_key` from the eagerly-loaded `user_row` relationship so downstream consumers no longer need to pass the access_key through separate channels. diff --git a/src/ai/backend/manager/api/adapters/session.py b/src/ai/backend/manager/api/adapters/session.py index 78ae9c9f807..69b59a5a03b 100644 --- a/src/ai/backend/manager/api/adapters/session.py +++ b/src/ai/backend/manager/api/adapters/session.py @@ -1011,8 +1011,8 @@ def _kernel_info_to_node(info: KernelInfo) -> KernelNode: session_type=info.session.session_type.value, ), user_info=KernelUserInfoGQLDTO( - user_id=info.user_permission.user_uuid, - access_key=info.user_permission.access_key, + user_id=info.user_permission.owner_id, + access_key=info.user_permission.main_access_key, domain_name=info.user_permission.domain_name, group_id=info.user_permission.group_id, ), diff --git a/src/ai/backend/manager/data/kernel/types.py b/src/ai/backend/manager/data/kernel/types.py index 075419d536a..ff495b2aae2 100644 --- a/src/ai/backend/manager/data/kernel/types.py +++ b/src/ai/backend/manager/data/kernel/types.py @@ -207,8 +207,8 @@ class ClusterConfig: @dataclass class UserPermission: - user_uuid: UUID - access_key: str + owner_id: UUID + main_access_key: str | None domain_name: str group_id: UUID uid: int | None diff --git a/src/ai/backend/manager/models/kernel/row.py b/src/ai/backend/manager/models/kernel/row.py index c5a570039a2..7332ad3ba0b 100644 --- a/src/ai/backend/manager/models/kernel/row.py +++ b/src/ai/backend/manager/models/kernel/row.py @@ -820,9 +820,8 @@ def set_status( else: self.status_data = dict(status_data) - def delegate_ownership(self, user_uuid: uuid.UUID, access_key: AccessKey) -> None: - self.user_uuid = user_uuid - self.access_key = access_key + def delegate_ownership(self, owner_id: uuid.UUID) -> None: + self.user_uuid = owner_id @classmethod async def set_kernel_status( @@ -945,8 +944,7 @@ def from_kernel_info(cls, info: KernelInfo) -> Self: agent_addr=info.resource.agent_addr, domain_name=info.user_permission.domain_name, group_id=info.user_permission.group_id, - user_uuid=info.user_permission.user_uuid, - access_key=info.user_permission.access_key, + user_uuid=info.user_permission.owner_id, image=info.image.identifier.canonical if info.image.identifier else None, architecture=info.image.identifier.architecture if info.image.identifier else None, registry=info.image.registry, @@ -1002,8 +1000,8 @@ def to_kernel_info(self) -> KernelInfo: session_type=self.session_type, ), user_permission=UserPermission( - user_uuid=self.user_uuid, - access_key=self.access_key or "", + owner_id=self.user_uuid, + main_access_key=self.user_row.main_access_key if self.user_row else None, domain_name=self.domain_name, group_id=self.group_id, uid=self.uid, @@ -1113,12 +1111,19 @@ async def recalc_concurrency_used( ) -> None: from ai.backend.manager.models.session import PRIVATE_SESSION_TYPES + # TODO(BA-5609 phase D): kernels.access_key is removed. Resolve the + # owner_id for this access_key (via users.main_access_key) and filter by + # KernelRow.user_uuid instead. The join below is a temporary shim that + # selects kernels whose owning user has main_access_key == access_key. + owner_id_subq = ( + sa.select(users.c.uuid).where(users.c.main_access_key == access_key).scalar_subquery() + ) async with db_sess.begin_nested(): result = await db_sess.execute( sa.select(sa.func.count()) .select_from(KernelRow) .where( - (KernelRow.access_key == access_key) + (KernelRow.user_uuid == owner_id_subq) & (KernelRow.status.in_(USER_RESOURCE_OCCUPYING_KERNEL_STATUSES)) & (KernelRow.session_type.not_in(PRIVATE_SESSION_TYPES)) ), @@ -1128,7 +1133,7 @@ async def recalc_concurrency_used( sa.select(sa.func.count()) .select_from(KernelRow) .where( - (KernelRow.access_key == access_key) + (KernelRow.user_uuid == owner_id_subq) & (KernelRow.status.in_(USER_RESOURCE_OCCUPYING_KERNEL_STATUSES)) & (KernelRow.session_type.in_(PRIVATE_SESSION_TYPES)) ), diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index 4f7beeaa5cd..926edf28afa 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -1288,7 +1288,7 @@ def delegate_ownership(self, user_uuid: UUID, access_key: AccessKey) -> None: self.user_uuid = user_uuid self.access_key = access_key for kernel_row in self.kernels: - kernel_row.delegate_ownership(user_uuid, access_key) + kernel_row.delegate_ownership(user_uuid) @staticmethod async def delete_by_user_id(user_uuid: UUID, *, db_session: SASession) -> None: diff --git a/src/ai/backend/manager/repositories/session/db_source/db_source.py b/src/ai/backend/manager/repositories/session/db_source/db_source.py index b544b94dd89..df5a75c24a7 100644 --- a/src/ai/backend/manager/repositories/session/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/session/db_source/db_source.py @@ -594,7 +594,7 @@ async def search_kernels( KernelListResult with items, total count, and pagination info """ async with self._db.begin_readonly_session() as db_sess: - query = sa.select(KernelRow) + query = sa.select(KernelRow).options(selectinload(KernelRow.user_row)) result = await execute_batch_querier( db_sess, diff --git a/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py b/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py index 746ec15b527..67e17f5b158 100644 --- a/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py +++ b/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py @@ -484,7 +484,7 @@ def _generate_slice_specs( spec = KernelUsageRecordCreatorSpec( kernel_id=UUID(str(kernel.id)), session_id=UUID(kernel.session.session_id), - user_uuid=kernel.user_permission.user_uuid, + user_uuid=kernel.user_permission.owner_id, project_id=kernel.user_permission.group_id, domain_name=kernel.user_permission.domain_name, resource_group=scaling_group, diff --git a/tests/unit/manager/api/compute_sessions/test_handler.py b/tests/unit/manager/api/compute_sessions/test_handler.py index c66a3db2431..2a3625d2ee7 100644 --- a/tests/unit/manager/api/compute_sessions/test_handler.py +++ b/tests/unit/manager/api/compute_sessions/test_handler.py @@ -123,8 +123,8 @@ def create_kernel_info( session_type=SessionTypes.INTERACTIVE, ), user_permission=UserPermission( - user_uuid=uuid4(), - access_key="TESTKEY", + owner_id=uuid4(), + main_access_key="TESTKEY", domain_name="default", group_id=uuid4(), uid=None, diff --git a/tests/unit/manager/services/session/test_session_service.py b/tests/unit/manager/services/session/test_session_service.py index 20999c0e168..c1cab977150 100644 --- a/tests/unit/manager/services/session/test_session_service.py +++ b/tests/unit/manager/services/session/test_session_service.py @@ -1712,8 +1712,8 @@ def sample_kernel_info(self) -> KernelInfo: session_type=SessionTypes.INTERACTIVE, ), user_permission=UserPermission( - user_uuid=user_id, - access_key="TESTKEY", + owner_id=user_id, + main_access_key="TESTKEY", domain_name="default", group_id=group_id, uid=1000, diff --git a/tests/unit/manager/sokovan/scheduler/handlers/conftest.py b/tests/unit/manager/sokovan/scheduler/handlers/conftest.py index 97cc68b2458..1f742e85e54 100644 --- a/tests/unit/manager/sokovan/scheduler/handlers/conftest.py +++ b/tests/unit/manager/sokovan/scheduler/handlers/conftest.py @@ -160,8 +160,8 @@ def _create_session( session_type=session_type, ), user_permission=UserPermission( - user_uuid=user_uuid, - access_key=access_key, + owner_id=user_uuid, + main_access_key=access_key, domain_name="default", group_id=group_id, uid=None, @@ -261,8 +261,8 @@ def _create_kernel( session_type=SessionTypes.INTERACTIVE, ), user_permission=UserPermission( - user_uuid=user_uuid, - access_key="test-access-key", + owner_id=user_uuid, + main_access_key="test-access-key", domain_name="default", group_id=group_id, uid=None, diff --git a/tests/unit/manager/sokovan/scheduler/terminator/conftest.py b/tests/unit/manager/sokovan/scheduler/terminator/conftest.py index f959b16b790..e33c6b66917 100644 --- a/tests/unit/manager/sokovan/scheduler/terminator/conftest.py +++ b/tests/unit/manager/sokovan/scheduler/terminator/conftest.py @@ -213,8 +213,8 @@ def _create_kernel_info( session_type=SessionTypes.INTERACTIVE, ), user_permission=UserPermission( - user_uuid=uuid4(), - access_key="test-access-key", + owner_id=uuid4(), + main_access_key="test-access-key", domain_name="default", group_id=uuid4(), uid=None,