upstream-archive/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-09-21 13:38:19 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix: keep the usage pool cleanup task alive across lock loss and Redis errors (#28834)

Browse Source 
With WEBSOCKET_MANAGER=redis on a multi-node deployment, the usage pool cleanup task could stop permanently for the whole cluster. Nodes that lost the startup lock race gave up for good after three attempts, and the winner died on a single failed renew or on any Redis connection error, releasing the lock with nobody left to take it over. From then on expired entries accumulated in the usage pool until a node restarted, so /api/usage over-reported models in use and every disconnect handler walked an ever-growing pool.

The task now retries lock acquisition forever like the session pool cleanup does, and any error is logged and answered by releasing the lock and returning to acquisition, so a transient failure costs one cleanup cycle and every node stays a takeover candidate. The delete of an emptied model entry is KeyError-guarded because a disconnect handler on another node can remove the same key between the sweep's snapshot and its delete; unguarded, that race was a permanent task killer that needed nothing rarer than a chat finishing while its tab closed.
This commit is contained in:
Classic298
2026-08-20 12:59:38 -07:00
committed by GitHub
parent b0fdc00452
commit 5586964bb2
1 changed files with 35 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/socket/main.py
+35 -34
View File
@@ -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(