ray-embedding 0.12.2__py3-none-any.whl → 0.12.4__py3-none-any.whl

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/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.dist-info → ray_embedding-0.12.4.dist-info}/METADATA

@@ -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.4.dist-info/RECORD

@@ -0,0 +1,9 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=xE-NznVlYftTPIfN3aAqCF0DpFIvTuC4vMTNrGfjkxI,2627
+ray_embedding/dto.py,sha256=rzPEB-R7XDYlTqeaXGGgfOjTWyTeRinnJ6LbI1oOWGY,1463
+ray_embedding/embedding_model.py,sha256=j8jPhfVqS_x11oZrlXDbjo6z1NHXxOk8CnGIQemGfUQ,4040
+ray_embedding/model_router.py,sha256=BsOEz24ttvpDD4LZsDVg9rLhn26FxgUsDAvcjI0Feao,5917
+ray_embedding-0.12.4.dist-info/METADATA,sha256=e7i2fUgJhRYIvIrzLVlGYVr3FvPD1Eaoj7L08Slbxp4,1094
+ray_embedding-0.12.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.12.4.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.12.4.dist-info/RECORD,,

ray_embedding-0.12.2.dist-info/RECORD

@@ -1,9 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=xE-NznVlYftTPIfN3aAqCF0DpFIvTuC4vMTNrGfjkxI,2627
-ray_embedding/dto.py,sha256=rzPEB-R7XDYlTqeaXGGgfOjTWyTeRinnJ6LbI1oOWGY,1463
-ray_embedding/embedding_model.py,sha256=wHyDgDCR11VcV2by4bCCp6CSi0mI9zcwkmSbnL3WdRY,3519
-ray_embedding/model_router.py,sha256=ELmD9gU6Tn8AUChN2CTaP3Puvnl8S5EcljnpRwsRI1Y,5701
-ray_embedding-0.12.2.dist-info/METADATA,sha256=VHUh1jHfc47rgef1RuAmlepj_uS081sEp6JkxaGMWUw,1094
-ray_embedding-0.12.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.12.2.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.12.2.dist-info/RECORD,,