ray-embedding 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

ray_embedding/deploy.py

@@ -5,6 +5,10 @@ import torch
 
 
 def deploy_model(args: Dict[str, Any]) -> Application:
+    """Builds and deploys a SentenceTransformer embedding model.
+    :arg args: arguments for initializing a SentenceTransformer model
+    :returns: a Ray Serve Application
+    """
     assert args
     deployment_name: str = args.pop("deployment", "")
     assert deployment_name
@@ -17,13 +21,15 @@ def deploy_model(args: Dict[str, Any]) -> Application:
     trust_remote_code: Optional[bool] = args.pop("trust_remote_code", False)
     model_kwargs: Dict[str, Any] = args.pop("model_kwargs", {})
     if "torch_dtype" in model_kwargs:
-        model_kwargs["torch_dtype"] = model_kwargs["torch_dtype"].strip()
-        if model_kwargs["torch_dtype"] == "float16":
+        torch_dtype = model_kwargs["torch_dtype"].strip()
+        if torch_dtype == "float16":
             model_kwargs["torch_dtype"] = torch.float16
-        elif model_kwargs["torch_dtype"] == "bfloat16":
+        elif torch_dtype == "bfloat16":
             model_kwargs["torch_dtype"] = torch.bfloat16
+        elif torch_dtype == "float32":
+            model_kwargs["torch_dtype"] = torch.float32
         else:
-            del model_kwargs["torch_dtype"]
+            del model_kwargs["torch_dtype"] # Remove
 
     deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
                                                                    backend=backend,

ray_embedding/embedding_model.py

@@ -1,7 +1,7 @@
 import logging
 import os.path
 import time
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
 
 import torch
 from fastapi import FastAPI, HTTPException
@@ -14,12 +14,22 @@ web_api = FastAPI(title=f"Ray Embeddings - OpenAI-compatible API")
 
 
 @serve.deployment(
-    ray_actor_options={"num_gpus": 1},
+    num_replicas="auto",
+    ray_actor_options={
+        "num_cpus": 1,
+        "num_gpus": 0,
+        "max_restarts": -1,
+        "max_task_retries": -1
+    },
     autoscaling_config={
         "target_ongoing_requests": 2,
         "min_replicas": 0,
         "initial_replicas": 1,
         "max_replicas": 1,
+    },
+    user_config={
+        "max_batch_size": 8,
+        "batch_wait_timeout_s": 0.25,
     }
 )
 @serve.ingress(web_api)
@@ -46,7 +56,13 @@ class EmbeddingModel:
              "owned_by": "openai",
              "permission": []}
         ]
-        self.logger.info(f"Successfully initialized embedding model {self.model}")
+        self.logger.info(f"Successfully initialized embedding model {self.model} using device {self.torch_device}")
+
+    def reconfigure(self, user_config: Dict):
+        assert "max_batch_size" in user_config and "batch_wait_timeout_s" in user_config, "Invalid user config"
+        self.logger.info(f"Reconfiguring dynamic batching parameters: {user_config}")
+        self.create_embeddings_batch.set_max_batch_size(user_config["max_batch_size"])
+        self.create_embeddings_batch.set_batch_wait_timeout_s(user_config["batch_wait_timeout_s"])
 
     @web_api.post("/v1/embeddings", response_model=EmbeddingResponse)
     async def create_embeddings(self, request: EmbeddingRequest):
@@ -55,34 +71,46 @@ class EmbeddingModel:
             assert request.model == self.served_model_name, (
                 f"Model '{request.model}' is not supported. Use '{self.served_model_name}' instead."
             )
+            return await self.create_embeddings_batch(request)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
 
+    @serve.batch(max_batch_size=8, batch_wait_timeout_s=0.25)
+    async def create_embeddings_batch(self, requests: List[EmbeddingRequest]) -> List[EmbeddingResponse]:
+        # Batch the text inputs
+        inputs = [], truncate_dims = []
+        for request in requests:
             if isinstance(request.input, str):
                 request.input = [request.input]
