ray-embedding 0.12.9__tar.gz → 0.12.11__tar.gz

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/PKG-INFO

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

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/ray_embedding/deploy.py

@@ -58,7 +58,9 @@ def build_app(args: AppConfig) -> Application:
     assert model_router.path_prefix
 
     node_health_check_interval_s = args.node_health_check_interval_s or DEFAULT_NODE_HEALTH_CHECK_INTERVAL_S
-    node_health_tracker = NodeHealthTracker.options(health_check_period_s=node_health_check_interval_s).bind()
+    tracked_model_deployments = [model_config.deployment for model_config in models]
+    node_health_tracker = (NodeHealthTracker.options(health_check_period_s=node_health_check_interval_s)
+                           .bind(tracked_model_deployments=tracked_model_deployments))
     deployed_models = {model_config.served_model_name: build_model(model_config, node_health_tracker=node_health_tracker)
                        for model_config in models}
     router = (ModelRouter.options(name=model_router.deployment)

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/ray_embedding/dto.py

@@ -38,13 +38,6 @@ class ModelDeploymentConfig(BaseModel):
     deployment: str
 
 
-class ReplicaState(BaseModel):
-    deployment_name: str
-    replica_actor_name: str
-    node_ip: str
-    state: Optional[str] = None
-
-
 class AppConfig(BaseModel):
     model_router: ModelRouterConfig
     node_health_check_interval_s: Optional[int] = 30

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/ray_embedding/embedding_model.py

@@ -3,6 +3,7 @@ import os.path
 import time
 from typing import Optional, Dict, Any, List, Union
 
+import ray
 import torch
 from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
 from ray import serve
@@ -76,16 +77,11 @@ class EmbeddingModel:
     def wait_for_cuda(self, wait: int = 10):
         if self.init_device == "cuda" and not torch.cuda.is_available():
             time.sleep(wait)
-        self.check_health()
+        self.check_cuda()
 
-    async def check_health(self):
+    def check_cuda(self) -> Any:
         if self.init_device != "cuda":
-            return
-
-        if self.node_health_tracker:
-            if await self.node_health_tracker.is_bad_node.remote(self.node_ip):
-                raise RuntimeError(f"The node {self.node_ip} is marked bad.")
-
+            return None
         try:
             # Even though CUDA was available at init time,
             # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
@@ -96,12 +92,22 @@ class EmbeddingModel:
 
             # replicas only have access to GPU 0
             handle = nvmlDeviceGetHandleByIndex(0)
-            mem_info = nvmlDeviceGetMemoryInfo(handle)
+            return handle
         except Exception as e:
-            await self.node_health_tracker.report_bad_node.remote(self.node_ip, self.deployment_name, self.replica_actor_name)
-            raise RuntimeError(f"CUDA health check failed for deployment: "
-                               f"{self.deployment_name}, replica: {self.replica_actor_name}, node: {self.node_ip}.\n{e}")
+            error_msg = f"CUDA health check failed for deployment: " \
+                        f"{self.deployment_name}, replica: {self.replica_actor_name}, node: {self.node_ip}.\n{e}"
+            self.logger.error(error_msg)
+            if self.node_health_tracker:
+                self.node_health_tracker.report_bad_gpu_node.remote(self.node_ip, self.deployment_name, self.replica_actor_name)
+            raise RuntimeError(error_msg)
+
+    async def check_health(self):
+        if self.node_health_tracker:
+            if await self.node_health_tracker.is_bad_gpu_node.remote(self.node_ip):
+                raise RuntimeError(f"The node {self.node_ip} is marked bad.")
 
+        handle = self.check_cuda()  # Raises an exception if CUDA is unavailable
+        mem_info = nvmlDeviceGetMemoryInfo(handle)
         reserved = torch.cuda.memory_reserved()  # bytes currently reserved by CUDA cache
         threshold_bytes = self.cuda_memory_flush_threshold * mem_info.total
 

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/ray_embedding/model_router.py

@@ -124,5 +124,5 @@ class ModelRouter:
 
     async def check_health(self):
         if self.node_health_tracker:
-            if await self.node_health_tracker.is_bad_node.remote(self.node_ip):
-                raise RuntimeError(f"The node {self.node_ip} is marked bad.")
+            if await self.node_health_tracker.is_bad_gpu_or_no_tracked_model_on_node.remote(self.node_ip):
+                raise RuntimeError(f"The node {self.node_ip} is marked bad, or there are no tracked models running on it.")

ray_embedding-0.12.11/ray_embedding/node_health.py

@@ -0,0 +1,90 @@
+import logging
+import threading
+from typing import Set, List
+
+import ray
+from networkx.algorithms.boundary import node_boundary
+from networkx.algorithms.cuts import node_expansion
+from ray import serve
+from ray._private.services import get_node_ip_address
+from ray.util.state import list_actors
+
+
+@serve.deployment(autoscaling_config=dict(min_replicas=1, max_replicas=1),
+                  ray_actor_options=dict(num_cpus=0.1))
+class NodeHealthTracker:
+    """Maintains a list of bad nodes, as reported by replicas that call the report_bad_node func.
+    Bad nodes are those that fail GPU/CUDA health check.
+    What's the purpose? Because when an embedding model replica becomes unhealthy
+    (due to GPU/CUDA issues), we want Ray to kill all replicas running on the node.
+    When Ray detects that there are no running replicas on a node, the node is stopped
+    and replaced with a new one.
+    """
+    def __init__(self, tracked_model_deployments: List[str] = None):
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.tracked_model_deployments = tracked_model_deployments or []
+        self.bad_gpu_node_ips: Set[str] = set()
+        self.lock = threading.RLock()
+        replica_context = serve.get_replica_context()
+        self.deployment_name = replica_context.deployment
+        self.replica_actor_name = replica_context.replica_id.to_full_id_str()
+        self.node_ip = get_node_ip_address()
+        self.logger.info(f"Successfully initialized NodeHealthTracker. Tracked model deployments: {self.tracked_model_deployments}")
+
+    async def report_bad_gpu_node(self, node_ip: str, deployment_name: str, replica_actor_name: str):
+        with self.lock:
+            if node_ip not in self.bad_gpu_node_ips:
+                self.bad_gpu_node_ips.add(node_ip)
+                self.logger.warning(
+                    f"[Bad GPU node reported] Deployment: {deployment_name}, Replica: {replica_actor_name}, Node IP: {node_ip}"
+                )
+
+    async def is_bad_gpu_node(self, node_ip: str) -> bool:
+        with self.lock:
+            return node_ip in self.bad_gpu_node_ips
+
+    async def is_bad_gpu_or_no_tracked_model_on_node(self, node_ip: str):
+        return (await self.is_bad_gpu_node(node_ip) or
+                not await self.is_tracked_model_running_on_node(node_ip))
+
+    async def check_health(self):
+        """Called periodically by Ray Serve. Used here to clean up stale node IDs."""
+        try:
+            current_node_ips = {node["NodeManagerAddress"] for node in ray.nodes() if node["Alive"]}
+            with self.lock:
+                stale_nodes = self.bad_gpu_node_ips - current_node_ips
+                if stale_nodes:
+                    self.logger.info(f"Removing stale bad node_ips: {stale_nodes}")
+                self.bad_gpu_node_ips.intersection_update(current_node_ips)
+            self.logger.info(f"Current nodes: {current_node_ips}. Bad GPU nodes: {self.bad_gpu_node_ips}.")
+        except Exception as e:
+            raise RuntimeError(f"An error occurred in check_health during bad node cleanup: {e}")
+
+        if not await self.is_tracked_model_running_on_node(node_ip=self.node_ip):
+            raise RuntimeError(f"There are no model replicas running on node {self.node_ip}.")
+
+
+    async def is_tracked_model_running_on_node(self, node_ip: str) -> bool:
+        """
+        Return True if there is at least one replica of any of the self.tracked_model_deployments
+        running on the specified node_ip.
+        """
+        try:
+            target_node_id = next(node["NodeID"] for node in ray.nodes() if node["Alive"] and node["NodeManagerAddress"] == node_ip)
+            assert target_node_id, f"No node found with IP {node_ip}"
+            prefixes = tuple(f"SERVE_REPLICA::{d}" for d in self.tracked_model_deployments)
+
+            for actor in list_actors(detail=False):
+                if (actor.state in ["DEPENDENCIES_UNREADY", 'PENDING_CREATION', 'ALIVE', 'RESTARTING'] and
+                        actor.node_id == target_node_id and
+                            actor.name.startswith(prefixes)):
+                    self.logger.info(f"There is at least one replica of tracked_deployments={self.tracked_model_deployments} "
+                        f"running on node {node_ip}")
+                    return True
+
+            self.logger.info(f"No replicas of tracked deployments={self.tracked_model_deployments} running on node: {node_ip}.")
+            return False
+        except Exception as e:
+            self.logger.error(f"An error occurred while checking replicas on node {node_ip}: {e}")
+            return False

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/ray_embedding.egg-info/PKG-INFO

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

{ray_embedding-0.12.9 → ray_embedding-0.12.11}/setup.cfg

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

ray_embedding-0.12.9/ray_embedding/node_health.py

@@ -1,48 +0,0 @@
-import logging
-import threading
-from typing import Set
-
-import ray
-from ray import serve
-
-
-@serve.deployment(autoscaling_config=dict(min_replicas=1, max_replicas=1))
-class NodeHealthTracker:
-    """Maintains a list of bad nodes, as reported by replicas that call the report_bad_node func.
-    Bad nodes are those that fail GPU/CUDA health check.
-    What's the purpose? Because when an embedding model replica becomes unhealthy
-    (due to GPU/CUDA issues), we want Ray to kill all replicas running on the node.
-    When Ray detects that there are no running replicas on a node, the node is stopped
-    and replaced with a new one.
-    """
-    def __init__(self):
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(self.__class__.__name__)
-        self.bad_node_ips: Set[str] = set()
-        self.lock = threading.RLock()
-        self.logger.info(f"Successfully initialized NodeHealthTracker.")
-
-    async def report_bad_node(self, node_ip: str, deployment_name: str, replica_actor_name: str):
-        with self.lock:
-            if node_ip not in self.bad_node_ips:
-                self.bad_node_ips.add(node_ip)
-                self.logger.warning(
-                    f"[Bad Node Reported] Deployment: {deployment_name}, Replica: {replica_actor_name}, Node IP: {node_ip}"
-                )
-
-    async def is_bad_node(self, node_ip: str) -> bool:
-        with self.lock:
-            return node_ip in self.bad_node_ips
-
-    async def check_health(self):
-        """Called periodically by Ray Serve. Used here to clean up stale node IDs."""
-        try:
-            current_node_ips = {node["NodeManagerAddress"] for node in ray.nodes() if node["Alive"]}
-            with self.lock:
-                stale_nodes = self.bad_node_ips - current_node_ips
-                if stale_nodes:
-                    self.logger.info(f"Removing stale bad node_ips: {stale_nodes}")
-                self.bad_node_ips.intersection_update(current_node_ips)
-            self.logger.info(f"Current nodes: {current_node_ips}. Bad nodes: {self.bad_node_ips}.")
-        except Exception as e:
-            raise RuntimeError(f"Exception in check_health during bad node cleanup: {e}")