From 0915ad05a86ec75ec2d51da5c59c99ed8cbae9b8 Mon Sep 17 00:00:00 2001 From: Osamaali313 Date: Wed, 15 Jul 2026 22:52:48 +0300 Subject: [PATCH 1/2] fix(rate_limiter): use a threading.Lock for the synchronous methods RateLimiter creates self._lock = asyncio.Lock() and correctly uses it as `async with` in the async _wait_if_needed path. But the synchronous methods get_stats/reset and AdaptiveRateLimiter.record_success/record_error/get_stats acquire the same asyncio lock with a plain `with self._lock:`. asyncio.Lock implements only __aenter__/__aexit__, so `with` raises TypeError: "'asyncio.locks.Lock' object does not support the context manager protocol". Every one of those sync calls fails -- e.g. the findata fetcher calls record_success()/record_error() around each adaptive-rate-limited fetch, so the first success throws and the except handler's record_error throws again. Add a dedicated threading.Lock (self._state_lock) for the sync bookkeeping methods; the async request-pacing path keeps the asyncio lock. --- ms_agent/utils/rate_limiter.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/ms_agent/utils/rate_limiter.py b/ms_agent/utils/rate_limiter.py index 3bb94ec52..ee69d08ac 100644 --- a/ms_agent/utils/rate_limiter.py +++ b/ms_agent/utils/rate_limiter.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio +import threading import time from collections import deque from typing import Callable @@ -47,6 +48,10 @@ def __init__( # Concurrency control self._semaphore = asyncio.Semaphore(max_concurrent) self._lock = asyncio.Lock() + # Separate synchronous lock for the non-async bookkeeping methods + # (get_stats/reset/record_success/record_error); an asyncio.Lock only + # supports `async with`, so using it under a plain `with` raises. + self._state_lock = threading.Lock() logger.info( f'RateLimiter initialized: {max_requests_per_second} req/s, ' @@ -137,7 +142,7 @@ def get_stats(self) -> dict: Returns: Dictionary containing current state """ - with self._lock: + with self._state_lock: now = time.time() cutoff_time = now - 1.0 recent_requests = sum(1 for t in self._request_times @@ -161,7 +166,7 @@ def get_stats(self) -> dict: def reset(self): """Reset rate limiter state""" - with self._lock: + with self._state_lock: self._request_times.clear() self._last_request_time = 0.0 logger.info('RateLimiter reset') @@ -219,7 +224,7 @@ def __init__( def record_success(self): """Record successful request""" - with self._lock: + with self._state_lock: self._total_requests += 1 self._consecutive_successes += 1 self._consecutive_errors = 0 @@ -243,7 +248,7 @@ def record_error(self, is_rate_limit_error: bool = False): Args: is_rate_limit_error: Whether this is a rate limit error """ - with self._lock: + with self._state_lock: self._total_requests += 1 self._total_errors += 1 self._consecutive_errors += 1 @@ -279,7 +284,7 @@ def record_error(self, is_rate_limit_error: bool = False): def get_stats(self) -> dict: """Get extended statistics""" stats = super().get_stats() - with self._lock: + with self._state_lock: stats.update({ 'total_requests': self._total_requests, From 18e97f560d5a73aa317f336dc56612c126ac50ee Mon Sep 17 00:00:00 2001 From: Osamaali313 Date: Wed, 15 Jul 2026 23:46:32 +0300 Subject: [PATCH 2/2] fix(rate_limiter): close state-lock gap in _wait_if_needed; resize adaptive deque Addresses review feedback: - _wait_if_needed mutated _request_times/_last_request_time under only the asyncio lock, which does not synchronize with the threading _state_lock used by the synchronous bookkeeping methods (callable from other threads). Guard those state accesses with _state_lock too, in brief critical sections that are never held across an await; _lock still serializes coroutines. - AdaptiveRateLimiter adjusts max_requests_per_second, but _request_times was a deque with a fixed maxlen, so raising the rate never took effect (len could never reach the new higher limit). Recreate the deque with the new maxlen (preserving recent history) whenever the rate changes. --- ms_agent/utils/rate_limiter.py | 70 +++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 23 deletions(-) diff --git a/ms_agent/utils/rate_limiter.py b/ms_agent/utils/rate_limiter.py index ee69d08ac..786028f92 100644 --- a/ms_agent/utils/rate_limiter.py +++ b/ms_agent/utils/rate_limiter.py @@ -74,8 +74,10 @@ async def _wait_if_needed(self): now = time.time() # 1. Check minimum request interval - if self._last_request_time > 0: - elapsed = now - self._last_request_time + with self._state_lock: + last_request_time = self._last_request_time + if last_request_time > 0: + elapsed = now - last_request_time if elapsed < self.min_request_interval: wait_time = self.min_request_interval - elapsed logger.debug( @@ -83,33 +85,43 @@ async def _wait_if_needed(self): await asyncio.sleep(wait_time) now = time.time() - # 2. Check requests per second limit + # 2. Check requests per second limit. _lock serializes coroutines; + # _state_lock (a threading.Lock shared with the synchronous bookkeeping + # methods) guards the actual shared state and is never held across an + # await. async with self._lock: - # 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: - self._request_times.popleft() - - # If rate limit reached, wait until oldest request expires - if len(self._request_times) >= self.max_requests_per_second: - oldest_request = self._request_times[0] - wait_time = 1.0 - (now - - oldest_request) + 0.01 # Add 10ms margin - if wait_time > 0: - logger.debug( - f'Rate limit reached ({self.max_requests_per_second} req/s): ' - f'waiting {wait_time:.3f}s') - await asyncio.sleep(wait_time) - now = time.time() - # Clean up expired records + with self._state_lock: + # 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: + self._request_times.popleft() + + # If rate limit reached, wait until oldest request expires + if len(self._request_times) >= self.max_requests_per_second: + oldest_request = self._request_times[0] + wait_time = 1.0 - (now - oldest_request + ) + 0.01 # Add 10ms margin + else: + wait_time = 0.0 + + if wait_time > 0: + logger.debug( + f'Rate limit reached ({self.max_requests_per_second} req/s): ' + f'waiting {wait_time:.3f}s') + await asyncio.sleep(wait_time) + now = time.time() + with self._state_lock: + # Clean up expired records after waiting cutoff_time = now - 1.0 while self._request_times and self._request_times[ 0] < cutoff_time: self._request_times.popleft() - # Record this request time - self._request_times.append(now) - self._last_request_time = now + with self._state_lock: + # Record this request time + self._request_times.append(now) + self._last_request_time = now async def execute(self, func: Callable, *args, **kwargs): """ @@ -236,6 +248,10 @@ def record_success(self): round(old_rps * self._recovery_factor), self._max_rps) if new_rps > old_rps: self.max_requests_per_second = new_rps + # deque maxlen is immutable, so resize it to the new limit + # (preserving recent history) or the rps change has no effect. + self._request_times = deque( + self._request_times, maxlen=new_rps) logger.info( f'Rate limit increased: {old_rps} → {new_rps} req/s ' f'(after {self._consecutive_successes} successes)') @@ -261,6 +277,10 @@ def record_error(self, is_rate_limit_error: bool = False): int(old_rps * self._backoff_factor), self._min_rps) if new_rps < old_rps: self.max_requests_per_second = new_rps + # deque maxlen is immutable, so resize it to the new limit + # (preserving recent history) or the rps change has no effect. + self._request_times = deque( + self._request_times, maxlen=new_rps) # Also increase minimum request interval self.min_request_interval = min( self.min_request_interval * 1.5, 2.0) @@ -276,6 +296,10 @@ def record_error(self, is_rate_limit_error: bool = False): int(old_rps * self._backoff_factor), self._min_rps) if new_rps < old_rps: self.max_requests_per_second = new_rps + # deque maxlen is immutable, so resize it to the new limit + # (preserving recent history) or the rps change has no effect. + self._request_times = deque( + self._request_times, maxlen=new_rps) logger.warning( f'Multiple errors detected! Reducing rate: {old_rps} → {new_rps} req/s ' f'(after {self._consecutive_errors} errors)')