-
-            truncate_dim = request.dimensions or self.matryoshka_dim
-
-            # Compute embeddings and convert to a PyTorch tensor on the GPU
-            embeddings = self.embedding_model.encode(
-                request.input, convert_to_tensor=True, normalize_embeddings=True
-            ).to(self.torch_device)
+            inputs.extend(request.input)
+            truncate_dims.append(request.dimensions or self.matryoshka_dim)
+
+        # Compute embeddings for the batch of text inputs
+        embeddings = self.embedding_model.encode(
+            inputs, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
+        ).to(self.torch_device)
+
+        # Truncate the embeddings; note that the truncate_dim can be different for each request
+        # so we need to this step one by one
+        results = []
+        ix = 0
+        for truncate_dim, request in zip(truncate_dims, requests):
+            num_inputs = len(request.input)
+            batch_embeddings = embeddings[ix: ix + num_inputs]
+            ix += num_inputs
 
             if truncate_dim is not None:
-                # Truncate and re-normalize the embeddings
-                embeddings = embeddings[:, :truncate_dim]
-                embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
-
-            # Move all embeddings to CPU at once before conversion
-            embeddings = embeddings.cpu().tolist()
+                batch_embeddings = batch_embeddings[:, :truncate_dim]
+                batch_embeddings = batch_embeddings / torch.norm(batch_embeddings, dim=1, keepdim=True)
 
-            # Convert embeddings to list format for response
+            batch_embeddings = batch_embeddings.cpu().tolist()
             response_data = [
-                {"index": idx, "embedding": emb}
-                for idx, emb in enumerate(embeddings)
+                {"index": emb_ix, "embedding": emb}
+                for emb_ix, emb in enumerate(batch_embeddings)
             ]
-            return EmbeddingResponse(object="list", data=response_data, model=request.model)
+            results.append(EmbeddingResponse(object="list", data=response_data, model=request.model))
 
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
+        return results
 
     @web_api.get("/v1/models")
     async def list_models(self):

{ray_embedding-0.9.0.dist-info → ray_embedding-0.9.1.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.9.0
-Summary: Deploy SentenceTransformers models to a ray cluster
+Version: 0.9.1
+Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: crispin.almodovar@docorto.ai
 Classifier: Programming Language :: Python :: 3
@@ -16,9 +16,15 @@ A tool for deploying SentenceTransformers models to a ray cluster.
 
 ### Supports the following backends
 
-- sbert-pytorch-gpu
-- sbert-pytorch-cpu
-- sbert-onnx-gpu
-- sbert-onnx-cpu
-- sbert-openvino-cpu
+- pytorch-gpu
+- pytorch-cpu
+
+### Planned:
+- onnx-gpu
+- onnx-cpu
+- openvino-cpu
 - fastembed-onnx-cpu
+ 
+- spot instances
+- grpc
+

ray_embedding-0.9.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=YD_udSm13QbFPgSAkCrTQso15DmtIn0QEhErOFNg7jM,1841
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=QvJRwQYgqMs99F8Nbo3aC0N5oMsMEtNB6W2C4-zSfPs,5131
+ray_embedding-0.9.1.dist-info/METADATA,sha256=w1Khp-pKHk6AwTzljTU_Z-JQG3kqImS6ICEMzHQHF1c,657
+ray_embedding-0.9.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+ray_embedding-0.9.1.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.9.1.dist-info/RECORD,,

{ray_embedding-0.9.0.dist-info → ray_embedding-0.9.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (77.0.3)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ray_embedding-0.9.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=E79J0bcVNXAFlFMVrjZTaSXLnrMZ6LvtRdD4d1DKu1w,1598
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=1sx5jXo61UgLLL8BtFBPflsbBay6J3yG1u2UMIPgAtk,3768
-ray_embedding-0.9.0.dist-info/METADATA,sha256=zfS00PtWu6mTcHloTDJDqEQ9bwbLJWUSgiv21degWNc,636
-ray_embedding-0.9.0.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-ray_embedding-0.9.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.9.0.dist-info/RECORD,,