refactor: Change global id generator to support type safety

Author: seedspirit

src/ai/backend/manager/api/gql/base.py

@@ -5,6 +5,7 @@
 from enum import StrEnum
 from typing import TYPE_CHECKING, Any, Optional, Type, cast
 
+import graphene
 import orjson
 import strawberry
 from graphql import StringValueNode
@@ -160,7 +161,17 @@ def from_resource_slot(resource_slot: ResourceSlot) -> JSONString:
         return JSONString.serialize(resource_slot.to_json())
 
 
-def to_global_id(type_: Type[Any], local_id: uuid.UUID | str) -> str:
+def to_global_id(
+    type_: Type[Any], local_id: uuid.UUID | str, is_target_graphene_object: bool = False
+) -> str:
+    if is_target_graphene_object:
+        # For compatibility with existing Graphene-based global IDs
+        if not issubclass(type_, graphene.ObjectType):
+            raise TypeError(
+                "type_ must be a graphene ObjectType when is_target_graphene_object is True."
+            )
+        typename = type_.__name__
+        return base64(f"{typename}:{local_id}")
     if not has_object_definition(type_):
         raise TypeError("type_ must be a Strawberry object type (Node or Edge).")
     typename = get_object_definition(type_, strict=True).name

src/ai/backend/manager/api/gql/model_deployment/model_deployment.py

@@ -17,7 +17,12 @@
     ModelDeploymentStatus as CommonDeploymentStatus,
 )
 from ai.backend.common.exception import ModelDeploymentUnavailableError
-from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, resolve_global_id
+from ai.backend.manager.api.gql.base import (
+    OrderDirection,
+    StringFilter,
+    resolve_global_id,
+    to_global_id,
+)
 from ai.backend.manager.api.gql.domain import Domain
 from ai.backend.manager.api.gql.model_deployment.access_token import (
     AccessToken,
@@ -46,7 +51,9 @@
     ReplicaSpec,
     ReplicaStateData,
 )
-from ai.backend.manager.models.gql_relay import AsyncNode
+from ai.backend.manager.models.gql_models.domain import DomainNode
+from ai.backend.manager.models.gql_models.group import GroupNode
+from ai.backend.manager.models.gql_models.user import UserNode
 from ai.backend.manager.services.deployment.actions.auto_scaling_rule.get_auto_scaling_rule_by_deployment_id import (
     GetAutoScalingRulesByDeploymentIdAction,
 )
@@ -185,12 +192,16 @@ class ModelDeploymentMetadata:
 
     @strawberry.field
     async def project(self, info: Info[StrawberryGQLContext]) -> Project:
-        project_global_id = AsyncNode.to_global_id("GroupNode", self._project_id)
+        project_global_id = to_global_id(
+            GroupNode, self._project_id, is_target_graphene_object=True
+        )
         return Project(id=ID(project_global_id))
 
     @strawberry.field
     async def domain(self, info: Info[StrawberryGQLContext]) -> Domain:
-        domain_global_id = AsyncNode.to_global_id("DomainNode", self._domain_name)
+        domain_global_id = to_global_id(
+            DomainNode, self._domain_name, is_target_graphene_object=True
+        )
         return Domain(id=ID(domain_global_id))
 
     @classmethod
@@ -266,7 +277,9 @@ class ModelDeployment(Node):
 
     @strawberry.field
     async def created_user(self, info: Info[StrawberryGQLContext]) -> User:
-        user_global_id = AsyncNode.to_global_id("UserNode", self._created_user_id)
+        user_global_id = to_global_id(
+            UserNode, self._created_user_id, is_target_graphene_object=True
+        )
         return User(id=strawberry.ID(user_global_id))
 
     @strawberry.field

src/ai/backend/manager/api/gql/model_deployment/model_replica.py

@@ -14,11 +14,16 @@
 from ai.backend.common.data.model_deployment.types import LivenessStatus as CommonLivenessStatus
 from ai.backend.common.data.model_deployment.types import ReadinessStatus as CommonReadinessStatus
 from ai.backend.common.exception import ModelDeploymentUnavailableError
-from ai.backend.manager.api.gql.base import JSONString, OrderDirection, resolve_global_id
+from ai.backend.manager.api.gql.base import (
+    JSONString,
+    OrderDirection,
+    resolve_global_id,
+    to_global_id,
+)
 from ai.backend.manager.api.gql.session import Session
 from ai.backend.manager.api.gql.types import StrawberryGQLContext
 from ai.backend.manager.data.deployment.types import ModelReplicaData
-from ai.backend.manager.models.gql_relay import AsyncNode
+from ai.backend.manager.models.gql_models.session import ComputeSessionNode
 from ai.backend.manager.services.deployment.actions.get_replicas_by_deployment_id import (
     GetReplicasByDeploymentIdAction,
 )
@@ -119,7 +124,9 @@ class ModelReplica(Node):
         description="The session ID associated with the replica. This can be null right after replica creation."
     )
     async def session(self, info: Info[StrawberryGQLContext]) -> "Session":
-        session_global_id = AsyncNode.to_global_id("ComputeSessionNode", self._session_id)
+        session_global_id = to_global_id(
+            ComputeSessionNode, self._session_id, is_target_graphene_object=True
+        )
         return Session(id=ID(session_global_id))
 
     @strawberry.field

src/ai/backend/manager/api/gql/model_deployment/model_revision.py

@@ -22,6 +22,7 @@
     OrderDirection,
     StringFilter,
     resolve_global_id,
+    to_global_id,
 )
 from ai.backend.manager.api.gql.image import (
     Image,
@@ -55,7 +56,9 @@
     ResourceSpec,
 )
 from ai.backend.manager.data.image.types import ImageIdentifier
-from ai.backend.manager.models.gql_relay import AsyncNode
+from ai.backend.manager.models.gql_models.image import ImageNode
+from ai.backend.manager.models.gql_models.scaling_group import ScalingGroupNode
+from ai.backend.manager.models.gql_models.vfolder import VirtualFolderNode
 from ai.backend.manager.services.deployment.actions.model_revision.add_model_revision import (
     AddModelRevisionAction,
 )
@@ -91,7 +94,9 @@ class ModelMountConfig:
 
     @strawberry.field
     async def vfolder(self, info: Info[StrawberryGQLContext]) -> VFolder:
-        vfolder_global_id = AsyncNode.to_global_id("VirtualFolderNode", self._vfolder_id)
+        vfolder_global_id = to_global_id(
+            VirtualFolderNode, self._vfolder_id, is_target_graphene_object=True
+        )
         return VFolder(id=ID(vfolder_global_id))
 
     @classmethod
@@ -135,7 +140,9 @@ class ResourceConfig:
     @strawberry.field
     def resource_group(self) -> "ResourceGroup":
         """Resolves the federated ResourceGroup."""
-        global_id = AsyncNode.to_global_id("ScalingGroupNode", self._resource_group_name)
+        global_id = to_global_id(
+            ScalingGroupNode, self._resource_group_name, is_target_graphene_object=True
+        )
         return ResourceGroup(id=ID(global_id))
 
     @classmethod
@@ -174,7 +181,7 @@ class ModelRevision(Node):
 
     @strawberry.field
     async def image(self, info: Info[StrawberryGQLContext]) -> Image:
-        image_global_id = AsyncNode.to_global_id("ImageNode", self._image_id)
+        image_global_id = to_global_id(ImageNode, self._image_id, is_target_graphene_object=True)
         return Image(id=ID(image_global_id))
 
     @classmethod