ray-embedding 0.13.17__py3-none-any.whl → 0.14.5__py3-none-any.whl

ray_embedding/deploy.py

@@ -1,13 +1,13 @@
 import os
+from typing import Any, Dict
 
 import torch
 from ray.serve import Application
 
 from ray_embedding.dto import AppConfig, ModelDeploymentConfig, DeployedModel, NodeReaperConfig
 from ray_embedding.embedding_model import EmbeddingModel
-from ray_embedding.node_reaper import NodeReaper, NODE_REAPER_DEPLOYMENT_NAME
-from ray_embedding.utils import node_affinity_for_head, node_affinity_for_worker
 from ray_embedding.model_router import ModelRouter
+from ray_embedding.node_reaper import NodeReaper, NODE_REAPER_DEPLOYMENT_NAME
 
 
 def build_model(model_config: ModelDeploymentConfig, node_reaper):
@@ -56,7 +56,7 @@ def build_app(args: AppConfig) -> Application:
 
     node_reaper_config = args.node_reaper or NodeReaperConfig()
 
-    node_reaper_kwargs = {
+    node_reaper_kwargs: Dict[str, Any] = {
         "ssh_user": node_reaper_config.ssh_user,
         "ssh_private_key": node_reaper_config.ssh_private_key,
     }
@@ -67,13 +67,20 @@ def build_app(args: AppConfig) -> Application:
 
     node_reaper = NodeReaper.options( 
         name=NODE_REAPER_DEPLOYMENT_NAME,
-        ray_actor_options={"num_cpus": 0.25, f"node:{HEAD_NODE_IP}": 1},
+        ray_actor_options={"num_cpus": 0.25, "resources": {"head_node": 1}},
         autoscaling_config={"initial_replicas": 1, "min_replicas": 1, "max_replicas": 1}
     ).bind(**node_reaper_kwargs)
-    
+
     deployed_models = {model_config.served_model_name: build_model(model_config, node_reaper) for model_config in models}
+    model_router_kwargs = {
+        "deployed_models": deployed_models,
+        "path_prefix": model_router.path_prefix,
+        "max_concurrency": model_router.max_concurrency,
+        "node_reaper": node_reaper
+    }
     router = ModelRouter.options(
-        name=model_router.deployment
-    ).bind(deployed_models, model_router.path_prefix, node_reaper)
+        name=model_router.deployment,
+        ray_actor_options={"num_cpus": 0.25, "resources": {"worker_node": 1}}
+    ).bind(**model_router_kwargs)
     
     return router

ray_embedding/embedding_model.py

@@ -25,16 +25,17 @@ class EmbeddingModel:
         self.served_model_name = served_model_name or os.path.basename(self.model)
         self.init_device = device
         self.cuda_memory_flush_threshold = cuda_memory_flush_threshold
-        if self.init_device is None or self.init_device == "auto":
-            self.init_device = "cuda" if torch.cuda.is_available() else "cpu"
-        if self.init_device == "cuda":
-            self.wait_for_cuda()
+        self.node_reaper = node_reaper
         self.torch_device = torch.device(self.init_device)
         self.backend = backend or "torch"
         self.matryoshka_dim = matryoshka_dim
         self.trust_remote_code = trust_remote_code or False
         self.model_kwargs = model_kwargs or {}
