refactor(BA-5650-E): collapse scheduler signatures to owner_id

Scheduler / predicates / scheduler-type collapse of the owner key:

- ``scheduler/predicates.py``: predicates now take SessionRow and
  resolve ``main_access_key`` from the owner via a helper when a
  keypair-scoped lookup (Redis concurrency, keypair resource policy)
  is required. Renames ``SessionRow.user_uuid`` references throughout.
- ``scheduler/drf.py``: per-user fairness tracking keyed by
  ``owner_id``/``main_access_key`` pair.
- ``repositories/scheduler/options.py``: drop the duplicated
  ``by_access_key_*`` factories — session filters go through
  ``SessionConditions`` helpers instead.
- ``repositories/scheduler/types/*``: rename ``access_key`` to
  ``main_access_key`` on ``ScheduledSessionData``,
  ``TerminatingSessionData``, ``SweptSessionInfo``, ``KernelEnqueueData``
  and ``SessionEnqueueData``.
- ``repositories/events/db_source/db_source.py`` and
  ``repositories/stream/db_source/db_source.py``: resolve the owner
  UUID from ``main_access_key`` via a sub-select shim while the schema
  still exposes ``sessions.access_key``.

Author: jopemachine

changes/BA-5650-E.misc.md

@@ -0,0 +1 @@
+Collapse scheduler / predicates / scheduler options signatures to take `owner_id: UUID`. Rename `access_key` field to `main_access_key` on `ScheduledSessionData` / `TerminatingSessionData` / `SweptSessionInfo` / scheduler types.

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

@@ -6,6 +6,7 @@
 from ai.backend.manager.errors.resource import ProjectNotFound
 from ai.backend.manager.models.group import groups
 from ai.backend.manager.models.session import SessionRow
+from ai.backend.manager.models.user import UserRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
 
 
@@ -21,8 +22,13 @@ async def match_sessions_by_name(
         access_key: AccessKey,
     ) -> list[SessionRow]:
         async with self._db.begin_readonly_session(isolation_level="READ COMMITTED") as db_sess:
+            owner_id = await db_sess.scalar(
+                sa.select(UserRow.uuid).where(UserRow.main_access_key == access_key)
+            )
+            if owner_id is None:
+                return []
             return await SessionRow.match_sessions(
-                db_sess, session_name, access_key, allow_prefix=False
+                db_sess, session_name, owner_id=owner_id, allow_prefix=False
             )
 
     async def resolve_group_id(self, group_name: str) -> uuid.UUID:

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

@@ -107,58 +107,6 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:
 
         return inner
 
-    @staticmethod
-    def by_access_key_contains(spec: StringMatchSpec) -> QueryCondition:
-        def inner() -> sa.sql.expression.ColumnElement[bool]:
-            if spec.case_insensitive:
-                condition = SessionRow.access_key.ilike(f"%{spec.value}%")
-            else:
-                condition = SessionRow.access_key.like(f"%{spec.value}%")
-            if spec.negated:
-                condition = sa.not_(condition)
-            return condition
-
-        return inner
-
-    @staticmethod
-    def by_access_key_equals(spec: StringMatchSpec) -> QueryCondition:
-        def inner() -> sa.sql.expression.ColumnElement[bool]:
-            if spec.case_insensitive:
-                condition = sa.func.lower(SessionRow.access_key) == spec.value.lower()
-            else:
-                condition = SessionRow.access_key == spec.value
-            if spec.negated:
-                condition = sa.not_(condition)
-            return condition
-
-        return inner
-
-    @staticmethod
-    def by_access_key_starts_with(spec: StringMatchSpec) -> QueryCondition:
-        def inner() -> sa.sql.expression.ColumnElement[bool]:
-            if spec.case_insensitive:
-                condition = SessionRow.access_key.ilike(f"{spec.value}%")
-            else:
-                condition = SessionRow.access_key.like(f"{spec.value}%")
-            if spec.negated:
-                condition = sa.not_(condition)
-            return condition
-
-        return inner
-
-    @staticmethod
-    def by_access_key_ends_with(spec: StringMatchSpec) -> QueryCondition:
-        def inner() -> sa.sql.expression.ColumnElement[bool]:
-            if spec.case_insensitive:
-                condition = SessionRow.access_key.ilike(f"%{spec.value}")
-            else:
-                condition = SessionRow.access_key.like(f"%{spec.value}")
-            if spec.negated:
-                condition = sa.not_(condition)
-            return condition
-
-        return inner
-
     @staticmethod
     def by_domain_name_contains(spec: StringMatchSpec) -> QueryCondition:
         def inner() -> sa.sql.expression.ColumnElement[bool]:
@@ -413,8 +361,8 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:
         return inner
 
     @staticmethod
-    def by_user_uuid_filter_equals(spec: UUIDEqualMatchSpec) -> QueryCondition:
-        """Factory for user UUID equality filter."""
+    def by_owner_id_filter_equals(spec: UUIDEqualMatchSpec) -> QueryCondition:
+        """Factory for owner_id equality filter."""
 
         def inner() -> sa.sql.expression.ColumnElement[bool]:
             if spec.negated:
@@ -424,8 +372,8 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:
         return inner
 
     @staticmethod
-    def by_user_uuid_filter_in(spec: UUIDInMatchSpec) -> QueryCondition:
-        """Factory for user UUID IN filter."""
+    def by_owner_id_filter_in(spec: UUIDInMatchSpec) -> QueryCondition:
+        """Factory for owner_id IN filter."""
 
         def inner() -> sa.sql.expression.ColumnElement[bool]:
             if spec.negated:

src/ai/backend/manager/repositories/scheduler/types/allocation.py

@@ -69,7 +69,7 @@ class SessionAllocation:
     kernel_allocations: list[KernelAllocation]
     # List of agent allocations for this session
     agent_allocations: list[AgentAllocation]
-    # Keypair associated with the session
+    # Resolved main_access_key of the owner; required for keypair-scoped concurrency tracking and resource policy lookups.
     access_key: AccessKey
     # Phases that passed during scheduling
     passed_phases: list[SchedulingPredicate] = field(default_factory=list)

src/ai/backend/manager/repositories/scheduler/types/results.py

@@ -13,5 +13,5 @@ class ScheduledSessionData:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
     reason: str

src/ai/backend/manager/repositories/scheduler/types/session.py

@@ -44,9 +44,10 @@ class PendingSessionData:
     """Pending session data for scheduling."""
 
     id: SessionId
+    # Resolved main_access_key of the owner; required for keypair-scoped concurrency tracking and resource policy lookups.
     access_key: AccessKey
     requested_slots: ResourceSlot
-    user_uuid: UUID
+    owner_id: UUID
     group_id: UUID
     domain_name: str
     scaling_group_name: str
@@ -64,9 +65,9 @@ def to_session_workload(self) -> SessionWorkload:
         kernel_workloads = [k.to_kernel_workload() for k in self.kernels]
         return SessionWorkload(
             session_id=self.id,
-            access_key=self.access_key,
+            main_access_key=self.access_key,
             requested_slots=self.requested_slots,
-            user_uuid=self.user_uuid,
+            owner_id=self.owner_id,
             group_id=self.group_id,
             domain_name=self.domain_name,
             scaling_group=self.scaling_group_name,
@@ -90,12 +91,12 @@ class PendingSessions:
     @cached_property
     def access_keys(self) -> set[AccessKey]:
         """Extract unique access keys from pending sessions."""
-        return {s.access_key for s in self.sessions}
+        return {s.main_access_key for s in self.sessions}
 
     @cached_property
-    def user_uuids(self) -> set[UUID]:
-        """Extract unique user UUIDs from pending sessions."""
-        return {s.user_uuid for s in self.sessions}
+    def owner_ids(self) -> set[UUID]:
+        """Extract unique owner (user) UUIDs from pending sessions."""
+        return {s.owner_id for s in self.sessions}
 
     @cached_property
     def group_ids(self) -> set[UUID]:
@@ -125,7 +126,7 @@ class TerminatingSessionData:
     """Data for a session that needs to be terminated."""
 
     session_id: SessionId
-    access_key: AccessKey
+    main_access_key: AccessKey
     creation_id: str
     status: SessionStatus
     status_info: str
@@ -161,7 +162,7 @@ class SweptSessionInfo:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
 
 
 @dataclass

src/ai/backend/manager/repositories/scheduler/types/session_creation.py

@@ -222,8 +222,8 @@ class KernelEnqueueData:
     scaling_group: str
     domain_name: str
     group_id: UUID
-    user_uuid: UUID
-    access_key: AccessKey
+    owner_id: UUID
+    main_access_key: AccessKey
     image: str  # Canonical image name
     architecture: str
     registry: str
@@ -268,8 +268,8 @@ class SessionEnqueueData:
     id: SessionId
     creation_id: str
     name: str
-    access_key: AccessKey
-    user_uuid: UUID
+    main_access_key: AccessKey
+    owner_id: UUID
     group_id: UUID
     domain_name: str
     scaling_group_name: str

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

@@ -1,5 +1,9 @@
+import sqlalchemy as sa
+
 from ai.backend.common.types import AccessKey
+from ai.backend.manager.errors.kernel import SessionNotFound
 from ai.backend.manager.models.session import KernelLoadingStrategy, SessionRow
+from ai.backend.manager.models.user import UserRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
 
 
@@ -15,9 +19,14 @@ async def get_streaming_session(
         access_key: AccessKey,
     ) -> SessionRow:
         async with self._db.begin_readonly_session() as db_sess:
+            owner_id = await db_sess.scalar(
+                sa.select(UserRow.uuid).where(UserRow.main_access_key == access_key)
+            )
+            if owner_id is None:
+                raise SessionNotFound(f"Unknown access_key: {access_key}")
             return await SessionRow.get_session(
                 db_sess,
                 session_name,
-                access_key,
+                owner_id=owner_id,
                 kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY,
             )

src/ai/backend/manager/scheduler/drf.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import uuid
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from decimal import Decimal
@@ -9,7 +10,6 @@
 import trafaret as t
 
 from ai.backend.common.types import (
-    AccessKey,
     ResourceSlot,
     SessionId,
 )
@@ -24,7 +24,7 @@
 
 
 class DRFScheduler(AbstractScheduler):
-    per_user_dominant_share: dict[AccessKey, Decimal]
+    per_user_dominant_share: dict[uuid.UUID, Decimal]
     total_capacity: ResourceSlot
 
     def __init__(
@@ -60,31 +60,30 @@ def pick_session(
                 slot_share = Decimal(value) / slot_cap
                 if dominant_share < slot_share:
                     dominant_share = slot_share
-            raw_access_key = existing_sess.access_key
-            if raw_access_key is not None:
-                access_key = AccessKey(raw_access_key)
-                if self.per_user_dominant_share[access_key] < dominant_share:
-                    self.per_user_dominant_share[access_key] = dominant_share
+            owner_id = existing_sess.owner_id
+            if owner_id is not None:
+                if self.per_user_dominant_share[owner_id] < dominant_share:
+                    self.per_user_dominant_share[owner_id] = dominant_share
         log.debug("per-user dominant share: {}", dict(self.per_user_dominant_share))
 
         # Find who has the least dominant share among the pending session.
-        users_with_pending_session: set[AccessKey] = {
-            AccessKey(pending_sess.access_key)
+        users_with_pending_session: set[uuid.UUID] = {
+            pending_sess.user_uuid
             for pending_sess in pending_sessions
-            if pending_sess.access_key is not None
+            if pending_sess.user_uuid is not None
         }
         if not users_with_pending_session:
             return None
         least_dominant_share_user, dshare = min(
-            ((akey, self.per_user_dominant_share[akey]) for akey in users_with_pending_session),
+            ((oid, self.per_user_dominant_share[oid]) for oid in users_with_pending_session),
             key=lambda item: item[1],
         )
         log.debug("least dominant share user: {} ({})", least_dominant_share_user, dshare)
 
         # Pick the first pending session of the user
         # who has the lowest dominant share.
         for pending_sess in pending_sessions:
-            if pending_sess.access_key == least_dominant_share_user:
+            if pending_sess.user_uuid == least_dominant_share_user:
                 return SessionId(pending_sess.id)
 
         return None
@@ -96,10 +95,7 @@ def update_allocation(
     ) -> None:
         # In such case, we just skip updating self.per_user_dominant_share state
         # and the scheduler continues to pick another session within the same scaling group.
-        raw_access_key = scheduled_session_or_kernel.access_key
-        if raw_access_key is None:
-            return
-        access_key = AccessKey(raw_access_key)
+        owner_id = scheduled_session_or_kernel.user_uuid
         requested_slots = scheduled_session_or_kernel.requested_slots
 
         # Update the dominant share.
@@ -114,5 +110,5 @@ def update_allocation(
             slot_share = Decimal(value) / slot_cap
             if dominant_share_from_request < slot_share:
                 dominant_share_from_request = slot_share
-        if self.per_user_dominant_share[access_key] < dominant_share_from_request:
-            self.per_user_dominant_share[access_key] = dominant_share_from_request
+        if self.per_user_dominant_share[owner_id] < dominant_share_from_request:
+            self.per_user_dominant_share[owner_id] = dominant_share_from_request