fix: Add missing resolver of Agent.compute_containers GraphQL field #…

…3011 (#3120)

changes/3011.fix.md

@@ -0,0 +1 @@
+Fix `Agent.compute_containers` GraphQL field by adding missing resolver

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

@@ -33,7 +33,13 @@
     simple_db_mutate,
 )
 from .group import association_groups_users
-from .kernel import AGENT_RESOURCE_OCCUPYING_KERNEL_STATUSES, KernelRow, kernels
+from .kernel import (
+    AGENT_RESOURCE_OCCUPYING_KERNEL_STATUSES,
+    ComputeContainer,
+    KernelRow,
+    KernelStatus,
+    kernels,
+)
 from .keypair import keypairs
 from .minilang.ordering import OrderSpecItem, QueryOrderParser
 from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter
@@ -157,9 +163,7 @@ class Meta:
     cpu_cur_pct = graphene.Float()
     mem_cur_bytes = graphene.Float()
 
-    compute_containers = graphene.List(
-        "ai.backend.manager.models.ComputeContainer", status=graphene.String()
-    )
+    compute_containers = graphene.List(ComputeContainer, status=graphene.String())
 
     @classmethod
     def from_row(
@@ -195,6 +199,18 @@ def from_row(
             used_tpu_slots=float(row["occupied_slots"].get("tpu.device", 0)),
         )
 
+    async def resolve_compute_containers(
+        self, info: graphene.ResolveInfo, *, status: Optional[str] = None
+    ) -> list[ComputeContainer]:
+        ctx: GraphQueryContext = info.context
+        _status = KernelStatus[status] if status is not None else None
+        loader = ctx.dataloader_manager.get_loader(
+            ctx,
+            "ComputeContainer.by_agent_id",
+            status=_status,
+        )
+        return await loader.load(self.id)
+
     async def resolve_live_stat(self, info: graphene.ResolveInfo) -> Any:
         ctx: GraphQueryContext = info.context
         loader = ctx.dataloader_manager.get_loader(ctx, "Agent.live_stat")

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

@@ -35,6 +35,7 @@
 from ai.backend.common.logging import BraceStyleAdapter
 from ai.backend.common.types import (
     AccessKey,
+    AgentId,
     BinarySize,
     ClusterMode,
     KernelId,
@@ -70,6 +71,7 @@
     StructuredJSONObjectListColumn,
     URLColumn,
     batch_multiresult,
+    batch_multiresult_in_scalar_stream,
     batch_result,
     mapper_registry,
 )
@@ -1010,6 +1012,31 @@ async def batch_load_by_session(
                 lambda row: row.session_id,
             )
 
+    @classmethod
+    async def batch_load_by_agent_id(
+        cls,
+        ctx: GraphQueryContext,
+        agent_ids: Sequence[AgentId],
+        *,
+        status: Optional[KernelStatus] = None,
+    ) -> Sequence[Sequence[ComputeContainer]]:
+        query_stmt = (
+            sa.select(KernelRow)
+            .where(KernelRow.agent.in_(agent_ids))
+            .options(selectinload(KernelRow.image_row).options(selectinload(ImageRow.aliases)))
+        )
+        if status is not None:
+            query_stmt = query_stmt.where(KernelRow.status == status)
+        async with ctx.db.begin_readonly_session() as db_session:
+            return await batch_multiresult_in_scalar_stream(
+                ctx,
+                db_session,
+                query_stmt,
+                cls,
+                agent_ids,
+                lambda row: row.agent,
+            )
+
     @classmethod
     async def batch_load_detail(
         cls,