diff --git a/changes/2013.misc.md b/changes/2013.misc.md new file mode 100644 index 0000000000..da16eec4bd --- /dev/null +++ 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 index 6cd6b7182c..ab1837274d 100644 --- 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 index f3f5b2c9f0..cf6e3ef3e4 100644 --- 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)