embeddings: serialize per-model loads — concurrent bursts stacked duplicate copies

A bulk indexer's burst of first-requests all saw model_obj None and EACH
loaded its own copy of the embedder (observed pool_instances=3 for the
4B model + 6 for DINOv2 ≈ the whole card) — this duplicate stacking,
not a single load, was what exhausted VRAM and poisoned the measured
footprint (26.8 GB recorded for an 8 GB model, measured across
overlapping loads). Loads now take a per-model asyncio lock with a
re-check, so one request loads and the rest reuse the loaded model.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014S8VtAvG499SsCbeESRK7V

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.1.36"
+__version__ = "0.1.37"
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
 # expandable_segments lets the allocator return freed pages to the driver even

codai/api/embeddings.py

@@ -500,22 +500,24 @@ async def create_embeddings(request: EmbeddingsRequest, http_request: Request =
         raise
 
 
-async def _run_embeddings(request: EmbeddingsRequest, http_request: Request = None):
-    """Core embeddings logic; registered as a task by create_embeddings()."""
-    model_info = await asyncio.to_thread(
-        multi_model_manager.request_model, request.model, model_type="embedding")
-    model_name = model_info.get('model_name')
-    if not model_name:
-        err = model_info.get('error', f"Model '{request.model}' not found")
-        raise HTTPException(status_code=404, detail=err)
+# Per-model-key asyncio locks serializing embedding model loads (see the
+# comment at the acquire site in _run_embeddings).
+_load_locks: dict = {}
 
-    model_key = model_info['model_key']
-    model_obj = model_info.get('model_object')
 
-    _emb_cfg = (multi_model_manager.config.get(f"embedding:{model_name}")
-                or multi_model_manager.config.get(model_name) or {})
+async def _load_embedding_locked(request, model_key: str, model_name: str,
+                                 _emb_cfg: dict):
+    """Load an embedding model — called with the per-model load lock HELD.
+
+    Re-checks the registry first (another request may have finished the load
+    while we waited on the lock), then loads with the standard evict-and-retry
+    contract every other model type follows."""
+    # Re-check under the lock: the request that held the lock before us has
+    # usually just loaded the model.
+    model_obj = multi_model_manager.models.get(model_key)
+    if model_obj is not None:
+        return model_obj
 
-    if model_obj is None:
     device = _derive_device()
     from codai.tasks import loading_task
     # Snapshot VRAM around the load so the model's real footprint is measured
@@ -570,12 +572,41 @@ async def _run_embeddings(request: EmbeddingsRequest, http_request: Request = No
                 model_obj = await asyncio.get_event_loop().run_in_executor(
                     None, _load_embedding_model, model_name, device, _emb_cfg)
     except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to load embedding model: {e}")
+        raise HTTPException(status_code=500,
+                            detail=f"Failed to load embedding model: {e}")
     # Register through add_model (pool + models_in_vram bookkeeping) rather than
     # a bare dict assignment, so eviction/unload treat it like every other model.
     multi_model_manager.add_model(model_key, model_obj)
     multi_model_manager.current_model_key = model_key
     multi_model_manager.record_vram_delta(model_key, _snap)
+    return model_obj
+
+
+async def _run_embeddings(request: EmbeddingsRequest, http_request: Request = None):
+    """Core embeddings logic; registered as a task by create_embeddings()."""
+    model_info = await asyncio.to_thread(
+        multi_model_manager.request_model, request.model, model_type="embedding")
+    model_name = model_info.get('model_name')
+    if not model_name:
+        err = model_info.get('error', f"Model '{request.model}' not found")
+        raise HTTPException(status_code=404, detail=err)
+
+    model_key = model_info['model_key']
+    model_obj = model_info.get('model_object')
+
+    _emb_cfg = (multi_model_manager.config.get(f"embedding:{model_name}")
+                or multi_model_manager.config.get(model_name) or {})
+
+    if model_obj is None:
+        # Serialize loads per model: a burst of first-requests (a bulk indexer)
+        # otherwise ALL see model_obj None and EACH loads its own copy — several
+        # 8 GB instances of the same embedder stacked on the card (observed as
+        # pool_instances=3/6), which is what actually filled the GPU. One
+        # request loads; the rest wait on the lock and reuse the loaded model.
+        _lock = _load_locks.setdefault(model_key, asyncio.Lock())
+        async with _lock:
+            model_obj = await _load_embedding_locked(
+                request, model_key, model_name, _emb_cfg)
 
     texts: List[str] = []
     if request.input is not None: