refactor(BA-5650-C): rename SessionData user_uuid to owner_id

Data-layer rename of SessionData / SessionMetadata user_uuid to
owner_id, plus the matching SessionRow adapters
(to_dataclass/from_dataclass/to_session_info/from_session_info).
ComputeSessionNode.from_dataclass becomes async and resolves
main_access_key from the owner via UserRepository when owner is not
eagerly loaded. models/resource_usage.py sources user_id from
SessionRow.user_uuid through the session relationship.

SessionRow._build_session_fetch_query / _match_sessions_by_* also
rename the filter parameter to owner_id; repository and db_source
callers are updated in the next slice so intermediate builds may
fail in isolation.

Author: jopemachine

changes/BA-5650-C.misc.md

@@ -0,0 +1 @@
+Rename `SessionData.user_uuid` / `SessionMetadata.user_uuid` to `owner_id` and drop the redundant `access_key` snapshot fields. `SessionRow.to_dataclass` / `from_dataclass`, `ComputeSessionNode.from_dataclass`, and `resource_usage` queries now resolve the owner's `main_access_key` from the eager-loaded `user` relationship.

src/ai/backend/manager/api/gql_legacy/session.py

@@ -23,7 +23,6 @@
 
 from ai.backend.common import validators as tx
 from ai.backend.common.defs.session import SESSION_PRIORITY_MAX, SESSION_PRIORITY_MIN
