Concurrency Limits and Rate Limiting

from __future__ import annotations

import asyncio
from collections import defaultdict


class KeyedLimiter:
    """Separate concurrency cap per key (host, provider, tenant)."""

    def __init__(self, default: int, overrides: dict[str, int] | None = None) -> None:
        overrides = overrides or {}
        self._sems: dict[str, asyncio.Semaphore] = defaultdict(lambda: asyncio.Semaphore(default))
        for key, n in overrides.items():
            self._sems[key] = asyncio.Semaphore(n)

    def __call__(self, key: str) -> asyncio.Semaphore:
        return self._sems[key]


async def main() -> None:
    limits = KeyedLimiter(default=4, overrides={"slow-site.example": 1})
    active: dict[str, int] = defaultdict(int)
    peak: dict[str, int] = defaultdict(int)

    async def fetch(host: str) -> None:
        async with limits(host):
            active[host] += 1
            peak[host] = max(peak[host], active[host])
            await asyncio.sleep(0.01)
            active[host] -= 1

    hosts = ["slow-site.example"] * 5 + ["fast.example"] * 10
    await asyncio.gather(*(fetch(h) for h in hosts))
    print(dict(peak))


asyncio.run(main())
# Output: {'slow-site.example': 1, 'fast.example': 4}

import asyncio
from typing import AsyncIterator, Awaitable, Callable, Iterable, Tuple, TypeVar

T = TypeVar("T")
R = TypeVar("R")


async def bounded_map(fn: Callable[[T], Awaitable[R]], items: Iterable[T], limit: int) -> AsyncIterator[Tuple[T, R]]:
    """Yield (item, result) as they complete, with at most `limit` in flight. Errors propagate."""
    it = iter(items)
    pending: dict[asyncio.Task, T] = {}

    def refill() -> None:
        while len(pending) < limit:
            try:
                item = next(it)
            except StopIteration:
                return
            pending[asyncio.ensure_future(fn(item))] = item

    refill()
    try:
        while pending:
            done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                item = pending.pop(task)
                yield item, task.result()
            refill()
    finally:
        for task in pending:                       # consumer stopped early or an error escaped
            task.cancel()
        await asyncio.gather(*pending, return_exceptions=True)


async def main() -> None:
    peak = active = 0

    async def embed(i: int) -> int:
        nonlocal peak, active
        active += 1
        peak = max(peak, active)
        await asyncio.sleep(0.001)
        active -= 1
        return i * 2

    total = 0
    async for _, vec in bounded_map(embed, range(1000), limit=8):
        total += vec
    print(total, "peak", peak)


asyncio.run(main())
# Output: 999000 peak 8

import asyncio
import time


class TokenBucket:
    def __init__(self, rate: float, capacity: float) -> None:
        if rate <= 0 or capacity <= 0:
            raise ValueError("rate and capacity must be positive")
        self.rate, self.capacity = rate, capacity
        self._tokens = capacity
        self._updated = time.monotonic()
        self._lock = asyncio.Lock()          # FIFO fairness: waiters are served in arrival order

    def _refill(self) -> None:
        now = time.monotonic()
        self._tokens = min(self.capacity, self._tokens + (now - self._updated) * self.rate)
        self._updated = now

    async def acquire(self, cost: float = 1.0) -> None:
        if cost > self.capacity:
            raise ValueError(f"cost {cost} exceeds bucket capacity {self.capacity}")
        async with self._lock:
            while True:
                self._refill()
                if self._tokens >= cost:
                    self._tokens -= cost
                    return
                await asyncio.sleep((cost - self._tokens) / self.rate)

    def refund(self, amount: float) -> None:
        self._refill()
        self._tokens = min(self.capacity, self._tokens + amount)


async def main() -> None:
    bucket = TokenBucket(rate=50, capacity=5)          # 50 req/s, burst of 5
    start = time.monotonic()
    await asyncio.gather(*(bucket.acquire() for _ in range(30)))
    elapsed = time.monotonic() - start
    print("30 requests:", 0.45 <= elapsed <= 0.8)      # 5 immediately, 25 more at 50/s = 0.5 s


