feat: Add GraphQL field rolling_count to check current rate limit w…

…ithout increasing (#2050)

Co-authored-by: Joongi Kim <[email protected]>

changes/2050.feature.md

@@ -0,0 +1 @@
+Introduce the `rolling_count` GraphQL field to provide the current rate limit counter for a keypair within the designated time window slice

src/ai/backend/client/cli/image.py

@@ -2,7 +2,6 @@
 
 import click
 
-# from ai.backend.client.output.fields import image_fields
 from ai.backend.cli.main import main
 from ai.backend.cli.types import ExitCode
 from ai.backend.client.exceptions import BackendAPIError

src/ai/backend/manager/api/schema.graphql

@@ -516,6 +516,9 @@ type KeyPair implements Item {
   last_used: DateTime
   rate_limit: Int
   num_queries: Int
+
+  """Added in 24.09.0."""
+  rolling_count: Int
   user: UUID
   projects: [String]
   ssh_public_key: String

src/ai/backend/manager/models/keypair.py

@@ -13,12 +13,14 @@
 from cryptography.hazmat.primitives.asymmetric import rsa
 from dateutil.parser import parse as dtparse
 from graphene.types.datetime import DateTime as GQLDateTime
+from redis.asyncio import Redis
 from sqlalchemy.engine.row import Row
 from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 from sqlalchemy.orm import relationship
 from sqlalchemy.sql.expression import false
 
 from ai.backend.common import msgpack, redis_helper
+from ai.backend.common.defs import REDIS_RLIM_DB
 from ai.backend.common.types import AccessKey, SecretKey
 
 if TYPE_CHECKING:
@@ -167,6 +169,7 @@ class Meta:
     last_used = GQLDateTime()
     rate_limit = graphene.Int()
     num_queries = graphene.Int()
+    rolling_count = graphene.Int(description="Added in 24.09.0.")
     user = graphene.UUID()
     projects = graphene.List(lambda: graphene.String)
 
@@ -228,6 +231,18 @@ async def resolve_num_queries(self, info: graphene.ResolveInfo) -> int:
             return n
         return 0
 
+    async def resolve_rolling_count(self, info: graphene.ResolveInfo) -> int:
+        ctx: GraphQueryContext = info.context
+        redis_rlim = redis_helper.get_redis_object(
+            ctx.shared_config.data["redis"], name="ratelimit", db=REDIS_RLIM_DB
+        )
+
+        async def _zcard(r: Redis):
+            return await r.zcard(self.access_key)
+
+        ret = await redis_helper.execute(redis_rlim, _zcard)
+        return int(ret) if ret is not None else 0
+
     async def resolve_vfolders(self, info: graphene.ResolveInfo) -> Sequence[VirtualFolder]:
         ctx: GraphQueryContext = info.context
         loader = ctx.dataloader_manager.get_loader(ctx, "VirtualFolder")