refactor(BA-5710): thread main_access_key through UserPermission

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.

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,
             ),

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

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))
             ),

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:

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,

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,

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,

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,

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,

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,