fix: Regression of AgentSummary GQL resolver (#3045) (#3172)

changes/3045.fix.md

@@ -0,0 +1 @@
+Fix regression of the `AgentSummary` resolver caused by an incorrect `batch_load_func` assignment.

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

@@ -3,7 +3,18 @@
 import enum
 import uuid
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence, TypeAlias, cast, override
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Mapping,
+    Optional,
+    Self,
+    Sequence,
+    TypeAlias,
+    cast,
+    override,
+)
 
 import graphene
 import sqlalchemy as sa
@@ -539,7 +550,7 @@ def from_row(
         cls,
         ctx: GraphQueryContext,
         row: Row,
-    ) -> Agent:
+    ) -> Self:
         return cls(
             id=row["id"],
             status=row["status"].name,
@@ -572,11 +583,11 @@ async def batch_load(
         graph_ctx: GraphQueryContext,
         agent_ids: Sequence[AgentId],
         *,
+        access_key: AccessKey,
         domain_name: str | None,
         raw_status: Optional[str] = None,
         scaling_group: Optional[str] = None,
-        access_key: str,
-    ) -> Sequence[Agent | None]:
+    ) -> Sequence[Optional[Self]]:
         query = (
             sa.select([agents])
             .select_from(agents)
@@ -638,7 +649,7 @@ async def load_slice(
         raw_status: Optional[str] = None,
         filter: Optional[str] = None,
         order: Optional[str] = None,
-    ) -> Sequence[Agent]:
+    ) -> Sequence[Self]:
         query = sa.select([agents]).select_from(agents).limit(limit).offset(offset)
         query = await _append_sgroup_from_clause(
             graph_ctx, query, access_key, domain_name, scaling_group

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

@@ -18,6 +18,7 @@
     TYPE_CHECKING,
     Any,
     ClassVar,
+    Concatenate,
     Final,
     Generic,
     NamedTuple,
@@ -752,7 +753,12 @@ def _get_func_key(
     def get_loader_by_func(
         self,
         context: ContextT,
-        batch_load_func: Callable[[ContextT, Sequence[LoaderKeyT]], Awaitable[LoaderResultT]],
+        batch_load_func: Callable[
+            Concatenate[ContextT, Sequence[LoaderKeyT], ...], Awaitable[LoaderResultT]
+        ],
+        # Using kwargs-only to prevent argument position confusion
+        # when DataLoader calls `batch_load_func(keys)` which is `partial(batch_load_func, **kwargs)(keys)`.
+        **kwargs,
     ) -> DataLoader:
         key = self._get_func_key(batch_load_func)
         loader = self.cache.get(key)
@@ -761,6 +767,7 @@ def get_loader_by_func(
                 functools.partial(
                     batch_load_func,
                     context,
+                    **kwargs,
                 ),
                 max_batch_size=128,
             )

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

@@ -892,9 +892,9 @@ async def resolve_agent_summary(
         if ctx.local_config["manager"]["hide-agents"]:
             raise ObjectNotFound(object_name="agent")
 
-        loader = ctx.dataloader_manager.get_loader(
+        loader = ctx.dataloader_manager.get_loader_by_func(
             ctx,
-            "Agent",
+            AgentSummary.batch_load,
             raw_status=None,
             scaling_group=scaling_group,
             domain_name=domain_name,