fix(core): reuse a single embedding provider per process (#903)

Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/basic_memory/repository/embedding_provider_factory.py

@@ -3,9 +3,19 @@
 import os
 from threading import Lock
 
+from loguru import logger
+
 from basic_memory.config import BasicMemoryConfig, default_fastembed_cache_dir
 from basic_memory.repository.embedding_provider import EmbeddingProvider
 
+# Cache key fields are limited to values that change the *identity* of the loaded
+# model (provider, model_name, dimensions, LiteLLM role/input-type/forward-dimension
+# settings, batch/request knobs that affect the LiteLLM identity, and the resolved
+# cache dir). Thread/parallel knobs are deliberately excluded — they change ONNX
+# *execution* only, not the loaded weights. Including them caused #872: in a
+# container/cgroup the CPU-derived thread count can drift between calls, producing
+# a fresh cache key and reloading the ~2.3GB model into a CPU arena that never
+# returns memory to the OS.
 type ProviderCacheKey = tuple[
     str,
     str,
@@ -16,8 +26,6 @@ type ProviderCacheKey = tuple[
     str | None,
     str | None,
     str,
-    int | None,
-    int | None,
 ]
 
 _EMBEDDING_PROVIDER_CACHE: dict[ProviderCacheKey, EmbeddingProvider] = {}
@@ -78,13 +86,16 @@ def _resolve_fastembed_runtime_knobs(
 
 
 def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey:
-    """Build a stable cache key from provider-relevant semantic embedding config.
+    """Build a stable cache key from model-identity semantic embedding config.
 
     Uses the *resolved* cache dir — not the raw config field — so different
     FASTEMBED_CACHE_PATH values produce distinct cache keys even when the
     config field itself is unset.
+
+    Deliberately excludes the FastEmbed thread/parallel knobs: they tune ONNX
+    execution, not which model weights are loaded, and resolving them from the
+    runtime CPU budget makes the key drift between calls in a container (#872).
     """
-    resolved_threads, resolved_parallel = _resolve_fastembed_runtime_knobs(app_config)
     return (
         app_config.semantic_embedding_provider.strip().lower(),
         app_config.semantic_embedding_model,
@@ -95,8 +106,6 @@ def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey:
         app_config.semantic_embedding_document_input_type,
         app_config.semantic_embedding_query_input_type,
         _resolve_cache_dir(app_config),
-        resolved_threads,
-        resolved_parallel,
     )
 
 
@@ -116,6 +125,14 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide
     embedding response is available.
     """
     cache_key = _provider_cache_key(app_config)
+    # Trigger: two threads miss the cache for the same key concurrently.
+    # Why: provider construction loads the ~2.3GB ONNX model and is slow, so we
+    # deliberately build it *outside* the lock to avoid serializing every caller
+    # behind a single cold start. This opens a by-design TOCTOU window where both
+    # threads may construct a provider.
+    # Outcome: the second check-and-set below resolves the race — the first writer
+    # wins and the loser's redundant provider is discarded, so the cache still
+    # yields a single process-wide singleton per key.
     with _EMBEDDING_PROVIDER_CACHE_LOCK:
         if cached_provider := _EMBEDDING_PROVIDER_CACHE.get(cache_key):
             return cached_provider
@@ -190,5 +207,19 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide
     with _EMBEDDING_PROVIDER_CACHE_LOCK:
         if cached_provider := _EMBEDDING_PROVIDER_CACHE.get(cache_key):
             return cached_provider
+        # Trigger: a distinct cache key is being inserted while the cache already
+        # holds entries for other keys.
+        # Why: the provider is meant to be a process-wide singleton (#872). A second
+        # key means something bypassed reuse — a real config change, or a regression
+        # that reintroduces volatile fields into the key — and each new key reloads
+        # the ~2.3GB ONNX model into a CPU arena that never releases memory.
+        # Outcome: surface the bypass so future leaks are diagnosable from logs.
+        if _EMBEDDING_PROVIDER_CACHE:
+            logger.warning(
+                "Creating a second distinct embedding provider in this process; "
+                "the model will be loaded again. existing_keys={existing} new_key={new}",
+                existing=list(_EMBEDDING_PROVIDER_CACHE.keys()),
+                new=cache_key,
+            )
         _EMBEDDING_PROVIDER_CACHE[cache_key] = provider
         return provider

src/basic_memory/repository/fastembed_provider.py

@@ -87,17 +87,20 @@ class FastEmbedEmbeddingProvider(EmbeddingProvider):
                 "pip install -U basic-memory"
             ) from exc
         resolved_model_name = self._resolved_model_name()
-        if self.cache_dir is not None and self.threads is not None:
-            return TextEmbedding(
-                model_name=resolved_model_name,
-                cache_dir=self.cache_dir,
-                threads=self.threads,
-            )
+        # Constraint: onnxruntime's CPU memory arena grows to fit peak usage and never
+        # returns that memory to the OS. If a model is ever loaded more than once in a
+        # long-running process it leaks tens of GB (#872). FastEmbed exposes
+        # enable_cpu_mem_arena via its session-option kwargs, so we disable the arena to
+        # let any transient extra load free memory.
+        model_kwargs: dict = {
+            "model_name": resolved_model_name,
+            "enable_cpu_mem_arena": False,
+        }
         if self.cache_dir is not None:
-            return TextEmbedding(model_name=resolved_model_name, cache_dir=self.cache_dir)
+            model_kwargs["cache_dir"] = self.cache_dir
         if self.threads is not None:
-            return TextEmbedding(model_name=resolved_model_name, threads=self.threads)
-        return TextEmbedding(model_name=resolved_model_name)
+            model_kwargs["threads"] = self.threads
+        return TextEmbedding(**model_kwargs)
 
     def _model_cache_candidates(self) -> list[tuple[Path, str]]:
         """Resolve ``(snapshot_dir, model_file)`` pairs for this model under ``cache_dir``.

src/basic_memory/repository/search_repository.py

@@ -13,6 +13,7 @@ from sqlalchemy import Result
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
 from basic_memory.config import BasicMemoryConfig, ConfigManager, DatabaseBackend
+from basic_memory.repository.embedding_provider_factory import create_embedding_provider
 from basic_memory.repository.postgres_search_repository import PostgresSearchRepository
 from basic_memory.repository.search_index_row import SearchIndexRow
 from basic_memory.repository.search_repository_base import VectorSyncBatchResult
@@ -120,19 +121,37 @@ def create_search_repository(
     Returns:
         SearchRepository: Backend-appropriate search repository instance
     """
-    # Prefer explicit parameter; fall back to ConfigManager for backwards compatibility
+    # Resolve config once so backend detection and the shared embedding provider
+    # come from the same source. Prefer the explicit arg; fall back to ConfigManager
+    # for backwards compatibility.
+    config = app_config or ConfigManager().config
     if database_backend is None:
-        config = app_config or ConfigManager().config
         database_backend = config.database_backend
 
+    # Trigger: every request, sync batch, and project builds its own search repo.
+    # Why: each repo __init__ would otherwise call create_embedding_provider(), and
+    # the process-wide cache can be bypassed if its key ever drifts (#872), reloading
+    # the ~2.3GB ONNX model and leaking memory in onnxruntime's CPU arena.
+    # Outcome: resolve the cached singleton here once and inject it, so the provider
+    # is the single source of truth across all callers of this factory.
+    embedding_provider = None
+    if config.semantic_search_enabled:
+        embedding_provider = create_embedding_provider(config)
+
     if database_backend == DatabaseBackend.POSTGRES:  # pragma: no cover
         return PostgresSearchRepository(  # pragma: no cover
             session_maker,
             project_id=project_id,
             app_config=app_config,
+            embedding_provider=embedding_provider,
         )
     else:
-        return SQLiteSearchRepository(session_maker, project_id=project_id, app_config=app_config)
+        return SQLiteSearchRepository(
+            session_maker,
+            project_id=project_id,
+            app_config=app_config,
+            embedding_provider=embedding_provider,
+        )
 
 
 __all__ = [

src/basic_memory/sync/sync_service.py

@@ -1744,7 +1744,9 @@ async def get_sync_service(project: Project) -> SyncService:  # pragma: no cover
     entity_repository = EntityRepository(session_maker, project_id=project.id)
     observation_repository = ObservationRepository(session_maker, project_id=project.id)
     relation_repository = RelationRepository(session_maker, project_id=project.id)
-    search_repository = create_search_repository(session_maker, project_id=project.id)
+    search_repository = create_search_repository(
+        session_maker, project_id=project.id, app_config=app_config
+    )
     project_repository = ProjectRepository(session_maker)
 
     # Initialize services

test-int/semantic/test_embedding_provider_reuse.py

@@ -0,0 +1,123 @@
+"""Integration tests for process-wide embedding provider reuse (#872).
+
+A long-running ``basic-memory mcp`` server constructs a new search repository per
+request, per sync batch, and per project. Each repository used to derive its own
+embedding provider via ``create_embedding_provider()``. If the provider cache key
+ever drifted, that reloaded the ~2.3GB FastEmbed/ONNX model and leaked memory in
+onnxruntime's CPU arena (which never returns memory to the OS).
+
+These tests use the *real* composition paths — ``create_embedding_provider``, the
+``create_search_repository`` factory, and the FastAPI deps function
+``get_search_repository`` — with a real FastEmbed provider. FastEmbed loads the
+ONNX model lazily on first embed, so constructing providers/repositories here is
+cheap and never touches the native model.
+
+They deliberately use the semantic suite's ``sqlite_engine_factory`` (not the
+parent ``engine_factory``): the parent fixture depends on ``postgres_engine``,
+which this directory overrides to spin up a Docker testcontainer gated only by a
+``docker`` binary on PATH. On Windows CI that binary exists without a usable
+daemon, so requesting the parent factory aborts these SQLite-only tests with a
+testcontainers Ryuk error (#872 follow-up). Staying on the SQLite factory keeps
+them backend-correct and Docker-free.
+"""
+
+from __future__ import annotations
+
+from typing import cast
+
+import pytest
+
+from basic_memory.config import BasicMemoryConfig, DatabaseBackend, ProjectEntry
+from basic_memory.deps.repositories import get_search_repository
+from basic_memory.repository.embedding_provider_factory import (
+    create_embedding_provider,
+    reset_embedding_provider_cache,
+)
+from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider
+from basic_memory.repository.search_repository import create_search_repository
+from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
+
+
+def _semantic_config(project_path) -> BasicMemoryConfig:
+    """Build a semantic-enabled FastEmbed config rooted at the test path."""
+    return BasicMemoryConfig(
+        env="test",
+        projects={"test-project": ProjectEntry(path=str(project_path))},
+        default_project="test-project",
+        database_backend=DatabaseBackend.SQLITE,
+        semantic_search_enabled=True,
+        semantic_embedding_provider="fastembed",
+    )
+
+
+@pytest.fixture(autouse=True)
+def _reset_provider_cache():
+    reset_embedding_provider_cache()
+    yield
+    reset_embedding_provider_cache()
+
+
+@pytest.mark.asyncio
+async def test_factory_resolves_single_provider_across_repositories(
+    tmp_path, sqlite_engine_factory
+):
+    """The factory must inject one cached provider into every search repository."""
+    _engine, session_maker = sqlite_engine_factory
+    config = _semantic_config(tmp_path)
+
+    expected_provider = create_embedding_provider(config)
+    assert isinstance(expected_provider, FastEmbedEmbeddingProvider)
+
+    # Two repositories, mimicking per-request / per-sync construction.
+    repo_a = cast(
+        SQLiteSearchRepository,
+        create_search_repository(session_maker, project_id=1, app_config=config),
+    )
+    repo_b = cast(
+        SQLiteSearchRepository,
+        create_search_repository(session_maker, project_id=2, app_config=config),
+    )
+
+    # Both repos reuse the exact same cached provider object — no second model load.
+    assert repo_a._embedding_provider is expected_provider
+    assert repo_b._embedding_provider is expected_provider
+
+
+@pytest.mark.asyncio
+async def test_deps_path_reuses_cached_provider(tmp_path, sqlite_engine_factory):
+    """The real FastAPI deps function must reuse the cached provider, not rebuild it."""
+    _engine, session_maker = sqlite_engine_factory
+    config = _semantic_config(tmp_path)
+
+    expected_provider = create_embedding_provider(config)
+
+    repo = cast(
+        SQLiteSearchRepository,
+        await get_search_repository(
+            session_maker=session_maker,
+            project_id=1,
+            app_config=config,
+        ),
+    )
+
+    assert repo._embedding_provider is expected_provider
+
+
+@pytest.mark.asyncio
+async def test_factory_skips_provider_when_semantic_disabled(tmp_path, sqlite_engine_factory):
+    """With semantic search off, no provider is created and none is injected."""
+    _engine, session_maker = sqlite_engine_factory
+    config = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": ProjectEntry(path=str(tmp_path))},
+        default_project="test-project",
+        database_backend=DatabaseBackend.SQLITE,
+        semantic_search_enabled=False,
+    )
+
+    repo = cast(
+        SQLiteSearchRepository,
+        create_search_repository(session_maker, project_id=1, app_config=config),
+    )
+
+    assert repo._embedding_provider is None

tests/repository/test_fastembed_provider.py

@@ -24,13 +24,20 @@ class _StubTextEmbedding:
     last_init_kwargs: dict = {}
     last_embed_kwargs: dict = {}
 
-    def __init__(self, model_name: str, cache_dir: str | None = None, threads: int | None = None):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        enable_cpu_mem_arena: bool | None = None,
+    ):
         self.model_name = model_name
         self.embed_calls = 0
         _StubTextEmbedding.last_init_kwargs = {
             "model_name": model_name,
             "cache_dir": cache_dir,
             "threads": threads,
+            "enable_cpu_mem_arena": enable_cpu_mem_arena,
         }
         _StubTextEmbedding.init_count += 1
 
@@ -120,10 +127,36 @@ async def test_fastembed_provider_passes_runtime_knobs_to_fastembed(monkeypatch)
         "model_name": "stub-model",
         "cache_dir": "/tmp/fastembed-cache",
         "threads": 3,
+        # onnxruntime CPU mem arena is disabled so transient extra loads free memory (#872)
+        "enable_cpu_mem_arena": False,
     }
     assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 8, "parallel": 2}
 
 
