scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-28 05:27:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

perf: write response streams to Redis incrementally

Browse source 
Persisting an in-flight response stream rewrote the whole accumulated
message into one global hash field on every streamed delta: quadratic
write volume, about 2 GB for a 32k-token response, about 16 GB for a
tool-heavy 8 MB turn, 43000x amplification observed on a production fleet.

Streams now live in per-stream keys with a TTL from creation: one hash
field per output item, an append-only journal following the actively
growing text lane and an append-only content string. Unchanged items are
never rewritten and slowly drifting lanes batch under a geometric budget,
keeping bytes proportional to new content while the text the UI renders
stays byte-exact at every save, so reloading mid-response never shows gaps.
Measured: 2263x fewer bytes on a 4000-delta message, 356x on a tool-heavy
turn; the worst case, two items alternating growth on every delta, stays at
roughly the old writer's cost, and the writer holds about twice the
in-flight message in memory. Throttling the saves instead was rejected
because it keeps the quadratic total and serves stale reconnect overlays.
Redis Cluster, Sentinel and the non-Redis fallback behave as before,
rolling upgrades keep old-format entries readable and this supersedes the
best-effort HEXPIRE of the previous fix.
This commit is contained in:
Classic298 2026-08-20 19:28:07 +02:00
parent 35fbde0a3f
commit 91d9bd8f1d
2 changed files with 499 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
backend/open_webui/main.py
View file

