diff --git a/backend/open_webui/tasks.py b/backend/open_webui/tasks.py index f60a5b2c86..04bb6ce06f 100644 --- a/backend/open_webui/tasks.py +++ b/backend/open_webui/tasks.py @@ -1,6 +1,7 @@ # tasks.py import asyncio import logging +from contextlib import suppress from uuid import uuid4 from redis.asyncio import Redis @@ -20,25 +21,46 @@ REDIS_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks' REDIS_ITEM_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks:item' REDIS_RESPONSE_STREAMS_KEY = f'{REDIS_KEY_PREFIX}:tasks:response_streams' REDIS_PUBSUB_CHANNEL = f'{REDIS_KEY_PREFIX}:tasks:commands' +REDIS_PUBSUB_RECONNECT_INTERVAL = 1.0 +REDIS_PUBSUB_MAX_RECONNECT_INTERVAL = 30.0 async def redis_task_command_listener(app): redis: Redis = app.state.redis - pubsub = redis.pubsub() - await pubsub.subscribe(REDIS_PUBSUB_CHANNEL) + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL - async for message in pubsub.listen(): - if message['type'] != 'message': - continue + while True: + pubsub = None try: - command = JSONCodec.loads(message['data']) - if command.get('action') == 'stop': - task_id = command.get('task_id') - local_task = tasks.get(task_id) - if local_task: - local_task.cancel() + pubsub = redis.pubsub() + await pubsub.subscribe(REDIS_PUBSUB_CHANNEL) + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL + + async for message in pubsub.listen(): + if message['type'] != 'message': + continue + try: + command = JSONCodec.loads(message['data']) + if command.get('action') != 'stop': + continue + + local_task = tasks.get(command.get('task_id')) + if local_task: + local_task.cancel() + except Exception as e: + log.exception(f'Error handling distributed task command: {e}') + log.warning('Redis task command listener stopped. Retrying.') + except asyncio.CancelledError: + raise except Exception as e: - log.exception(f'Error handling distributed task command: {e}') + log.exception(f'Redis task command listener failed. Retrying: {e}') + finally: + if pubsub: + with suppress(Exception): + await pubsub.aclose() + + await asyncio.sleep(reconnect_interval) + reconnect_interval = min(reconnect_interval * 2, REDIS_PUBSUB_MAX_RECONNECT_INTERVAL) ### ------------------------------