+@pytest.mark.asyncio
+async def test_fastembed_provider_disables_cpu_mem_arena_by_default(monkeypatch):
+    """Even with no cache_dir/threads set, the ONNX CPU memory arena must be disabled.
+
+    onnxruntime's CPU arena never returns memory to the OS, so a duplicate model
+    load would leak tens of GB in a long-running process (#872). The provider must
+    always pass enable_cpu_mem_arena=False to FastEmbed regardless of other knobs.
+    """
+    module = type(sys)("fastembed")
+    setattr(module, "TextEmbedding", _StubTextEmbedding)
+    monkeypatch.setitem(sys.modules, "fastembed", module)
+    _StubTextEmbedding.last_init_kwargs = {}
+
+    provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4)
+    await provider.embed_documents(["arena default"])
+
+    assert _StubTextEmbedding.last_init_kwargs == {
+        "model_name": "stub-model",
+        "cache_dir": None,
+        "threads": None,
+        "enable_cpu_mem_arena": False,
+    }
+
+
 @pytest.mark.asyncio
 async def test_fastembed_provider_parallel_one_disables_multiprocessing(monkeypatch):
     """parallel=1 should not pass FastEmbed multiprocessing kwargs."""
@@ -245,7 +278,13 @@ class _SelfHealStubTextEmbedding:
     RESOLVED_MODEL = "stub-model"
     MODEL_FILE = "model_optimized.onnx"
 
