Fix OOM crash causing missing jobs by limiting PyTorch concurrency

agents/crew.py

@@ -285,12 +285,13 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
     if not retrieval_is_empty:
         try:
             reranker = _get_reranker()
-            
-            # Prepare pairs of (query, chunk_text)
-            pairs = [[query, chunk["text"]] for chunk in unique_chunks]
-            
-            # Predict scores using the CrossEncoder
-            scores = reranker.predict(pairs)
+            import threading
+            # Lock the PyTorch inference block to prevent OOM
+            if not hasattr(_get_reranker, '_lock'):
+                _get_reranker._lock = threading.Lock()
+            with _get_reranker._lock:
+                # Predict scores using the CrossEncoder
+                scores = reranker.predict([(query, chunk["text"]) for chunk in unique_chunks])
             
             # Assign scores back to the chunks
             for i, chunk in enumerate(unique_chunks):

agents/embed_llm.py

@@ -82,14 +82,18 @@ _embed_dim  = _get_dim()
 log.info("SentenceTransformer model ready — dim=%d", _embed_dim)
 
 
+import threading
+_embed_lock = threading.Lock()
+
 def _embed_sentences(sentences: list[str]) -> np.ndarray:
     """Embed a list of sentences and return dense vectors as ndarray (N, dim)."""
-    vecs = _st_model.encode(
-        sentences,
-        batch_size=BATCH_SIZE,
-        show_progress_bar=False,
-        normalize_embeddings=True,
-    )
+    with _embed_lock:
+        vecs = _st_model.encode(
+            sentences,
+            batch_size=BATCH_SIZE,
+            show_progress_bar=False,
+            normalize_embeddings=True,
+        )
     return vecs if isinstance(vecs, np.ndarray) else np.array(vecs)
 
 

app.py

@@ -133,7 +133,7 @@ _session_uploads: dict[str, int] = {}
 
 # Global thread pool to allow concurrent RAG execution
 from concurrent.futures import ThreadPoolExecutor
-_query_executor = ThreadPoolExecutor(max_workers=8)
+_query_executor = ThreadPoolExecutor(max_workers=2)
 
 # Auto-ingest background progress tracker
 _auto_ingest_status: dict = {
@@ -889,7 +889,7 @@ def query_stream(job_id):
                     return
                 idx += 1
 
-            if job["error"] or job["done"]:
+            if job.get("error") or job.get("done"):
                 break
                 
             time.sleep(0.5)

start.sh

@@ -41,6 +41,12 @@ if [ "$HF_MODE_FLAG" -eq 1 ]; then
   export TOP_K_GRAPH=3
   export EMBED_FP16=false          # CPU only — FP16 unsupported
   export TORCH_COMPILE_SKIP=1      # Skip torch.compile() on CPU (no benefit, adds 30s startup)
+  
+  # Crucial to prevent CPU thrashing and PyTorch OOM when multiple threads query
+  export OMP_NUM_THREADS=2
+  export MKL_NUM_THREADS=2
+  export OPENBLAS_NUM_THREADS=2
+  
   MODE_LABEL="HuggingFace / CPU"
 else
   export GEN_MODEL_ID="${GEN_MODEL_ID:-Jackrong/Qwen3.5-2B-Claude-4.6-Opus-Reasoning-Distilled-GGUF}"