-from ai.backend.common.exception import SessionWithInvalidStateError
 from ai.backend.common.types import (
     ClusterMode,
     KernelId,
@@ -391,7 +390,7 @@ def from_row(
         return result
 
     @classmethod
-    def from_dataclass(
+    async def from_dataclass(
         cls,
         ctx: GraphQueryContext,
         session_data: SessionData,
@@ -405,8 +404,16 @@ def from_dataclass(
         else:
             vfolder_mounts = [vf.vfid.folder_id for vf in session_data.vfolder_mounts]
 
-        if session_data.owner is None:
-            raise SessionWithInvalidStateError()
+        # access_key on this node is the owner's main_access_key. When the
+        # caller hasn't eagerly loaded session_data.owner, fall back to the
+        # user repository so the required field is always populated from the
+        # always-present owner_id.
+        if session_data.owner is not None:
+            main_access_key = session_data.owner.main_access_key
+        else:
+            main_access_key = await ctx.user_repository.get_main_access_key_by_id(
+                session_data.owner_id
+            )
 
         result = cls(
             # identity
@@ -422,9 +429,9 @@ def from_dataclass(
             # ownership
             domain_name=session_data.domain_name,
             project_id=session_data.group_id,
-            user_id=session_data.user_uuid,
-            access_key=session_data.access_key,
-            owner=UserNode.from_dataclass(ctx, session_data.owner),
+            user_id=session_data.owner_id,
+            access_key=main_access_key,
+            owner=UserNode.from_dataclass(ctx, session_data.owner) if session_data.owner else None,
             # status
             status=session_data.status.name,
             # status_changed=row.status_changed,  # FIXME: generated attribute
@@ -919,7 +926,7 @@ async def mutate_and_get_payload(
         )
 
         return ModifyComputeSession(
-            ComputeSessionNode.from_dataclass(graph_ctx, result.session_data),
+            await ComputeSessionNode.from_dataclass(graph_ctx, result.session_data),
             input.get("client_mutation_id"),
         )
 
@@ -970,7 +977,7 @@ async def mutate(
                 )
             )
             session_nodes.append(
-                ComputeSessionNode.from_dataclass(graph_ctx, action_result.session_data)
+                await ComputeSessionNode.from_dataclass(graph_ctx, action_result.session_data)
             )
 
         return CheckAndTransitStatus(session_nodes, input.get("client_mutation_id"))

src/ai/backend/manager/data/session/types.py

@@ -12,7 +12,6 @@
 
 from ai.backend.common.data.vfolder.types import VFolderMountData
 from ai.backend.common.types import (
-    AccessKey,
     CIStrEnum,
     ClusterMode,
     ResourceSlot,
@@ -155,7 +154,7 @@ class SessionData:
     cluster_size: int
     domain_name: str
     group_id: UUID
-    user_uuid: UUID
+    owner_id: UUID
     occupying_slots: Any  # TODO: ResourceSlot?
     requested_slots: Any
     use_host_network: bool
@@ -165,7 +164,6 @@ class SessionData:
     num_queries: int
     creation_id: str | None
     name: str | None
-    access_key: AccessKey | None
     agent_ids: list[str] | None
     images: list[str] | None
     tag: str | None
@@ -206,8 +204,7 @@ class SessionMetadata:
     name: str
     domain_name: str
     group_id: UUID
-    user_uuid: UUID
-    access_key: str
+    owner_id: UUID
     session_type: SessionTypes
     priority: int
     created_at: datetime | None

src/ai/backend/manager/models/resource_usage.py

@@ -524,7 +524,11 @@ async def parse_resource_usage_groups(
             last_stat=stat_map.get(kern.id),
             user_id=kern.session.user_uuid,
             user_email=kern.session.user.email if kern.session.user is not None else None,
-            access_key=kern.session.access_key,
+            # TODO(BA-5609 phase D): resolve access_key from owner via
+            # users.main_access_key. SessionRow.access_key has been removed.
+            access_key=(
+                kern.session.user.main_access_key if kern.session.user is not None else None
+            ),
             project_id=kern.session.group.id if kern.session.group is not None else None,
             project_name=kern.session.group.name if kern.session.group is not None else None,
             kernel_id=kern.id,
@@ -553,7 +557,8 @@ async def parse_resource_usage_groups(
     SessionRow.domain_name,
     SessionRow.id,
     SessionRow.group_id,
-    SessionRow.access_key,
+    # TODO(BA-5609 phase D): SessionRow.access_key removed. Callers should
+    # join UserRow and read users.main_access_key when an access_key is needed.
     SessionRow.images,
     SessionRow.cluster_mode,
     SessionRow.status_history,

src/ai/backend/manager/models/session/row.py

@@ -123,7 +123,6 @@
 
 if TYPE_CHECKING:
     from ai.backend.manager.models.domain import DomainRow
-    from ai.backend.manager.models.keypair import KeyPairRow
     from ai.backend.manager.models.scaling_group import ScalingGroupRow
     from ai.backend.manager.models.user import UserRow
 
@@ -494,17 +493,17 @@ async def handle_session_exception(
 
 def _build_session_fetch_query(
     base_cond: Any,
-    access_key: AccessKey | None = None,
     *,
+    owner_id: UUID | None = None,
     allow_stale: bool = True,
     for_update: bool = False,
     do_ordering: bool = False,
     max_matches: int | None = None,
     eager_loading_op: Sequence[_AbstractLoad] | None = None,
 ) -> sa.sql.Select[Any]:
     cond = base_cond
-    if access_key:
-        cond = cond & (SessionRow.access_key == access_key)
+    if owner_id is not None:
+        cond = cond & (SessionRow.user_uuid == owner_id)
     if not allow_stale:
         cond = cond & (~SessionRow.status.in_(DEAD_SESSION_STATUSES))
     query = (
@@ -528,8 +527,8 @@ def _build_session_fetch_query(
 async def _match_sessions_by_id(
     db_session: SASession,
     session_id_or_list: SessionId | list[SessionId],
-    access_key: AccessKey | None = None,
     *,
+    owner_id: UUID | None = None,
     allow_prefix: bool = False,
     allow_stale: bool = True,
     for_update: bool = False,
@@ -546,7 +545,7 @@ async def _match_sessions_by_id(
             cond = SessionRow.id == session_id_or_list
     query = _build_session_fetch_query(
         cond,
-        access_key,
+        owner_id=owner_id,
         max_matches=max_matches,
         allow_stale=allow_stale,
         for_update=for_update,
@@ -560,8 +559,8 @@ async def _match_sessions_by_id(
 async def _match_sessions_by_name(
     db_session: SASession,
     session_name: str,
-    access_key: AccessKey,
     *,
+    owner_id: UUID | None = None,
     allow_prefix: bool = False,
     allow_stale: bool = True,
     for_update: bool = False,
@@ -575,7 +574,7 @@ async def _match_sessions_by_name(
         cond = SessionRow.name == session_name
     query = _build_session_fetch_query(
         cond,
-        access_key,
+        owner_id=owner_id,
         max_matches=max_matches,
         allow_stale=allow_stale,
         for_update=for_update,
@@ -595,20 +594,6 @@ class ConcurrencyUsed:
     compute_session_ids: set[SessionId] = field(default_factory=set)
     system_session_ids: set[SessionId] = field(default_factory=set)
 
-    @property
-    def compute_concurrency_used_key(self) -> str:
-        return f"{COMPUTE_CONCURRENCY_USED_KEY_PREFIX}{self.access_key}"
-
-    @property
-    def system_concurrency_used_key(self) -> str:
-        return f"{SYSTEM_CONCURRENCY_USED_KEY_PREFIX}{self.access_key}"
-
-    def to_cnt_map(self) -> Mapping[str, int]:
-        return {
-            self.compute_concurrency_used_key: len(self.compute_session_ids),
-            self.system_concurrency_used_key: len(self.system_session_ids),
-        }
-
 
 class SessionOp(enum.StrEnum):
     CREATE = "create_session"
@@ -637,13 +622,6 @@ class KernelLoadingStrategy(enum.StrEnum):
 }
 
 
-# Defined for avoiding circular import
-def _get_keypair_row_join_condition() -> sa.sql.elements.ColumnElement[Any]:
-    from ai.backend.manager.models.keypair import KeyPairRow
-
-    return KeyPairRow.access_key == foreign(SessionRow.access_key)
-
-
 def _get_user_row_join_condition() -> sa.sql.elements.ColumnElement[Any]:
     from ai.backend.manager.models.user import UserRow
 
@@ -731,14 +709,7 @@ class SessionRow(Base):  # type: ignore[misc]
         back_populates="sessions",
         foreign_keys=[user_uuid],
     )
-
     access_key: Mapped[str | None] = mapped_column("access_key", sa.String(length=20))
-    access_key_row: Mapped[KeyPairRow | None] = relationship(
-        "KeyPairRow",
-        primaryjoin=_get_keypair_row_join_condition,
-        back_populates="sessions",
-        foreign_keys=[access_key],
-    )
 
     # `image` column is identical to kernels `image` column.
     images: Mapped[list[str] | None] = mapped_column("images", sa.ARRAY(sa.String), nullable=True)
@@ -884,7 +855,7 @@ class SessionRow(Base):  # type: ignore[misc]
         sa.Index("ix_session_status_with_priority", "status", "priority"),
         # Unique index for session names per user excluding terminal statuses
         sa.Index(
-            "ix_sessions_unique_name_per_user_nonterminal",
+            "ix_sessions_unique_name_per_owner_nonterminal",
             "name",
             "user_uuid",
             unique=True,
@@ -923,8 +894,7 @@ def from_dataclass(cls, session_data: SessionData) -> SessionRow:
             target_sgroup_names=session_data.target_sgroup_names,
             domain_name=session_data.domain_name,
             group_id=session_data.group_id,
-            user_uuid=session_data.user_uuid,
-            access_key=session_data.access_key,
+            user_uuid=session_data.owner_id,
             images=session_data.images,
             tag=session_data.tag,
             occupying_slots=session_data.occupying_slots,
@@ -968,8 +938,7 @@ def to_dataclass(self, owner: UserData | None = None) -> SessionData:
             target_sgroup_names=self.target_sgroup_names,
             domain_name=self.domain_name,
             group_id=self.group_id,
-            user_uuid=self.user_uuid,
-            access_key=AccessKey(self.access_key) if self.access_key else None,
+            owner_id=self.user_uuid,
             images=self.images,
             tag=self.tag,
             occupying_slots=self.occupying_slots,
@@ -1017,8 +986,7 @@ def from_session_info(cls, info: SessionInfo) -> Self:
             target_sgroup_names=info.resource.target_sgroup_names,
             domain_name=info.metadata.domain_name,
             group_id=info.metadata.group_id,
-            user_uuid=info.metadata.user_uuid,
-            access_key=info.metadata.access_key,
+            user_uuid=info.metadata.owner_id,
             images=info.image.images,
             tag=info.image.tag or info.metadata.tag,
             occupying_slots=info.resource.occupying_slots,
@@ -1059,8 +1027,7 @@ def to_session_info(self) -> SessionInfo:
                 name=self.name or "",
                 domain_name=self.domain_name,
                 group_id=self.group_id,
-                user_uuid=self.user_uuid,
-                access_key=self.access_key or "",
+                owner_id=self.user_uuid,
                 session_type=self.session_type,
                 priority=self.priority,
                 created_at=self.created_at,
@@ -1284,11 +1251,10 @@ def set_status(
         if _status_info is not None:
             self.status_info = _status_info
 
-    def delegate_ownership(self, user_uuid: UUID, access_key: AccessKey) -> None:
-        self.user_uuid = user_uuid
-        self.access_key = access_key
+    def delegate_ownership(self, owner_id: UUID) -> None:
+        self.user_uuid = owner_id
         for kernel_row in self.kernels:
-            kernel_row.delegate_ownership(user_uuid, access_key)
+            kernel_row.delegate_ownership(owner_id)
 
     @staticmethod
     async def delete_by_user_id(user_uuid: UUID, *, db_session: SASession) -> None:
@@ -1357,8 +1323,8 @@ async def match_sessions(
         cls,
         db_session: SASession,
         session_reference: str | UUID | list[UUID],
-        access_key: AccessKey | None,
         *,
+        owner_id: UUID | None = None,
         allow_prefix: bool = False,
         allow_stale: bool = True,
         for_update: bool = False,
@@ -1412,7 +1378,7 @@ async def match_sessions(
         for fetch_func in query_list:
             rows = await fetch_func(
                 db_session,
-                access_key=access_key,
+                owner_id=owner_id,
                 allow_stale=allow_stale,
                 for_update=for_update,
                 max_matches=max_matches,
@@ -1428,8 +1394,8 @@ async def get_session(
         cls,
         db_session: SASession,
         session_name_or_id: str | UUID,
-        access_key: AccessKey | None = None,
         *,
+        owner_id: UUID | None = None,
         allow_stale: bool = False,
         for_update: bool = False,
         kernel_loading_strategy: KernelLoadingStrategy = KernelLoadingStrategy.NONE,
@@ -1474,7 +1440,7 @@ async def get_session(
         session_list = await cls.match_sessions(
             db_session,
             session_name_or_id,
-            access_key,
+            owner_id=owner_id,
             allow_stale=allow_stale,
             for_update=for_update,
             eager_loading_op=_eager_loading_op,
@@ -1499,8 +1465,8 @@ async def list_sessions(
         cls,
         db_session: SASession,
         session_ids: list[UUID],
-        access_key: AccessKey | None = None,
         *,
+        owner_id: UUID | None = None,
         allow_stale: bool = False,
         for_update: bool = False,
         kernel_loading_strategy: KernelLoadingStrategy = KernelLoadingStrategy.NONE,
@@ -1531,7 +1497,7 @@ async def list_sessions(
         session_list = await cls.match_sessions(
             db_session,
             session_ids,
-            access_key,
+            owner_id=owner_id,
             allow_stale=allow_stale,
             for_update=for_update,
             eager_loading_op=_eager_loading_op,
@@ -1547,8 +1513,8 @@ async def get_session_by_id(
         cls,
         db_session: SASession,
         session_id: SessionId,
-        access_key: AccessKey | None = None,
         *,
+        owner_id: UUID | None = None,
         max_matches: int | None = None,
         allow_stale: bool = True,
         for_update: bool = False,
@@ -1557,7 +1523,7 @@ async def get_session_by_id(
         sessions = await _match_sessions_by_id(
             db_session,
             session_id,
-            access_key,
+            owner_id=owner_id,
             max_matches=max_matches,
             allow_stale=allow_stale,
             for_update=for_update,
@@ -1586,7 +1552,6 @@ async def get_sgroup_managed_sessions(
                 noload("*"),
                 selectinload(SessionRow.group).options(noload("*")),
                 selectinload(SessionRow.domain).options(noload("*")),
-                selectinload(SessionRow.access_key_row).options(noload("*")),
                 selectinload(SessionRow.kernels).options(noload("*")),
             )
         )