-    def __init__(self, model_name: str, cache_dir: str | None = None, threads: int | None = None):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **_kwargs,
+    ):
         type(self).construct_count += 1
         if type(self).construct_count <= type(self).fail_first_n:
             raise RuntimeError(
@@ -555,3 +594,80 @@ async def test_fastembed_provider_fails_fast_without_cache_dir(monkeypatch):
         await provider.embed_documents(["no cache dir"])
 
     assert _SelfHealStubTextEmbedding.construct_count == 1
+
+@pytest.mark.asyncio
+async def test_factory_loads_native_model_once_across_repo_constructions(monkeypatch):
+    """The native ONNX model must load exactly once per process despite reuse (#872).
+
+    Counting native model loads requires a stub TextEmbedding that increments a
+    counter on construction — there is no way to observe real ONNX loads otherwise.
+    This is the justified mock: the rest of the path (factory cache, repository
+    injection) is exercised with real implementations.
+
+    The test resolves the provider several times with a *drifting* CPU budget — the
+    exact condition that previously produced a fresh cache key and a second model
+    load — then builds multiple search repositories and embeds across all of them,
+    asserting the stub was constructed only once.
+    """
+    from typing import Any, cast
+
+    from basic_memory.config import BasicMemoryConfig, DatabaseBackend, ProjectEntry
+    from basic_memory.repository import embedding_provider_factory as factory_module
+    from basic_memory.repository.embedding_provider_factory import (
+        create_embedding_provider,
+        reset_embedding_provider_cache,
+    )
+    from basic_memory.repository.search_repository import create_search_repository
+    from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
+
+    module = type(sys)("fastembed")
+    setattr(module, "TextEmbedding", _StubTextEmbedding)
+    monkeypatch.setitem(sys.modules, "fastembed", module)
+    _StubTextEmbedding.init_count = 0
+    reset_embedding_provider_cache()
+
+    config = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": ProjectEntry(path="/tmp/basic-memory-test")},
+        default_project="test-project",
+        database_backend=DatabaseBackend.SQLITE,
+        semantic_search_enabled=True,
+        semantic_embedding_provider="fastembed",
+        semantic_embedding_model="stub-model",
+        semantic_embedding_dimensions=4,
+        semantic_embedding_threads=None,
+        semantic_embedding_parallel=None,
+    )
+
+    try:
+        # First resolution under one CPU budget.
+        monkeypatch.setattr(factory_module.os, "process_cpu_count", lambda: 8)
+        monkeypatch.setattr(factory_module.os, "cpu_count", lambda: 8)
+        provider_first = create_embedding_provider(config)
+
+        # CPU budget drifts (cgroup throttling) — used to force a second model load.
+        monkeypatch.setattr(factory_module.os, "process_cpu_count", lambda: 4)
+        monkeypatch.setattr(factory_module.os, "cpu_count", lambda: 4)
+
+        # Build several repositories the way per-request/per-sync code does. Each
+        # one is injected with the cached provider rather than deriving its own.
+        # session_maker is unused during construction and during embed_documents,
+        # so None is sufficient for this provider-identity assertion.
+        repos = [
+            cast(
+                SQLiteSearchRepository,
+                create_search_repository(cast(Any, None), project_id=project_id, app_config=config),
+            )
+            for project_id in (1, 2, 3)
+        ]
+        for repo in repos:
+            assert repo._embedding_provider is provider_first
+
+        # Embed several times to trigger lazy model loads; because every repo shares
+        # provider_first, repeated embeds must still construct the stub model once.
+        for _ in repos:
+            await provider_first.embed_documents(["auth token session"])
+
+        assert _StubTextEmbedding.init_count == 1
+    finally:
+        reset_embedding_provider_cache()

