From c1020fbff3a088d3837c5e7d8322fec2a741d0fb Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Sun, 16 Aug 2026 02:48:22 +0900 Subject: [PATCH 1/5] refactor(BA-7226): replace AgentRow.agent_resource_rows with explicit queries Add fetch_actual_occupied_slots() to the agent repository query module and pass the resolved slots into AgentRow.to_data(), so agent occupancy no longer depends on the agent_resource_rows relationship or its slot_type_row chain. The scheduler db_source loads per-slot capacity/reserved/used through its own batched query. Co-Authored-By: Claude Opus 5 (1M context) --- .../backend/manager/api/gql_legacy/agent.py | 20 ++++++---- src/ai/backend/manager/models/agent/row.py | 14 +------ .../repositories/agent/db_source/db_source.py | 25 ++++++------ .../manager/repositories/agent/query.py | 28 +++++++++++++- .../manager/repositories/agent/repository.py | 15 ++++---- .../scheduler/db_source/db_source.py | 38 +++++++++++++------ .../test_resource_slot_tables.py | 18 ++------- 7 files changed, 90 insertions(+), 68 deletions(-) diff --git a/src/ai/backend/manager/api/gql_legacy/agent.py b/src/ai/backend/manager/api/gql_legacy/agent.py index 9895d6cf258..eb39723cdfe 100644 --- a/src/ai/backend/manager/api/gql_legacy/agent.py +++ b/src/ai/backend/manager/api/gql_legacy/agent.py @@ -46,10 +46,13 @@ ScopeType, ) from ai.backend.manager.models.rbac.context import ClientContext -from ai.backend.manager.models.resource_slot import AgentResourceRow from ai.backend.manager.models.scaling_group import ScalingGroupRow from ai.backend.manager.models.user import UserRole, users -from ai.backend.manager.repositories.agent.query import QueryConditions, QueryOrders +from ai.backend.manager.repositories.agent.query import ( + QueryConditions, + QueryOrders, + fetch_actual_occupied_slots, +) from ai.backend.manager.services.agent.actions.update_resource_group import ( UpdateAgentResourceGroupAction, ) @@ -783,11 +786,6 @@ async def batch_load( query = ( sa.select(AgentRow) .where(AgentRow.id.in_(agent_ids)) - .options( - sa.orm.selectinload(AgentRow.agent_resource_rows).joinedload( - AgentResourceRow.slot_type_row - ) - ) .order_by( AgentRow.id, ) @@ -800,7 +798,13 @@ async def batch_load( async with graph_ctx.db.begin_readonly_session() as session: result = await session.scalars(query) agent_list = result.unique().all() - return [cls.from_data(agent.to_data()) for agent in agent_list] + occupied_slots = await fetch_actual_occupied_slots( + session, [AgentId(agent.id) for agent in agent_list] + ) + return [ + cls.from_data(agent.to_data(occupied_slots[AgentId(agent.id)])) + for agent in agent_list + ] @classmethod async def load_count( diff --git a/src/ai/backend/manager/models/agent/row.py b/src/ai/backend/manager/models/agent/row.py index 90ff834ccbc..c805bab076f 100644 --- a/src/ai/backend/manager/models/agent/row.py +++ b/src/ai/backend/manager/models/agent/row.py @@ -14,7 +14,6 @@ joinedload, load_only, mapped_column, - relationship, selectinload, ) from sqlalchemy.sql.expression import false, true @@ -132,16 +131,7 @@ class AgentRow(Base): default=False, ) - agent_resource_rows: Mapped[list[AgentResourceRow]] = relationship("AgentResourceRow") - - def actual_occupied_slots(self) -> ResourceSlot: - occupied = ResourceSlot() - sorted_rows = sorted(self.agent_resource_rows, key=lambda r: r.slot_type_row.rank) - for resource_row in sorted_rows: - occupied[resource_row.slot_name] = resource_row.used - return occupied - - def to_data(self) -> AgentData: + def to_data(self, actual_occupied_slots: ResourceSlot) -> AgentData: return AgentData( id=AgentId(self.id), status=self.status, @@ -151,7 +141,7 @@ def to_data(self) -> AgentData: schedulable=self.schedulable, available_slots=self.available_slots, cached_occupied_slots=self.occupied_slots, - actual_occupied_slots=self.actual_occupied_slots(), + actual_occupied_slots=actual_occupied_slots, addr=self.addr, public_host=self.public_host, first_contact=self.first_contact, diff --git a/src/ai/backend/manager/repositories/agent/db_source/db_source.py b/src/ai/backend/manager/repositories/agent/db_source/db_source.py index 5ba1c835abb..9fa1baf60b0 100644 --- a/src/ai/backend/manager/repositories/agent/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/agent/db_source/db_source.py @@ -31,6 +31,7 @@ from ai.backend.manager.models.resource_slot import AgentResourceRow from ai.backend.manager.models.scaling_group import ScalingGroupRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.repositories.agent.query import fetch_actual_occupied_slots from ai.backend.manager.repositories.agent.updaters import AgentStatusUpdaterSpec from ai.backend.manager.repositories.base import BulkUpserter, execute_bulk_upserter from ai.backend.manager.repositories.base.querier import BatchQuerier, execute_batch_querier @@ -85,18 +86,13 @@ async def get_images_by_digest(self, digests: list[str]) -> dict[ImageID, ImageD async def get_by_id(self, agent_id: AgentId) -> AgentData: async with self._db.begin_readonly_session_read_committed() as db_session: agent_row: AgentRow | None = await db_session.scalar( - sa.select(AgentRow) - .where(AgentRow.id == agent_id) - .options( - selectinload(AgentRow.agent_resource_rows).joinedload( - AgentResourceRow.slot_type_row - ) - ) + sa.select(AgentRow).where(AgentRow.id == agent_id) ) if agent_row is None: log.error("Agent with id {} not found", agent_id) raise AgentNotFound(f"Agent with id {agent_id} not found") - return agent_row.to_data() + occupied_slots = await fetch_actual_occupied_slots(db_session, [agent_id]) + return agent_row.to_data(occupied_slots[agent_id]) async def upsert_agent_with_state(self, upsert_data: AgentHeartbeatUpsert) -> UpsertResult: async with self._db.begin_session_read_committed() as session: @@ -260,11 +256,7 @@ async def search_agents( """Searches agents with total count.""" async with self._db.begin_readonly_session() as db_sess: - query = sa.select(AgentRow).options( - selectinload(AgentRow.agent_resource_rows).joinedload( - AgentResourceRow.slot_type_row - ), - ) + query = sa.select(AgentRow) result = await execute_batch_querier( db_sess, @@ -272,7 +264,12 @@ async def search_agents( querier, ) agent_rows: list[AgentRow] = [row.AgentRow for row in result.rows] - items = [agent_row.to_data() for agent_row in agent_rows] + occupied_slots = await fetch_actual_occupied_slots( + db_sess, [AgentId(row.id) for row in agent_rows] + ) + items = [ + agent_row.to_data(occupied_slots[AgentId(agent_row.id)]) for agent_row in agent_rows + ] admin_permissions = list(ADMIN_AGENT_PERMISSIONS) agents_with_permissions = [ AgentDetailData(agent=agent_data, permissions=admin_permissions) diff --git a/src/ai/backend/manager/repositories/agent/query.py b/src/ai/backend/manager/repositories/agent/query.py index 3ff877ddd15..2e7a29032bd 100644 --- a/src/ai/backend/manager/repositories/agent/query.py +++ b/src/ai/backend/manager/repositories/agent/query.py @@ -1,13 +1,39 @@ from collections.abc import Collection import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncSession as SASession from ai.backend.common.data.filter_specs import StringMatchSpec -from ai.backend.common.types import AgentId +from ai.backend.common.types import AgentId, ResourceSlot from ai.backend.manager.data.agent.types import AgentStatus from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.clauses import QueryCondition, QueryOrder from ai.backend.manager.models.condition_utils import make_string_in_factory +from ai.backend.manager.models.resource_slot import AgentResourceRow, ResourceSlotTypeRow + + +async def fetch_actual_occupied_slots( + db_session: SASession, + agent_ids: Collection[AgentId], +) -> dict[AgentId, ResourceSlot]: + """Load per-agent occupied slots from ``agent_resources``, keyed in slot type rank order.""" + occupied: dict[AgentId, ResourceSlot] = { + AgentId(agent_id): ResourceSlot() for agent_id in agent_ids + } + if not occupied: + return occupied + stmt = ( + sa.select(AgentResourceRow.agent_id, AgentResourceRow.slot_name, AgentResourceRow.used) + .join( + ResourceSlotTypeRow, + AgentResourceRow.slot_name == ResourceSlotTypeRow.slot_name, + ) + .where(AgentResourceRow.agent_id.in_(occupied.keys())) + .order_by(ResourceSlotTypeRow.rank) + ) + for row in await db_session.execute(stmt): + occupied[AgentId(row.agent_id)][row.slot_name] = row.used + return occupied class QueryConditions: diff --git a/src/ai/backend/manager/repositories/agent/repository.py b/src/ai/backend/manager/repositories/agent/repository.py index e74aa3279ca..6363628d03e 100644 --- a/src/ai/backend/manager/repositories/agent/repository.py +++ b/src/ai/backend/manager/repositories/agent/repository.py @@ -27,10 +27,10 @@ from ai.backend.manager.data.kernel.types import KernelInfo from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.clauses import QueryCondition, QueryOrder -from ai.backend.manager.models.resource_slot import AgentResourceRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.repositories.agent.cache_source.cache_source import AgentCacheSource from ai.backend.manager.repositories.agent.db_source.db_source import AgentDBSource +from ai.backend.manager.repositories.agent.query import fetch_actual_occupied_slots from ai.backend.manager.repositories.agent.stateful_source.stateful_source import ( AgentStatefulSource, ) @@ -212,11 +212,7 @@ async def list_data( conditions: Sequence[QueryCondition], order_by: Sequence[QueryOrder] = tuple(), ) -> list[AgentData]: - stmt: sa.sql.Select[Any] = sa.select(AgentRow).options( - sa.orm.selectinload(AgentRow.agent_resource_rows).joinedload( - AgentResourceRow.slot_type_row - ), - ) + stmt: sa.sql.Select[Any] = sa.select(AgentRow) for cond in conditions: stmt = stmt.where(cond()) @@ -226,7 +222,12 @@ async def list_data( async with self._db_source._db.begin_readonly_session() as db_session: result = await db_session.scalars(stmt) agent_rows = cast(list[AgentRow], result.unique().all()) - return [agent_row.to_data() for agent_row in agent_rows] + occupied_slots = await fetch_actual_occupied_slots( + db_session, [AgentId(row.id) for row in agent_rows] + ) + return [ + agent_row.to_data(occupied_slots[AgentId(agent_row.id)]) for agent_row in agent_rows + ] @agent_repository_resilience.apply() async def update_gpu_alloc_map(self, agent_id: AgentId, alloc_map: Mapping[str, Any]) -> None: diff --git a/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py b/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py index d1ea5e753bf..038935c6bd4 100644 --- a/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py @@ -626,9 +626,7 @@ async def _fetch_agents( from the normalized ``agent_resources`` table.""" agent_rows = ( await db_sess.scalars( - sa.select(AgentRow) - .options(selectinload(AgentRow.agent_resource_rows)) - .where( + sa.select(AgentRow).where( sa.and_( AgentRow.status == AgentStatus.ALIVE, AgentRow.resource_group_id == resource_group_id, @@ -638,19 +636,14 @@ async def _fetch_agents( ) ).all() container_counts = await self._fetch_agent_container_counts(db_sess, resource_group_id) + slots_by_agent = await self._fetch_agent_slot_resources( + db_sess, [AgentId(agent_row.id) for agent_row in agent_rows] + ) agents = [] for agent_row in agent_rows: - slots = { - ResourceSlotName(ar.slot_name): SlotResource( - capacity=ar.capacity, - # Advance reservations occupy capacity for scheduling - reserved=ar.reserved + ar.prereserved, - used=ar.used, - ) - for ar in agent_row.agent_resource_rows - } agent_id = AgentId(agent_row.id) + slots = slots_by_agent.get(agent_id, {}) agents.append( AgentMeta( id=agent_id, @@ -662,6 +655,27 @@ async def _fetch_agents( ) return agents + async def _fetch_agent_slot_resources( + self, db_sess: SASession, agent_ids: Sequence[AgentId] + ) -> dict[AgentId, dict[ResourceSlotName, SlotResource]]: + """Load per-agent, per-slot capacity/reserved/used from ``agent_resources``.""" + if not agent_ids: + return {} + rows = ( + await db_sess.scalars( + sa.select(AgentResourceRow).where(AgentResourceRow.agent_id.in_(agent_ids)) + ) + ).all() + slots_by_agent: dict[AgentId, dict[ResourceSlotName, SlotResource]] = defaultdict(dict) + for ar in rows: + slots_by_agent[AgentId(ar.agent_id)][ResourceSlotName(ar.slot_name)] = SlotResource( + capacity=ar.capacity, + # Advance reservations occupy capacity for scheduling + reserved=ar.reserved + ar.prereserved, + used=ar.used, + ) + return dict(slots_by_agent) + async def _fetch_agent_container_counts( self, db_sess: SASession, resource_group_id: ResourceGroupID ) -> dict[AgentId, int]: diff --git a/tests/unit/manager/models/resource_slot/test_resource_slot_tables.py b/tests/unit/manager/models/resource_slot/test_resource_slot_tables.py index 97359faad15..6d3cbc073d9 100644 --- a/tests/unit/manager/models/resource_slot/test_resource_slot_tables.py +++ b/tests/unit/manager/models/resource_slot/test_resource_slot_tables.py @@ -13,9 +13,8 @@ import pytest import sqlalchemy as sa -from sqlalchemy.orm import selectinload -from ai.backend.manager.models.agent import AgentRow +from ai.backend.common.types import AgentId from ai.backend.manager.models.resource_slot import ( AgentResourceRow, NumberFormat, @@ -23,6 +22,7 @@ ResourceSlotTypeRow, ) from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.repositories.agent.query import fetch_actual_occupied_slots class TestResourceSlotTypeRow: @@ -532,16 +532,6 @@ async def test_occupied_slots_sorted_by_rank( await db_sess.flush() async with database_with_resource_slot_tables.begin_readonly_session() as db_sess: - agent_row = await db_sess.scalar( - sa.select(AgentRow) - .where(AgentRow.id == agent_id) - .options( - selectinload(AgentRow.agent_resource_rows).joinedload( - AgentResourceRow.slot_type_row - ) - ) - ) - assert agent_row is not None - occupied = agent_row.actual_occupied_slots() + occupied_by_agent = await fetch_actual_occupied_slots(db_sess, [AgentId(agent_id)]) - assert list(occupied.keys()) == expected_order + assert list(occupied_by_agent[AgentId(agent_id)].keys()) == expected_order From e381abd11d2e0bf55abf57c5610bb7276a872cd8 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Sun, 16 Aug 2026 02:48:29 +0900 Subject: [PATCH 2/5] refactor(BA-7226): replace KernelRow scalar relationship usages with explicit queries - ComputeContainer.imageObject resolves through a new id-keyed ImageNode dataloader instead of the KernelRow.image_row join; group_row/user_row were eager-loaded but never read, so those loader options are gone. - KernelRow.get_kernel() filters live agents with an explicit join on agents. - get_direct_access_info() reads the agent public host through the session repository rather than walking main_kernel.agent_row. - Drop the now-dead agent_row eager loads in the registry, the stream/session db_sources and the session kernel loading strategies. Co-Authored-By: Claude Opus 5 (1M context) --- .../backend/manager/api/gql_legacy/image.py | 23 +++++++++ .../backend/manager/api/gql_legacy/kernel.py | 35 ++++++------- src/ai/backend/manager/models/kernel/row.py | 50 +++---------------- src/ai/backend/manager/models/session/row.py | 20 ++------ src/ai/backend/manager/registry.py | 5 +- .../session/db_source/db_source.py | 13 +++-- .../repositories/session/repository.py | 6 ++- .../stream/db_source/db_source.py | 6 +-- .../manager/services/session/service.py | 5 +- .../session/test_session_lifecycle_service.py | 8 +-- 10 files changed, 76 insertions(+), 95 deletions(-) diff --git a/src/ai/backend/manager/api/gql_legacy/image.py b/src/ai/backend/manager/api/gql_legacy/image.py index fba9a947bfa..3d4d202f0ff 100644 --- a/src/ai/backend/manager/api/gql_legacy/image.py +++ b/src/ai/backend/manager/api/gql_legacy/image.py @@ -103,6 +103,7 @@ ResourceLimit, ResourceLimitInput, batch_multiresult_in_scalar_stream, + batch_result_in_scalar_stream, extract_object_uuid, generate_sql_info_for_gql_connection, ) @@ -500,6 +501,28 @@ async def batch_load_by_name_and_arch( lambda row: (row.name, row.architecture), ) + @classmethod + async def batch_load_by_ids( + cls, + graph_ctx: GraphQueryContext, + image_ids: Sequence[ImageID], + ) -> Sequence[ImageNode | None]: + """Load images by row id, regardless of status.""" + query = ( + sa.select(ImageRow) + .where(ImageRow.id.in_(image_ids)) + .options(selectinload(ImageRow.aliases)) + ) + async with graph_ctx.db.begin_readonly_session() as db_session: + return await batch_result_in_scalar_stream( + graph_ctx, + db_session, + query, + cls, + image_ids, + lambda row: row.id, + ) + @classmethod async def batch_load_by_image_identifier( cls, diff --git a/src/ai/backend/manager/api/gql_legacy/kernel.py b/src/ai/backend/manager/api/gql_legacy/kernel.py index a37eeb70d42..8605bb74ce4 100644 --- a/src/ai/backend/manager/api/gql_legacy/kernel.py +++ b/src/ai/backend/manager/api/gql_legacy/kernel.py @@ -15,12 +15,13 @@ from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime from sqlalchemy.engine.row import Row -from sqlalchemy.orm import noload, selectinload +from sqlalchemy.orm import noload from ai.backend.common.types import ( AccessKey, AgentId, BinarySize, + ImageID, KernelId, SessionId, ) @@ -28,7 +29,6 @@ from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.models.group import groups -from ai.backend.manager.models.image import ImageRow from ai.backend.manager.models.kernel import ( AGENT_RESOURCE_OCCUPYING_KERNEL_STATUSES, DEFAULT_KERNEL_ORDERING, @@ -220,6 +220,8 @@ class ComputeContainer(graphene.ObjectType): # type: ignore[misc] class Meta: interfaces = (Item,) + _image_id: ImageID | None = None + # identity idx = graphene.Int() # legacy role = graphene.String() # legacy @@ -286,7 +288,6 @@ def parse_row(cls, ctx: GraphQueryContext, row: KernelRow) -> Mapping[str, Any]: "session_id": row.session_id, # image "image": row.image, - "image_object": ImageNode.from_row(ctx, row.image_row), "architecture": row.architecture, "registry": row.registry, # status @@ -314,7 +315,18 @@ def from_row(cls, ctx: GraphQueryContext, row: KernelRow | None) -> ComputeConta if row is None: return None props = cls.parse_row(ctx, row) - return cls(**props) + obj = cls(**props) + obj._image_id = row.image_id + return obj + + async def resolve_image_object(self, info: graphene.ResolveInfo) -> ImageNode | None: + if self._image_id is None: + return None + graph_ctx: GraphQueryContext = info.context + loader = graph_ctx.dataloader_manager.get_loader_by_func( + graph_ctx, ImageNode.batch_load_by_ids + ) + return cast(ImageNode | None, await loader.load(self._image_id)) # last_stat also fetches data from Redis, meaning that # both live_stat and last_stat will reference same data from same source @@ -425,7 +437,6 @@ async def load_slice( .where(KernelRow.session_id == session_id) .limit(limit) .offset(offset) - .options(selectinload(KernelRow.image_row).options(selectinload(ImageRow.aliases))) ) if cluster_role is not None: query = query.where(KernelRow.cluster_role == cluster_role) @@ -456,7 +467,6 @@ async def batch_load_by_session( sa.select(KernelRow) # TODO: use "owner session ID" when we implement multi-container session .where(KernelRow.session_id.in_(session_ids)) - .options(selectinload(KernelRow.image_row).options(selectinload(ImageRow.aliases))) ) async with ctx.db.begin_readonly_session() as conn: return await batch_multiresult( @@ -476,11 +486,7 @@ async def batch_load_by_agent_id( *, status: KernelStatus | None = None, ) -> Sequence[Sequence[ComputeContainer]]: - query_stmt = ( - sa.select(KernelRow) - .where(KernelRow.agent.in_(agent_ids)) - .options(selectinload(KernelRow.image_row).options(selectinload(ImageRow.aliases))) - ) + query_stmt = sa.select(KernelRow).where(KernelRow.agent.in_(agent_ids)) kernel_status: tuple[KernelStatus, ...] if status is not None: kernel_status = (status,) @@ -511,12 +517,7 @@ async def batch_load_detail( .where( (KernelRow.id.in_(container_ids)), ) - .options( - noload("*"), - selectinload(KernelRow.group_row), - selectinload(KernelRow.user_row), - selectinload(KernelRow.image_row), - ) + .options(noload("*")) ) if domain_name is not None: query = query.where(KernelRow.domain_name == domain_name) diff --git a/src/ai/backend/manager/models/kernel/row.py b/src/ai/backend/manager/models/kernel/row.py index e80264da690..5f187544ce6 100644 --- a/src/ai/backend/manager/models/kernel/row.py +++ b/src/ai/backend/manager/models/kernel/row.py @@ -15,11 +15,9 @@ from sqlalchemy.ext.asyncio import AsyncSession as SASession from sqlalchemy.orm import ( Mapped, - foreign, mapped_column, noload, relationship, - selectinload, ) from ai.backend.common.identifier.image import ImageID @@ -52,11 +50,7 @@ ) if TYPE_CHECKING: - from ai.backend.manager.models.agent import AgentRow - from ai.backend.manager.models.group import GroupRow - from ai.backend.manager.models.image import ImageRow from ai.backend.manager.models.session import SessionRow - from ai.backend.manager.models.user import UserRow from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.errors.kernel import SessionNotFound @@ -131,13 +125,6 @@ def default_hostname(context: Any) -> str: return f"{params['cluster_role']}{params['cluster_idx']}" -# Defined for avoiding circular import -def _get_user_row_join_condition() -> sa.sql.elements.ColumnElement[Any]: - from ai.backend.manager.models.user import UserRow - - return UserRow.uuid == foreign(KernelRow.user_uuid) - - class KernelRow(CreatedAtMixin, Base): __tablename__ = "kernels" @@ -431,17 +418,6 @@ class KernelRow(CreatedAtMixin, Base): ) session: Mapped[SessionRow] = relationship("SessionRow", back_populates="kernels") - image_row: Mapped[ImageRow | None] = relationship( - "ImageRow", - foreign_keys="KernelRow.image_id", - ) - agent_row: Mapped[AgentRow | None] = relationship("AgentRow") - group_row: Mapped[GroupRow] = relationship("GroupRow") - user_row: Mapped[UserRow] = relationship( - "UserRow", - primaryjoin=_get_user_row_join_condition, - foreign_keys="KernelRow.user_uuid", - ) @property def used_time(self) -> str | None: @@ -473,28 +449,18 @@ async def batch_load_main_kernels_by_session_id( async def get_kernel( db: ExtendedAsyncSAEngine, kern_id: uuid.UUID, allow_stale: bool = False ) -> KernelRow: - from ai.backend.manager.models.agent import AgentStatus + from ai.backend.manager.models.agent import AgentRow, AgentStatus async def _query() -> KernelRow: async with db.begin_readonly_session() as db_sess: - query = ( - sa.select(KernelRow) - .where(KernelRow.id == kern_id) - .options( - noload("*"), - selectinload(KernelRow.agent_row).options(noload("*")), - ) - ) - result = (await db_sess.execute(query)).scalars().all() - - cand = result + query = sa.select(KernelRow).where(KernelRow.id == kern_id).options(noload("*")) if not allow_stale: - cand = [ - k - for k in result - if (k.status not in DEAD_KERNEL_STATUSES) - and (k.agent_row is not None and k.agent_row.status == AgentStatus.ALIVE) - ] + # An inner join drops kernels with no agent assigned yet. + query = query.join(AgentRow, KernelRow.agent == AgentRow.id).where( + KernelRow.status.not_in(DEAD_KERNEL_STATUSES), + AgentRow.status == AgentStatus.ALIVE, + ) + cand = (await db_sess.execute(query)).scalars().all() if not cand: raise SessionNotFound return cand[0] diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index e33fd6560e4..52979644662 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -1039,20 +1039,14 @@ async def get_session( case KernelLoadingStrategy.ALL_KERNELS: _eager_loading_op.extend([ noload("*"), - selectinload(SessionRow.kernels).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(SessionRow.kernels).options(noload("*")), ]) case KernelLoadingStrategy.MAIN_KERNEL_ONLY: kernel_rel = SessionRow.kernels kernel_rel.and_(KernelRow.cluster_role == DEFAULT_ROLE) _eager_loading_op.extend([ noload("*"), - selectinload(kernel_rel).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(kernel_rel).options(noload("*")), ]) _eager_loading_op.append(joinedload(SessionRow.user)) @@ -1097,20 +1091,14 @@ async def list_sessions( case KernelLoadingStrategy.ALL_KERNELS: _eager_loading_op.extend([ noload("*"), - selectinload(SessionRow.kernels).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(SessionRow.kernels).options(noload("*")), ]) case KernelLoadingStrategy.MAIN_KERNEL_ONLY: kernel_rel = SessionRow.kernels kernel_rel.and_(KernelRow.cluster_role == DEFAULT_ROLE) _eager_loading_op.extend([ noload("*"), - selectinload(kernel_rel).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(kernel_rel).options(noload("*")), ]) session_list = await cls.match_sessions( diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py index 05862bae612..7336e10fab0 100644 --- a/src/ai/backend/manager/registry.py +++ b/src/ai/backend/manager/registry.py @@ -1456,10 +1456,7 @@ async def _fetch_session() -> tuple[SessionRow, str | None]: session_id, eager_loading_op=( noload("*"), - selectinload(SessionRow.kernels).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(SessionRow.kernels).options(noload("*")), ), ) network_ref_name = await sess.get_network_ref(db_sess) 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 329a78e46e4..b002dec0cf5 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 @@ -37,6 +37,7 @@ SessionNotFound, TooManySessionsMatched, ) +from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.group import groups from ai.backend.manager.models.image import ImageRow @@ -167,6 +168,13 @@ async def get_session_validated( eager_loading_op=list(eager_loading_op) if eager_loading_op else None, ) + async def get_agent_public_host(self, agent_id: AgentId) -> str | None: + """Look up the public host of an agent, or None when the agent is unknown.""" + async with self._db.begin_readonly_session_read_committed() as db_sess: + return await db_sess.scalar( + sa.select(AgentRow.public_host).where(AgentRow.id == agent_id) + ) + async def match_sessions( self, id_or_name_prefix: str, @@ -579,10 +587,7 @@ async def get_session_with_routing_minimal( noload("*"), selectinload( SessionRow.kernels.and_(KernelRow.cluster_role == DEFAULT_ROLE) - ).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + ).options(noload("*")), joinedload(SessionRow.user), ) ) diff --git a/src/ai/backend/manager/repositories/session/repository.py b/src/ai/backend/manager/repositories/session/repository.py index bcfba4d0380..c6663629a9e 100644 --- a/src/ai/backend/manager/repositories/session/repository.py +++ b/src/ai/backend/manager/repositories/session/repository.py @@ -14,7 +14,7 @@ from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy from ai.backend.common.resilience.resilience import Resilience -from ai.backend.common.types import AccessKey, ImageAlias, KernelId, SessionId +from ai.backend.common.types import AccessKey, AgentId, ImageAlias, KernelId, SessionId from ai.backend.manager.data.image.types import ImageIdentifier from ai.backend.manager.data.kernel.types import KernelListResult from ai.backend.manager.data.resource_slot.types import ResourceAllocationAggregate @@ -99,6 +99,10 @@ async def get_session_validated( eager_loading_op, ) + @session_repository_resilience.apply() + async def get_agent_public_host(self, agent_id: AgentId) -> str | None: + return await self._db_source.get_agent_public_host(agent_id) + @session_repository_resilience.apply() async def match_sessions( self, diff --git a/src/ai/backend/manager/repositories/stream/db_source/db_source.py b/src/ai/backend/manager/repositories/stream/db_source/db_source.py index 1b88a2c5b97..7c0e2bedb2b 100644 --- a/src/ai/backend/manager/repositories/stream/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/stream/db_source/db_source.py @@ -5,7 +5,6 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.errors.kernel import SessionNotFound, TooManySessionsMatched -from ai.backend.manager.models.kernel import KernelRow from ai.backend.manager.models.session import SessionRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine @@ -31,10 +30,7 @@ async def get_streaming_session( ) .options( noload("*"), - selectinload(SessionRow.kernels).options( - noload("*"), - selectinload(KernelRow.agent_row).noload("*"), - ), + selectinload(SessionRow.kernels).options(noload("*")), joinedload(SessionRow.user), ) .execution_options(populate_existing=True) diff --git a/src/ai/backend/manager/services/session/service.py b/src/ai/backend/manager/services/session/service.py index ee9cecc6ee3..497305810f5 100644 --- a/src/ai/backend/manager/services/session/service.py +++ b/src/ai/backend/manager/services/session/service.py @@ -1181,11 +1181,12 @@ async def get_direct_access_info( resp = {} sess_type = sess.session_type if sess_type in PRIVATE_SESSION_TYPES: - if sess.main_kernel.agent_row is None: + agent_id = sess.main_kernel.agent + if agent_id is None: raise KernelNotReady( f"Kernel of the session has no agent info yet (kernel: {sess.main_kernel.id}, kernel status: {sess.main_kernel.status.name})" ) - public_host = sess.main_kernel.agent_row.public_host + public_host = await self._session_repository.get_agent_public_host(AgentId(agent_id)) found_ports: dict[str, list[str]] = {} service_ports = sess.main_kernel.service_ports if service_ports is None: diff --git a/tests/unit/manager/services/session/test_session_lifecycle_service.py b/tests/unit/manager/services/session/test_session_lifecycle_service.py index 2ee6de25c6b..7c2d24ce6d8 100644 --- a/tests/unit/manager/services/session/test_session_lifecycle_service.py +++ b/tests/unit/manager/services/session/test_session_lifecycle_service.py @@ -286,8 +286,7 @@ def _make_mock_session( session.main_kernel.id = kernel_id session.main_kernel.status = MagicMock() session.main_kernel.status.name = "RUNNING" - session.main_kernel.agent_row = MagicMock() - session.main_kernel.agent_row.public_host = "10.0.0.1" + session.main_kernel.agent = "i-agent01" session.main_kernel.agent_addr = "tcp://10.0.0.1:6001" session.main_kernel.kernel_host = "10.0.0.1" session.main_kernel.service_ports = [ @@ -1784,6 +1783,7 @@ async def test_system_session_returns_ports( session_type=SessionTypes.SYSTEM, ) mock_session_repository.get_session_validated = AsyncMock(return_value=mock_session) + mock_session_repository.get_agent_public_host = AsyncMock(return_value="10.0.0.1") action = GetDirectAccessInfoAction( session_name="system-session", @@ -1823,7 +1823,7 @@ async def test_interactive_session_returns_empty_dict( assert result.result == {} - async def test_agent_row_none_raises_kernel_not_ready( + async def test_unassigned_agent_raises_kernel_not_ready( self, session_service: SessionService, mock_session_repository: MagicMock, @@ -1841,7 +1841,7 @@ async def test_agent_row_none_raises_kernel_not_ready( sample_kernel_id, session_type=SessionTypes.SYSTEM, ) - mock_session.main_kernel.agent_row = None + mock_session.main_kernel.agent = None mock_session_repository.get_session_validated = AsyncMock(return_value=mock_session) action = GetDirectAccessInfoAction( From 0553db335e92a7ee92ae501164aaa8f961f66d30 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Sun, 16 Aug 2026 02:48:38 +0900 Subject: [PATCH 3/5] refactor(BA-7226): replace SessionRow.user/group usages with explicit queries - ComputeSessionNode.owner resolves through a uuid-keyed UserNode dataloader. - gql_legacy session queries join users/groups explicitly so the full_name/user_email/group_name filters keep working. - The resource-usage report runs one explicit join and returns ResourceUsageRecord instead of walking kernel -> session -> user/group. - convert_session_to_image reads the project container registry through the session repository; get_session_with_group is gone. - join_by_related_field had no callers left and is removed. Co-Authored-By: Claude Opus 5 (1M context) --- .../backend/manager/api/gql_legacy/schema.py | 4 +- .../backend/manager/api/gql_legacy/session.py | 50 ++++--- src/ai/backend/manager/api/gql_legacy/user.py | 18 +++ .../backend/manager/models/resource_usage.py | 123 ++++++++++-------- src/ai/backend/manager/models/session/row.py | 27 +--- src/ai/backend/manager/models/types.py | 4 - .../repositories/group/db_source/db_source.py | 4 +- .../manager/repositories/group/repository.py | 4 +- .../session/db_source/db_source.py | 22 +--- .../repositories/session/repository.py | 13 +- .../stream/db_source/db_source.py | 3 +- .../repositories/user/db_source/db_source.py | 10 +- .../manager/services/session/service.py | 13 +- 13 files changed, 136 insertions(+), 159 deletions(-) diff --git a/src/ai/backend/manager/api/gql_legacy/schema.py b/src/ai/backend/manager/api/gql_legacy/schema.py index b57ce39465d..93156852da5 100644 --- a/src/ai/backend/manager/api/gql_legacy/schema.py +++ b/src/ai/backend/manager/api/gql_legacy/schema.py @@ -17,7 +17,7 @@ from graphql.type import GraphQLField, get_named_type, is_leaf_type from opentelemetry import trace from opentelemetry.trace import StatusCode -from sqlalchemy.orm import joinedload, selectinload +from sqlalchemy.orm import selectinload from ai.backend.common.clients.valkey_client.valkey_image.client import ValkeyImageClient from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient @@ -2615,7 +2615,7 @@ async def resolve_session_pending_queue( stmt = ( sa.select(SessionRow) .where(SessionRow.id.in_(pending_sessions)) - .options(selectinload(SessionRow.kernels), joinedload(SessionRow.user)) + .options(selectinload(SessionRow.kernels)) ) query_result = await db_session.scalars(stmt) for row in query_result: diff --git a/src/ai/backend/manager/api/gql_legacy/session.py b/src/ai/backend/manager/api/gql_legacy/session.py index 2d54e8cea66..b9b62cb292d 100644 --- a/src/ai/backend/manager/api/gql_legacy/session.py +++ b/src/ai/backend/manager/api/gql_legacy/session.py @@ -20,7 +20,7 @@ from dateutil.parser import parse as dtparse from graphene.types.datetime import DateTime as GQLDateTime from sqlalchemy.engine.row import Row -from sqlalchemy.orm import joinedload, selectinload +from sqlalchemy.orm import selectinload from ai.backend.common import validators as tx from ai.backend.common.defs.session import SESSION_PRIORITY_MAX, SESSION_PRIORITY_MIN @@ -67,7 +67,6 @@ from ai.backend.manager.models.types import ( QueryCondition, QueryOption, - join_by_related_field, load_related_field, ) from ai.backend.manager.models.user import UserRole, UserRow @@ -194,6 +193,13 @@ def parse_value(value: str) -> ComputeSessionPermission: return ComputeSessionPermission(value) +def _join_owner_rows(stmt: sa.sql.Select[Any]) -> sa.sql.Select[Any]: + """Join the owning user and project so their columns are filterable/orderable.""" + return stmt.join(UserRow, SessionRow.user_uuid == UserRow.uuid).join( + GroupRow, SessionRow.group_id == GroupRow.id + ) + + class _HasVFID(Protocol): vfid: VFolderID @@ -310,7 +316,7 @@ async def get_node( stmt = ( sa.select(SessionRow) .where(SessionRow.id == uuid.UUID(raw_session_id)) - .options(selectinload(SessionRow.kernels), joinedload(SessionRow.user)) + .options(selectinload(SessionRow.kernels)) ) query_result = await db_session.scalar(stmt) if query_result is None: @@ -322,21 +328,21 @@ async def get_node( def _add_basic_options_to_query( cls, stmt: sa.sql.Select[Any], is_count: bool = False ) -> sa.sql.Select[Any]: - options = [ - join_by_related_field(SessionRow.user), - join_by_related_field(SessionRow.group), - ] + # The user and project joins back the `full_name`/`user_email`/`group_name` filters. + stmt = _join_owner_rows(stmt) if not is_count: - options = [ - *options, - load_related_field(SessionRow.kernel_load_option()), - load_related_field(SessionRow.user_load_option(already_joined=True)), - load_related_field(SessionRow.project_load_option(already_joined=True)), - ] - for option in options: - stmt = option(stmt) + stmt = stmt.options(SessionRow.kernel_load_option()) return stmt + async def resolve_owner(self, info: graphene.ResolveInfo) -> UserNode | None: + if self.user_id is None: + return None + graph_ctx: GraphQueryContext = info.context + loader = graph_ctx.dataloader_manager.get_loader_by_func( + graph_ctx, UserNode.batch_load_by_uuids + ) + return cast(UserNode | None, await loader.load(self.user_id)) + async def resolve_queue_position(self, info: graphene.ResolveInfo) -> int | None: if self.status != SessionStatus.PENDING: return None @@ -377,7 +383,6 @@ def from_row( project_id=row.group_id, user_id=row.user_uuid, access_key=row.access_key, - owner=UserNode.from_row(ctx, row.user), # status status=row.status.name, # status_changed=row.status_changed, # FIXME: generated attribute @@ -592,13 +597,7 @@ async def resolve_graph( query = sa.select(dependency_cte.c.id) session_ids = (await db_sess.execute(query)).scalars().all() # Get the session rows in the graph - query = ( - sa.select(SessionRow) - .where(SessionRow.id.in_(session_ids)) - .options( - selectinload(SessionRow.user), - ) - ) + query = sa.select(SessionRow).where(SessionRow.id.in_(session_ids)) session_rows = list((await db_sess.execute(query)).scalars().all()) await batch_populate_session_occupied_slots(db_sess, session_rows) @@ -653,7 +652,6 @@ async def batch_load_by_dependee_id( .where(SessionRow.id.in_(dependent_ids)) .options( selectinload(SessionRow.kernels.and_(KernelRow.cluster_role == DEFAULT_ROLE)), - joinedload(SessionRow.user), ) ) rows = list((await db_sess.execute(sess_query)).unique().scalars().all()) @@ -698,7 +696,6 @@ async def batch_load_by_dependent_id( .where(SessionRow.id.in_(dependee_ids)) .options( selectinload(SessionRow.kernels.and_(KernelRow.cluster_role == DEFAULT_ROLE)), - joinedload(SessionRow.user), ) ) rows = list((await db_sess.execute(sess_query)).unique().scalars().all()) @@ -848,8 +845,7 @@ async def get_data( query_conditions.append(by_resource_group_name(resource_group_name)) query_options: list[QueryOption] = [ load_related_field(SessionRow.kernel_load_option()), - join_by_related_field(SessionRow.user), - join_by_related_field(SessionRow.group), + _join_owner_rows, ] session_rows = await SessionRow.list_session_by_condition( query_conditions, query_options, db=ctx.db diff --git a/src/ai/backend/manager/api/gql_legacy/user.py b/src/ai/backend/manager/api/gql_legacy/user.py index 7584dff210b..3d524177748 100644 --- a/src/ai/backend/manager/api/gql_legacy/user.py +++ b/src/ai/backend/manager/api/gql_legacy/user.py @@ -82,6 +82,7 @@ PaginatedList, batch_multiresult, batch_result, + batch_result_in_scalar_stream, generate_sql_info_for_gql_connection, ) from .gql_relay import AsyncNode, Connection, ConnectionResolverResult @@ -174,6 +175,23 @@ class Meta: description="Added in 25.5.0.", ) + @classmethod + async def batch_load_by_uuids( + cls, + ctx: GraphQueryContext, + user_uuids: Sequence[UUID], + ) -> Sequence[Self | None]: + query = sa.select(UserRow).where(UserRow.uuid.in_(user_uuids)) + async with ctx.db.begin_readonly_session() as db_session: + return await batch_result_in_scalar_stream( + ctx, + db_session, + query, + cls, + user_uuids, + lambda row: row.uuid, + ) + @classmethod def from_row(cls, ctx: GraphQueryContext, row: UserRow) -> Self: return cls( diff --git a/src/ai/backend/manager/models/resource_usage.py b/src/ai/backend/manager/models/resource_usage.py index a4454bf1ac0..67b007101ac 100644 --- a/src/ai/backend/manager/models/resource_usage.py +++ b/src/ai/backend/manager/models/resource_usage.py @@ -10,7 +10,7 @@ import attrs import msgpack import sqlalchemy as sa -from sqlalchemy.orm import joinedload, load_only +from sqlalchemy.orm import Load from sqlalchemy.sql.elements import ColumnElement from ai.backend.common.types import SlotName @@ -30,6 +30,7 @@ "BaseResourceUsageGroup", "ResourceGroupUnit", "ResourceUsage", + "ResourceUsageRecord", "fetch_resource_usage", "parse_resource_usage", "parse_resource_usage_groups", @@ -44,6 +45,17 @@ class ResourceGroupUnit(StrEnum): TOTAL = "total" +@attrs.define(slots=True, kw_only=True) +class ResourceUsageRecord: + """A kernel joined with the session, project and user attributes its usage report needs.""" + + kernel: KernelRow + session: SessionRow + project: GroupRow + user_email: str | None + user_full_name: str | None + + @attrs.define(slots=True) class ResourceUsage: nfs: set[Any] = attrs.field(factory=set) @@ -493,12 +505,12 @@ def parse_resource_usage( async def parse_resource_usage_groups( - kernels: list[KernelRow], + records: list[ResourceUsageRecord], valkey_stat_client: ValkeyStatClient, local_tz: tzinfo, ) -> list[BaseResourceUsageGroup]: - stat_map = {k.id: k.last_stat for k in kernels} - stat_empty_kerns = [k.id for k in kernels if not k.last_stat] + stat_map = {r.kernel.id: r.kernel.last_stat for r in records} + stat_empty_kerns = [r.kernel.id for r in records if not r.kernel.last_stat] kernel_ids_str = [str(kern_id) for kern_id in stat_empty_kerns] raw_stats = await valkey_stat_client.get_user_kernel_statistics_batch(kernel_ids_str) @@ -509,40 +521,40 @@ async def parse_resource_usage_groups( return [ BaseResourceUsageGroup( - kernel_row=kern, - project_row=kern.session.group, - session_row=kern.session, - created_at=kern.created_at, - terminated_at=kern.terminated_at, + kernel_row=record.kernel, + project_row=record.project, + session_row=record.session, + created_at=record.kernel.created_at, + terminated_at=record.kernel.terminated_at, scheduled_at=( - kern.status_history.get(KernelStatus.SCHEDULED.name) - if kern.status_history + record.kernel.status_history.get(KernelStatus.SCHEDULED.name) + if record.kernel.status_history else None ), - used_time=kern.used_time, - used_days=kern.get_used_days(local_tz), - 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, - 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, - container_ids={kern.container_id} if kern.container_id else set(), - session_id=kern.session_id, - session_name=kern.session.name, - domain_name=kern.session.domain_name, - full_name=kern.session.user.full_name if kern.session.user is not None else None, - images={kern.image} if kern.image else set(), - agents={kern.agent} if kern.agent else set(), - status=kern.status.name, - status_history=kern.status_history, - cluster_mode=kern.cluster_mode, - status_info=kern.status_info, + used_time=record.kernel.used_time, + used_days=record.kernel.get_used_days(local_tz), + last_stat=stat_map.get(record.kernel.id), + user_id=record.session.user_uuid, + user_email=record.user_email, + access_key=record.session.access_key, + project_id=record.project.id, + project_name=record.project.name, + kernel_id=record.kernel.id, + container_ids={record.kernel.container_id} if record.kernel.container_id else set(), + session_id=record.kernel.session_id, + session_name=record.session.name, + domain_name=record.session.domain_name, + full_name=record.user_full_name, + images={record.kernel.image} if record.kernel.image else set(), + agents={record.kernel.agent} if record.kernel.agent else set(), + status=record.kernel.status.name, + status_history=record.kernel.status_history, + cluster_mode=record.kernel.cluster_mode, + status_info=record.kernel.status_info, group_unit=ResourceGroupUnit.KERNEL, - total_usage=parse_resource_usage(kern, stat_map.get(kern.id)), + total_usage=parse_resource_usage(record.kernel, stat_map.get(record.kernel.id)), ) - for kern in kernels + for record in records ] @@ -594,25 +606,25 @@ def _parse_query( session_cond: ColumnElement[bool] | None = None, project_cond: ColumnElement[bool] | None = None, ) -> sa.sql.Select[Any]: - session_load = joinedload(KernelRow.session) - if session_cond is not None: - session_load = joinedload(KernelRow.session.and_(session_cond)) - - project_load = joinedload(SessionRow.group) - if project_cond is not None: - project_load = joinedload(SessionRow.group.and_(project_cond)) - query = sa.select(KernelRow).options( - load_only(*KERNEL_RESOURCE_SELECT_COLS), - session_load.options( - load_only(*SESSION_RESOURCE_SELECT_COLS), - joinedload(SessionRow.user).options( - load_only(UserRow.email, UserRow.username, UserRow.full_name) - ), - project_load.options(load_only(*PROJECT_RESOURCE_SELECT_COLS)), - ), + query = ( + sa.select(KernelRow, SessionRow, GroupRow, UserRow.email, UserRow.full_name) + .select_from(KernelRow) + .join(SessionRow, KernelRow.session_id == SessionRow.id) + .join(GroupRow, SessionRow.group_id == GroupRow.id) + # Sessions of a purged user keep their user_uuid, so the user side may be missing. + .outerjoin(UserRow, SessionRow.user_uuid == UserRow.uuid) + .options( + Load(KernelRow).load_only(*KERNEL_RESOURCE_SELECT_COLS), + Load(SessionRow).load_only(*SESSION_RESOURCE_SELECT_COLS), + Load(GroupRow).load_only(*PROJECT_RESOURCE_SELECT_COLS), + ) ) if kernel_cond is not None: query = query.where(kernel_cond) + if session_cond is not None: + query = query.where(session_cond) + if project_cond is not None: + query = query.where(project_cond) return query @@ -622,7 +634,7 @@ async def fetch_resource_usage( end_date: datetime, session_ids: Sequence[UUID] | None = None, project_ids: Sequence[UUID] | None = None, -) -> list[KernelRow]: +) -> list[ResourceUsageRecord]: project_cond = None if project_ids: project_cond = GroupRow.id.in_(project_ids) @@ -646,4 +658,13 @@ async def fetch_resource_usage( ) async with db_engine.begin_readonly_session() as db_sess: result = await db_sess.execute(query) - return list(result.scalars().all()) + return [ + ResourceUsageRecord( + kernel=row.KernelRow, + session=row.SessionRow, + project=row.GroupRow, + user_email=row.email, + user_full_name=row.full_name, + ) + for row in result.all() + ] diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index 52979644662..0c75234ebd9 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -23,8 +23,6 @@ from sqlalchemy.orm import ( Mapped, contains_eager, - foreign, - joinedload, load_only, mapped_column, noload, @@ -115,7 +113,7 @@ ) if TYPE_CHECKING: - from ai.backend.manager.models.user import UserRow + pass log = BraceStyleAdapter(logging.getLogger(__spec__.name)) @@ -348,13 +346,6 @@ class KernelLoadingStrategy(enum.StrEnum): } -# Defined for avoiding circular import -def _get_user_row_join_condition() -> sa.sql.elements.ColumnElement[Any]: - from ai.backend.manager.models.user import UserRow - - return UserRow.uuid == foreign(SessionRow.user_uuid) - - class SessionRow(CreatedAtMixin, Base): __tablename__ = "sessions" id: Mapped[SessionId] = mapped_column( @@ -474,16 +465,9 @@ class SessionRow(CreatedAtMixin, Base): group_id: Mapped[UUID] = mapped_column( "group_id", GUID, sa.ForeignKey("groups.id"), nullable=False ) - group: Mapped[GroupRow] = relationship("GroupRow") user_uuid: Mapped[UUID] = mapped_column( "user_uuid", GUID, server_default=sa.text("uuid_generate_v4()"), nullable=False ) - user: Mapped[UserRow] = relationship( - "UserRow", - primaryjoin=_get_user_row_join_condition, - foreign_keys=[user_uuid], - ) - access_key: Mapped[str | None] = mapped_column("access_key", sa.String(length=20)) # `images` stores canonical image name strings for historical audit. @@ -664,14 +648,6 @@ class SessionRow(CreatedAtMixin, Base): def kernel_load_option(cls, already_joined: bool = False) -> _AbstractLoad: return selectinload(cls.kernels) if not already_joined else contains_eager(cls.kernels) - @classmethod - def user_load_option(cls, already_joined: bool = False) -> _AbstractLoad: - return joinedload(cls.user) if not already_joined else contains_eager(cls.user) - - @classmethod - def project_load_option(cls, already_joined: bool = False) -> _AbstractLoad: - return joinedload(cls.group) if not already_joined else contains_eager(cls.group) - @classmethod def from_dataclass(cls, session_data: SessionData) -> SessionRow: vfolder_mounts = [] @@ -1048,7 +1024,6 @@ async def get_session( noload("*"), selectinload(kernel_rel).options(noload("*")), ]) - _eager_loading_op.append(joinedload(SessionRow.user)) session_list = await cls.match_sessions( db_session, diff --git a/src/ai/backend/manager/models/types.py b/src/ai/backend/manager/models/types.py index 51500efda33..a4c88134ff3 100644 --- a/src/ai/backend/manager/models/types.py +++ b/src/ai/backend/manager/models/types.py @@ -13,7 +13,3 @@ def load_related_field(field: _AbstractLoad) -> QueryOption: return lambda stmt: stmt.options(field) - - -def join_by_related_field(field: sa.orm.attributes.InstrumentedAttribute[Any]) -> QueryOption: - return lambda stmt: stmt.join(field) diff --git a/src/ai/backend/manager/repositories/group/db_source/db_source.py b/src/ai/backend/manager/repositories/group/db_source/db_source.py index fbc014b8464..602a768a6d4 100644 --- a/src/ai/backend/manager/repositories/group/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/group/db_source/db_source.py @@ -60,7 +60,7 @@ kernels, ) from ai.backend.manager.models.rbac_models.role import RoleRow -from ai.backend.manager.models.resource_usage import fetch_resource_usage +from ai.backend.manager.models.resource_usage import ResourceUsageRecord, fetch_resource_usage from ai.backend.manager.models.routing import RoutingRow from ai.backend.manager.models.specs.pagination import NoPagination from ai.backend.manager.models.user import UserRow, users @@ -472,7 +472,7 @@ async def fetch_project_resource_usage( start_date: datetime, end_date: datetime, project_ids: Sequence[UUID] | None = None, - ) -> list[KernelRow]: + ) -> list[ResourceUsageRecord]: """Fetch resource usage data for projects.""" return await fetch_resource_usage(self._db, start_date, end_date, project_ids=project_ids) diff --git a/src/ai/backend/manager/repositories/group/repository.py b/src/ai/backend/manager/repositories/group/repository.py index 483bec805f3..5b090c976fa 100644 --- a/src/ai/backend/manager/repositories/group/repository.py +++ b/src/ai/backend/manager/repositories/group/repository.py @@ -22,7 +22,7 @@ from ai.backend.manager.data.user.types import UserData from ai.backend.manager.errors.resource import InvalidUserUpdateMode from ai.backend.manager.models.group.row import GroupRow -from ai.backend.manager.models.kernel import KernelRow +from ai.backend.manager.models.resource_usage import ResourceUsageRecord from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.repositories.base.creator import Creator from ai.backend.manager.repositories.base.querier import BatchQuerier @@ -116,7 +116,7 @@ async def fetch_project_resource_usage( start_date: datetime, end_date: datetime, project_ids: Sequence[UUID] | None = None, - ) -> list[KernelRow]: + ) -> list[ResourceUsageRecord]: """Fetch resource usage data for projects.""" return await self._db_source.fetch_project_resource_usage(start_date, end_date, project_ids) 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 b002dec0cf5..b57d597601d 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 @@ -39,7 +39,7 @@ ) from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.container_registry import ContainerRegistryRow -from ai.backend.manager.models.group import groups +from ai.backend.manager.models.group import GroupRow, groups from ai.backend.manager.models.image import ImageRow from ai.backend.manager.models.kernel import KernelRow from ai.backend.manager.models.keypair import KeyPairRow @@ -552,22 +552,11 @@ async def find_dependency_sessions( access_key, ) - async def get_session_with_group( - self, - session_name_or_id: str | SessionId, - owner_access_key: AccessKey, - kernel_loading_strategy: KernelLoadingStrategy = KernelLoadingStrategy.MAIN_KERNEL_ONLY, - allow_stale: bool = False, - ) -> SessionRow: - """Get session with group information eagerly loaded""" + async def get_project_container_registry(self, group_id: uuid.UUID) -> dict[str, Any] | None: + """Look up the container registry configuration of a project.""" async with self._db.begin_readonly_session_read_committed() as db_sess: - return await SessionRow.get_session( - db_sess, - session_name_or_id, - owner_access_key, - kernel_loading_strategy=kernel_loading_strategy, - allow_stale=allow_stale, - eager_loading_op=[selectinload(SessionRow.group)], + return await db_sess.scalar( + sa.select(GroupRow.container_registry).where(GroupRow.id == group_id) ) async def get_session_with_routing_minimal( @@ -588,7 +577,6 @@ async def get_session_with_routing_minimal( selectinload( SessionRow.kernels.and_(KernelRow.cluster_role == DEFAULT_ROLE) ).options(noload("*")), - joinedload(SessionRow.user), ) ) async with self._ops.read_ops() as r: diff --git a/src/ai/backend/manager/repositories/session/repository.py b/src/ai/backend/manager/repositories/session/repository.py index c6663629a9e..dce986f4095 100644 --- a/src/ai/backend/manager/repositories/session/repository.py +++ b/src/ai/backend/manager/repositories/session/repository.py @@ -260,17 +260,8 @@ async def find_dependency_sessions( return await self._db_source.find_dependency_sessions(session_name_or_id, access_key) @session_repository_resilience.apply() - async def get_session_with_group( - self, - session_name_or_id: str | SessionId, - owner_access_key: AccessKey, - kernel_loading_strategy: KernelLoadingStrategy = KernelLoadingStrategy.MAIN_KERNEL_ONLY, - allow_stale: bool = False, - ) -> SessionRow: - """Get session with group information eagerly loaded""" - return await self._db_source.get_session_with_group( - session_name_or_id, owner_access_key, kernel_loading_strategy, allow_stale - ) + async def get_project_container_registry(self, group_id: uuid.UUID) -> dict[str, Any] | None: + return await self._db_source.get_project_container_registry(group_id) @session_repository_resilience.apply() async def get_session_with_routing_minimal( diff --git a/src/ai/backend/manager/repositories/stream/db_source/db_source.py b/src/ai/backend/manager/repositories/stream/db_source/db_source.py index 7c0e2bedb2b..568043eef70 100644 --- a/src/ai/backend/manager/repositories/stream/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/stream/db_source/db_source.py @@ -1,7 +1,7 @@ import uuid import sqlalchemy as sa -from sqlalchemy.orm import joinedload, noload, selectinload +from sqlalchemy.orm import noload, selectinload from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.errors.kernel import SessionNotFound, TooManySessionsMatched @@ -31,7 +31,6 @@ async def get_streaming_session( .options( noload("*"), selectinload(SessionRow.kernels).options(noload("*")), - joinedload(SessionRow.user), ) .execution_options(populate_existing=True) ) diff --git a/src/ai/backend/manager/repositories/user/db_source/db_source.py b/src/ai/backend/manager/repositories/user/db_source/db_source.py index fdb69cfda7b..7fc5702c0d4 100644 --- a/src/ai/backend/manager/repositories/user/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/user/db_source/db_source.py @@ -72,13 +72,11 @@ from ai.backend.manager.models.session import ( AGENT_RESOURCE_OCCUPYING_SESSION_STATUSES, QueryCondition, - QueryOption, SessionRow, by_status, by_user_id, ) from ai.backend.manager.models.specs.pagination import NoPagination -from ai.backend.manager.models.types import join_by_related_field from ai.backend.manager.models.user import UserRole, UserRow, UserStatus, users from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.models.vfolder import ( @@ -558,13 +556,7 @@ async def retrieve_active_sessions(self, user_uuid: UUID) -> list[SessionRow]: by_status(AGENT_RESOURCE_OCCUPYING_SESSION_STATUSES), ] - query_options: list[QueryOption] = [ - join_by_related_field(SessionRow.user), - ] - - return await SessionRow.list_session_by_condition( - query_conditions, query_options, db=self._db - ) + return await SessionRow.list_session_by_condition(query_conditions, [], db=self._db) async def delegate_endpoint_ownership( self, diff --git a/src/ai/backend/manager/services/session/service.py b/src/ai/backend/manager/services/session/service.py index 497305810f5..f289448570c 100644 --- a/src/ai/backend/manager/services/session/service.py +++ b/src/ai/backend/manager/services/session/service.py @@ -93,7 +93,6 @@ ) from ai.backend.manager.errors.storage import VFolderBadRequest from ai.backend.manager.idle import IdleCheckerHost -from ai.backend.manager.models.group import GroupRow from ai.backend.manager.models.session import ( DEAD_SESSION_STATUSES, PRIVATE_SESSION_TYPES, @@ -442,20 +441,22 @@ async def convert_session_to_image( }, ) - session = await self._session_repository.get_session_with_group( + session = await self._session_repository.get_session_validated( session_name, owner_access_key, kernel_loading_strategy=KernelLoadingStrategy.MAIN_KERNEL_ONLY, ) - project: GroupRow = session.group - if not project.container_registry: + container_registry = await self._session_repository.get_project_container_registry( + session.group_id + ) + if not container_registry: raise InvalidAPIParameters( "Project not ready to convert session image (registry configuration not populated)" ) - registry_hostname = project.container_registry["registry"] - registry_project = project.container_registry["project"] + registry_hostname = container_registry["registry"] + registry_project = container_registry["project"] registry_conf = await self._session_repository.get_container_registry( registry_hostname, registry_project From c4fdf5fc6c907bf5209857364c524410b5443925 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Sun, 16 Aug 2026 02:49:58 +0900 Subject: [PATCH 4/5] refactor(BA-7226): drop KernelRow.session and AgentResourceRow.slot_type_row Both relationships are now unreferenced: slot_type_row lost its last reader when agent occupancy moved to an explicit join, and nothing reads KernelRow.session. SessionRow.kernels drops its back_populates and becomes one-directional; no code writes through the relationship, so delete behaviour still comes from the database foreign keys. Co-Authored-By: Claude Opus 5 (1M context) --- src/ai/backend/manager/models/kernel/row.py | 12 +----------- src/ai/backend/manager/models/resource_slot/row.py | 6 +----- src/ai/backend/manager/models/session/row.py | 2 +- 3 files changed, 3 insertions(+), 17 deletions(-) diff --git a/src/ai/backend/manager/models/kernel/row.py b/src/ai/backend/manager/models/kernel/row.py index 5f187544ce6..a338eed7956 100644 --- a/src/ai/backend/manager/models/kernel/row.py +++ b/src/ai/backend/manager/models/kernel/row.py @@ -4,10 +4,7 @@ import uuid from collections.abc import Sequence from datetime import datetime, tzinfo -from typing import ( - TYPE_CHECKING, - Any, -) +from typing import Any import sqlalchemy as sa import yarl @@ -17,7 +14,6 @@ Mapped, mapped_column, noload, - relationship, ) from ai.backend.common.identifier.image import ImageID @@ -48,10 +44,6 @@ RuntimeConfig, UserPermission, ) - -if TYPE_CHECKING: - from ai.backend.manager.models.session import SessionRow - from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.errors.kernel import SessionNotFound from ai.backend.manager.models.base import ( @@ -417,8 +409,6 @@ class KernelRow(CreatedAtMixin, Base): ), ) - session: Mapped[SessionRow] = relationship("SessionRow", back_populates="kernels") - @property def used_time(self) -> str | None: if self.terminated_at is not None and self.created_at is not None: diff --git a/src/ai/backend/manager/models/resource_slot/row.py b/src/ai/backend/manager/models/resource_slot/row.py index be99244326a..51f53120e62 100644 --- a/src/ai/backend/manager/models/resource_slot/row.py +++ b/src/ai/backend/manager/models/resource_slot/row.py @@ -13,7 +13,7 @@ from decimal import Decimal import sqlalchemy as sa -from sqlalchemy.orm import Mapped, mapped_column, relationship +from sqlalchemy.orm import Mapped, mapped_column from ai.backend.common.identifier.resource_slot import ResourceSlotTypeUUID from ai.backend.manager.data.resource_slot.types import ( @@ -153,10 +153,6 @@ class AgentResourceRow(LifecycleTimestampsMixin, Base): "used", sa.Numeric(precision=24, scale=6), nullable=False, server_default=sa.text("0") ) - slot_type_row: Mapped[ResourceSlotTypeRow] = relationship( - "ResourceSlotTypeRow", foreign_keys=[slot_name], lazy="raise" - ) - __table_args__ = ( sa.ForeignKeyConstraint( ["agent_id"], diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index 0c75234ebd9..d4b900117fd 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -432,7 +432,7 @@ class SessionRow(CreatedAtMixin, Base): nullable=True, default=None, ) - kernels: Mapped[list[KernelRow]] = relationship("KernelRow", back_populates="session") + kernels: Mapped[list[KernelRow]] = relationship("KernelRow") # Resource ownership resource_group_id: Mapped[ResourceGroupID] = mapped_column( From 36d1a94a1f7225364600d9bef3e4e0c6071c9e67 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Sun, 16 Aug 2026 02:54:30 +0900 Subject: [PATCH 5/5] changelog: add news fragment for PR #13786 --- changes/13786.enhance.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/13786.enhance.md diff --git a/changes/13786.enhance.md b/changes/13786.enhance.md new file mode 100644 index 00000000000..df200e3c3ba --- /dev/null +++ b/changes/13786.enhance.md @@ -0,0 +1 @@ +Replace the session, kernel and agent ORM relationship usages in the manager with explicit queries and remove the relationship definitions