fix: lint

lablup · Nov 28, 2024 · 211f6e7 · 211f6e7
1 parent 66b0a99
commit 211f6e7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/src/ai/backend/manager/models/agent.py b/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
+from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Self, Sequence
 
 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,
@@ -510,7 +510,7 @@ def from_row(
         cls,
         ctx: GraphQueryContext,
         row: Row,
-    ) -> Agent:
+    ) -> Self:
         return cls(
             id=row["id"],
             status=row["status"].name,
@@ -543,11 +543,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)
@@ -609,7 +609,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

diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py
@@ -742,9 +742,9 @@ async def resolve_agent_summary(
         if ctx.local_config["manager"]["hide-agents"]:
             raise ObjectNotFound(object_name="agent")
 
-        loader = ctx.dataloader_manager.get_loader_by_func(
+        loader = ctx.dataloader_manager.get_loader(
             ctx,
-            AgentSummary.batch_load,
+            "AgentSummary.by_id",
             raw_status=None,
             scaling_group=scaling_group,
             domain_name=domain_name,