OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-05-17 12:03:41 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

fix: running models

Browse Source
This commit is contained in:
Timothy J. Baek 2024-06-04 11:38:31 -07:00
parent 2be9c25ba7
commit 92d056c982
2 changed files with 26 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
backend/apps/socket/main.py
View File

@ -34,6 +34,7 @@ async def connect(sid, environ, auth):
print(len(set(USER_POOL))) print(len(set(USER_POOL)))
await sio.emit("user-count", {"count": len(set(USER_POOL))}) await sio.emit("user-count", {"count": len(set(USER_POOL))})
await sio.emit("usage", {"models": get_models_in_use()})
@sio.on("user-join") @sio.on("user-join")
@ -64,10 +65,9 @@ async def user_count(sid):
def get_models_in_use(): def get_models_in_use():
# Aggregate all models in use # Aggregate all models in use
models_in_use = [] models_in_use = []
for sid, data in USAGE_POOL.items(): for model_id, data in USAGE_POOL.items():
models_in_use.extend(data["models"]) models_in_use.append(model_id)
print(f"Models in use: {models_in_use}") print(f"Models in use: {models_in_use}")
return models_in_use return models_in_use
@ -77,46 +77,45 @@ def get_models_in_use():
async def usage(sid, data): async def usage(sid, data):
print(f'Received "usage" event from {sid}: {data}') print(f'Received "usage" event from {sid}: {data}')
# Cancel previous task if there is one
if sid in USAGE_POOL:
USAGE_POOL[sid]["task"].cancel()
# Store the new usage data and task
model_id = data["model"] model_id = data["model"]
if sid in USAGE_POOL and "models" in USAGE_POOL[sid]: # Cancel previous callback if there is one
if model_id in USAGE_POOL:
USAGE_POOL[model_id]["callback"].cancel()
print(USAGE_POOL[sid]) # Store the new usage data and task
models = USAGE_POOL[sid]["models"] if model_id in USAGE_POOL:
if model_id not in models: USAGE_POOL[model_id]["sids"].append(sid)
models.append(model_id) USAGE_POOL[model_id]["sids"] = list(set(USAGE_POOL[model_id]["sids"]))
USAGE_POOL[sid] = {"models": models}
else: else:
USAGE_POOL[sid] = {"models": [model_id]} USAGE_POOL[model_id] = {"sids": [sid]}
# Schedule a task to remove the usage data after TIMEOUT_DURATION # Schedule a task to remove the usage data after TIMEOUT_DURATION
USAGE_POOL[sid]["task"] = asyncio.create_task(remove_after_timeout(sid, model_id)) USAGE_POOL[model_id]["callback"] = asyncio.create_task(
remove_after_timeout(sid, model_id)
)
models_in_use = get_models_in_use()
# Broadcast the usage data to all clients # Broadcast the usage data to all clients
await sio.emit("usage", {"models": models_in_use}) await sio.emit("usage", {"models": get_models_in_use()})
async def remove_after_timeout(sid, model_id): async def remove_after_timeout(sid, model_id):
try: try:
print("remove_after_timeout", sid, model_id)
await asyncio.sleep(TIMEOUT_DURATION) await asyncio.sleep(TIMEOUT_DURATION)
if sid in USAGE_POOL: if model_id in USAGE_POOL:
if model_id in USAGE_POOL[sid]["models"]: print(USAGE_POOL[model_id]["sids"])
USAGE_POOL[sid]["models"].remove(model_id) USAGE_POOL[model_id]["sids"].remove(sid)
if len(USAGE_POOL[sid]["models"]) == 0: USAGE_POOL[model_id]["sids"] = list(set(USAGE_POOL[model_id]["sids"]))
del USAGE_POOL[sid]
print(f"Removed usage data for {sid} due to timeout")
models_in_use = get_models_in_use() if len(USAGE_POOL[model_id]["sids"]) == 0:
del USAGE_POOL[model_id]
print(f"Removed usage data for {model_id} due to timeout")
# Broadcast the usage data to all clients # Broadcast the usage data to all clients
await sio.emit("usage", {"models": models_in_use}) await sio.emit("usage", {"models": get_models_in_use()})
except asyncio.CancelledError: except asyncio.CancelledError:
# Task was cancelled due to new 'usage' event # Task was cancelled due to new 'usage' event
pass pass

2
src/lib/components/layout/Sidebar/UserMenu.svelte
View File

@ -145,7 +145,7 @@
<Tooltip <Tooltip
content={$USAGE_POOL && $USAGE_POOL.length > 0 content={$USAGE_POOL && $USAGE_POOL.length > 0
? `Running: ${$USAGE_POOL.join(',')} ✨` ? `Running: ${$USAGE_POOL.join(', ')} ✨`
: ''} : ''}
> >
<div class="flex rounded-md py-1.5 px-3 text-xs gap-2.5 items-center"> <div class="flex rounded-md py-1.5 px-3 text-xs gap-2.5 items-center">