ray-embedding 0.12.3__tar.gz → 0.12.5__tar.gz

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.12.3
+Version: 0.12.5
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/ray_embedding/deploy.py

@@ -17,6 +17,7 @@ def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
     matryoshka_dim = model_config.matryoshka_dim
     trust_remote_code = model_config.trust_remote_code or False
     model_kwargs = model_config.model_kwargs or {}
+    cuda_memory_flush_threshold = model_config.cuda_memory_flush_threshold or 0.8
 
     if "torch_dtype" in model_kwargs:
         torch_dtype = model_kwargs["torch_dtype"].strip()

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/ray_embedding/dto.py

@@ -25,7 +25,6 @@ class ModelRouterConfig(BaseModel):
 
 
 class ModelDeploymentConfig(BaseModel):
-    deployment: str
     model: str
     served_model_name: str
     batch_size: Optional[int] = 8
@@ -35,6 +34,8 @@ class ModelDeploymentConfig(BaseModel):
     matryoshka_dim: Optional[int] = 768
     trust_remote_code: Optional[bool] = False
     model_kwargs: Optional[Dict[str, Any]] = {}
+    cuda_memory_flush_threshold: Optional[float] = 0.8
+    deployment: str
 
 
 class AppConfig(BaseModel):

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/ray_embedding/embedding_model.py

@@ -4,7 +4,7 @@ import time
 from typing import Optional, Dict, Any, List, Union
 
 import torch
-from pynvml import nvmlInit, nvmlDeviceGetCount
+from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
 from ray import serve
 from sentence_transformers import SentenceTransformer
 
@@ -14,7 +14,7 @@ class EmbeddingModel:
     def __init__(self, model: str, served_model_name: Optional[str] = None,
                  device: Optional[str] = None, backend: Optional[str] = "torch",
                  matryoshka_dim: Optional[int] = None, trust_remote_code: Optional[bool] = False,
-                 model_kwargs: Dict[str, Any] = None):
+                 model_kwargs: Dict[str, Any] = None, cuda_memory_flush_threshold: Optional[float] = 0.8):
         logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(self.__class__.__name__)
         self.model = model
@@ -29,6 +29,7 @@ class EmbeddingModel:
         self.matryoshka_dim = matryoshka_dim
         self.trust_remote_code = trust_remote_code or False
         self.model_kwargs = model_kwargs or {}
+        self.cuda_memory_flush_threshold = cuda_memory_flush_threshold
         self.logger.info(f"Initializing embedding model: {self.model}")
         self.embedding_model = SentenceTransformer(self.model, device=self.init_device, backend=self.backend,
                                                    trust_remote_code=self.trust_remote_code,
@@ -56,6 +57,10 @@ class EmbeddingModel:
 
         # Move all embeddings to CPU at once before conversion
         embeddings_list = embeddings.cpu().tolist()
+
+        # don't wait for GC
+        del embeddings
+
         return embeddings_list
 
     def wait_for_cuda(self, wait: int = 10):
@@ -64,12 +69,36 @@ class EmbeddingModel:
         self.check_health()
 
     def check_health(self):
-        if self.init_device == "cuda":
+        if self.init_device != "cuda":
+            return
+
+        try:
             # Even though CUDA was available at init time,
             # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
             # https://github.com/ray-project/ray/issues/49594
-            try:
-                nvmlInit()
-                assert nvmlDeviceGetCount() >= 1
-            except:
-                raise RuntimeError("CUDA device is not available")
+            nvmlInit()
+            count = nvmlDeviceGetCount()
+            assert count >= 1, "No CUDA devices found"
+
+            # replicas only have access to GPU 0
+            handle = nvmlDeviceGetHandleByIndex(0)
+            mem_info = nvmlDeviceGetMemoryInfo(handle)
+        except Exception as e:
+            raise RuntimeError(f"CUDA health check failed: {e}")
+
+        reserved = torch.cuda.memory_reserved()  # bytes currently reserved by CUDA cache
+        threshold_bytes = self.cuda_memory_flush_threshold * mem_info.total
+
+        if reserved > threshold_bytes:
+            # flush only when cache exceeds the percentage threshold
+            torch.cuda.empty_cache()
+
+    def __del__(self):
+        # Clean up and free any remaining GPU memory
+        try:
+            if hasattr(self, 'embedding_model'):
+                del self.embedding_model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception as e:
+            self.logger.warning(f"Error during cleanup: {e}")

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/ray_embedding/model_router.py

@@ -46,7 +46,7 @@ class ModelRouter:
                              f"to {len(batches)} mini-batches, each with max length {batch_size}.")
 
         # Call embedding model replicas in parallel (rate-limited)
-        tasks = [self._rate_limited_embedding_call(model_handle, batch, dimensions) for batch in batches]
+        tasks = [self._compute_embeddings_rate_limited(model_handle, batch, dimensions) for batch in batches]
         all_results = await asyncio.gather(*tasks, return_exceptions=True)
 
         # Retry any failed model calls
@@ -64,7 +64,7 @@ class ModelRouter:
         self.logger.info(f"Successfully computed embeddings from {len(batches)} mini-batches")
         return [emb for result in all_results for emb in result]
 
-    async def _rate_limited_embedding_call(self, model_handle: DeploymentHandle, batch: List[str], dimensions: int):
+    async def _compute_embeddings_rate_limited(self, model_handle: DeploymentHandle, batch: List[str], dimensions: int):
         async with self.rate_limiter:
             return await model_handle.remote(batch, dimensions)
 

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/ray_embedding.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.12.3
+Version: 0.12.5
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

{ray_embedding-0.12.3 → ray_embedding-0.12.5}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = ray-embedding
-version = 0.12.3
+version = 0.12.5
 author = Crispin Almodovar
 author_email = 
 description = Deploy SentenceTransformers embedding models to a ray cluster