chore: Add max_pending_session_count field to keypair_resource_poli…

…cy GQL schema (#2013) Backported-from: main (24.09) Backported-to: 24.03
lablup · May 13, 2024 · 0ad399b · 0ad399b
1 parent 05fd6fd
commit 0ad399b
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/changes/2013.misc.md b/changes/2013.misc.md
@@ -0,0 +1 @@
+Add max_pending_session_count field to Keypair resource policy GQL schema
diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql
@@ -596,6 +596,12 @@ type KeyPairResourcePolicy {
   max_vfolder_count: Int @deprecated(reason: "Deprecated since 23.09.4")
   max_vfolder_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
   max_quota_scope_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
+
+  """Added in 24.03.4."""
+  max_pending_session_count: Int
+
+  """Added in 24.03.4."""
+  max_pending_session_resource_slots: JSONString
 }
 
 type UserResourcePolicy {
@@ -1387,6 +1393,12 @@ input CreateKeyPairResourcePolicyInput {
   max_vfolder_count: Int @deprecated(reason: "Deprecated since 23.09.4")
   max_vfolder_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
   max_quota_scope_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
+
+  """Added in 24.03.4."""
+  max_pending_session_count: Int
+
+  """Added in 24.03.4."""
+  max_pending_session_resource_slots: JSONString
 }
 
 type ModifyKeyPairResourcePolicy {
@@ -1406,6 +1418,12 @@ input ModifyKeyPairResourcePolicyInput {
   max_vfolder_count: Int @deprecated(reason: "Deprecated since 23.09.4")
   max_vfolder_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
   max_quota_scope_size: BigInt @deprecated(reason: "Deprecated since 23.09.4")
+
+  """Added in 24.03.4."""
+  max_pending_session_count: Int
+
+  """Added in 24.03.4."""
+  max_pending_session_resource_slots: JSONString
 }
 
 type DeleteKeyPairResourcePolicy {

diff --git a/src/ai/backend/manager/models/resource_policy.py b/src/ai/backend/manager/models/resource_policy.py
@@ -159,6 +159,8 @@ class KeyPairResourcePolicy(graphene.ObjectType):
     max_vfolder_count = graphene.Int(deprecation_reason="Deprecated since 23.09.4")
     max_vfolder_size = BigInt(deprecation_reason="Deprecated since 23.09.4")
     max_quota_scope_size = BigInt(deprecation_reason="Deprecated since 23.09.4")
+    max_pending_session_count = graphene.Int(description="Added in 24.03.4.")
+    max_pending_session_resource_slots = graphene.JSONString(description="Added in 24.03.4.")
 
     @classmethod
     def from_row(
@@ -168,6 +170,11 @@ def from_row(
     ) -> KeyPairResourcePolicy | None:
         if row is None:
             return None
+
+        if row["max_pending_session_resource_slots"] is not None:
+            max_pending_session_resource_slots = row["max_pending_session_resource_slots"].to_json()
+        else:
+            max_pending_session_resource_slots = None
         return cls(
             name=row["name"],
             created_at=row["created_at"],
@@ -178,6 +185,8 @@ def from_row(
             max_containers_per_session=row["max_containers_per_session"],
             idle_timeout=row["idle_timeout"],
             allowed_vfolder_hosts=row["allowed_vfolder_hosts"].to_json(),
+            max_pending_session_count=row["max_pending_session_count"],
+            max_pending_session_resource_slots=max_pending_session_resource_slots,
         )
 
     @classmethod
@@ -311,6 +320,8 @@ class CreateKeyPairResourcePolicyInput(graphene.InputObjectType):
     max_vfolder_count = graphene.Int(required=False, deprecation_reason="Deprecated since 23.09.4")
     max_vfolder_size = BigInt(required=False, deprecation_reason="Deprecated since 23.09.4")
     max_quota_scope_size = BigInt(required=False, deprecation_reason="Deprecated since 23.09.4")
+    max_pending_session_count = graphene.Int(description="Added in 24.03.4.")
+    max_pending_session_resource_slots = graphene.JSONString(description="Added in 24.03.4.")
 
 
 class ModifyKeyPairResourcePolicyInput(graphene.InputObjectType):
@@ -325,6 +336,8 @@ class ModifyKeyPairResourcePolicyInput(graphene.InputObjectType):
     max_vfolder_count = graphene.Int(required=False, deprecation_reason="Deprecated since 23.09.4")
     max_vfolder_size = BigInt(required=False, deprecation_reason="Deprecated since 23.09.4")
     max_quota_scope_size = BigInt(required=False, deprecation_reason="Deprecated since 23.09.4")
+    max_pending_session_count = graphene.Int(description="Added in 24.03.4.")
+    max_pending_session_resource_slots = graphene.JSONString(description="Added in 24.03.4.")
 
 
 class CreateKeyPairResourcePolicy(graphene.Mutation):
@@ -357,6 +370,13 @@ async def mutate(
             "idle_timeout": props.idle_timeout,
             "allowed_vfolder_hosts": props.allowed_vfolder_hosts,
         }
+        set_if_set(props, data, "max_pending_session_count")
+        set_if_set(
+            props,
+            data,
+            "max_pending_session_resource_slots",
+            clean_func=lambda v: ResourceSlot.from_user_input(v, None),
+        )
         insert_query = sa.insert(keypair_resource_policies).values(data)
         return await simple_db_mutate_returning_item(
             cls,
@@ -403,6 +423,13 @@ async def mutate(
         set_if_set(props, data, "max_containers_per_session")
         set_if_set(props, data, "idle_timeout")
         set_if_set(props, data, "allowed_vfolder_hosts")
+        set_if_set(props, data, "max_pending_session_count")
+        set_if_set(
+            props,
+            data,
+            "max_pending_session_resource_slots",
+            clean_func=lambda v: ResourceSlot.from_user_input(v, None),
+        )
         update_query = (
             sa.update(keypair_resource_policies)
             .values(data)