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

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

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

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

@@ -56,6 +56,12 @@ class EmbeddingModel:
 
         # Move all embeddings to CPU at once before conversion
         embeddings_list = embeddings.cpu().tolist()
+
+        # free GPU memory now, don't wait for GC
+        del embeddings
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
         return embeddings_list
 
     def wait_for_cuda(self, wait: int = 10):
@@ -73,3 +79,13 @@ class EmbeddingModel:
                 assert nvmlDeviceGetCount() >= 1
             except:
                 raise RuntimeError("CUDA device is not available")
+
+    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.4}/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.4}/ray_embedding.egg-info/PKG-INFO

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

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

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