ray-embedding 0.11.7__py3-none-any.whl → 0.11.8__py3-none-any.whl

ray_embedding/deploy.py

@@ -1,18 +1,19 @@
+import os
 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.dto import AppConfig, ModelDeploymentConfig, DeployedModel
 from ray_embedding.embedding_model import EmbeddingModel
 import torch
 
 from ray_embedding.model_router import ModelRouter
 
 
-def build_model(model_config: ModelDeploymentConfig) -> DeploymentHandle:
+def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
     deployment_name = model_config.deployment
     model = model_config.model
-    served_model_name = model_config.served_model_name
+    served_model_name = model_config.served_model_name or os.path.basename(model)
     device = model_config.device
     backend = model_config.backend or "torch"
     matryoshka_dim = model_config.matryoshka_dim
@@ -37,12 +38,17 @@ def build_model(model_config: ModelDeploymentConfig) -> DeploymentHandle:
                                                                    trust_remote_code=trust_remote_code,
                                                                    model_kwargs=model_kwargs
                                                                    )
-    return deployment
+    return DeployedModel(model=served_model_name,
+                         deployment_handle=deployment,
+                         batch_size=model_config.batch_size,
+                         num_retries=model_config.num_retries
+                         )
+
 
 def build_app(args: AppConfig) -> Application:
     model_router, models = args.model_router, args.models
     assert model_router and models
 
-    served_models = {model_config.served_model_name: build_model(model_config) for model_config in models}
-    router = ModelRouter.options(name=model_router.deployment).bind(served_models)
+    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)
     return router

ray_embedding/model_router.py

@@ -13,33 +13,32 @@ web_api = FastAPI(title="Ray Embeddings - OpenAI-compatible API")
 @serve.deployment
 @serve.ingress(web_api)
 class ModelRouter:
-    def __init__(self, served_models: Dict[str, DeployedModel]):
-        assert served_models, "models cannot be empty"
+    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.served_models = served_models
+        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.served_models.keys()
+             "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.served_models
-        self.logger.info(f"Retrieving model handle for: {model} ")
-        model_handle = self.served_models[model].deployment_handle
-        self.logger.info(f"Model handle: {model_handle} ")
-        batch_size = self.served_models[model].batch_size
-        num_retries = self.served_models[model].num_retries
+        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 is greater than {batch_size}. "
-                             f"It was resized to {len(batches)} mini-batches of size {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)
 
@@ -68,7 +67,7 @@ class ModelRouter:
     async def compute_embeddings(self, request: EmbeddingRequest):
         try:
             inputs = request.input if isinstance(request.input, list) else [request.input]
-            self.logger.info(f"Received input of size {len(inputs)} text chunks; model: {request.model}")
+            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}

{ray_embedding-0.11.7.dist-info → ray_embedding-0.11.8.dist-info}/METADATA

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

ray_embedding-0.11.8.dist-info/RECORD

@@ -0,0 +1,9 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=fHyxSGOcNTD_EFmPC2QYugvR77Ulr92bNU6Xk2Mwwf4,2649
+ray_embedding/dto.py,sha256=QlduDoqkFHaeF_KgsFeUKq2XWiPMmrgRPy_QjCTSCRE,1399
+ray_embedding/embedding_model.py,sha256=2wLk54BZIhHMCnwx5vneU0z4Y7EQs220BFNeLf-UQh4,3387
+ray_embedding/model_router.py,sha256=q7SgfOEyslbzKe-_gbIq02T2KLpUI2DmDivdmPf-_ts,3906
+ray_embedding-0.11.8.dist-info/METADATA,sha256=jRwJmh8p85_h7w8zuyLqYvS3nBE1f_YYzaR56rNOmXc,1094
+ray_embedding-0.11.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.11.8.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.11.8.dist-info/RECORD,,

ray_embedding-0.11.7.dist-info/RECORD

@@ -1,9 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=2lnHmSONHWyRQdZSXJrU8UGtXVQWgyFZGwUYdD7gzOM,2356
-ray_embedding/dto.py,sha256=QlduDoqkFHaeF_KgsFeUKq2XWiPMmrgRPy_QjCTSCRE,1399
-ray_embedding/embedding_model.py,sha256=2wLk54BZIhHMCnwx5vneU0z4Y7EQs220BFNeLf-UQh4,3387
-ray_embedding/model_router.py,sha256=QBf188zZXsCrYMtb06PWlRSB0DEiAK3UojSLhKMMEEM,3997
-ray_embedding-0.11.7.dist-info/METADATA,sha256=_b15ndqKOJUd7G9zgk31wEnYDASKE8D3QOSUpFpAyOU,1094
-ray_embedding-0.11.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.11.7.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.11.7.dist-info/RECORD,,