asyncio.run(main())
# Output: 30 requests: True

import asyncio


class DualRateLimiter:
    """Waits until both the request-per-minute and token-per-minute budgets allow the call."""

    def __init__(self, rpm: int, tpm: int, burst_seconds: float = 6.0) -> None:
        self.requests = TokenBucket(rate=rpm / 60, capacity=max(1.0, rpm / 60 * burst_seconds))
        self.tokens = TokenBucket(rate=tpm / 60, capacity=max(1.0, tpm / 60 * burst_seconds))

    async def acquire(self, tokens: int) -> None:
        # Take tokens first: it is the scarcer budget for long prompts, and holding a request
        # slot while waiting for tokens would starve short requests.
        await self.tokens.acquire(min(tokens, self.tokens.capacity))
        await self.requests.acquire(1)

    def reconcile(self, estimated: int, actual: int) -> None:
        """Give back over-estimated tokens once the response reports real usage."""
        if actual < estimated:
            self.tokens.refund(estimated - actual)

async def chat(limiter: DualRateLimiter, sem: asyncio.Semaphore, client, messages, max_tokens: int = 512):
    estimate = estimate_tokens(messages) + max_tokens
    async with sem:                                   # concurrency (sockets, memory)
        await limiter.acquire(estimate)               # rate (RPM + TPM)
        resp = await with_retry(lambda: client.complete(messages, max_tokens=max_tokens))
    limiter.reconcile(estimate, resp.usage.total_tokens)
    return resp

class PausableLimiter:
    def __init__(self) -> None:
        self._resume_at = 0.0

    def pause_for(self, seconds: float) -> None:
        loop = asyncio.get_running_loop()
        self._resume_at = max(self._resume_at, loop.time() + seconds)

    async def wait(self) -> None:
        loop = asyncio.get_running_loop()
        delay = self._resume_at - loop.time()
        if delay > 0:
            await asyncio.sleep(delay)

import asyncio


class RateLimited(Exception):
    pass


class AdaptiveLimiter:
    def __init__(self, initial: int = 8, minimum: int = 1, maximum: int = 128) -> None:
        self.limit, self.minimum, self.maximum = initial, minimum, maximum
        self.in_flight = 0
        self._cond = asyncio.Condition()

    async def __aenter__(self) -> "AdaptiveLimiter":
        async with self._cond:
            await self._cond.wait_for(lambda: self.in_flight < self.limit)
            self.in_flight += 1
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        async with self._cond:
            self.in_flight -= 1
            if exc_type is None:
                self.limit = min(self.maximum, self.limit + 1 / self.limit)   # about +1 per window
            elif issubclass(exc_type, (RateLimited, TimeoutError)):
                self.limit = max(self.minimum, int(self.limit * 0.5))         # halve on overload
            self._cond.notify_all()


async def main() -> None:
    limiter = AdaptiveLimiter(initial=16)
    for _ in range(3):
        try:
            async with limiter:
                raise RateLimited("429")
        except RateLimited:
            pass
    print("limit after three 429s:", limiter.limit)


asyncio.run(main())
# Output: limit after three 429s: 2

from dataclasses import dataclass, field


@dataclass
class ModelBudget:
    rpm: int
    tpm: int
    max_concurrency: int
    limiter: "DualRateLimiter" = field(init=False)
    sem: asyncio.Semaphore = field(init=False)

    def __post_init__(self) -> None:
        self.limiter = DualRateLimiter(self.rpm, self.tpm)
        self.sem = asyncio.Semaphore(self.max_concurrency)


BUDGETS = {
    "provider-a/large": ModelBudget(rpm=500, tpm=200_000, max_concurrency=32),
    "provider-a/small": ModelBudget(rpm=5_000, tpm=2_000_000, max_concurrency=128),
    "provider-b/embed": ModelBudget(rpm=3_000, tpm=1_000_000, max_concurrency=64),
}