# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Embedder concurrency tests: the fast tokenizer isn't thread-safe, so encode and token counting must be serialized (else threads panic "Already borrowed").""" import os import sys import threading import time from types import SimpleNamespace import numpy as np import pytest from core.rag import config, embeddings @pytest.fixture(autouse = True) def _pin_st_backend(monkeypatch): # Tests patch ST internals (_get), so force the ST backend. monkeypatch.setattr(config, "EMBED_BACKEND", "sentence-transformers") embeddings._reset_backend() yield embeddings._reset_backend() class _ConcurrencyProbe: """Records whether two callers were in the guarded body at once.""" def __init__(self): self.inside = 0 self.saw_overlap = False self._g = threading.Lock() def enter(self): with self._g: self.inside += 1 if self.inside > 1: self.saw_overlap = True time.sleep(0.005) # widen the race window with self._g: self.inside -= 1 class _FakeModel: def __init__(self, probe): self._probe = probe self.tokenizer = _FakeTokenizer(probe) def encode(self, texts, **_kw): self._probe.enter() return np.zeros((len(texts), 4), dtype = np.float32) class _FakeTokenizer: def __init__(self, probe): self._probe = probe def encode(self, text, **_kw): self._probe.enter() return list(range(len(text.split()))) def _hammer(fn, n = 8): errors: list[Exception] = [] def worker(): try: fn() except Exception as exc: # noqa: BLE001 errors.append(exc) threads = [threading.Thread(target = worker) for _ in range(n)] for t in threads: t.start() for t in threads: t.join() return errors def test_encode_is_serialized(monkeypatch): probe = _ConcurrencyProbe() monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _FakeModel(probe)) errors = _hammer(lambda: embeddings.encode(["alpha beta", "gamma"])) assert errors == [] assert probe.saw_overlap is False # compute lock serialized encode() def test_token_counter_is_serialized(monkeypatch): probe = _ConcurrencyProbe() monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _FakeModel(probe)) count = embeddings.token_counter() errors = _hammer(lambda: count("one two three four")) assert errors == [] assert probe.saw_overlap is False # counting shares the tokenizer lock def test_encode_enables_parallelism_only_during_call(monkeypatch): seen = {} class _M: tokenizer = None def encode(self, texts, **_kw): seen["during"] = os.environ.get("TOKENIZERS_PARALLELISM") return np.zeros((len(texts), 4), dtype = np.float32) monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _M()) os.environ["TOKENIZERS_PARALLELISM"] = "false" embeddings.encode(["alpha", "beta"]) assert seen["during"] == "true" # rayon batch tokenization enabled in-call assert os.environ.get("TOKENIZERS_PARALLELISM") == "false" # restored after def test_token_counter_enables_parallelism_only_during_call(monkeypatch): seen = {} class _Tok: def encode(self, text, **_kw): seen["during"] = os.environ.get("TOKENIZERS_PARALLELISM") return list(range(len(text.split()))) class _M: tokenizer = _Tok() monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _M()) os.environ["TOKENIZERS_PARALLELISM"] = "false" count = embeddings.token_counter() count("alpha beta gamma") assert seen["during"] == "true" # rayon enabled in-call, like _st_encode assert os.environ.get("TOKENIZERS_PARALLELISM") == "false" # restored after def test_sentence_transformer_load_uses_live_cache(monkeypatch, tmp_path): observed = {} class FakeSentenceTransformer: def __init__(self, name, **kwargs): observed["name"] = name observed.update(kwargs) monkeypatch.setitem( sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer = FakeSentenceTransformer), ) monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) monkeypatch.setattr(embeddings, "_guard_model_security", lambda *_a, **_k: None) monkeypatch.setattr(embeddings, "_device", lambda: "cpu") monkeypatch.setattr( "utils.hf_cache_settings.active_hf_hub_cache", lambda: str(tmp_path / "selected-hub"), ) embeddings._model = None embeddings._name = None embeddings._get("Org/Embedder") assert observed["name"] == "Org/Embedder" assert observed["cache_folder"] == str(tmp_path / "selected-hub") class _SentinelLlamaBackend: """Stand-in for LlamaServerBackend; never spawns a real server.""" def _force_st_load_failure(monkeypatch): """Make the ST warm-probe raise.""" def _boom(model_name = None): raise RuntimeError("torch is broken on this machine") monkeypatch.setattr(embeddings, "_get", _boom) def _patch_llama_backend(monkeypatch, *, binary): from core.inference.llama_cpp import LlamaCppBackend from core.rag import embed_llama_server monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary)) monkeypatch.setattr(embed_llama_server, "LlamaServerBackend", _SentinelLlamaBackend) def test_st_failure_falls_back_to_llama_server(monkeypatch): # ST can't load but llama-server is available -> use it. _force_st_load_failure(monkeypatch) _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") embeddings._reset_backend() backend = embeddings._get_backend() assert isinstance(backend, _SentinelLlamaBackend) def test_st_failure_without_llama_binary_reraises(monkeypatch): # No llama-server binary -> surface the failure, don't degrade to nothing. _force_st_load_failure(monkeypatch) _patch_llama_backend(monkeypatch, binary = None) embeddings._reset_backend() with pytest.raises(RuntimeError, match = "torch is broken"): embeddings._get_backend() def test_st_success_keeps_sentence_transformers(monkeypatch): # Clean ST probe -> ST backend stays selected, no fallback. monkeypatch.setattr(embeddings, "_get", lambda model_name = None: object()) _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") embeddings._reset_backend() backend = embeddings._get_backend() assert isinstance(backend, embeddings._SentenceTransformersBackend) class _BoomOnEncodeModel: """Loads fine (init probe passes) but raises when encoding.""" tokenizer = None def encode(self, texts, **_kw): raise RuntimeError("CUDA error during encode") def test_st_encode_runtime_failure_switches_to_llama(monkeypatch): # encode() blows up mid-run -> switch to llama-server and stay switched. monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) _patch_llama_backend(monkeypatch, binary = "/fake/llama-server") calls = {} def _sentinel_encode( self, texts, *, model_name = None, normalize = True, ): calls["used"] = True return np.zeros((len(texts), 4), dtype = np.float32) monkeypatch.setattr(_SentinelLlamaBackend, "encode", _sentinel_encode, raising = False) embeddings._reset_backend() out = embeddings.encode(["alpha", "beta"]) assert calls.get("used") is True # retried on the llama fallback assert out.shape == (2, 4) # Switch is process-wide: later calls keep using llama, not ST. assert isinstance(embeddings._get_backend(), _SentinelLlamaBackend) def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): # No llama-server binary -> surface the encode error. monkeypatch.setattr(embeddings, "_get", lambda model_name = None: _BoomOnEncodeModel()) _patch_llama_backend(monkeypatch, binary = None) embeddings._reset_backend() with pytest.raises(RuntimeError, match = "CUDA error during encode"): embeddings.encode(["alpha", "beta"])