feat(BA-6644): Add domain id, resource group id in session row (#12452)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: seedspirit

changes/12452.feature.md

@@ -0,0 +1 @@
+Add domain_id and resource_group_id columns to session rows and populate them at session creation

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

@@ -11,6 +11,8 @@
 from uuid import UUID
 
 from ai.backend.common.docker import ImageRef
+from ai.backend.common.identifier.domain import DomainID
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import AccessKey
 
 
@@ -62,7 +64,9 @@ class DeploymentContext:
     created_user: UserContext
     session_owner: UserContext
     container_user: ContainerUserContext
+    domain_id: DomainID
     group_id: UUID
+    resource_group_id: ResourceGroupID
     resource_policy: dict[str, Any]
     image: ImageContext
     resolved_presets: ResolvedPresetValues | None = None

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

@@ -37,10 +37,10 @@
 import yarl
 from pydantic import ConfigDict, Field
 
-from ai.backend.common.identifier.domain import DomainName
+from ai.backend.common.identifier.domain import DomainID, DomainName
 from ai.backend.common.identifier.image import ImageID
 from ai.backend.common.identifier.project import ProjectID
-from ai.backend.common.identifier.resource_group import ResourceGroupName
+from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName
 from ai.backend.common.identifier.session import SessionID
 from ai.backend.common.types import (
     AccessKey,
@@ -179,8 +179,10 @@ class SessionNetworkDraft(_DraftBaseModel):
 class SessionScopeDraft(_DraftBaseModel):
     """Optional-heavy mirror of ``SessionScope``."""
 
+    domain_id: DomainID | None = None
     domain_name: DomainName | None = None
     project_id: ProjectID | None = None
+    resource_group_id: ResourceGroupID | None = None
     resource_group_name: ResourceGroupName | None = None
 
 

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

@@ -24,9 +24,9 @@
 from pydantic import ConfigDict, Field
 
 from ai.backend.common.exception import BackendAIError
-from ai.backend.common.identifier.domain import DomainName
+from ai.backend.common.identifier.domain import DomainID, DomainName
 from ai.backend.common.identifier.project import ProjectID
-from ai.backend.common.identifier.resource_group import ResourceGroupName
+from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName
 from ai.backend.common.identifier.session import SessionID
 from ai.backend.common.types import (
     AccessKey,
@@ -69,8 +69,10 @@ class SessionIdentity(_SpecBaseModel):
 class SessionScope(_SpecBaseModel):
     """Ownership / placement scope of the session."""
 
+    domain_id: DomainID
     domain_name: DomainName
     project_id: ProjectID
+    resource_group_id: ResourceGroupID
     resource_group_name: ResourceGroupName
 
 

src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py

@@ -0,0 +1,74 @@
+"""add session scope id columns
+
+Revision ID: ada41cb881bb
+Revises: a379b72f1206
+Create Date: 2026-06-30
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+from ai.backend.manager.models.base import GUID
+
+# revision identifiers, used by Alembic.
+revision = "ada41cb881bb"
+down_revision = "a379b72f1206"
+# Part of: 26.7.0
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column("sessions", sa.Column("domain_id", GUID(), nullable=True))
+    op.add_column("sessions", sa.Column("resource_group_id", GUID(), nullable=True))
+    op.create_index("ix_sessions_domain_id", "sessions", ["domain_id"])
+    op.create_index("ix_sessions_resource_group_id", "sessions", ["resource_group_id"])
+
+    op.execute(
+        sa.text("""
+        UPDATE sessions
+        SET domain_id = domains.id
+        FROM domains
+        WHERE sessions.domain_name = domains.name
+          AND sessions.domain_id IS NULL
+        """)
+    )
+    op.execute(
+        sa.text("""
+        UPDATE sessions
+        SET resource_group_id = scaling_groups.id
+        FROM scaling_groups
+        WHERE sessions.scaling_group_name = scaling_groups.name
+          AND sessions.resource_group_id IS NULL
+        """)
+    )
+    op.alter_column("sessions", "domain_id", nullable=False)
+    op.alter_column("sessions", "resource_group_id", nullable=False)
+    op.create_foreign_key(
+        op.f("fk_sessions_domain_id_domains"),
+        "sessions",
+        "domains",
+        ["domain_id"],
+        ["id"],
+    )
+    op.create_foreign_key(
+        op.f("fk_sessions_resource_group_id_scaling_groups"),
+        "sessions",
+        "scaling_groups",
+        ["resource_group_id"],
+        ["id"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_constraint(
+        op.f("fk_sessions_resource_group_id_scaling_groups"),
+        "sessions",
+        type_="foreignkey",
+    )
+    op.drop_constraint(op.f("fk_sessions_domain_id_domains"), "sessions", type_="foreignkey")
+    op.drop_index("ix_sessions_resource_group_id", table_name="sessions")
+    op.drop_index("ix_sessions_domain_id", table_name="sessions")
+    op.drop_column("sessions", "resource_group_id")
+    op.drop_column("sessions", "domain_id")

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

@@ -135,7 +135,11 @@ class DomainRow(Base):  # type: ignore[misc]
         "dotfiles", sa.LargeBinary(length=MAXIMUM_DOTFILE_SIZE), nullable=False, default=b"\x90"
     )
 
-    sessions: Mapped[list[SessionRow]] = relationship("SessionRow", back_populates="domain")
+    sessions: Mapped[list[SessionRow]] = relationship(
+        "SessionRow",
+        back_populates="domain",
+        foreign_keys="[SessionRow.domain_name]",
+    )
     users: Mapped[list[UserRow]] = relationship("UserRow", back_populates="domain")
     groups: Mapped[list[GroupRow]] = relationship("GroupRow", back_populates="domain")
     sgroup_for_domains_rows: Mapped[list[ScalingGroupForDomainRow]] = relationship(

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

@@ -330,7 +330,11 @@ class ScalingGroupRow(Base):  # type: ignore[misc]
         default=DefaultSessionOptions,
     )
 
-    sessions: Mapped[list[SessionRow]] = relationship("SessionRow", back_populates="scaling_group")
+    sessions: Mapped[list[SessionRow]] = relationship(
+        "SessionRow",
+        back_populates="scaling_group",
+        foreign_keys="[SessionRow.scaling_group_name]",
+    )
     agents: Mapped[list[AgentRow]] = relationship("AgentRow", back_populates="scaling_group_row")
 
     sgroup_for_domains_rows: Mapped[list[ScalingGroupForDomainRow]] = relationship(

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

@@ -39,6 +39,8 @@
 
 from ai.backend.common.defs.session import SESSION_PRIORITY_DEFAULT
 from ai.backend.common.exception import BackendAIError
+from ai.backend.common.identifier.domain import DomainID
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import (
     AccessKey,
     ClusterMode,
@@ -453,11 +455,20 @@ class SessionRow(Base):  # type: ignore[misc]
     kernels: Mapped[list[KernelRow]] = relationship("KernelRow", back_populates="session")
 
     # Resource ownership
+    resource_group_id: Mapped[ResourceGroupID] = mapped_column(
+        "resource_group_id",
+        GUID(ResourceGroupID),
+        sa.ForeignKey("scaling_groups.id"),
+        index=True,
+        nullable=False,
+    )
     scaling_group_name: Mapped[str] = mapped_column(
         "scaling_group_name", sa.ForeignKey("scaling_groups.name"), index=True, nullable=False
     )
     scaling_group: Mapped[ScalingGroupRow] = relationship(
-        "ScalingGroupRow", back_populates="sessions"
+        "ScalingGroupRow",
+        back_populates="sessions",
+        foreign_keys=[scaling_group_name],
     )
     target_sgroup_names: Mapped[list[str] | None] = mapped_column(
         "target_sgroup_names",
@@ -469,7 +480,18 @@ class SessionRow(Base):  # type: ignore[misc]
     domain_name: Mapped[str] = mapped_column(
         "domain_name", sa.String(length=64), sa.ForeignKey("domains.name"), nullable=False
     )
-    domain: Mapped[DomainRow] = relationship("DomainRow", back_populates="sessions")
+    domain_id: Mapped[DomainID] = mapped_column(
+        "domain_id",
+        GUID(DomainID),
+        sa.ForeignKey("domains.id"),
+        index=True,
+        nullable=False,
+    )
+    domain: Mapped[DomainRow] = relationship(
+        "DomainRow",
+        back_populates="sessions",
+        foreign_keys=[domain_name],
+    )
     group_id: Mapped[UUID] = mapped_column(
         "group_id", GUID, sa.ForeignKey("groups.id"), nullable=False
     )

src/ai/backend/manager/registry.py

@@ -1119,8 +1119,13 @@ def _build_execution_spec(
         if not groups_by_role:
             raise InvalidAPIParameters("No kernel groups resolved from the enqueue request.")
 
+        domain_name = DomainName(user_scope.domain_name)
+        domain_id = await self._scheduler_repository.get_domain_id_by_name(domain_name)
         if scaling_group:
             resource_group_name = ResourceGroupName(scaling_group)
+            resource_group_id = await self._scheduler_repository.get_resource_group_id_by_name(
+                resource_group_name
+            )
         else:
             resource_group_id = await self._scheduler_repository.pick_default_resource_group(
                 access_key=access_key,
@@ -1140,8 +1145,10 @@ def _build_execution_spec(
                 user_uuid=user_scope.user_uuid,
             ),
             scope=SessionScopeDraft(
-                domain_name=DomainName(user_scope.domain_name),
+                domain_id=domain_id,
+                domain_name=domain_name,
                 project_id=ProjectID(user_scope.group_id),
+                resource_group_id=resource_group_id,
                 resource_group_name=resource_group_name,
             ),
             classification=SessionClassificationDraft(

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

@@ -29,10 +29,11 @@
 from ai.backend.common.identifier.deployment import DeploymentID
 from ai.backend.common.identifier.deployment_preset import DeploymentPresetID
 from ai.backend.common.identifier.deployment_revision import DeploymentRevisionID
+from ai.backend.common.identifier.domain import DomainID
 from ai.backend.common.identifier.image import ImageID
 from ai.backend.common.identifier.replica import ReplicaID
 from ai.backend.common.identifier.replica_group import ReplicaGroupID
-from ai.backend.common.identifier.resource_group import ResourceGroupName
+from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName
 from ai.backend.common.identifier.runtime_variant import RuntimeVariantID
 from ai.backend.common.identifier.vfolder import VFolderUUID
 from ai.backend.common.types import (
@@ -105,6 +106,7 @@
     UserNotFoundInDeployment,
 )
 from ai.backend.manager.errors.resource import (
+    DomainNotFound,
     ProjectNotFound,
     RuntimeVariantNotFound,
     ScalingGroupNotFound,
@@ -124,6 +126,7 @@
 from ai.backend.manager.models.deployment_revision_preset.row import (
     DeploymentRevisionPresetRow,
 )
+from ai.backend.manager.models.domain import DomainRow
 from ai.backend.manager.models.endpoint import (
     EndpointAutoScalingRuleRow,
     EndpointRow,
@@ -2099,6 +2102,29 @@ async def _resolve_user_and_active_access_key(
             raise UserNotFoundInDeployment(f"{user_uuid} not found")
         raise NoActiveKeypairForDeployment(f"{user_uuid} has no active keypair")
 
+    async def _resolve_deployment_scope_ids(
+        self,
+        db_sess: SASession,
+        deployment_info: DeploymentInfo,
+    ) -> tuple[DomainID, ResourceGroupID]:
+        domain_id = await db_sess.scalar(
+            sa.select(DomainRow.id).where(DomainRow.name == deployment_info.metadata.domain)
+        )
+        if domain_id is None:
+            raise DomainNotFound(deployment_info.metadata.domain)
+
+        resource_group_id = await db_sess.scalar(
+            sa.select(ScalingGroupRow.id).where(
+                ScalingGroupRow.name == deployment_info.metadata.resource_group
+            )
+        )
+        if resource_group_id is None:
+            raise ScalingGroupNotFound(
+                f"Resource group {deployment_info.metadata.resource_group!r} not found"
+            )
+
+        return DomainID(domain_id), ResourceGroupID(resource_group_id)
+
     async def fetch_deployment_context(
         self,
         deployment_info: DeploymentInfo,
@@ -2147,6 +2173,10 @@ async def fetch_deployment_context(
             group_id = user_info.group_id
             resource_policy = user_info.resource_policy
 
+            domain_id, resource_group_id = await self._resolve_deployment_scope_ids(
+                db_sess, deployment_info
+            )
+
             revision_query = (
                 sa.select(DeploymentRevisionRow)
                 .where(DeploymentRevisionRow.id == revision_id)
@@ -2212,7 +2242,9 @@ async def fetch_deployment_context(
                     main_gid=session_owner_user.container_main_gid,
                     supplementary_gids=session_owner_user.container_gids or [],
                 ),
+                domain_id=domain_id,
                 group_id=group_id,
+                resource_group_id=resource_group_id,
                 resource_policy=dict(resource_policy),
                 image=ImageContext(
                     ref=image_row.image_ref,