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:
Classic298 2026-09-06 21:20:41 +02:00 committed by GitHub
parent 1bfa59acbd
commit 68a74da70b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 24 additions and 12 deletions

View file

@ -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
####################################

View file

@ -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_-].

View file

@ -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 {}