diff --git a/ms_agent/utils/rate_limiter.py b/ms_agent/utils/rate_limiter.py index 3bb94ec52..20ce91d82 100644 --- a/ms_agent/utils/rate_limiter.py +++ b/ms_agent/utils/rate_limiter.py @@ -66,20 +66,20 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): async def _wait_if_needed(self): """Check and wait to satisfy rate limits""" - now = time.time() - - # 1. Check minimum request interval - if self._last_request_time > 0: - elapsed = now - self._last_request_time - if elapsed < self.min_request_interval: - wait_time = self.min_request_interval - elapsed - logger.debug( - f'Enforcing min interval: waiting {wait_time:.3f}s') - await asyncio.sleep(wait_time) - now = time.time() - - # 2. Check requests per second limit async with self._lock: + now = time.time() + + # 1. Check minimum request interval + if self._last_request_time > 0: + elapsed = now - self._last_request_time + if elapsed < self.min_request_interval: + wait_time = self.min_request_interval - elapsed + logger.debug( + f'Enforcing min interval: waiting {wait_time:.3f}s') + await asyncio.sleep(wait_time) + now = time.time() + + # 2. Check requests per second limit # Clean up expired request records (older than 1 second) cutoff_time = now - 1.0 while self._request_times and self._request_times[0] < cutoff_time: diff --git a/tests/utils/test_rate_limiter_pacing.py b/tests/utils/test_rate_limiter_pacing.py new file mode 100644 index 000000000..bac027611 --- /dev/null +++ b/tests/utils/test_rate_limiter_pacing.py @@ -0,0 +1,75 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Concurrent callers must observe the same request-start spacing.""" +import asyncio +import time + +import pytest +from ms_agent.utils.rate_limiter import RateLimiter + + +@pytest.mark.parametrize('max_concurrent', [1, 4]) +def test_minimum_interval_between_concurrent_requests(max_concurrent): + async def run(): + limiter = RateLimiter( + max_requests_per_second=100, + min_request_interval=0.05, + max_concurrent=max_concurrent, + ) + starts = [] + + async def request(): + async with limiter: + starts.append(time.monotonic()) + + await asyncio.wait_for( + asyncio.gather(*(request() for _ in range(4))), timeout=3) + gaps = [later - earlier for earlier, later in zip(starts, starts[1:])] + assert all(gap >= 0.045 for gap in gaps), gaps + + asyncio.run(run()) + + +def test_spacing_does_not_serialize_request_bodies(): + async def run(): + limiter = RateLimiter( + max_requests_per_second=100, + min_request_interval=0.05, + max_concurrent=3, + ) + starts = [] + all_entered = asyncio.Event() + + async def request(): + async with limiter: + starts.append(time.monotonic()) + if len(starts) == 3: + all_entered.set() + await all_entered.wait() + + await asyncio.wait_for( + asyncio.gather(*(request() for _ in range(3))), timeout=3) + gaps = [later - earlier for earlier, later in zip(starts, starts[1:])] + assert all(gap >= 0.045 for gap in gaps), gaps + + asyncio.run(run()) + + +def test_requests_per_second_still_limits_concurrent_requests(): + async def run(): + limiter = RateLimiter( + max_requests_per_second=2, + min_request_interval=0.05, + max_concurrent=3, + ) + starts = [] + + async def request(): + async with limiter: + starts.append(time.monotonic()) + + await asyncio.wait_for( + asyncio.gather(*(request() for _ in range(3))), timeout=3) + assert starts[1] - starts[0] >= 0.045 + assert starts[2] - starts[0] >= 1.0 + + asyncio.run(run())