@ -194,6 +194,7 @@ from open_webui.socket.main import (
)
from open_webui.tasks import (
cleanup_task,
clear_response_stream,
create_task,
has_active_tasks,
list_task_ids_by_item_id,
@ -1750,6 +1751,9 @@ async def chat_completion(
)
except asyncio.CancelledError:
pass
elif chat_id:
# mirrors middleware's fallback: with no task the stream is keyed by message id
await clear_response_stream(request.app.state.redis, metadata.get('message_id'))
except Exception:
pass

530
backend/open_webui/tasks.py
View file

@ -1,13 +1,16 @@
# tasks.py
import asyncio
import logging
import time
from contextlib import suppress
from copy import deepcopy
from uuid import uuid4
from redis.asyncio import Redis
from open_webui.env import REDIS_KEY_PREFIX, REDIS_RESPONSE_STREAM_TTL
from open_webui.utils.json_codec import JSONCodec, dumps_bytes
from open_webui.utils.misc import get_output_text, sanitize_text_for_db
log = logging.getLogger(__name__)
@ -20,10 +23,67 @@ response_streams: dict[str, dict] = {}
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'
# lane before id: a client-supplied id must not collide with another stream's key
REDIS_RESPONSE_STREAM_KEY_PREFIX = f'{REDIS_KEY_PREFIX}:tasks:response_stream'
REDIS_PUBSUB_CHANNEL = f'{REDIS_KEY_PREFIX}:tasks:commands'
REDIS_PUBSUB_RECONNECT_INTERVAL = 1.0
REDIS_PUBSUB_MAX_RECONNECT_INTERVAL = 30.0
RESPONSE_STREAM_EXPIRE_SECONDS = REDIS_RESPONSE_STREAM_TTL if REDIS_RESPONSE_STREAM_TTL > 0 else None
# TTL refresh cadence; also bounds how long an evicted ledger key goes unnoticed
RESPONSE_STREAM_EXPIRE_REFRESH_SECONDS = (
min(10, RESPONSE_STREAM_EXPIRE_SECONDS / 2) if RESPONSE_STREAM_EXPIRE_SECONDS else 10
)
# Append growth outside the journaled leaf is rewritten in batches of at least this size
RESPONSE_STREAM_DRIFT_REWRITE_CHARS = 4096
# Per-process incremental-writer state for in-flight response streams, keyed by task id
_stream_states: dict[str, dict] = {}
# Sentinel distinguishing "structures differ" from "no difference" in _grown_leaves
_DIFFERS = object()
# Marks a slot whose field does not exist in the hash yet; None is not usable
# as the marker because output items themselves may legally be None
_UNWRITTEN = object()
def _ledger_key(task_id: str) -> str:
return f'{REDIS_RESPONSE_STREAM_KEY_PREFIX}:ledger:{task_id}'
def _content_key(task_id: str) -> str:
return f'{REDIS_RESPONSE_STREAM_KEY_PREFIX}:content:{task_id}'
def _journal_key(task_id: str) -> str:
return f'{REDIS_RESPONSE_STREAM_KEY_PREFIX}:journal:{task_id}'
def _queue_stream_key_deletes(pipe, task_id: str):
# separate deletes: the three keys hash to different cluster slots
pipe.delete(_ledger_key(task_id))
pipe.delete(_content_key(task_id))
pipe.delete(_journal_key(task_id))
# previous-release writers still use the global hash; drop with the next release
pipe.hdel(REDIS_RESPONSE_STREAMS_KEY, task_id)
def _new_stream_state() -> dict:
return {
'slots': [], # one slot per output item
'next_field': 0,
'saved_order': None, # fkey list as last written into the meta field
'saved_journal': None, # [fkey, leaf path, nonce] of the active journal
'journal_bytes': 0, # server-side byte length of the journal value
'content': '',
'content_bytes': 0, # server-side byte length of the content value
# keys are armed at creation, so the first probe is due one refresh later
'expire_after': time.monotonic() + RESPONSE_STREAM_EXPIRE_REFRESH_SECONDS,
}
async def redis_task_command_listener(app):
redis: Redis = app.state.redis
@ -79,7 +139,7 @@ async def redis_save_task(redis: Redis, task_id: str, item_id: str | None):
async def redis_cleanup_task(redis: Redis, task_id: str, item_id: str | None):
pipe = redis.pipeline()
pipe.hdel(REDIS_TASKS_KEY, task_id)
pipe.hdel(REDIS_RESPONSE_STREAMS_KEY, task_id)
_queue_stream_key_deletes(pipe, task_id)
if item_id:
pipe.srem(f'{REDIS_ITEM_TASKS_KEY}:{item_id}', task_id)
await pipe.execute()
@ -112,11 +172,10 @@ async def cleanup_task(redis, task_id: str, id=None):
"""
Remove a completed or canceled task from the global `tasks` dictionary.
"""
if redis:
await redis_cleanup_task(redis, task_id, id)
# local state first: it must go even when Redis is unreachable
tasks.pop(task_id, None) # Remove the task if it exists
response_streams.pop(task_id, None)
_stream_states.pop(task_id, None)
# If an ID is provided, remove the task from the item_tasks dictionary
if id and task_id in item_tasks.get(id, []):
@ -124,6 +183,9 @@ async def cleanup_task(redis, task_id: str, id=None):
if not item_tasks[id]: # If no tasks left for this ID, remove the entry
item_tasks.pop(id, None)
if redis:
await redis_cleanup_task(redis, task_id, id)
async def create_task(redis, coroutine, id=None, task_id=None):
"""
@ -166,6 +228,305 @@ async def list_task_ids_by_item_id(redis, id):
return item_tasks.get(id, [])
def _rewrite_budget(item_json_chars: int) -> int:
# scales with the item's serialized size so batch rewrites amortize to O(new content)
return max(RESPONSE_STREAM_DRIFT_REWRITE_CHARS, item_json_chars // 4)
def _grown_leaves(previous, current, path=()): # noqa: C901
"""Compare two JSON structures; return None when identical, a list of
(path, grown chars) when every difference is a string leaf growing by
appending, _DIFFERS otherwise."""
if isinstance(previous, str) and isinstance(current, str):
if previous == current:
return None
if len(current) > len(previous) and current.startswith(previous):
return [(path, len(current) - len(previous))]
return _DIFFERS
if isinstance(previous, dict) and isinstance(current, dict):
if list(previous.keys()) != list(current.keys()):
return _DIFFERS
pairs = [(previous[key], current[key], key) for key in current]
elif isinstance(previous, list) and isinstance(current, list):
if len(previous) != len(current):
return _DIFFERS
pairs = [(p, c, index) for index, (p, c) in enumerate(zip(previous, current))]
else:
# type-sensitive equality: True == 1 must not count as identical
if type(previous) is type(current) and previous == current:
return None
return _DIFFERS
found = []
for child_previous, child_current, step in pairs:
result = _grown_leaves(child_previous, child_current, path + (step,))
if result is None:
continue
if result is _DIFFERS:
return _DIFFERS
found.extend(result)
return found or None
def _set_leaf(node, path, value):
for step in path[:-1]:
node = node[step]
node[path[-1]] = value
def _leaf_value(node, path):
for step in path:
node = node[step]
return node
def _total_growth(grown, skip=None):
return sum(growth for path, growth in grown if path != skip)
def _item_id(item):
# output is user-editable via the chats API: items may not be dicts, ids may not be str
item_id = item.get('id') if isinstance(item, dict) else None
return item_id if isinstance(item_id, str) else None
def _choose_journal_slot(slots, diffs, appendable, journal_ref):
"""The journal follows growth, not position: it stays on its item while that
item keeps appending and moves to the biggest grower once it goes quiet.
A demoted slot's diff is set to _DIFFERS in place so it rewrites in full
(its field lags the journal). Returns (journal_ref, journaled fkey)."""
if journal_ref:
for index, slot in enumerate(slots):
if slot['fkey'] != journal_ref[0]:
continue
bound_diff = diffs[index]
if bound_diff is _DIFFERS or (bound_diff is None and appendable):
diffs[index] = _DIFFERS
break
return journal_ref, journal_ref[0]
# falling out of the loop means the bound item was realigned away
if appendable:
return None, max(appendable, key=lambda fkey: _total_growth(appendable[fkey]))
return None, None
def _collect_changed_fields(slots, output, journal_ref):
"""Serialize every item whose content changed. Returns the dirty fields,
how many of them are new, the journal op and the updated journal ref."""
diffs = [
_DIFFERS if slot['snapshot'] is _UNWRITTEN else _grown_leaves(slot['snapshot'], item)
for slot, item in zip(slots, output)
]
growing = {slot['fkey']: diff for slot, diff in zip(slots, diffs) if isinstance(diff, list)}
# a root-level string leaf has an empty path and cannot be journaled
appendable = {fkey: diff for fkey, diff in growing.items() if all(path for path, _ in diff)}
journal_ref, target_fkey = _choose_journal_slot(slots, diffs, appendable, journal_ref)
fields = {}
new_fields = 0
journal = None
for slot, item, diff in zip(slots, output, diffs):
fkey = slot['fkey']
if diff is None:
continue
if fkey == target_fkey:
write_field, journal, journal_ref = _diff_journal_item(slot, journal_ref, item, appendable[fkey])
if not write_field:
continue
elif fkey in growing and _total_growth(growing[fkey]) <= _rewrite_budget(slot['item_json_chars']):
# non-journaled append growth defers under the geometric budget
continue
if slot['snapshot'] is _UNWRITTEN:
new_fields += 1
fields[fkey] = JSONCodec.dumps(item)
slot['snapshot'] = deepcopy(item)
slot['item_json_chars'] = len(fields[fkey])
return fields, new_fields, journal, journal_ref
def _is_rendered(path):
# the lanes structuredOutput.ts renders as message text
return len(path) == 3 and path[0] in ('content', 'summary') and path[2] == 'text'
def _journal_leaf(grown, item):
"""Pick the leaf to journal: prefer what the UI renders (content and summary
text lanes) over larger unrendered twins such as reasoning_details."""
rendered = [leaf_path for leaf_path, _ in grown if _is_rendered(leaf_path)]
candidates = rendered or [leaf_path for leaf_path, _ in grown]
return max(candidates, key=lambda p: len(_leaf_value(item, p)))
def _diff_journal_item(slot, journal_ref, item, grown):
"""Decide how to persist the journaled item. Returns
(write_field, journal op, journal_ref)."""
preferred = _journal_leaf(grown, item)
journal = None
if journal_ref is None:
text = _leaf_value(item, preferred)
journal = ('set', text)
journal_ref = [slot['fkey'], preferred, uuid4().hex]
_set_leaf(slot['snapshot'], preferred, text)
else:
if _is_rendered(preferred) and not _is_rendered(journal_ref[1]):
# a rendered lane started growing while an unrendered one holds the
# journal: rewrite in full so the next save re-binds to it
return True, None, None
growth = next((growth for leaf_path, growth in grown if leaf_path == journal_ref[1]), 0)
if not growth and preferred != journal_ref[1]:
# the bound leaf idles while another grows: rewrite in full, re-bind next save
return True, None, None
# growth > 0 here: the idle case returned above
leaf = _leaf_value(item, journal_ref[1])
journal = ('append', leaf[len(leaf) - growth :])
_set_leaf(slot['snapshot'], journal_ref[1], leaf)
drift = _total_growth(grown, skip=journal_ref[1])
if drift > _rewrite_budget(slot['item_json_chars']):
if preferred != journal_ref[1]:
# a better leaf outranks the journaled one: drop the journal
# with the full write so the next save re-binds it
return True, None, None
return True, journal, journal_ref
return False, journal, journal_ref
def _realign_slots(state, output):
previous_slots = state['slots']
by_obj_id = {id(slot['item']): index for index, slot in enumerate(previous_slots)}
matched = set()
slots = [None] * len(output)
for position, item in enumerate(output):
index = by_obj_id.get(id(item))
if index is not None and index not in matched:
matched.add(index)
slots[position] = previous_slots[index]
by_item_id = {}
for index, slot in enumerate(previous_slots):
if index not in matched:
item_id = _item_id(slot['item'])
if item_id is not None:
by_item_id.setdefault(item_id, []).append(index)
for position, item in enumerate(output):
if slots[position] is not None:
continue
candidates = by_item_id.get(_item_id(item))
if candidates:
index = candidates.pop(0)
matched.add(index)
slots[position] = previous_slots[index]
else:
slots[position] = {
'fkey': f'i{state["next_field"]}',
'item': None,
'snapshot': _UNWRITTEN,
'item_json_chars': 0,
}
state['next_field'] += 1
state['slots'] = slots
return [slot['fkey'] for index, slot in enumerate(previous_slots) if index not in matched]
async def _refresh_ttls(redis, state, task_id) -> bool:
"""Periodic TTL refresh that doubles as the key-loss probe."""
if time.monotonic() < state['expire_after']:
return True
keys = [_ledger_key(task_id)]
if state['content']:
keys.append(_content_key(task_id))
if state['saved_journal']:
keys.append(_journal_key(task_id))
for key in keys: # separate calls: the keys hash to different cluster slots
if RESPONSE_STREAM_EXPIRE_SECONDS:
alive = await redis.expire(key, RESPONSE_STREAM_EXPIRE_SECONDS)
else:
alive = await redis.exists(key)
if not alive:
return False
state['expire_after'] = time.monotonic() + RESPONSE_STREAM_EXPIRE_REFRESH_SECONDS
return True
async def _write_stream(redis, state, task_id, chat_id, message_id, content, output) -> bool: # noqa: C901
"""One incremental persistence pass. False means the server state no longer
matches the local mirror and the stream must be rebuilt."""
ledger_key = _ledger_key(task_id)
content_key = _content_key(task_id)
journal_key = _journal_key(task_id)
orphans = []
if len(state['slots']) != len(output):
orphans = _realign_slots(state, output)
slots = state['slots']
fields, new_fields, journal, journal_ref = _collect_changed_fields(slots, output, state['saved_journal'])
for slot, item in zip(slots, output):
slot['item'] = item # pin: keeps id(item) unique while the slot lives
order = [slot['fkey'] for slot in slots]
if order != state['saved_order'] or journal_ref != state['saved_journal']:
if state['saved_order'] is None:
new_fields += 1
fields['meta'] = JSONCodec.dumps(
{'v': 1, 'chat_id': chat_id, 'message_id': message_id, 'order': order, 'journal': journal_ref}
)
# text lanes first: a torn read then sees newer text, never older;
# stripped like the DB convention, raw lanes would raise on strict UTF-8
if journal is not None:
kind, text = journal
if kind == 'set':
journal_value = f'{journal_ref[2]}:{sanitize_text_for_db(text)}'
await redis.set(journal_key, journal_value, ex=RESPONSE_STREAM_EXPIRE_SECONDS)
state['journal_bytes'] = len(journal_value.encode('utf-8'))
else:
journal_suffix = sanitize_text_for_db(text)
expected = state['journal_bytes'] + len(journal_suffix.encode('utf-8'))
if await redis.append(journal_key, journal_suffix) != expected:
return False
state['journal_bytes'] = expected
previous = state['content']
if content != previous:
state['content'] = content
if previous and content.startswith(previous):
# per-code-point stripping distributes over concatenation
suffix = sanitize_text_for_db(content[len(previous) :])
expected = state['content_bytes'] + len(suffix.encode('utf-8'))
if await redis.append(content_key, suffix) != expected:
return False
state['content_bytes'] = expected
else:
sanitized = sanitize_text_for_db(content)
await redis.set(content_key, sanitized, ex=RESPONSE_STREAM_EXPIRE_SECONDS)
state['content_bytes'] = len(sanitized.encode('utf-8'))
if fields:
if RESPONSE_STREAM_EXPIRE_SECONDS and state['saved_order'] is None:
# one packet: a kill between HSET and EXPIRE would leak an unexpiring key
pipe = redis.pipeline()
pipe.hset(ledger_key, mapping=fields)
pipe.expire(ledger_key, RESPONSE_STREAM_EXPIRE_SECONDS)
added = (await pipe.execute())[0]
else:
added = await redis.hset(ledger_key, mapping=fields)
if added != new_fields:
return False
state['saved_order'] = order
state['saved_journal'] = journal_ref
if orphans:
# after the HSET: the new meta no longer names these fields
await redis.hdel(ledger_key, *orphans)
return await _refresh_ttls(redis, state, task_id)
async def save_response_stream(
redis,
task_id: str | None,
@ -177,51 +538,150 @@ async def save_response_stream(
if not task_id or not chat_id or not message_id:
return
data = {
'chat_id': chat_id,
'message_id': message_id,
'content': content,
'output': output,
}
if not redis:
response_streams[task_id] = {
'chat_id': chat_id,
'message_id': message_id,
'content': content,
'output': output,
}
return
if redis:
await redis.hset(REDIS_RESPONSE_STREAMS_KEY, task_id, dumps_bytes(data))
if REDIS_RESPONSE_STREAM_TTL > 0:
with suppress(Exception):
await redis.hexpire(REDIS_RESPONSE_STREAMS_KEY, REDIS_RESPONSE_STREAM_TTL, task_id)
else:
response_streams[task_id] = data
state = _stream_states.get(task_id)
if state is None:
state = _stream_states[task_id] = _new_stream_state()
try:
if await _write_stream(redis, state, task_id, chat_id, message_id, content, output):
return
log.info(f'Response stream {task_id} for chat {chat_id} state out of sync, rebuilding')
state = _stream_states[task_id] = _new_stream_state()
pipe = redis.pipeline()
_queue_stream_key_deletes(pipe, task_id)
await pipe.execute()
if await _write_stream(redis, state, task_id, chat_id, message_id, content, output):
return
except BaseException:
# the local mirror may be ahead of Redis: drop it and re-prove next save
_stream_states.pop(task_id, None)
raise
# even the rebuild failed (e.g. a concurrent cleanup): start clean next save
_stream_states.pop(task_id, None)
log.warning(
f'Response stream {task_id} for chat {chat_id} was not persisted this save, retrying from scratch next save'
)
async def _load_stream(redis, task_id: str, fields: dict) -> dict | None:
for is_retry in (False, True):
try:
meta = JSONCodec.loads(fields['meta'])
if meta.get('v') != 1:
# written by a newer release during a rolling upgrade
log.debug(f'Response stream {task_id} skipped: unknown layout version')
return None
output = [JSONCodec.loads(fields[fkey]) for fkey in meta['order']]
journal_ref = meta.get('journal')
if journal_ref:
fkey, path, nonce = journal_ref
journal = await redis.get(_journal_key(task_id))
marker, _, text = (journal or '').partition(':')
if marker != nonce:
# the journal was rewritten between our HGETALL and GET: re-read once
if is_retry:
log.debug(f'Response stream {task_id} skipped: journal nonce mismatch')
return None
fields = await redis.hgetall(_ledger_key(task_id))
if not fields:
log.debug(f'Response stream {task_id} skipped: ledger gone during retry')
return None
continue
_set_leaf(output[meta['order'].index(fkey)], path, text)
content = await redis.get(_content_key(task_id))
return {
'chat_id': meta.get('chat_id'),
'message_id': meta.get('message_id'),
'content': content or get_output_text(output),
'output': output,
}
except (KeyError, ValueError, IndexError, TypeError, AttributeError):
# torn or partially cleaned-up write: skip rather than render a wrong snapshot
log.debug(f'Response stream {task_id} skipped: torn or partial write', exc_info=True)
return None
async def _gather_ledger_streams(redis, chat_id: str, task_ids: list[str]) -> tuple[list[dict], list[str]]:
"""Load the new-format overlays; returns (streams, task ids without a ledger)."""
ledgers = await asyncio.gather(
*(redis.hgetall(_ledger_key(task_id)) for task_id in task_ids), return_exceptions=True
)
with_ledger = []
legacy_task_ids = []
for task_id, fields in zip(task_ids, ledgers):
if isinstance(fields, BaseException):
log.warning(f'Response stream {task_id} overlay skipped for chat {chat_id}', exc_info=fields)
elif fields:
with_ledger.append((task_id, fields))
else:
legacy_task_ids.append(task_id)
loaded = await asyncio.gather(
*(_load_stream(redis, task_id, fields) for task_id, fields in with_ledger),
return_exceptions=True,
)
streams = []
for (task_id, _), stream in zip(with_ledger, loaded):
if isinstance(stream, BaseException):
log.warning(f'Response stream {task_id} overlay unreadable for chat {chat_id}', exc_info=stream)
elif stream and stream.get('chat_id') == chat_id:
streams.append(stream)
return streams, legacy_task_ids
async def get_response_streams_by_chat_id(redis, chat_id: str) -> list[dict]:
task_ids = await list_task_ids_by_item_id(redis, chat_id)
if not task_ids:
return []
if not redis:
return [
stream
for task_id in await list_task_ids_by_item_id(redis, chat_id)
if (stream := response_streams.get(task_id)) and stream.get('chat_id') == chat_id
]
if redis:
values = await redis.hmget(REDIS_RESPONSE_STREAMS_KEY, task_ids)
streams = []
for value in values:
if not value:
continue
try:
task_ids = await list_task_ids_by_item_id(redis, chat_id)
if not task_ids:
return []
streams, legacy_task_ids = await _gather_ledger_streams(redis, chat_id, task_ids)
if legacy_task_ids:
# streams written by processes still running the previous release
try:
data = JSONCodec.loads(value)
for value in await redis.hmget(REDIS_RESPONSE_STREAMS_KEY, legacy_task_ids):
if not value:
continue
try:
data = JSONCodec.loads(value)
except Exception:
continue
if data.get('chat_id') == chat_id:
streams.append(data)
except Exception:
continue
if data.get('chat_id') == chat_id:
streams.append(data)
# the legacy hash lives on its own slot: its loss must not drop the rest
log.warning(f'Legacy response stream overlay unavailable for chat {chat_id}', exc_info=True)
return streams
return [
stream for task_id in task_ids if (stream := response_streams.get(task_id)) and stream.get('chat_id') == chat_id
]
except Exception:
# the overlay is cosmetic: a Redis hiccup must not fail the whole chat load
log.warning(f'Response stream overlay unavailable for chat {chat_id}', exc_info=True)
return []
async def clear_response_stream(redis, task_id: str | None):
if not task_id:
return
_stream_states.pop(task_id, None)
if redis:
await redis.hdel(REDIS_RESPONSE_STREAMS_KEY, task_id)
pipe = redis.pipeline()
_queue_stream_key_deletes(pipe, task_id)
await pipe.execute()
else:
response_streams.pop(task_id, None)