mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-22 00:13:08 +00:00
fix: fix segfault in load model (#3879)
# What does this PR do? Fix segfault with load model The cc-vec integration failed with segfault when used with default embedding model on macOS `model_id: nomic-ai/nomic-embed-text-v1.5` and `provider_id: sentence-transformers` Checked crash report and see this is due to torch OPENMP settings. Constrainting to 1 thread works without crashes. ## Test Plan Tested with cc-vec integration 1. start server llama stack run starter 2. Do the setup in https://github.com/raghotham/cc-vec to set env variables and try `uv run cc-vec index --url-patterns "%.github.io" --vector-store-name "ml-research" --limit 50 --chunk-size 800 --overlap 400`
This commit is contained in:
parent
1ec7216c3f
commit
cb6a5e2687
1 changed files with 12 additions and 0 deletions
|
@ -6,9 +6,12 @@
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import platform
|
||||
import struct
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from llama_stack.log import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
@ -24,6 +27,8 @@ from llama_stack.apis.inference import (
|
|||
|
||||
EMBEDDING_MODELS = {}
|
||||
|
||||
DARWIN = "Darwin"
|
||||
|
||||
|
||||
log = get_logger(name=__name__, category="providers::utils")
|
||||
|
||||
|
@ -83,6 +88,13 @@ class SentenceTransformerEmbeddingMixin:
|
|||
def _load_model():
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
platform_name = platform.system()
|
||||
if platform_name == DARWIN:
|
||||
# PyTorch's OpenMP kernels can segfault on macOS when spawned from background
|
||||
# threads with the default parallel settings, so force a single-threaded CPU run.
|
||||
log.debug(f"Constraining torch threads on {platform_name} to a single worker")
|
||||
torch.set_num_threads(1)
|
||||
|
||||
return SentenceTransformer(model, trust_remote_code=True)
|
||||
|
||||
loaded_model = await asyncio.to_thread(_load_model)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue