ewilliams28/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-13 01:02:25 -06:00
Code Issues Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-02-15 18:41:16 -06:00
parent 911eecac85
commit b780d5c556
3 changed files with 110 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/main.py
+2
View File
@@ -64,6 +64,7 @@ from open_webui.socket.main import (
MODELS,
app as socket_app,
periodic_usage_pool_cleanup,
periodic_session_pool_cleanup,
get_event_emitter,
get_models_in_use,
)
@@ -635,6 +636,7 @@ async def lifespan(app: FastAPI):
limiter.total_tokens = THREAD_POOL_SIZE
asyncio.create_task(periodic_usage_pool_cleanup())
asyncio.create_task(periodic_session_pool_cleanup())
if app.state.config.ENABLE_BASE_MODELS_CACHE:
await get_all_models(
backend/open_webui/socket/main.py
+74 -18
View File
@@ -99,6 +99,7 @@ else:
# Timeout duration in seconds
TIMEOUT_DURATION = 3
SESSION_POOL_TIMEOUT = 120 # seconds without heartbeat before session is reaped
# Dictionary to maintain the user pool
@@ -147,6 +148,17 @@ if WEBSOCKET_MANAGER == "redis":
aquire_func = clean_up_lock.aquire_lock
renew_func = clean_up_lock.renew_lock
release_func = clean_up_lock.release_lock
session_cleanup_lock = RedisLock(
redis_url=WEBSOCKET_REDIS_URL,
lock_name=f"{REDIS_KEY_PREFIX}:session_cleanup_lock",
timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
redis_sentinels=redis_sentinels,
redis_cluster=WEBSOCKET_REDIS_CLUSTER,
)
session_aquire_func = session_cleanup_lock.aquire_lock
session_renew_func = session_cleanup_lock.renew_lock
session_release_func = session_cleanup_lock.release_lock
else:
MODELS = {}
@@ -154,6 +166,7 @@ else:
USAGE_POOL = {}
aquire_func = release_func = renew_func = lambda: True
session_aquire_func = session_release_func = session_renew_func = lambda: True
YDOC_MANAGER = YdocManager(
@@ -162,6 +175,31 @@ YDOC_MANAGER = YdocManager(
)
async def periodic_session_pool_cleanup():
"""Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance)."""
if not session_aquire_func():
log.debug("Session cleanup lock held by another node. Skipping.")
return
try:
while True:
if not session_renew_func():
log.error("Unable to renew session cleanup lock. Exiting.")
return
now = int(time.time())
for sid in list(SESSION_POOL.keys()):
entry = SESSION_POOL.get(sid)
if entry and now - entry.get("last_seen_at", 0) > SESSION_POOL_TIMEOUT:
log.warning(
f"Reaping orphaned session {sid} (user {entry.get('id')})"
)
del SESSION_POOL[sid]
await asyncio.sleep(SESSION_POOL_TIMEOUT)
finally:
session_release_func()
async def periodic_usage_pool_cleanup():
max_retries = 2
retry_delay = random.uniform(
@@ -313,15 +351,18 @@ async def connect(sid, environ, auth):
user = Users.get_user_by_id(data["id"])
if user:
SESSION_POOL[sid] = user.model_dump(
exclude=[
"profile_image_url",
"profile_banner_image_url",
"date_of_birth",
"bio",
"gender",
]
)
SESSION_POOL[sid] = {
**user.model_dump(
exclude=[
"profile_image_url",
"profile_banner_image_url",
"date_of_birth",
"bio",
"gender",
]
),
"last_seen_at": int(time.time()),
}
await sio.enter_room(sid, f"user:{user.id}")
@@ -340,15 +381,18 @@ async def user_join(sid, data):
if not user:
return
SESSION_POOL[sid] = user.model_dump(
exclude=[
"profile_image_url",
"profile_banner_image_url",
"date_of_birth",
"bio",
"gender",
]
)
SESSION_POOL[sid] = {
**user.model_dump(
exclude=[
"profile_image_url",
"profile_banner_image_url",
"date_of_birth",
"bio",
"gender",
]
),
"last_seen_at": int(time.time()),
}
await sio.enter_room(sid, f"user:{user.id}")
@@ -366,6 +410,7 @@ async def user_join(sid, data):
async def heartbeat(sid, data):
user = SESSION_POOL.get(sid)
if user:
SESSION_POOL[sid] = {**user, "last_seen_at": int(time.time())}
Users.update_last_active_by_id(user["id"])
@@ -709,6 +754,17 @@ async def disconnect(sid):
if sid in SESSION_POOL:
user = SESSION_POOL[sid]
del SESSION_POOL[sid]
# Clean up USAGE_POOL entries for this session
for model_id in list(USAGE_POOL.keys()):
connections = USAGE_POOL.get(model_id)
if connections and sid in connections:
del connections[sid]
if not connections:
del USAGE_POOL[model_id]
else:
USAGE_POOL[model_id] = connections
await YDOC_MANAGER.remove_user_from_all_documents(sid)
else:
pass
backend/open_webui/socket/utils.py
+34
View File
@@ -118,6 +118,8 @@ class RedisDict:
class YdocManager:
COMPACTION_THRESHOLD = 500
def __init__(
self,
redis=None,
@@ -133,10 +135,42 @@ class YdocManager:
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
await self._redis.rpush(redis_key, json.dumps(list(update)))
list_len = await self._redis.llen(redis_key)
if list_len >= self.COMPACTION_THRESHOLD:
await self._compact_updates_redis(document_id)
else:
if document_id not in self._updates:
self._updates[document_id] = []
self._updates[document_id].append(update)
if len(self._updates[document_id]) >= self.COMPACTION_THRESHOLD:
self._compact_updates_memory(document_id)
async def _compact_updates_redis(self, document_id: str):
"""Rolling compaction: squash oldest half into one snapshot."""
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
all_updates = await self._redis.lrange(redis_key, 0, -1)
if len(all_updates) <= 1:
return
mid = len(all_updates) // 2
ydoc = Y.Doc()
for raw in all_updates[:mid]:
ydoc.apply_update(bytes(json.loads(raw)))
snapshot = json.dumps(list(ydoc.get_update()))
pipe = self._redis.pipeline()
pipe.delete(redis_key)
pipe.rpush(redis_key, snapshot, *all_updates[mid:])
await pipe.execute()
def _compact_updates_memory(self, document_id: str):
"""Rolling compaction: squash oldest half into one snapshot."""
updates = self._updates.get(document_id, [])
if len(updates) <= 1:
return
mid = len(updates) // 2
ydoc = Y.Doc()
for update in updates[:mid]:
ydoc.apply_update(bytes(update))
self._updates[document_id] = [ydoc.get_update()] + updates[mid:]
async def get_updates(self, document_id: str) -> List[bytes]:
document_id = document_id.replace(":", "_")