-        self.node_reaper = node_reaper
+
+        if self.init_device is None or self.init_device == "auto":
+            self.init_device = "cuda" if torch.cuda.is_available() else "cpu"
+        if self.init_device == "cuda":
+            self.wait_for_cuda()
 
         self.logger.info(f"Initializing embedding model: {self.model}")
         self.embedding_model = SentenceTransformer(self.model, device=self.init_device, backend=self.backend,

ray_embedding/model_router.py

@@ -4,7 +4,6 @@ import time
 from typing import Optional, Dict, List, Tuple
 
 from fastapi import FastAPI, HTTPException
-import ray
 from ray import serve
 from ray.serve.handle import DeploymentHandle
 
@@ -116,12 +115,12 @@ class ModelRouter:
             raise HTTPException(status_code=400, detail=f"The API path prefix specified is invalid: '{path_prefix}'")
         return {"object": "list", "data": self.available_models}
 
-    def check_health(self):
+    async def check_health(self):
         if not self.node_reaper:
             return
 
         try:
-            unhealthy_node_ips = ray.get(self.node_reaper.get_unhealthy_node_ips.remote())
+            unhealthy_node_ips = await self.node_reaper.get_unhealthy_node_ips.remote()
         except Exception as exc:
             self.logger.warning(f"Unable to fetch node reaper data: {exc}")
             return

ray_embedding/node_reaper.py

@@ -5,17 +5,12 @@ from pathlib import Path
 from typing import Dict, Any, List, Optional, Set
 
 from ray import serve
-from ray_embedding.utils import HEAD_NODE_IP
 
 
 NODE_REAPER_DEPLOYMENT_NAME = "NodeReaper"
 
 
-@serve.deployment(
-    name=NODE_REAPER_DEPLOYMENT_NAME,
-    ray_actor_options={"num_cpus": 0.25, f"node:{HEAD_NODE_IP}": 1},
-    autoscaling_config={"initial_replicas": 1, "min_replicas": 1, "max_replicas": 1},
-)
+@serve.deployment
 class NodeReaper:
     def __init__(
         self,

ray_embedding/utils.py

@@ -1,15 +1,12 @@
 from typing import Optional, Tuple
 
-import ray
 from ray import serve
 from ray.serve.handle import DeploymentHandle
 from ray.util import get_node_ip_address, state
-from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, NotIn
 
 from ray_embedding.node_reaper import NODE_REAPER_DEPLOYMENT_NAME
 
 
-
 def get_head_node_id() -> Tuple[str, str]:
     try:
         nodes = state.list_nodes(filters=[("is_head_node", "=", True)])
@@ -21,17 +18,6 @@ def get_head_node_id() -> Tuple[str, str]:
         raise RuntimeError("Unable to locate the head node ID for NodeReaper deployment.") from exc
 
 
-HEAD_NODE_ID, HEAD_NODE_IP = get_head_node_id()
-
-
-def node_affinity_for_head() -> NodeAffinitySchedulingStrategy:
-    return NodeAffinitySchedulingStrategy(node_id=HEAD_NODE_ID, soft=False)
-
-
-def node_affinity_for_worker() -> NodeAffinitySchedulingStrategy:
-    return NodeAffinitySchedulingStrategy(node_id=NotIn(HEAD_NODE_ID), soft=False)
-
-
 def get_node_reaper_handle() -> DeploymentHandle:
     try:
         return serve.context.get_deployment_handle(NODE_REAPER_DEPLOYMENT_NAME)

{ray_embedding-0.13.17.dist-info → ray_embedding-0.14.5.dist-info}/METADATA

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

ray_embedding-0.14.5.dist-info/RECORD

@@ -0,0 +1,11 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=Q3J-zVVqGkJ0CiIL_cevDrPWhOzLf5i-5Oz2wjlPR48,4111
+ray_embedding/dto.py,sha256=6JuAcD6pLfzUL48HfyPnZI7Hb-o66KFM5UtYZOOgwc8,1739
+ray_embedding/embedding_model.py,sha256=_X4W4uozmcPgfFhSGaP2UDVMLQVoXlZjcFpNqFoIRxg,4906
+ray_embedding/model_router.py,sha256=0CKlT4AvnzY-7OjzGF8L5-HA4vvmdJ9MUfHmKdWp1Pw,6650
+ray_embedding/node_reaper.py,sha256=ISwSHnQs22B_f3PihND3KYTLkJSDbg1JWIAaKS-qCm0,4800
+ray_embedding/utils.py,sha256=FAV1xQWa9GeAIm5w2HrP-pgFZBkB0LYUcX6ly6zlYg0,1805
+ray_embedding-0.14.5.dist-info/METADATA,sha256=QBBj7TFoIqZB35c3hPsSKS7FnUzq6YTHujFP9zdADAw,1094
+ray_embedding-0.14.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.14.5.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.14.5.dist-info/RECORD,,

ray_embedding-0.13.17.dist-info/RECORD

@@ -1,11 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=xFPRUQVYkjBp9WbOV8Ag-l8I4Hjq7iZijXZzDl09GK0,3882
-ray_embedding/dto.py,sha256=6JuAcD6pLfzUL48HfyPnZI7Hb-o66KFM5UtYZOOgwc8,1739
-ray_embedding/embedding_model.py,sha256=Zr5lxVuy60y8-JgsOmKDD44FZlbTL1tiiY-3_72sTR4,4905
-ray_embedding/model_router.py,sha256=W2c0hvqwDe1iCfNx4ee2UT7wKduywMP8dY0Ggb8xBvU,6658
-ray_embedding/node_reaper.py,sha256=MBTf20RmXvW90w3ufwNCIiakWeVOhfXZey2GCnE816Q,5041
-ray_embedding/utils.py,sha256=NcxUbZVonN0FZIvRdJt3nuyKT7JSPM4x6i2-9Ncv3h4,2241
-ray_embedding-0.13.17.dist-info/METADATA,sha256=aoLmk2j7H9iDCSLy6eOZGr21guKA47Rbq4o0UmmouZA,1095
-ray_embedding-0.13.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.13.17.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.13.17.dist-info/RECORD,,