tests/repository/test_openai_provider.py

@@ -523,7 +523,37 @@ async def test_openai_provider_fails_fast_on_malformed_concurrent_batch(monkeypa
 
 
 def test_embedding_provider_factory_creates_new_provider_for_different_cache_key():
-    """Factory should create distinct providers when cache key fields differ."""
+    """Factory should create distinct providers when model-identity fields differ."""
+    config_a = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": "/tmp/basic-memory-test"},
+        default_project="test-project",
+        semantic_search_enabled=True,
+        semantic_embedding_provider="fastembed",
+        semantic_embedding_model="bge-small-en-v1.5",
+    )
+    config_b = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": "/tmp/basic-memory-test"},
+        default_project="test-project",
+        semantic_search_enabled=True,
+        semantic_embedding_provider="fastembed",
+        semantic_embedding_model="some-other-model",
+    )
+
+    provider_a = create_embedding_provider(config_a)
+    provider_b = create_embedding_provider(config_b)
+
+    assert provider_a is not provider_b
+
+
+def test_embedding_provider_factory_reuses_provider_when_only_thread_knobs_differ():
+    """Thread/parallel knobs tune ONNX execution, not model identity (#872).
+
+    A different CPU-derived thread count must not invalidate the process-wide
+    provider cache and reload the model. Two configs that differ only in the
+    FastEmbed thread/parallel knobs must return the exact same provider instance.
+    """
     config_a = BasicMemoryConfig(
         env="test",
         projects={"test-project": "/tmp/basic-memory-test"},
@@ -531,6 +561,7 @@ def test_embedding_provider_factory_creates_new_provider_for_different_cache_key
         semantic_search_enabled=True,
         semantic_embedding_provider="fastembed",
         semantic_embedding_threads=2,
+        semantic_embedding_parallel=1,
     )
     config_b = BasicMemoryConfig(
         env="test",
@@ -539,12 +570,42 @@ def test_embedding_provider_factory_creates_new_provider_for_different_cache_key
         semantic_search_enabled=True,
         semantic_embedding_provider="fastembed",
         semantic_embedding_threads=4,
+        semantic_embedding_parallel=2,
     )
 
     provider_a = create_embedding_provider(config_a)
     provider_b = create_embedding_provider(config_b)
 
