diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 8b335419..d9507787 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -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) diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index 066406fa..5402be7c 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -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() diff --git a/hindsight-api/tests/conftest.py b/hindsight-api/tests/conftest.py index 6fec5abc..b3ba72fe 100644 --- a/hindsight-api/tests/conftest.py +++ b/hindsight-api/tests/conftest.py @@ -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():