refactor(BA-3516): Integrate ModifyScalingGroup action with existing GQL resolver (#7524)

Author: jopemachine

changes/7524.enhance.md

@@ -0,0 +1 @@
+Integrate `ModifyScalingGroup` action with existing GQL resolver

src/ai/backend/manager/api/gql_legacy/scaling_group.py

@@ -15,6 +15,7 @@
 import graphene_federation
 import sqlalchemy as sa
 from graphene.types.datetime import DateTime as GQLDateTime
+from graphql import Undefined
 from sqlalchemy.engine.row import Row
 from sqlalchemy.orm import load_only
 
@@ -36,19 +37,31 @@
 from ai.backend.manager.models.user import UserRole
 from ai.backend.manager.repositories.base.creator import Creator
 from ai.backend.manager.repositories.base.purger import Purger
+from ai.backend.manager.repositories.base.updater import Updater
 from ai.backend.manager.repositories.scaling_group.creators import ScalingGroupCreatorSpec
+from ai.backend.manager.repositories.scaling_group.updaters import (
+    ScalingGroupDriverConfigUpdaterSpec,
+    ScalingGroupMetadataUpdaterSpec,
+    ScalingGroupNetworkConfigUpdaterSpec,
+    ScalingGroupSchedulerConfigUpdaterSpec,
+    ScalingGroupStatusUpdaterSpec,
+    ScalingGroupUpdaterSpec,
+)
 from ai.backend.manager.services.scaling_group.actions.create import (
     CreateScalingGroupAction,
 )
+from ai.backend.manager.services.scaling_group.actions.modify import (
+    ModifyScalingGroupAction,
+)
 from ai.backend.manager.services.scaling_group.actions.purge_scaling_group import (
     PurgeScalingGroupAction,
 )
+from ai.backend.manager.types import OptionalState, TriState
 
 from .base import (
     batch_multiresult,
     batch_multiresult_in_scalar_stream,
     batch_result,
-    set_if_set,
     simple_db_mutate,
 )
 from .gql_relay import (
@@ -606,6 +619,41 @@ class ModifyScalingGroupInput(graphene.InputObjectType):
     scheduler_opts = graphene.JSONString(required=False)
     use_host_network = graphene.Boolean(required=False)
 
+    def to_updater(self, name: str) -> Updater[ScalingGroupRow]:
+        """Convert GraphQL input to Updater for scaling group modification."""
+        status_spec = ScalingGroupStatusUpdaterSpec(
+            is_active=OptionalState.from_graphql(self.is_active),
+            is_public=OptionalState.from_graphql(self.is_public),
+        )
+        metadata_spec = ScalingGroupMetadataUpdaterSpec(
+            description=TriState.from_graphql(self.description),
+        )
+        network_spec = ScalingGroupNetworkConfigUpdaterSpec(
+            wsproxy_addr=TriState.from_graphql(self.wsproxy_addr),
+            wsproxy_api_token=TriState.from_graphql(self.wsproxy_api_token),
+            use_host_network=OptionalState.from_graphql(self.use_host_network),
+        )
+        driver_spec = ScalingGroupDriverConfigUpdaterSpec(
+            driver=OptionalState.from_graphql(self.driver),
+            driver_opts=OptionalState.from_graphql(self.driver_opts),
+        )
+        scheduler_spec = ScalingGroupSchedulerConfigUpdaterSpec(
+            scheduler=OptionalState.from_graphql(self.scheduler),
+            scheduler_opts=OptionalState.from_graphql(
+                ScalingGroupOpts.from_json(self.scheduler_opts)
+                if self.scheduler_opts is not None and self.scheduler_opts is not Undefined
+                else Undefined
+            ),
+        )
+        spec = ScalingGroupUpdaterSpec(
+            status=status_spec,
+            metadata=metadata_spec,
+            network=network_spec,
+            driver=driver_spec,
+            scheduler=scheduler_spec,
+        )
+        return Updater(spec=spec, pk_value=name)
+
 
 class CreateScalingGroup(graphene.Mutation):
     allowed_roles = (UserRole.SUPERADMIN,)
@@ -683,21 +731,11 @@ async def mutate(
         name: str,
         props: ModifyScalingGroupInput,
     ) -> ModifyScalingGroup:
-        data: dict[str, Any] = {}
-        set_if_set(props, data, "description")
-        set_if_set(props, data, "is_active")
-        set_if_set(props, data, "is_public")
-        set_if_set(props, data, "wsproxy_addr")
-        set_if_set(props, data, "wsproxy_api_token")
-        set_if_set(props, data, "driver")
-        set_if_set(props, data, "driver_opts")
-        set_if_set(props, data, "scheduler")
-        set_if_set(
-            props, data, "scheduler_opts", clean_func=lambda v: ScalingGroupOpts.from_json(v)
+        graph_ctx: GraphQueryContext = info.context
+        await graph_ctx.processors.scaling_group.modify_scaling_group.wait_for_complete(
+            ModifyScalingGroupAction(updater=props.to_updater(name))
         )
-        set_if_set(props, data, "use_host_network")
-        update_query = sa.update(scaling_groups).values(data).where(scaling_groups.c.name == name)
-        return await simple_db_mutate(cls, info.context, update_query)
+        return cls(ok=True, msg="success")
 
 
 class DeleteScalingGroup(graphene.Mutation):

src/ai/backend/manager/repositories/scaling_group/repository.py

@@ -84,6 +84,7 @@ async def purge_scaling_group(
         """
         return await self._db_source.purge_scaling_group(purger)
 
+    @scaling_group_repository_resilience.apply()
     async def update_scaling_group(
         self,
         updater: Updater[ScalingGroupRow],