mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-10 22:43:29 +00:00
Fix: mps inference evaluations (#29735)
* fix(retrieval): serialize local embedding and reranking on MPS On Apple Silicon the server process is killed outright (SIGSEGV or SIGTRAP, no traceback) partway through answering any question that retrieves from a knowledge base with hybrid search and a local reranking model. The client sees a dropped connection and the answer is lost. Hybrid search fans its queries out concurrently and every task calls the same shared local model on a worker thread. Torch's Metal shader cache is a process-wide singleton whose lookup tables have no lock, so two of those threads racing inside it corrupt the cache and take the process down with it. Guard the local SentenceTransformer and CrossEncoder calls with a shared lock that is only a real lock when the selected device is MPS. CPU and CUDA installs keep the concurrency they have today, and external reranking endpoints are untouched. Reranking several queries on a Mac now runs one at a time, which is the cost of the process staying alive. Verified by driving the real hybrid-search fan-out with 16 concurrent queries: peak simultaneous entries into the local model drops from 16 to 1 on MPS, stays at 16 on CPU, and the returned documents, scores and ordering are byte-identical in every case. Fixes #29722 * fix(evaluations): serialize the leaderboard embedder against retrieval on MPS The leaderboard's tag-similarity search builds its own SentenceTransformer, and on Apple Silicon sentence-transformers places it on the MPS device. It runs on a worker thread, so an admin running a leaderboard search while anyone queries a knowledge base puts two threads into torch's Metal backend at the same time, which kills the server process outright with no traceback. Move the lock added for the retrieval path into env.py, beside the device selection that decides whether MPS is used at all, and take it around the leaderboard's embedding calls as well. Sharing one lock between the two modules is the whole point, since two separate locks would still let a leaderboard search collide with a retrieval query. Only inference is guarded, matching the retrieval path. Model construction stays as it is here and in the retrieval routers. Verified by driving the leaderboard similarity path and retrieval reranking from six threads against one instrumented model: peak simultaneous entries drops from six to one on MPS, and the similarity scores are unchanged. Related to #29722.
This commit is contained in:
parent
1bfa59acbd
commit
68a74da70b
3 changed files with 24 additions and 12 deletions
|
|
@ -7,7 +7,9 @@ import pkgutil
|
|||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import threading
|
||||
import traceback
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
|
|
@ -66,6 +68,9 @@ if sys.platform == 'darwin' and DEVICE_TYPE == 'cpu':
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
# Torch MPS inference is not thread-safe and a concurrent call kills the whole process.
|
||||
MPS_INFERENCE_LOCK = threading.Lock() if DEVICE_TYPE == 'mps' else nullcontext()
|
||||
|
||||
####################################
|
||||
# LOGGING
|
||||
####################################
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from open_webui.env import (
|
|||
BYPASS_RETRIEVAL_ACCESS_CONTROL,
|
||||
ENABLE_FORWARD_USER_INFO_HEADERS,
|
||||
ENABLE_RETRIEVAL_UNSCOPED_COLLECTIONS,
|
||||
MPS_INFERENCE_LOCK,
|
||||
OFFLINE_MODE,
|
||||
)
|
||||
from open_webui.models.access_grants import AccessGrants
|
||||
|
|
@ -1116,17 +1117,16 @@ def get_embedding_function(
|
|||
'SentenceTransformer model name, or configure an external '
|
||||
'RAG_EMBEDDING_ENGINE (ollama, openai, azure_openai).'
|
||||
)
|
||||
return await asyncio.to_thread(
|
||||
(
|
||||
lambda query, prefix=None: embedding_function.encode(
|
||||
|
||||
def encode():
|
||||
with MPS_INFERENCE_LOCK:
|
||||
return embedding_function.encode(
|
||||
query,
|
||||
batch_size=int(embedding_batch_size),
|
||||
**({'prompt': prefix} if prefix else {}),
|
||||
).tolist()
|
||||
),
|
||||
query,
|
||||
prefix,
|
||||
)
|
||||
|
||||
return await asyncio.to_thread(encode)
|
||||
|
||||
return async_embedding_function
|
||||
elif embedding_engine in ['ollama', 'openai', 'azure_openai']:
|
||||
|
|
@ -1250,9 +1250,14 @@ def get_reranking_function(reranking_engine, reranking_model, reranking_function
|
|||
[(query, doc.page_content) for doc in documents], user=user
|
||||
)
|
||||
else:
|
||||
return lambda query, documents, user=None: reranking_function.predict(
|
||||
[(query, doc.page_content) for doc in documents], batch_size=int(reranking_batch_size)
|
||||
)
|
||||
|
||||
def predict(query, documents, user=None):
|
||||
with MPS_INFERENCE_LOCK:
|
||||
return reranking_function.predict(
|
||||
[(query, doc.page_content) for doc in documents], batch_size=int(reranking_batch_size)
|
||||
)
|
||||
|
||||
return predict
|
||||
|
||||
|
||||
# UUIDs, SHA-256 digests, and prefixed variants thereof all fit [A-Za-z0-9_-].
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import Optional
|
|||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import MPS_INFERENCE_LOCK
|
||||
from open_webui.events import EVENTS, publish_event
|
||||
from open_webui.internal.db import get_async_session
|
||||
from open_webui.models.config import Config
|
||||
|
|
@ -179,8 +180,9 @@ def _compute_similarities(feedbacks: list[LeaderboardFeedbackData], query: str)
|
|||
return {}
|
||||
|
||||
try:
|
||||
tag_embeddings = embedding_model.encode(all_tags)
|
||||
query_embedding = embedding_model.encode([query])[0]
|
||||
with MPS_INFERENCE_LOCK:
|
||||
tag_embeddings = embedding_model.encode(all_tags)
|
||||
query_embedding = embedding_model.encode([query])[0]
|
||||
except Exception as e:
|
||||
log.error(f'Embedding error: {e}')
|
||||
return {}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue