Compare commits

...
3 Commits
Author SHA1 Message Date
Nicolò Boschi 44f7940f47 fix: add MPS support for macOS Apple Silicon
Extend GPU detection to include Apple MPS backend in addition to CUDA.
This ensures macOS users with Apple Silicon use MPS acceleration
instead of being incorrectly routed to the CPU fallback path.
2026-01-20 14:06:36 +01:00
Nicolò Boschi a03fd32214 fix: add filelock for model initialization in parallel tests
When pytest-xdist runs multiple workers in parallel, they all try to
load models from the HuggingFace cache simultaneously, causing race
conditions and intermittent meta tensor errors.

Added filelock around embeddings and cross_encoder initialization in
conftest.py, similar to how pg0 database setup is serialized. Models
are now pre-initialized in the fixture before being passed to tests.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-01-20 13:54:01 +01:00
Nicolò Boschi c935c576f1 fix: prevent meta tensor issues when accelerate is installed without GPU
When accelerate is installed but no GPU is available, transformers can
incorrectly use lazy loading (meta tensors) which fails when
sentence-transformers tries to move the model to a device.

The fix checks hardware and installed packages to determine the right
loading strategy:
- GPU available: device=None, device_map=None (auto-detect GPU)
- No GPU + accelerate: device='cpu', device_map='cpu' (force CPU loading)
- No GPU + no accelerate: device='cpu', device_map=None (normal CPU)

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-01-20 13:41:03 +01:00
3 changed files with 111 additions and 15 deletions
@@ -130,17 +130,38 @@ class LocalSTCrossEncoder(CrossEncoderModel):
"Install it with: pip install sentence-transformers"
)
# Note: We use CPU even when GPU/MPS is available because:
# 1. The reranker model (MiniLM) is tiny (~22M params)
# 2. Batch sizes are small (~100-200 pairs)
# 3. Data transfer overhead to GPU outweighs compute benefit
# 4. CPU inference is actually faster for this workload
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate.
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized.
# Determine device and device_map based on hardware and installed packages.
# When accelerate is installed but no GPU/MPS is available, transformers can
# incorrectly use lazy loading (meta tensors) which fails on .to(device).
# We use device_map="cpu" in that case to force direct CPU loading.
import torch
try:
import accelerate # type: ignore[import-not-found] # noqa: F401
accelerate_available = True
except ImportError:
accelerate_available = False
# Check for GPU (CUDA) or Apple Silicon (MPS)
has_gpu = torch.cuda.is_available() or (hasattr(torch.backends, "mps") and torch.backends.mps.is_available())
if has_gpu:
device = None # Let sentence-transformers auto-detect GPU/MPS
device_map = None
elif accelerate_available:
device = "cpu"
device_map = "cpu" # Force direct CPU loading to avoid meta tensors
else:
device = "cpu"
device_map = None
self._model = CrossEncoder(
self.model_name,
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
device=device,
model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map},
)
# Initialize shared executor (limited workers naturally limits concurrency)
@@ -128,11 +128,37 @@ class LocalSTEmbeddings(Embeddings):
)
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
# Determine device and device_map based on hardware and installed packages.
# When accelerate is installed but no GPU/MPS is available, transformers can
# incorrectly use lazy loading (meta tensors) which fails on .to(device).
# We use device_map="cpu" in that case to force direct CPU loading.
import torch
try:
import accelerate # type: ignore[import-not-found] # noqa: F401
accelerate_available = True
except ImportError:
accelerate_available = False
# Check for GPU (CUDA) or Apple Silicon (MPS)
has_gpu = torch.cuda.is_available() or (hasattr(torch.backends, "mps") and torch.backends.mps.is_available())
if has_gpu:
device = None # Let sentence-transformers auto-detect GPU/MPS
device_map = None
elif accelerate_available:
device = "cpu"
device_map = "cpu" # Force direct CPU loading to avoid meta tensors
else:
device = "cpu"
device_map = None
self._model = SentenceTransformer(
self.model_name,
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
device=device,
model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map},
)
self._dimension = self._model.get_sentence_embedding_dimension()
+53 -4
View File
@@ -116,16 +116,65 @@ def llm_config():
@pytest.fixture(scope="session")
def embeddings():
def embeddings(tmp_path_factory, worker_id):
"""
Session-scoped embeddings fixture with filelock to prevent race conditions.
return LocalSTEmbeddings()
When pytest-xdist runs multiple workers in parallel, they all try to load
models from the HuggingFace cache simultaneously, which can cause race
conditions and meta tensor errors. We use a filelock to serialize model
initialization across workers.
"""
# Get shared temp dir for coordination between xdist workers
if worker_id == "master":
root_tmp_dir = tmp_path_factory.getbasetemp()
else:
root_tmp_dir = tmp_path_factory.getbasetemp().parent
lock_file = root_tmp_dir / "embeddings_init.lock"
emb = LocalSTEmbeddings()
# Serialize model initialization across workers
with filelock.FileLock(str(lock_file)):
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(emb.initialize())
finally:
loop.close()
return emb
@pytest.fixture(scope="session")
def cross_encoder():
def cross_encoder(tmp_path_factory, worker_id):
"""
Session-scoped cross-encoder fixture with filelock to prevent race conditions.
return LocalSTCrossEncoder()
When pytest-xdist runs multiple workers in parallel, they all try to load
models from the HuggingFace cache simultaneously, which can cause race
conditions and meta tensor errors. We use a filelock to serialize model
initialization across workers.
"""
# Get shared temp dir for coordination between xdist workers
if worker_id == "master":
root_tmp_dir = tmp_path_factory.getbasetemp()
else:
root_tmp_dir = tmp_path_factory.getbasetemp().parent
lock_file = root_tmp_dir / "cross_encoder_init.lock"
ce = LocalSTCrossEncoder()
# Serialize model initialization across workers
with filelock.FileLock(str(lock_file)):
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(ce.initialize())
finally:
loop.close()
return ce
@pytest.fixture(scope="session")
def query_analyzer():