ray-embedding 0.10.12__py3-none-any.whl → 0.11.0__py3-none-any.whl

ray_embedding/__init__.py

@@ -1,2 +1,2 @@
-from ray_embedding.deploy import deploy_model
+from ray_embedding.deploy import build_app
 

ray_embedding/deploy.py

@@ -1,26 +1,22 @@
 from typing import Dict, Any, Optional
 from ray.serve import Application
+from ray.serve.handle import DeploymentHandle
+
+from ray_embedding.dto import AppConfig, ModelDeploymentConfig
 from ray_embedding.embedding_model import EmbeddingModel
 import torch
 
+from ray_embedding.embedding_service import EmbeddingService
 
-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
-
-    model: str = args.pop("model", "")
-    assert model
 
-    device: Optional[str] = args.pop("device", None)
-    backend: Optional[str] = args.pop("backend", "torch")
-    matryoshka_dim: Optional[int] = args.pop("matryoshka_dim", None)
-    trust_remote_code: Optional[bool] = args.pop("trust_remote_code", False)
-    model_kwargs: Dict[str, Any] = args.pop("model_kwargs", {})
+def build_model(model_config: ModelDeploymentConfig) -> DeploymentHandle:
+    deployment_name = model_config.deployment_name
+    model = model_config.model
+    device = model_config.device
+    backend = model_config.backend or "torch"
+    matryoshka_dim = model_config.matryoshka_dim
+    trust_remote_code = model_config.trust_remote_code or False
+    model_kwargs = model_config.model_kwargs or {}
     if "torch_dtype" in model_kwargs:
         torch_dtype = model_kwargs["torch_dtype"].strip()
         if torch_dtype == "float16":
@@ -39,4 +35,12 @@ def deploy_model(args: Dict[str, Any]) -> Application:
                                                                    trust_remote_code=trust_remote_code,
                                                                    model_kwargs=model_kwargs
                                                                    )
-    return deployment
+    return deployment
+
+def build_app(args: AppConfig) -> Application:
+    model_router, models = args.model_router, args.models
+    assert model_router and models
+
+    served_models = {model_config.model: build_model(model_config) for model_config in models}
+    app = EmbeddingService.options(name=model_router.deployment).bind(served_models)
+    return app

ray_embedding/dto.py

@@ -1,5 +1,7 @@
-from typing import Union, List, Optional
+import dataclasses
+from typing import Union, List, Optional, Dict, Any
 from pydantic import BaseModel
+from ray.serve.handle import DeploymentHandle
 
 
 class EmbeddingRequest(BaseModel):
@@ -13,4 +15,34 @@ class EmbeddingResponse(BaseModel):
     """Schema of embedding response (compatible with OpenAI)"""
     object: str
     data: List[dict]  # Embedding data including index and vector
-    model: str  # Model name used for embedding
+    model: str  # Model name used for embedding
+
+
+class ModelRouterConfig(BaseModel):
+    deployment: str
+
+
+class ModelDeploymentConfig(BaseModel):
+    deployment: str
+    model: str
+    served_model_name: str
+    batch_size: Optional[int] = 8
+    num_retries: Optional[int] = 2
+    device: Optional[str] = None
+    backend: Optional[str] = None
+    matryoshka_dim: Optional[int] = 768
+    trust_remote_code: Optional[bool] = False
+    model_kwargs: Optional[Dict[str, Any]] = {}
+
+
+class AppConfig(BaseModel):
+    model_router: ModelRouterConfig
+    models: List[ModelDeploymentConfig]
+
+
+@dataclasses.dataclass
+class DeployedModel:
+    model: str
+    deployment_handle: DeploymentHandle
+    batch_size: int
+    num_retries: Optional[int] = 2

ray_embedding/embedding_model.py

@@ -1,34 +1,18 @@
 import logging
 import os.path
 import time
-from typing import Optional, Dict, Any, List
+from typing import Optional, Dict, Any, List, Union
 
 import torch
-from fastapi import FastAPI, HTTPException
+from pynvml import nvmlInit, nvmlDeviceGetCount
 from ray import serve
 from sentence_transformers import SentenceTransformer
 
-from ray_embedding.dto import EmbeddingResponse, EmbeddingRequest
 
