ray-embedding 0.12.6__tar.gz → 0.12.8__tar.gz

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.12.6
+Version: 0.12.8
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 
@@ -31,6 +31,6 @@ to see how this library is used.
 - onnx-gpu
 - onnx-cpu
 - openvino-cpu
-- fastembed-onnx-cpu
+
 
 

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/README.md

@@ -19,6 +19,6 @@ to see how this library is used.
 - onnx-gpu
 - onnx-cpu
 - openvino-cpu
-- fastembed-onnx-cpu
+
 
 

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/ray_embedding/deploy.py

@@ -1,14 +1,19 @@
 import os
+from typing import 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
 from ray_embedding.model_router import ModelRouter
+from ray_embedding.node_health import NodeHealthTracker
 
+DEFAULT_NODE_HEALTH_CHECK_INTERVAL_S = 30
 
-def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
+
+def build_model(model_config: ModelDeploymentConfig, node_health_tracker: Optional[DeploymentHandle] = None) -> DeployedModel:
     deployment_name = model_config.deployment
     model = model_config.model
     served_model_name = model_config.served_model_name or os.path.basename(model)
@@ -37,7 +42,8 @@ def build_model(model_config: ModelDeploymentConfig) -> DeployedModel:
                                                                    matryoshka_dim=matryoshka_dim,
                                                                    trust_remote_code=trust_remote_code,
                                                                    model_kwargs=model_kwargs,
-                                                                   cuda_memory_flush_threshold=cuda_memory_flush_threshold
+                                                                   cuda_memory_flush_threshold=cuda_memory_flush_threshold,
+                                                                   node_health_tracker=node_health_tracker
                                                                    )
     return DeployedModel(model=served_model_name,
                          deployment_handle=deployment,
@@ -51,6 +57,12 @@ def build_app(args: AppConfig) -> Application:
     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, 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()
+    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)
+              .bind(deployed_models=deployed_models,
+                    path_prefix=model_router.path_prefix,
+                    node_health_tracker=node_health_tracker))
     return router

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/ray_embedding/dto.py

@@ -38,8 +38,16 @@ 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
     models: List[ModelDeploymentConfig]
 
 

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/ray_embedding/embedding_model.py

@@ -6,6 +6,8 @@ from typing import Optional, Dict, Any, List, Union
 import torch
 from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
 from ray import serve
+from ray.util import get_node_ip_address
+from ray.serve.handle import DeploymentHandle
 from sentence_transformers import SentenceTransformer
 
 
@@ -14,12 +16,14 @@ class EmbeddingModel:
     def __init__(self, model: str, served_model_name: Optional[str] = None,
                  device: Optional[str] = None, backend: Optional[str] = "torch",
                  matryoshka_dim: Optional[int] = None, trust_remote_code: Optional[bool] = False,
-                 model_kwargs: Dict[str, Any] = None, cuda_memory_flush_threshold: Optional[float] = 0.8):
+                 model_kwargs: Dict[str, Any] = None, cuda_memory_flush_threshold: Optional[float] = 0.8,
+                 node_health_tracker: Optional[DeploymentHandle] = None):
         logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(self.__class__.__name__)
         self.model = model
         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":
@@ -29,13 +33,19 @@ class EmbeddingModel:
         self.matryoshka_dim = matryoshka_dim
         self.trust_remote_code = trust_remote_code or False
         self.model_kwargs = model_kwargs or {}
-        self.cuda_memory_flush_threshold = cuda_memory_flush_threshold
+
         self.logger.info(f"Initializing embedding model: {self.model}")
         self.embedding_model = SentenceTransformer(self.model, device=self.init_device, backend=self.backend,
                                                    trust_remote_code=self.trust_remote_code,
                                                    model_kwargs=self.model_kwargs)
 
-        self.logger.info(f"Successfully initialized model {self.model} using device {self.torch_device}")
+        self.node_health_tracker = node_health_tracker
+        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 model {self.model} using device {self.torch_device}. "
+                         f"Deployment name: {self.deployment_name}, Replica actor name: {self.replica_actor_name}, Node IP: {self.node_ip}")
 
     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."""
@@ -68,10 +78,14 @@ class EmbeddingModel:
             time.sleep(wait)
         self.check_health()
 
-    def check_health(self):
+    async def check_health(self):
         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.")
+
         try:
             # Even though CUDA was available at init time,
             # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
@@ -84,7 +98,9 @@ class EmbeddingModel:
             handle = nvmlDeviceGetHandleByIndex(0)
             mem_info = nvmlDeviceGetMemoryInfo(handle)
         except Exception as e:
-            raise RuntimeError(f"CUDA health check failed: {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}")
 
         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.6 → ray_embedding-0.12.8}/ray_embedding/model_router.py

@@ -6,6 +6,7 @@ from typing import Optional, Dict, List, Tuple
 from fastapi import FastAPI, HTTPException
 from ray import serve
 from ray.serve.handle import DeploymentHandle
+from ray.util import get_node_ip_address
 
 from ray_embedding.dto import DeployedModel, EmbeddingRequest, EmbeddingResponse
 
@@ -14,7 +15,9 @@ 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):
+    def __init__(self, deployed_models: Dict[str, DeployedModel],
+                 path_prefix: List[str], max_concurrency: Optional[int] = 32,
+                 node_health_tracker: Optional[DeploymentHandle] = None):
         assert deployed_models, "models cannot be empty"
         assert path_prefix, "path_prefix cannot be empty"
 
@@ -32,6 +35,13 @@ class ModelRouter:
              "permission": []} for item in self.deployed_models.keys()
         ]
         self.logger.info(f"Successfully registered models: {self.available_models}")
+        self.node_health_tracker = node_health_tracker
+        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 model router. "
+                         f"Deployment name: {self.deployment_name}, Replica actor name: {self.replica_actor_name}, Node IP: {self.node_ip}")
 
     async def _compute_embeddings_from_resized_batches(self, model: str, inputs: List[str], dimensions: Optional[int] = None):
         deployed_model = self.deployed_models[model]
@@ -111,3 +121,8 @@ class ModelRouter:
         if path_prefix not in self.path_prefix:
             raise HTTPException(status_code=400, detail=f"The API path prefix specified is invalid: '{path_prefix}'")
         return {"object": "list", "data": self.available_models}
+
+    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.")

ray_embedding-0.12.8/ray_embedding/node_health.py

@@ -0,0 +1,48 @@
+import logging
+import threading
+from typing import Set
+
+import ray
+from ray import serve
+
+
+@serve.deployment(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}")

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/ray_embedding.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.12.6
+Version: 0.12.8
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 
@@ -31,6 +31,6 @@ to see how this library is used.
 - onnx-gpu
 - onnx-cpu
 - openvino-cpu
-- fastembed-onnx-cpu
+
 
 

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/ray_embedding.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ ray_embedding/deploy.py
 ray_embedding/dto.py
 ray_embedding/embedding_model.py
 ray_embedding/model_router.py
+ray_embedding/node_health.py
 ray_embedding.egg-info/PKG-INFO
 ray_embedding.egg-info/SOURCES.txt
 ray_embedding.egg-info/dependency_links.txt

{ray_embedding-0.12.6 → ray_embedding-0.12.8}/setup.cfg

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