diff --git a/changes/2075.feature.md b/changes/2075.feature.md new file mode 100644 index 0000000000..1da9a5f9f1 --- /dev/null +++ b/changes/2075.feature.md @@ -0,0 +1 @@ +Implement RateLimit API and `get_hot_anonymous_clients`. \ No newline at end of file diff --git a/src/ai/backend/client/func/ratelimit.py b/src/ai/backend/client/func/ratelimit.py new file mode 100644 index 0000000000..0449abe3ef --- /dev/null +++ b/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() diff --git a/src/ai/backend/client/session.py b/src/ai/backend/client/session.py index 312b626fbd..76e3e47698 100644 --- a/src/ai/backend/client/session.py +++ b/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: diff --git a/src/ai/backend/manager/api/ratelimit.py b/src/ai/backend/manager/api/ratelimit.py index c07541d838..a6779dbda0 100644 --- a/src/ai/backend/manager/api/ratelimit.py +++ b/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)