-web_api = FastAPI(title=f"Ray Embeddings - OpenAI-compatible API")
-
-
-@serve.deployment(
-    num_replicas="auto",
-    ray_actor_options={
-        "num_cpus": 1,
-        "num_gpus": 0
-    },
-    autoscaling_config={
-        "target_ongoing_requests": 2,
-        "min_replicas": 0,
-        "initial_replicas": 1,
-        "max_replicas": 1,
-    }
-)
-@serve.ingress(web_api)
+@serve.deployment
 class EmbeddingModel:
-    def __init__(self, model: str, device: Optional[str] = None, backend: Optional[str] = "torch",
+    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):
         logging.basicConfig(level=logging.INFO)
@@ -49,56 +33,29 @@ class EmbeddingModel:
                                                    trust_remote_code=self.trust_remote_code,
                                                    model_kwargs=self.model_kwargs)
 
-        self.served_model_name = os.path.basename(self.model)
-        self.available_models = [
-            {"id": self.served_model_name,
-             "object": "model",
-             "created": int(time.time()),
-             "owned_by": "openai",
-             "permission": []}
-        ]
+        self.served_model_name = served_model_name or os.path.basename(self.model)
         self.logger.info(f"Successfully initialized embedding model {self.model} using device {self.torch_device}")
 
-    @web_api.post("/v1/embeddings", response_model=EmbeddingResponse)
-    async def create_embeddings(self, request: EmbeddingRequest):
-        """Generate embeddings for the input text using the specified model."""
-        try:
-            assert request.model == self.served_model_name, (
-                f"Model '{request.model}' is not supported. Use '{self.served_model_name}' instead."
-            )
-            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, show_progress_bar=False,
-            ).to(self.torch_device)
+    async def __call__(self, text: Union[str, List[str]], dimensions: Optional[int] = None):
+        """Compute embeddings for the input text using the loaded model."""
+        if isinstance(text, str):
+            text = [text]
+        truncate_dim = dimensions or self.matryoshka_dim
 
-            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)
+        # Compute embeddings and convert to a PyTorch tensor on the GPU
+        embeddings = self.embedding_model.encode(
+            text, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
+        ).to(self.torch_device)
 
-            # Move all embeddings to CPU at once before conversion
-            embeddings = embeddings.cpu().tolist()
+        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)
 
-            # Convert embeddings to list format for response
-            response_data = [
-                {"index": idx, "embedding": emb}
-                for idx, emb in enumerate(embeddings)
-            ]
-            return EmbeddingResponse(object="list", data=response_data, model=request.model)
+        # Move all embeddings to CPU at once before conversion
+        embeddings = embeddings.cpu().tolist()
+        return embeddings
 
-        except Exception as e:
-            self.logger.error(e)
-            raise HTTPException(status_code=500, detail=str(e))
-
-    @web_api.get("/v1/models")
-    async def list_models(self):
-        """Returns the list of available models in OpenAI-compatible format."""
-        return {"object": "list", "data": self.available_models}
 
     def wait_for_cuda(self, wait: int = 10):
         if self.init_device == "cuda" and not torch.cuda.is_available():
@@ -106,8 +63,12 @@ class EmbeddingModel:
         self.check_health()
 
     def check_health(self):
-        if self.init_device == "cuda" and not torch.cuda.is_available():
+        if self.init_device == "cuda":
             # Even though CUDA was available at init time,
             # CUDA can become unavailable - this is a known problem in AWS EC2
             # https://github.com/ray-project/ray/issues/49594
-            raise RuntimeError("CUDA device is not available")
+            try:
+                nvmlInit()
+                assert nvmlDeviceGetCount() >= 1
+            except:
+                raise RuntimeError("CUDA device is not available")

ray_embedding/embedding_service.py

