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

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

@@ -2,7 +2,7 @@
 
 import enum
 import uuid
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence, cast
+from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Self, Sequence, cast
 
 import graphene
 import sqlalchemy as sa
@@ -17,7 +17,7 @@
 from sqlalchemy.sql.expression import false, true
 
 from ai.backend.common import msgpack, redis_helper
-from ai.backend.common.types import AgentId, BinarySize, HardwareMetadata, ResourceSlot
+from ai.backend.common.types import AccessKey, AgentId, BinarySize, HardwareMetadata, ResourceSlot
 
 from .base import (
     Base,
@@ -528,7 +528,7 @@ def from_row(
         cls,
         ctx: GraphQueryContext,
         row: Row,
-    ) -> Agent:
+    ) -> Self:
         return cls(
             id=row["id"],
             status=row["status"].name,
@@ -561,11 +561,11 @@ async def batch_load(
         graph_ctx: GraphQueryContext,
         agent_ids: Sequence[AgentId],
         *,
-        domain_name: str | None,
+        access_key: AccessKey,
+        domain_name: Optional[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)
@@ -627,7 +627,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 @@
     Awaitable,
     Callable,
     ClassVar,
+    Concatenate,
     Coroutine,
     Generic,
     List,
@@ -712,7 +713,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)
@@ -721,6 +727,7 @@ def get_loader_by_func(
                 functools.partial(
                     batch_load_func,
                     context,
+                    **kwargs,
                 ),
                 max_batch_size=128,
             )

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

@@ -819,9 +819,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,