feat: Implement get_hot_anonymous_clients to RateLimit API

changes/2075.feature.md

@@ -0,0 +1 @@
+Implement RateLimit API and `get_hot_anonymous_clients`.

src/ai/backend/client/func/ratelimit.py

@@ -0,0 +1,18 @@
+from ..request import Request
+from .base import BaseFunction, api_function
+
+__all__ = ("RateLimit",)
+
+
+class RateLimit(BaseFunction):
+    """
+    Provides RateLimiting API functions.
+    """
+
+    @api_function
+    @classmethod
+    async def get_hot_anonymous_clients(cls):
+        """ """
+        rqst = Request("GET", "/ratelimit/hot_anonymous_clients")
+        async with rqst.fetch() as resp:
+            return await resp.json()

src/ai/backend/client/session.py

@@ -272,6 +272,7 @@ class BaseSession(metaclass=abc.ABCMeta):
         "Service",
         "Model",
         "QuotaScope",
+        "RateLimit",
     )
 
     aiohttp_session: aiohttp.ClientSession
@@ -307,6 +308,7 @@ def __init__(
         from .func.manager import Manager
         from .func.model import Model
         from .func.quota_scope import QuotaScope
+        from .func.ratelimit import RateLimit
         from .func.resource import Resource
         from .func.scaling_group import ScalingGroup
         from .func.server_log import ServerLog
@@ -344,6 +346,7 @@ def __init__(
         self.Service = Service
         self.Model = Model
         self.QuotaScope = QuotaScope
+        self.RateLimit = RateLimit
 
     @property
     def proxy_mode(self) -> bool:

src/ai/backend/manager/api/ratelimit.py

@@ -5,6 +5,7 @@
 from decimal import Decimal
 from typing import Final, Iterable, Tuple
 
+import aiohttp_cors
 import attrs
 from aiohttp import web
 from aiotools import apartial
@@ -14,6 +15,8 @@
 from ai.backend.common.logging import BraceStyleAdapter
 from ai.backend.common.networking import get_client_ip
 from ai.backend.common.types import RedisConnectionInfo
+from ai.backend.manager.api.auth import superadmin_required
+from ai.backend.manager.api.manager import READ_ALLOWED, server_status_required
 
 from .context import RootContext
 from .exceptions import RateLimitExceeded
@@ -145,6 +148,24 @@ async def init(app: web.Application) -> None:
     )
 
 
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+async def get_hot_anonymous_clients(request: web.Request) -> web.Response:
+    """
+    Retrieve a dictionary of anonymous client IP addresses and their corresponding suspicion scores.
+    suspicion scores are based on the number of requests made by the client.
+    """
+    log.info("RATELIMIT.GET_HOT_ANONYMOUS_CLIENTS ()")
+    rlimit_ctx: RateLimitContext = request.app["ratelimit.context"]
+    rr = rlimit_ctx.redis_rlim
+    result: list[tuple[bytes, float]] = await redis_helper.execute(
+        rr, lambda r: r.zrange("suspicious_ips", 0, -1, withscores=True)
+    )
+    suspicious_ips = {k.decode(): v for k, v in dict(result).items()}
+
+    return web.json_response(suspicious_ips, status=200)
+
+
 async def shutdown(app: web.Application) -> None:
     app_ctx: RateLimitContext = app["ratelimit.context"]
     await redis_helper.execute(app_ctx.redis_rlim, lambda r: r.flushdb())
@@ -157,6 +178,10 @@ def create_app(
     app = web.Application()
     app["api_versions"] = (1, 2, 3, 4)
     app["ratelimit.context"] = RateLimitContext()
+    app["prefix"] = "ratelimit"
+    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
+    add_route = app.router.add_route
+    cors.add(add_route("GET", "/hot_anonymous_clients", get_hot_anonymous_clients))
 
     app.on_startup.append(init)
     app.on_shutdown.append(shutdown)