ray-embedding 0.13.15__py3-none-any.whl → 0.14.0__py3-none-any.whl

ray_embedding/deploy.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, Dict
 
 import torch
 from ray.serve import Application
@@ -6,7 +7,7 @@ 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.utils import get_head_node_id
 from ray_embedding.model_router import ModelRouter
 
 
@@ -56,7 +57,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,
     }
@@ -65,15 +66,22 @@ def build_app(args: AppConfig) -> Application:
     if node_reaper_config.reap_interval_seconds is not None:
         node_reaper_kwargs["reap_interval_seconds"] = node_reaper_config.reap_interval_seconds
 
-    node_reaper = NodeReaper.options(
+    node_reaper = NodeReaper.options( 
         name=NODE_REAPER_DEPLOYMENT_NAME,
-        scheduling_strategy=node_affinity_for_head()
+        ray_actor_options={"num_cpus": 0.25, "resources": {"node_type:head": 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,
-        scheduling_strategy=node_affinity_for_worker()
-    ).bind(deployed_models, model_router.path_prefix, node_reaper)
+        ray_actor_options={"num_cpus": 0.25, "resources": {"node_type:worker": 1}}
+    ).bind(**model_router_kwargs)
     
     return router

ray_embedding/node_reaper.py

@@ -10,11 +10,7 @@ from ray import serve
 NODE_REAPER_DEPLOYMENT_NAME = "NodeReaper"
 
 
-@serve.deployment(
-    name=NODE_REAPER_DEPLOYMENT_NAME,
-    ray_actor_options={"num_cpus": 0.25},
-    autoscaling_config={"initial_replicas": 1, "min_replicas": 1, "max_replicas": 1},
-)
+@serve.deployment
 class NodeReaper:
     def __init__(
         self,

ray_embedding/utils.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Tuple
 
 import ray
 from ray import serve
@@ -10,28 +10,17 @@ from ray_embedding.node_reaper import NODE_REAPER_DEPLOYMENT_NAME
 
 
 
-def get_head_node_id() -> str:
+def get_head_node_id() -> Tuple[str, str]:
     try:
         nodes = state.list_nodes(filters=[("is_head_node", "=", True)])
         if not nodes:
             raise RuntimeError("Unable to locate head node for NodeReaper deployment.")
         head_node = nodes[0]
-        return head_node["node_id"]
+        return head_node["node_id"], head_node["node_ip"]
     except Exception as exc:
         raise RuntimeError("Unable to locate the head node ID for NodeReaper deployment.") from exc
 
 
-HEAD_NODE_ID = 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.15.dist-info → ray_embedding-0.14.0.dist-info}/METADATA

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

ray_embedding-0.14.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=NYpGDGF8y1rh3Thts-NC4nb8anXQJDC0dFZC18_R2f8,4170
+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=ISwSHnQs22B_f3PihND3KYTLkJSDbg1JWIAaKS-qCm0,4800
+ray_embedding/utils.py,sha256=cbdI7q6xSvbl31ZthdM8mz55VrN8pubkoD6RqKGYLUc,1898
+ray_embedding-0.14.0.dist-info/METADATA,sha256=uJlMttPN4bYVZbvLi5g37dCQRyzEzHrj7eqK1fYiv_w,1094
+ray_embedding-0.14.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.14.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.14.0.dist-info/RECORD,,

ray_embedding-0.13.15.dist-info/RECORD

@@ -1,11 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=QWbDjApgVWaMy_uWOC67zSuwrpj44vZx1ws1ersq-lw,3824
-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=0lQvq94pdqZE21JeahIyc3LXA4xKtSEv8oELJ1y25QI,4969
-ray_embedding/utils.py,sha256=IE0uyBRIW9HVQjO1I46qvtxGQfdorPwU4FsAlYHjI9g,2186
-ray_embedding-0.13.15.dist-info/METADATA,sha256=je1cLWUJfUQ8skQDOTzktBgFpvot7nz8Dsy01gDo2mM,1095
-ray_embedding-0.13.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.13.15.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.13.15.dist-info/RECORD,,