fix: vram not release after indexing

Author: qduc

src/api.py

@@ -334,6 +334,9 @@ def index_codebase(
 
             metadata_index.mark_reindex_complete()
 
+        # Clear VRAM cache after heavy indexing
+        self.embedding_model.clear_cache()
+
         return total_chunks
 
     def incremental_reindex(
@@ -459,6 +462,9 @@ def incremental_reindex(
         # Mark reindex complete
         metadata_index.mark_reindex_complete()
 
+        # Clear VRAM cache after heavy indexing
+        self.embedding_model.clear_cache()
+
         return {
             'added_count': len(added),
             'modified_count': len(modified),

src/embedding_server.py

@@ -129,17 +129,23 @@ def _load_models(self):
                 self._embedding_model = OpenAIEmbedding(model_name)
             else:
                 from .embeddings.sentence_transformer_embedding import SentenceTransformerEmbedding
-                # No idle timeout - server manages lifecycle
-                self._embedding_model = SentenceTransformerEmbedding(model_name, lazy_load=False, idle_timeout=0)
+                # Use configured idle timeout (default 30 min) to release VRAM
+                idle_timeout = self.config.get_model_idle_timeout()
+                self._embedding_model = SentenceTransformerEmbedding(
+                    model_name,
+                    lazy_load=False,
+                    idle_timeout=idle_timeout
+                )
 
             # Load reranker if enabled
             if self.config.is_reranker_enabled():
                 print(f"Loading reranker model: {self.config.get_reranker_model()}", file=sys.stderr)
                 from .reranker.cross_encoder_reranker import CrossEncoderReranker
+                idle_timeout = self.config.get_model_idle_timeout()
                 self._reranker = CrossEncoderReranker(
                     self.config.get_reranker_model(),
                     lazy_load=False,
-                    idle_timeout=0
+                    idle_timeout=idle_timeout
                 )
 
             self._models_loaded = True
@@ -242,6 +248,23 @@ async def embedding_dimension(request: DimensionRequest):
             dim = self._embedding_model.get_embedding_dimension()
             return {"dimension": dim}
 
+        @app.post("/clear_cache")
+        async def clear_cache(request: DimensionRequest):
+            """Clear memory cache (CUDA) without unloading models."""
+            self._update_heartbeat(request.client_id)
+
+            with self._models_lock:
+                if self._embedding_model:
+                    if hasattr(self._embedding_model, "clear_cache"):
+                        self._embedding_model.clear_cache()
+
+                # Also clear reranker cache if loaded
+                if self._reranker:
+                    if hasattr(self._reranker, "clear_cache"):
+                        self._reranker.clear_cache()
+
+            return {"status": "ok"}
+
         @app.post("/rerank")
         async def rerank(request: RerankRequest):
             """Rerank documents."""

src/embeddings/embedding_interface.py

@@ -47,3 +47,20 @@ def embed_batch(self, texts: List[str]) -> List[List[float]]:
             A list of embedding vectors
         """
         pass
+
+    @abstractmethod
+    def get_embedding_dimension(self) -> int:
+        """
+        Get the dimension of the embedding vectors.
+
+        Returns:
+            The dimension of embeddings produced by this model
+        """
+        pass
+
+    def clear_cache(self):
+        """
+        Clear memory cache (e.g. CUDA memory) without unloading the model.
+        Useful after memory-intensive operations like indexing.
+        """
+        pass

src/embeddings/http_embedding.py

@@ -258,3 +258,15 @@ def start_background_loading(self):
     def unload_model(self):
         """No-op for HTTP client - server manages model lifecycle."""
         pass
+
+    def clear_cache(self):
+        """Request server to clear memory cache."""
+        try:
+            # We don't retry this if it fails, as it's an optimization only
+            requests.post(
+                f"{self.base_url}/clear_cache",
+                json={"client_id": self.client_id},
+                timeout=CONNECTION_TIMEOUT
+            )
+        except (requests.ConnectionError, requests.Timeout):
+            pass

src/embeddings/sentence_transformer_embedding.py

@@ -98,6 +98,12 @@ def _cleanup_loop(self):
                         self.model = None
                         # Force garbage collection to release GPU memory if applicable
                         gc.collect()
+                        try:
+                            import torch
+                            if torch.cuda.is_available():
+                                torch.cuda.empty_cache()
+                        except ImportError:
+                            pass
 
     def _ensure_model_loaded(self):
         """Ensure the model is loaded before use (blocks if still loading)."""
@@ -174,6 +180,12 @@ def unload_model(self):
                 del self.model
                 self.model = None
                 gc.collect()
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                except ImportError:
+                    pass
 
     def stop_cleanup_thread(self):
         """Stop the background cleanup thread."""
@@ -185,3 +197,15 @@ def __del__(self):
         """Cleanup when the object is destroyed."""
         self.stop_cleanup_thread()
         self.unload_model()
+
+    def clear_cache(self):
+        """Clear memory cache (CUDA memory) without unloading the model."""
+        with self._loading_lock:
+            # Force garbage collection
+            gc.collect()
+            try:
+                import torch
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except ImportError:
+                pass

src/reranker/cross_encoder_reranker.py

@@ -85,6 +85,12 @@ def _cleanup_loop(self):
                         self.model = None
                         # Force garbage collection to release GPU memory if applicable
                         gc.collect()
+                        try:
+                            import torch
+                            if torch.cuda.is_available():
+                                torch.cuda.empty_cache()
+                        except ImportError:
+                            pass
 
     def _ensure_model_loaded(self):
         """Ensure the model is loaded before use (blocks if still loading)."""
@@ -141,6 +147,12 @@ def unload_model(self):
                 del self.model
                 self.model = None
                 gc.collect()
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                except ImportError:
+                    pass
 
     def stop_cleanup_thread(self):
         """Stop the background cleanup thread."""