ray-embedding 0.11.8__tar.gz → 0.12.1__tar.gz

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/PKG-INFO

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

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/ray_embedding/deploy.py

@@ -1,12 +1,10 @@
 import os
-from typing import Dict, Any, Optional
+
+import torch
 from ray.serve import Application
-from ray.serve.handle import DeploymentHandle
 
 from ray_embedding.dto import AppConfig, ModelDeploymentConfig, DeployedModel
 from ray_embedding.embedding_model import EmbeddingModel
-import torch
-
 from ray_embedding.model_router import ModelRouter
 
 
@@ -19,6 +17,7 @@ def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
     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":
@@ -48,7 +47,8 @@ def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
 def build_app(args: AppConfig) -> Application:
     model_router, models = args.model_router, args.models
     assert model_router and models
+    assert model_router.path_prefix
 
     deployed_models = {model_config.served_model_name: build_model(model_config) for model_config in models}
-    router = ModelRouter.options(name=model_router.deployment).bind(deployed_models)
+    router = ModelRouter.options(name=model_router.deployment).bind(deployed_models, model_router.path_prefix)
     return router

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/ray_embedding/dto.py

@@ -20,6 +20,8 @@ class EmbeddingResponse(BaseModel):
 
 class ModelRouterConfig(BaseModel):
     deployment: str
+    path_prefix: List[str] = []
+    max_concurrency: int = 32
 
 
 class ModelDeploymentConfig(BaseModel):

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/ray_embedding/embedding_model.py

@@ -34,15 +34,17 @@ class EmbeddingModel:
                                                    trust_remote_code=self.trust_remote_code,
                                                    model_kwargs=self.model_kwargs)
 
-        self.logger.info(f"Successfully initialized embedding model {self.model} using device {self.torch_device}")
+        self.logger.info(f"Successfully initialized model {self.model} using device {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]
+    async def __call__(self, text: Union[str, List[str]], dimensions: Optional[int] = None) -> List[List[float]]:
+        """Compute embeddings for the input text using the current model."""
+        if not text or (isinstance(text, list) and not all(text)):
+            raise ValueError("Input text is empty or invalid")
+
+        text = [text] if isinstance(text, str) else text
         truncate_dim = dimensions or self.matryoshka_dim
 
-        # Compute embeddings and convert to a PyTorch tensor on the GPU
+        # Compute embeddings in PyTorch format
         embeddings = self.embedding_model.encode(
             text, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
         ).to(self.torch_device)
@@ -53,9 +55,8 @@ class EmbeddingModel:
             embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
 
         # Move all embeddings to CPU at once before conversion
-        embeddings = embeddings.cpu().tolist()
-        return embeddings
-
+        embeddings_list = embeddings.cpu().tolist()
+        return embeddings_list
 
     def wait_for_cuda(self, wait: int = 10):
         if self.init_device == "cuda" and not torch.cuda.is_available():
@@ -65,7 +66,7 @@ class EmbeddingModel:
     def check_health(self):
         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
+            # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
             # https://github.com/ray-project/ray/issues/49594
             try:
                 nvmlInit()

ray_embedding-0.12.1/ray_embedding/model_router.py

@@ -0,0 +1,111 @@
+import asyncio
+import logging
+import time
+from typing import Optional, Dict, List, Tuple
+
+from fastapi import FastAPI, HTTPException
+from ray import serve
+from ray.serve.handle import DeploymentHandle
+
+from ray_embedding.dto import DeployedModel, EmbeddingRequest, EmbeddingResponse
+
+web_api = FastAPI(title="Ray Embeddings - OpenAI-compatible API")
+
+@serve.deployment
+@serve.ingress(web_api)
+class ModelRouter:
+    def __init__(self, deployed_models: Dict[str, DeployedModel], path_prefix: List[str], max_concurrency: Optional[int] = 32):
+        assert deployed_models, "models cannot be empty"
+        assert path_prefix, "path_prefix cannot be empty"
+
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.deployed_models = deployed_models
+        self.path_prefix = [item.removeprefix("/").removesuffix("/") for item in path_prefix]
+        self.max_concurrency = max_concurrency
+        self.rate_limiter = asyncio.Semaphore(self.max_concurrency)
+        self.available_models = [
+            {"id": str(item),
+             "object": "model",
+             "created": int(time.time()),
+             "owned_by": "openai",
+             "permission": []} for item in self.deployed_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):
+        deployed_model = self.deployed_models[model]
+        model_handle = deployed_model.deployment_handle
+        batch_size = deployed_model.batch_size
+        num_retries = deployed_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)]
+        if len(inputs) > batch_size:
+            self.logger.info(f"Original input (length {len(inputs)} was resized "
+                             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]
+        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):
+                self.logger.warning(f"Retrying mini-batch {i} due to exception: {result}")
+                result_retried, retries = await self._retry_failed_embedding_call(model_handle, batches[i], dimensions,
+                                                                                  num_retries)
+                if retries >= num_retries and (isinstance(result_retried, Exception) or result_retried is None):
+                    raise result_retried or ValueError(f"Failed to compute `{model}` embeddings for mini-batch {i} after {num_retries} retries.")
+
+                all_results[i] = result_retried
+
+        # Flatten the results because `all_results` is a list of lists
+        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):
+        with self.rate_limiter:
+            return await model_handle.remote(batch, dimensions)
+
+    async def _retry_failed_embedding_call(self,  model_handle: DeploymentHandle, batch: List[str],
+                                           dimensions: Optional[int] = None, num_retries: Optional[int] = 2) \
+            -> Tuple[List[List[float]] | Exception, int]:
+
+        result_retried, retries = None, 0
+        while retries < num_retries:
+            try:
+                result_retried = await model_handle.remote(batch, dimensions)
+            except Exception as e:
+                result_retried = e
+                self.logger.warning(e)
+            finally:
+                retries += 1
+            if not isinstance(result_retried, Exception) and result_retried is not None:
+                break
+
+        return result_retried, retries
+
+    @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:
+            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)
+            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("/{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}"
+        return {"object": "list", "data": self.available_models}

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/ray_embedding.egg-info/PKG-INFO

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

{ray_embedding-0.11.8 → ray_embedding-0.12.1}/setup.cfg

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

ray_embedding-0.11.8/ray_embedding/model_router.py

@@ -1,84 +0,0 @@
-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 ModelRouter:
-    def __init__(self, deployed_models: Dict[str, DeployedModel]):
-        assert deployed_models, "models cannot be empty"
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(self.__class__.__name__)
-        self.deployed_models = deployed_models
-        self.available_models = [
-            {"id": str(item),
-             "object": "model",
-             "created": int(time.time()),
-             "owned_by": "openai",
-             "permission": []} for item in self.deployed_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.deployed_models
-        deployed_model = self.deployed_models[model]
-        model_handle = deployed_model.deployment_handle
-        batch_size = deployed_model.batch_size
-        num_retries = deployed_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)]
-        if len(inputs) > batch_size:
-            self.logger.info(f"Original input (length {len(inputs)} was resized "
-                             f"to {len(batches)} mini-batches of length {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
-        self.logger.info(f"Successfully computed embeddings from {len(batches)} mini-batches")
-        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]
-            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)
-            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}