diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 9224501767..7739636887 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -218,47 +218,48 @@ async def periodic_session_pool_cleanup(): async def periodic_usage_pool_cleanup(): - max_retries = 2 retry_delay = random.uniform(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT) - for attempt in range(max_retries + 1): - if aquire_func(): - break - else: - if attempt < max_retries: - log.debug('Cleanup lock already exists. Retry %s after %ss...', attempt + 1, retry_delay) + while True: + try: + if not aquire_func(): + log.debug('Usage cleanup lock held by another node. Retrying.') await asyncio.sleep(retry_delay) - else: - log.warning('Failed to acquire cleanup lock after retries. Skipping cleanup.') - return + continue - log.debug('Running periodic_cleanup') - try: - while True: - if not renew_func(): - log.error('Unable to renew cleanup lock. Exiting usage pool cleanup.') - raise Exception('Unable to renew usage pool cleanup lock.') + try: + while True: + if not renew_func(): + log.warning('Unable to renew usage cleanup lock. Retrying cleanup ownership.') + break - now = int(time.time()) - for model_id, connections in list(USAGE_POOL.items()): - # Creating a list of sids to remove if they have timed out - expired_sids = [ - sid for sid, details in connections.items() if now - details['updated_at'] > TIMEOUT_DURATION - ] + now = int(time.time()) + for model_id, connections in list(USAGE_POOL.items()): + expired_sids = [ + sid + for sid, details in connections.items() + if now - details['updated_at'] > TIMEOUT_DURATION + ] - if connections and not expired_sids: - continue + if connections and not expired_sids: + continue - for sid in expired_sids: - del connections[sid] + for sid in expired_sids: + del connections[sid] - if not connections: - log.debug('Cleaning up model %s from usage pool', model_id) - del USAGE_POOL[model_id] - else: - USAGE_POOL[model_id] = connections - await asyncio.sleep(TIMEOUT_DURATION) - finally: - release_func() + if not connections: + log.debug('Cleaning up model %s from usage pool', model_id) + try: + del USAGE_POOL[model_id] + except KeyError: + pass + else: + USAGE_POOL[model_id] = connections + await asyncio.sleep(TIMEOUT_DURATION) + finally: + release_func() + except Exception: + log.exception('Usage pool cleanup failed. Retrying.') + await asyncio.sleep(retry_delay) app = socketio.ASGIApp(