"""
title: GPU Scaling Filter
author: projectmoon
author_url: https://git.agnos.is/projectmoon/open-webui-filters
version: 0.2.2
license: AGPL-3.0+
required_open_webui_version: 0.3.32
"""
# Documentation: https://git.agnos.is/projectmoon/open-webui-filters
# System Imports
import chromadb
from chromadb import ClientAPI as ChromaAPI
from chromadb import Collection as ChromaCollection
from pydantic import BaseModel, Field
from typing import Callable, Awaitable, Any, Optional, Literal
import json
# OpenWebUI imports
from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT
from utils.misc import get_last_user_message, get_last_assistant_message
from apps.ollama.main import generate_chat_completion, GenerateChatCompletionForm
from apps.webui.models.users import UserModel
# Why refactor when you can janky monkey patch? This will NOT be fixed at
# some point, as this filter is deprecated.
CHROMA_CLIENT = VECTOR_DB_CLIENT.client
class GpuChatState:
"""
Get or set GPU layer count by base model for a given chat.
"""
collection_name = "gpu_layers_by_chat"
def __init__(self, chroma_client: ChromaAPI, chat_id: str):
self.chroma_client = chroma_client
self.chat_id = chat_id
self.gpu_layers = {}
def _get_collection(self) -> ChromaCollection:
return self.chroma_client.get_or_create_collection(
name=GpuChatState.collection_name
)
def _parse_results(self, results) -> dict:
if 'documents' in results:
doc = results['documents'][0] if len(results['documents']) > 0 else None
return json.loads(doc) if doc else {}
else:
return {}
def get_gpu_layers(self):
coll = self._get_collection()
if self.gpu_layers == {}:
self.gpu_layers = self._parse_results(
coll.get(ids=[self.chat_id], include=["documents"])
)
return self.gpu_layers
def get_gpu_layers_for_model(self, model_id: str) -> Optional[int]:
info = self.get_gpu_layers()
return info[model_id] if model_id in info else None
def set_gpu_layers(self, model: str, amount: int):
# set gpu layers for this chat.
self.gpu_layers[model] = amount
self._get_collection().upsert(
ids=[self.chat_id],
documents=[json.dumps(self.gpu_layers)]
)
self.gpu_layers = self.get_gpu_layers()
class SessionInfo(BaseModel):
chat_id: str
message_id: str
session_id: str
def dict_to_attributes(input_dict):
class AttrDict:
def __init__(self, attr_dict):
for key, value in attr_dict.items():
setattr(self, key, value)
return AttrDict(input_dict)
def extract_model_id(model: dict) -> Optional[str]:
model_id = None
if "info" in model:
if "base_model_id" in model["info"]:
model_id = model["info"]["base_model_id"]
else:
if "ollama" in model and "id" in model["ollama"]:
model_id = model["ollama"]["id"]
if not model_id:
model_id = model["id"]
return model_id
def extract_session_info(event_emitter) -> Optional[SessionInfo]:
"""The latest innovation in hacky workarounds."""
try:
info = event_emitter.__closure__[0].cell_contents
return SessionInfo(
chat_id=info["chat_id"],
message_id=info["message_id"],
session_id=info["session_id"]
)
except:
return None
class Filter:
class Valves(BaseModel):
reduction_start: int = Field(
default=20, description="Amount of GPU layers to reduce to immediately on failure"
)
scaling_step: int = Field(
default=5, description="Amount of GPU layers to reduce by on continued failures"
)
show_status: bool = Field(
default=True, description="Show status message when running downscaled model."
)
pass
def __init__(self):
self.valves = self.Valves()
pass
async def send_message_adjusting(self, done: bool, amount: int=0, steps: int=0):
if steps > 0:
steps_desc = f"reduced by {steps}"
else:
steps_desc = "initial reduction"
desc = (
"Downscaling GPU layers..." if not done
else f"GPU layers downscaled to {amount} ({steps_desc}). Please retry.")
await self.event_emitter(
{
"type": "status",
"data": {
"description": desc,
"done": done
},
}
)
async def send_message_downscaled(self):
await self.event_emitter(
{
"type": "status",
"data": {
"description": "Running at reduced GPU capacity. Responses will be slower.",
"done": True
},
}
)
def get_num_layers_for_model(
self,
gpu_layer_info: GpuChatState,
__model__: dict
) -> Optional[int]:
model_id = extract_model_id(__model__)
if model_id:
return gpu_layer_info.get_gpu_layers_for_model(model_id)
else:
return None
async def downscale(self, model):
"""Update tracked downscale GPU layers for this chat + model."""
# this logic is currently very basic. does not yet take into
# account the actual number of layers in a model. but it's
# better than nothing. if this is the first failure (no entry
# in gpu chat state), set number of layers to the valve
# parameter. if this is a subsequent failure (we have entry
# for this chat already), reduce by the step valve parameter,
# to a minimum of CPU (100% cpu).
model_id = extract_model_id(model)
if not model_id:
print("Could not extract model ID for GPU downscaling!")
return
await self.send_message_adjusting(False)
gpu_layer_info = GpuChatState(CHROMA_CLIENT, self.session_info.chat_id)
num_layers = self.get_num_layers_for_model(gpu_layer_info, model)
downscale_steps = 0
if num_layers:
print(f"Downscaling layers by {self.valves.scaling_step}")
num_layers -= self.valves.scaling_step
downscale_steps = self.valves.scaling_step
if num_layers < 0:
num_layers = 0
else:
num_layers = self.valves.reduction_start
gpu_layer_info.set_gpu_layers(model_id, num_layers)
await self.send_message_adjusting(True, amount=num_layers, steps=downscale_steps)
print(
f"Set GPU layers for chat {self.session_info.chat_id} to {num_layers}"
)
async def inlet(
self,
body: dict,
__event_emitter__: Callable[[Any], Awaitable[None]],
__model__: Optional[dict] = None,
) -> dict:
"""Intercept incoming messages and downscale if necessary."""
if not __model__ or __model__["owned_by"] != "ollama":
return body
self.event_emitter = __event_emitter__
self.session_info = extract_session_info(__event_emitter__)
if self.session_info:
gpu_layer_info = GpuChatState(CHROMA_CLIENT, self.session_info.chat_id)
num_layers = self.get_num_layers_for_model(gpu_layer_info, __model__)
if num_layers and "options" in body:
model_id = extract_model_id(__model__)
body["options"]["num_gpu"] = num_layers
if self.valves.show_status:
await self.send_message_downscaled()
print((
f"Downscaled GPU layers for incoming request for {model_id} "
f"to {num_layers}"
))
return body
async def outlet(
self,
body: dict,
__user__: dict,
__event_emitter__: Callable[[Any], Awaitable[None]],
__model__: Optional[dict] = None,
) -> dict:
"""On response failure, downscale the GPU layers for next try."""
if not __model__ or __model__["owned_by"] != "ollama":
return body
self.event_emitter = __event_emitter__
self.session_info = extract_session_info(__event_emitter__)
if not self.session_info or not __model__:
return body
if len(body["messages"]) == 0:
return body
last_reply = body["messages"][-1]
broke = last_reply["content"] == "" and last_reply["info"] == {}
if broke:
# while we could actually redo the message itself, it is
# useless, because open web ui does not currently have a
# way to clear error state when message content is
# replaced. so we just lower gpu layers and tell user to
# try again. the inlet will intercept the incoming request
# and lower the gpu layers.
await self.downscale(__model__)
return body