-    assert provider_a is not provider_b
+    assert provider_a is provider_b
+
+
+def test_embedding_provider_factory_reuses_provider_when_cpu_budget_drifts(monkeypatch):
+    """A drifting CPU budget between calls must not reload the model (#872).
+
+    Simulates a container/cgroup where the auto-tuned thread count changes between
+    two calls. Before the fix, this produced a fresh cache key and reloaded the
+    ~2.3GB ONNX model; now the cached singleton is reused.
+    """
+    config = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": "/tmp/basic-memory-test"},
+        default_project="test-project",
+        semantic_search_enabled=True,
+        semantic_embedding_provider="fastembed",
+        semantic_embedding_threads=None,
+        semantic_embedding_parallel=None,
+    )
+
+    monkeypatch.setattr(embedding_provider_factory_module.os, "process_cpu_count", lambda: 8)
+    monkeypatch.setattr(embedding_provider_factory_module.os, "cpu_count", lambda: 8)
+    provider_first = create_embedding_provider(config)
+
+    # CPU budget shrinks (e.g. cgroup throttling) → auto-tuned thread count changes.
+    monkeypatch.setattr(embedding_provider_factory_module.os, "process_cpu_count", lambda: 4)
+    monkeypatch.setattr(embedding_provider_factory_module.os, "cpu_count", lambda: 4)
+    provider_second = create_embedding_provider(config)
+
+    assert provider_first is provider_second
 
 
 def test_embedding_provider_factory_forwards_openai_request_concurrency():