From 7ded62b712f41978b4264791ea83a96c817d2544 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 26 Apr 2024 04:58:56 +0000 Subject: [PATCH 1/6] feat: Implement `get_hot_anonymous_clients` to RateLimit API --- src/ai/backend/client/func/ratelimit.py | 18 ++++++++++++++++++ src/ai/backend/client/session.py | 3 +++ src/ai/backend/manager/api/ratelimit.py | 22 ++++++++++++++++++++++ 3 files changed, 43 insertions(+) create mode 100644 src/ai/backend/client/func/ratelimit.py diff --git a/src/ai/backend/client/func/ratelimit.py b/src/ai/backend/client/func/ratelimit.py new file mode 100644 index 0000000000..6119463290 --- /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 interactions with Anonymous Ratelimit. + """ + + @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..e7e3e02474 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,21 @@ 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: + """ """ + log.info("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 +175,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) From 0eb185f647a09fa2eab4e499a23513962e0fb4d5 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 26 Apr 2024 05:33:03 +0000 Subject: [PATCH 2/6] chore: Add fragment --- changes/2075.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/2075.feature.md diff --git a/changes/2075.feature.md b/changes/2075.feature.md new file mode 100644 index 0000000000..4f1b47e67c --- /dev/null +++ b/changes/2075.feature.md @@ -0,0 +1 @@ +Implement `get_hot_anonymous_clients` to RateLimit API From 1e9f65ca97308934b996054c09c360300949f641 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 26 Apr 2024 08:49:11 +0000 Subject: [PATCH 3/6] chore: add prefix logging message --- src/ai/backend/manager/api/ratelimit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/api/ratelimit.py b/src/ai/backend/manager/api/ratelimit.py index e7e3e02474..9365a1dcb5 100644 --- a/src/ai/backend/manager/api/ratelimit.py +++ b/src/ai/backend/manager/api/ratelimit.py @@ -152,7 +152,7 @@ async def init(app: web.Application) -> None: @superadmin_required async def get_hot_anonymous_clients(request: web.Request) -> web.Response: """ """ - log.info("GET_HOT_ANONYMOUS_CLIENTS ()") + 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( From 63aeb6ac9c2e3aec703e6c9f85e873d8712dc5ab Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 26 Apr 2024 08:55:11 +0000 Subject: [PATCH 4/6] chore: Add comment for `get_hot_anonymous_clients` --- src/ai/backend/manager/api/ratelimit.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ai/backend/manager/api/ratelimit.py b/src/ai/backend/manager/api/ratelimit.py index 9365a1dcb5..a6779dbda0 100644 --- a/src/ai/backend/manager/api/ratelimit.py +++ b/src/ai/backend/manager/api/ratelimit.py @@ -151,7 +151,10 @@ 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 From f4423d7d0b8575687d2518d2a640685424b075eb Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Fri, 26 Apr 2024 09:32:12 +0000 Subject: [PATCH 5/6] chore: Update comment --- src/ai/backend/client/func/ratelimit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/client/func/ratelimit.py b/src/ai/backend/client/func/ratelimit.py index 6119463290..0449abe3ef 100644 --- a/src/ai/backend/client/func/ratelimit.py +++ b/src/ai/backend/client/func/ratelimit.py @@ -6,7 +6,7 @@ class RateLimit(BaseFunction): """ - Provides interactions with Anonymous Ratelimit. + Provides RateLimiting API functions. """ @api_function From d11c05090fdb404fbe46122146588638c824e832 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 5 Aug 2024 04:16:31 +0000 Subject: [PATCH 6/6] docs: rename news fragment --- changes/2075.feature.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changes/2075.feature.md b/changes/2075.feature.md index 4f1b47e67c..1da9a5f9f1 100644 --- a/changes/2075.feature.md +++ b/changes/2075.feature.md @@ -1 +1 @@ -Implement `get_hot_anonymous_clients` to RateLimit API +Implement RateLimit API and `get_hot_anonymous_clients`. \ No newline at end of file