Skip to content

Commit

Permalink
feat: Implement get_hot_anonymous_clients to RateLimit API
Browse files Browse the repository at this point in the history
  • Loading branch information
jopemachine committed Apr 26, 2024
1 parent 5bd661c commit 379b972
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 0 deletions.
18 changes: 18 additions & 0 deletions src/ai/backend/client/func/ratelimit.py
Original file line number Diff line number Diff line change
@@ -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()
3 changes: 3 additions & 0 deletions src/ai/backend/client/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,7 @@ class BaseSession(metaclass=abc.ABCMeta):
"Service",
"Model",
"QuotaScope",
"RateLimit",
)

aiohttp_session: aiohttp.ClientSession
Expand Down Expand Up @@ -306,6 +307,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
Expand Down Expand Up @@ -343,6 +345,7 @@ def __init__(
self.Service = Service
self.Model = Model
self.QuotaScope = QuotaScope
self.RateLimit = RateLimit

@property
def proxy_mode(self) -> bool:
Expand Down
22 changes: 22 additions & 0 deletions src/ai/backend/manager/api/ratelimit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -146,6 +149,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())
Expand All @@ -158,6 +176,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)
Expand Down

0 comments on commit 379b972

Please sign in to comment.