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__ = [