@@ -0,0 +1,77 @@
+import asyncio
+import logging
+import time
+from typing import Optional, Dict, List
+
+from fastapi import FastAPI, HTTPException
+from ray import serve
+
+from ray_embedding.dto import EmbeddingResponse, EmbeddingRequest, DeployedModel
+
+web_api = FastAPI(title="Ray Embeddings - OpenAI-compatible API")
+
+@serve.deployment
+@serve.ingress(web_api)
+class EmbeddingService:
+    def __init__(self, served_models: Dict[str, DeployedModel]):
+        self.logger = logging.getLogger(self.__class__.__name__)
+        assert served_models, "models cannot be empty"
+        self.served_models = served_models
+        self.available_models = [
+            {"id": str(item),
+             "object": "model",
+             "created": int(time.time()),
+             "owned_by": "openai",
+             "permission": []} for item in self.served_models.keys()
+        ]
+        self.logger.info(f"Successfully registered models: {self.available_models}")
+
+    async def _compute_embeddings_from_resized_batches(self, model: str, inputs: List[str], dimensions: Optional[int] = None):
+        assert model in self.served_models
+        model_handle = self.served_models[model].deployment_handle
+        batch_size = self.served_models[model].batch_size
+        num_retries = self.served_models[model].num_retries
+
+        # Resize the inputs into batch_size items, and dispatch in parallel
+        batches = [inputs[i:i+batch_size] for i in range(0, len(inputs), batch_size)]
+        tasks = [model_handle.remote(batch, dimensions) for batch in batches]
+        all_results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Retry any failed model calls
+        for i, result in enumerate(all_results):
+            if isinstance(result, Exception):
+                retries = 0
+                while retries < num_retries:
+                    try:
+                        all_results[i] = await model_handle.remote(batches[i], dimensions)
+                    except Exception as e:
+                        self.logger.warning(e)
+                    finally:
+                        retries += 1
+                    if not isinstance(all_results[i], Exception):
+                        break
+
+                if retries >= num_retries and isinstance(all_results[i], Exception):
+                    raise all_results[i]
+
+        # Flatten the results because all_results is a list of lists
+        return [emb for result in all_results for emb in result]
+
+    @web_api.post("/v1/embeddings", response_model=EmbeddingResponse)
+    async def compute_embeddings(self, request: EmbeddingRequest):
+        try:
+            inputs = request.input if isinstance(request.input, list) else [request.input]
+            embeddings = await self._compute_embeddings_from_resized_batches(request.model, inputs, request.dimensions)
+            response_data = [
+                {"index": idx, "embedding": emb}
+                for idx, emb in enumerate(embeddings)
+            ]
+            return EmbeddingResponse(object="list", data=response_data, model=request.model)
+        except Exception as e:
+            self.logger.error(f"Failed to create embeddings: {e}")
+            raise HTTPException(status_code=500, detail=str(e))
+
+    @web_api.get("/v1/models")
+    async def list_models(self):
+        """Returns the list of available models in OpenAI-compatible format."""
+        return {"object": "list", "data": self.available_models}

{ray_embedding-0.10.12.dist-info → ray_embedding-0.11.0.dist-info}/METADATA

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

ray_embedding-0.11.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=rqVJ8GVh3hGyjojJ9wzHosiPe6gRNVv5-8zkENKKqMc,2195
+ray_embedding/dto.py,sha256=QlduDoqkFHaeF_KgsFeUKq2XWiPMmrgRPy_QjCTSCRE,1399
+ray_embedding/embedding_model.py,sha256=d_gqqcKK3B2nF8qMDk7NKW8gHBbtPwvPFArKTBKNdFA,3372
+ray_embedding/embedding_service.py,sha256=HwIaVQFegKYF3sp6ST5G2vf0oW0Sl3W_SDp6eq79Tng,3412
+ray_embedding-0.11.0.dist-info/METADATA,sha256=VdFyMJDMJcaWinYq7hAhztPkCzjGtSmtD7UNibfdlPY,1094
+ray_embedding-0.11.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.11.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.11.0.dist-info/RECORD,,

{ray_embedding-0.10.12.dist-info → ray_embedding-0.11.0.dist-info}/WHEEL

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

ray_embedding-0.10.12.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=ZGxcG4589WcRtaM6H84YJarw0m1XqHNgfOf3PLAhM5M,1995
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=JfNt0rJYXGlbNjg5xmY14k4jQmNZKHnTQZnTi5SbNSc,4829
-ray_embedding-0.10.12.dist-info/METADATA,sha256=7EjNTy3P2zGYP6JzjrigmpbAujjaroJA-e1slz_Ra1E,1095
-ray_embedding-0.10.12.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-ray_embedding-0.10.12.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.10.12.dist-info/RECORD,,