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

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

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

{ray_embedding-0.12.2 → 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.2 → 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)
 
@@ -88,10 +88,10 @@ class ModelRouter:
 
     @web_api.post("/{path_prefix}/v1/embeddings", response_model=EmbeddingResponse)
     async def compute_embeddings(self, path_prefix: str, request: EmbeddingRequest):
-        assert path_prefix in self.path_prefix, f"Invalid path prefix: {path_prefix}"
-        assert request.model in self.deployed_models, f"Invalid model: {request.model}"
-
         try:
+            assert path_prefix in self.path_prefix, f"The API path prefix specified is invalid: '{path_prefix}'"
+            assert request.model in self.deployed_models, f"The model specified is invalid: {request.model}"
+
             inputs = request.input if isinstance(request.input, list) else [request.input]
             self.logger.info(f"Computing embeddings for a batch of {len(inputs)} texts using model: {request.model}")
             embeddings = await self._compute_embeddings_from_resized_batches(request.model, inputs, request.dimensions)
@@ -101,11 +101,13 @@ class ModelRouter:
             ]
             return EmbeddingResponse(object="list", data=response_data, model=request.model)
         except Exception as e:
+            status_code = 400 if isinstance(e, AssertionError) else 500
             self.logger.error(f"Failed to create embeddings: {e}")
-            raise HTTPException(status_code=500, detail=str(e))
+            raise HTTPException(status_code=status_code, detail=str(e))
 
     @web_api.get("/{path_prefix}/v1/models")
     async def list_models(self, path_prefix: str):
         """Returns the list of available models in OpenAI-compatible format."""
-        assert path_prefix in self.path_prefix, f"Invalid path prefix: {path_prefix}"
+        if path_prefix not in self.path_prefix:
+            raise HTTPException(status_code=400, detail=f"The API path prefix specified is invalid: '{path_prefix}'")
         return {"object": "list", "data": self.available_models}

{ray_embedding-0.12.2 → 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.2
+Version: 0.12.4
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

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

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