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

ray_embedding/deploy.py

@@ -1,70 +1,87 @@
-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, 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)
-    device = model_config.device
-    backend = model_config.backend or "torch"
-    matryoshka_dim = model_config.matryoshka_dim
-    trust_remote_code = model_config.trust_remote_code or False
-    model_kwargs = model_config.model_kwargs or {}
-    cuda_memory_flush_threshold = model_config.cuda_memory_flush_threshold or 0.8
-
-    if "torch_dtype" in model_kwargs:
-        torch_dtype = model_kwargs["torch_dtype"].strip()
-        if torch_dtype == "float16":
-            model_kwargs["torch_dtype"] = torch.float16
-        elif torch_dtype == "bfloat16":
-            model_kwargs["torch_dtype"] = torch.bfloat16
-        elif torch_dtype == "float32":
-            model_kwargs["torch_dtype"] = torch.float32
-        else:
-            raise ValueError(f"Invalid torch_dtype: '{torch_dtype}'")
-
-    deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
-                                                                   served_model_name=served_model_name,
-                                                                   device=device,
-                                                                   backend=backend,
-                                                                   matryoshka_dim=matryoshka_dim,
-                                                                   trust_remote_code=trust_remote_code,
-                                                                   model_kwargs=model_kwargs,
-                                                                   cuda_memory_flush_threshold=cuda_memory_flush_threshold,
-                                                                   node_health_tracker=node_health_tracker
-                                                                   )
-    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
-    assert model_router.path_prefix
-
-    node_health_check_interval_s = args.node_health_check_interval_s or DEFAULT_NODE_HEALTH_CHECK_INTERVAL_S
-    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)
-              .bind(deployed_models=deployed_models,
-                    path_prefix=model_router.path_prefix,
-                    node_health_tracker=node_health_tracker))
-    return router
+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 get_head_node_id
+from ray_embedding.model_router import ModelRouter
+
+
+def build_model(model_config: ModelDeploymentConfig, node_reaper):
+    deployment_name = model_config.deployment
+    model = model_config.model
+    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
+    trust_remote_code = model_config.trust_remote_code or False
+    model_kwargs = model_config.model_kwargs or {}
+    cuda_memory_flush_threshold = model_config.cuda_memory_flush_threshold or 0.8
+
+    if "torch_dtype" in model_kwargs:
+        torch_dtype = model_kwargs["torch_dtype"].strip()
+        if torch_dtype == "float16":
+            model_kwargs["torch_dtype"] = torch.float16
+        elif torch_dtype == "bfloat16":
+            model_kwargs["torch_dtype"] = torch.bfloat16
+        elif torch_dtype == "float32":
+            model_kwargs["torch_dtype"] = torch.float32
+        else:
+            raise ValueError(f"Invalid torch_dtype: '{torch_dtype}'")
+
+    deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
+                                                                   served_model_name=served_model_name,
+                                                                   device=device,
+                                                                   backend=backend,
+                                                                   matryoshka_dim=matryoshka_dim,
+                                                                   trust_remote_code=trust_remote_code,
+                                                                   model_kwargs=model_kwargs,
+                                                                   cuda_memory_flush_threshold=cuda_memory_flush_threshold,
+                                                                   node_reaper=node_reaper,
+                                                                   )
+    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
+    assert model_router.path_prefix
+
+    node_reaper_config = args.node_reaper or NodeReaperConfig()
+
+    node_reaper_kwargs: Dict[str, Any] = {
+        "ssh_user": node_reaper_config.ssh_user,
+        "ssh_private_key": node_reaper_config.ssh_private_key,
+    }
+    if node_reaper_config.retention_seconds is not None:
+        node_reaper_kwargs["retention_seconds"] = node_reaper_config.retention_seconds
+    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( 
+        name=NODE_REAPER_DEPLOYMENT_NAME,
+        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,
+        ray_actor_options={"num_cpus": 0.25, "resources": {"node_type:worker": 1}}
+    ).bind(**model_router_kwargs)
+    
+    return router

ray_embedding/dto.py

@@ -1,52 +1,59 @@
-import dataclasses
-from typing import Union, List, Optional, Dict, Any
-from pydantic import BaseModel
-from ray.serve.handle import DeploymentHandle
-
-
-class EmbeddingRequest(BaseModel):
-    """Schema of embedding requests (compatible with OpenAI)"""
-    model: str  # Model name (for compatibility; only one model is used here)
-    input: Union[str, List[str]]  # List of strings to embed
-    dimensions: Optional[int] = None
-
-
-class EmbeddingResponse(BaseModel):
-    """Schema of embedding response (compatible with OpenAI)"""
-    object: str
-    data: List[dict]  # Embedding data including index and vector
-    model: str  # Model name used for embedding
-
-
-class ModelRouterConfig(BaseModel):
-    deployment: str
-    path_prefix: List[str] = []
-    max_concurrency: int = 32
-
-
-class ModelDeploymentConfig(BaseModel):
-    model: str
-    served_model_name: str
-    batch_size: Optional[int] = 8
-    num_retries: Optional[int] = 2
-    device: Optional[str] = None
-    backend: Optional[str] = None
-    matryoshka_dim: Optional[int] = 768
-    trust_remote_code: Optional[bool] = False
-    model_kwargs: Optional[Dict[str, Any]] = {}
-    cuda_memory_flush_threshold: Optional[float] = 0.8
-    deployment: str
-
-
-class AppConfig(BaseModel):
-    model_router: ModelRouterConfig
-    node_health_check_interval_s: Optional[int] = 30
-    models: List[ModelDeploymentConfig]
-
-
-@dataclasses.dataclass
-class DeployedModel:
-    model: str
-    deployment_handle: DeploymentHandle
-    batch_size: int
-    num_retries: Optional[int] = 2
+import dataclasses
+from typing import Union, List, Optional, Dict, Any
+from pydantic import BaseModel
+from ray.serve.handle import DeploymentHandle
+
+
+class EmbeddingRequest(BaseModel):
+    """Schema of embedding requests (compatible with OpenAI)"""
+    model: str  # Model name (for compatibility; only one model is used here)
+    input: Union[str, List[str]]  # List of strings to embed
+    dimensions: Optional[int] = None
+
+
+class EmbeddingResponse(BaseModel):
+    """Schema of embedding response (compatible with OpenAI)"""
+    object: str
+    data: List[dict]  # Embedding data including index and vector
+    model: str  # Model name used for embedding
+
+
+class ModelRouterConfig(BaseModel):
+    deployment: str
+    path_prefix: List[str] = []
+    max_concurrency: int = 32
+
+
+class ModelDeploymentConfig(BaseModel):
+    model: str
+    served_model_name: str
+    batch_size: Optional[int] = 8
+    num_retries: Optional[int] = 2
+    device: Optional[str] = None
+    backend: Optional[str] = None
+    matryoshka_dim: Optional[int] = 768
+    trust_remote_code: Optional[bool] = False
+    model_kwargs: Optional[Dict[str, Any]] = {}
+    cuda_memory_flush_threshold: Optional[float] = 0.8
+    deployment: str
+
+
+class NodeReaperConfig(BaseModel):
+    ssh_user: str = "ubuntu"
+    ssh_private_key: str = "/home/ray/ray_bootstrap_key.pem"
+    retention_seconds: Optional[int] = 900
+    reap_interval_seconds: Optional[int] = 60
+
+
+class AppConfig(BaseModel):
+    model_router: ModelRouterConfig
+    node_reaper: Optional[NodeReaperConfig] = None
+    models: List[ModelDeploymentConfig]
+    
+
+@dataclasses.dataclass
+class DeployedModel:
+    model: str
+    deployment_handle: DeploymentHandle
+    batch_size: int
+    num_retries: Optional[int] = 2

ray_embedding/embedding_model.py

@@ -1,126 +1,112 @@
-import logging
-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
-from ray.util import get_node_ip_address
-from ray.serve.handle import DeploymentHandle
-from sentence_transformers import SentenceTransformer
-
-
-@serve.deployment
-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,
-                 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":
-            self.wait_for_cuda()
-        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.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.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."""
-        if not text or (isinstance(text, list) and not all(text)):
-            raise ValueError("Input text is empty or invalid")
-
-        text = [text] if isinstance(text, str) else text
-        truncate_dim = dimensions or self.matryoshka_dim
-
-        # Compute embeddings in PyTorch format
-        embeddings = self.embedding_model.encode(
-            text, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
-        ).to(self.torch_device)
-
-        if truncate_dim is not None:
-            # Truncate and re-normalize the embeddings
-            embeddings = embeddings[:, :truncate_dim]
-            embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
-
-        # Move all embeddings to CPU at once before conversion
-        embeddings_list = embeddings.cpu().tolist()
-
-        # don't wait for GC
-        del embeddings
-
-        return embeddings_list
-
-    def wait_for_cuda(self, wait: int = 10):
-        if self.init_device == "cuda" and not torch.cuda.is_available():
-            time.sleep(wait)
-        self.check_cuda()
-
-    def check_cuda(self) -> Any:
-        if self.init_device != "cuda":
-            return None
-        try:
-            # Even though CUDA was available at init time,
-            # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
-            # https://github.com/ray-project/ray/issues/49594
-            nvmlInit()
-            count = nvmlDeviceGetCount()
-            assert count >= 1, "No CUDA devices found"
-
-            # replicas only have access to GPU 0
-            handle = nvmlDeviceGetHandleByIndex(0)
-            return handle
-        except Exception as 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
-
-        if reserved > threshold_bytes:
-            # flush only when cache exceeds the percentage threshold
-            torch.cuda.empty_cache()
-
-    def __del__(self):
-        # Clean up and free any remaining GPU memory
-        try:
-            if hasattr(self, 'embedding_model'):
-                del self.embedding_model
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-        except Exception as e:
-            self.logger.warning(f"Error during cleanup: {e}")
+import logging
+import os.path
+import time
+from typing import Optional, Dict, Any, List, Union
+
+import torch
+from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
+from ray import serve
+from ray.serve.handle import DeploymentHandle
+from sentence_transformers import SentenceTransformer
+
+from ray_embedding.utils import report_unhealthy_replica
+
+
+@serve.deployment
+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,
+                 node_reaper: 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":
+            self.wait_for_cuda()
+        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
+
+        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}")
+
+    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."""
+        if not text or (isinstance(text, list) and not all(text)):
+            raise ValueError("Input text is empty or invalid")
+
+        text = [text] if isinstance(text, str) else text
+        truncate_dim = dimensions or self.matryoshka_dim
+
+        # Compute embeddings in PyTorch format
+        embeddings = self.embedding_model.encode(
+            text, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
+        ).to(self.torch_device)
+
+        if truncate_dim is not None:
+            # Truncate and re-normalize the embeddings
+            embeddings = embeddings[:, :truncate_dim]
+            embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
+
+        # Move all embeddings to CPU at once before conversion
+        embeddings_list = embeddings.cpu().tolist()
+
+        # don't wait for GC
+        del embeddings
+
+        return embeddings_list
+
+    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()
+
+    def check_health(self):
+        if self.init_device != "cuda":
+            return
+
+        try:
+            # Even though CUDA was available at init time,
+            # CUDA can become unavailable - this is a known problem in AWS EC2+Docker
+            # https://github.com/ray-project/ray/issues/49594
+            nvmlInit()
+            count = nvmlDeviceGetCount()
+            assert count >= 1, "No CUDA devices found"
+
+            # replicas only have access to GPU 0
+            handle = nvmlDeviceGetHandleByIndex(0)
+            mem_info = nvmlDeviceGetMemoryInfo(handle)
+        except Exception as e:
+            error_message = f"CUDA health check failed: {e}"
+            report_unhealthy_replica(error=error_message, node_reaper=self.node_reaper)
+            raise RuntimeError(error_message)
+
+        reserved = torch.cuda.memory_reserved()  # bytes currently reserved by CUDA cache
+        threshold_bytes = self.cuda_memory_flush_threshold * mem_info.total
+
+        if reserved > threshold_bytes:
+            # flush only when cache exceeds the percentage threshold
+            torch.cuda.empty_cache()
+
+    def __del__(self):
+        # Clean up and free any remaining GPU memory
+        try:
+            if hasattr(self, 'embedding_model'):
+                del self.embedding_model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception as e:
+            self.logger.warning(f"Error during cleanup: {e}")

ray_embedding/model_router.py

@@ -4,20 +4,20 @@ 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
-from ray.util import get_node_ip_address
 
 from ray_embedding.dto import DeployedModel, EmbeddingRequest, EmbeddingResponse
+from ray_embedding.utils import get_current_node_ip
 
 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,
-                 node_health_tracker: Optional[DeploymentHandle] = None):
+    def __init__(self, deployed_models: Dict[str, DeployedModel], path_prefix: List[str],
+                 max_concurrency: Optional[int] = 32, node_reaper: Optional[DeploymentHandle] = None):
         assert deployed_models, "models cannot be empty"
         assert path_prefix, "path_prefix cannot be empty"
 
@@ -35,13 +35,7 @@ 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}")
+        self.node_reaper = node_reaper
 
     async def _compute_embeddings_from_resized_batches(self, model: str, inputs: List[str], dimensions: Optional[int] = None):
         deployed_model = self.deployed_models[model]
@@ -122,7 +116,19 @@ 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}
 
-    async def check_health(self):
-        if self.node_health_tracker:
-            if await self.node_health_tracker.is_bad_gpu_or_no_model_replica_on_node.remote(self.node_ip):
-                raise RuntimeError(f"The node {self.node_ip} is marked bad, or no model replica running on the node.")
+    def check_health(self):
+        if not self.node_reaper:
+            return
+
+        try:
+            unhealthy_node_ips = ray.get(self.node_reaper.get_unhealthy_node_ips.remote())
+        except Exception as exc:
+            self.logger.warning(f"Unable to fetch node reaper data: {exc}")
+            return
+
+        if not unhealthy_node_ips:
+            return
+
+        node_ip = get_current_node_ip()
+        if node_ip and node_ip in unhealthy_node_ips:
+            raise RuntimeError("Model router replica is colocated with an unhealthy embedding replica node.")

ray_embedding/node_reaper.py

@@ -0,0 +1,124 @@
+import asyncio
+import logging
+import time
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Set
+
+from ray import serve
+
+
+NODE_REAPER_DEPLOYMENT_NAME = "NodeReaper"
+
+
+@serve.deployment
+class NodeReaper:
+    def __init__(
+        self,
+        ssh_user: str,
+        ssh_private_key: str,
+        retention_seconds: int = 900,
+        reap_interval_seconds: int = 60,
+    ):
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.ssh_user = ssh_user
+        key_path = Path(ssh_private_key).expanduser()
+        if not key_path.exists():
+            raise FileNotFoundError(f"SSH private key not found: {key_path}")
+        self.ssh_private_key = key_path.as_posix()
+        self.retention_seconds = retention_seconds
+        self.reap_interval_seconds = max(30, reap_interval_seconds)
+
+        self._unhealthy_replicas: Dict[str, Dict[str, Any]] = {}
+        self._nodes_marked_for_reap: Dict[str, float] = {}
+        self._nodes_inflight: Set[str] = set()
+
+        loop = asyncio.get_event_loop()
+        self._reaper_task = loop.create_task(self._reap_loop())
+        self.logger.info("NodeReaper initialized; monitoring unhealthy nodes for recycling")
+
+    def __del__(self):
+        if hasattr(self, "_reaper_task") and self._reaper_task and not self._reaper_task.done():
+            self._reaper_task.cancel()
+    
+    def report_failure(self, replica_id: str, node_ip: str, error: Optional[str] = None):
+        self._unhealthy_replicas[replica_id] = {
+            "node_ip": node_ip,
+            "error": error,
+            "timestamp": time.time(),
+        }
+        self._nodes_marked_for_reap[node_ip] = self._nodes_marked_for_reap.get(node_ip, time.time())
+        self.logger.warning(f"Replica {replica_id} on {node_ip} marked for reaping: {error}")
+        self._purge_stale()
+
+    def get_unhealthy_node_ips(self) -> List[str]:
+        self._purge_stale()
+        return list(self._nodes_marked_for_reap.keys())
+    
+    async def _reap_loop(self):
+        while True:
+            try:
+                await asyncio.sleep(self.reap_interval_seconds)
+                await self._reap_pending_nodes()
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                self.logger.warning(f"Unexpected error in reap loop: {exc}")
+
+    async def _reap_pending_nodes(self):
+        nodes = self.get_unhealthy_node_ips()
+        for node_ip in nodes:
+            if node_ip in self._nodes_inflight:
+                continue
+            self._nodes_inflight.add(node_ip)
+            try:
+                await self._reap_node(node_ip)
+                self._clear_node(node_ip)
+                self.logger.info(f"Successfully reaped node {node_ip}")
+            except Exception as exc:
+                self.logger.error(f"Failed to reap node {node_ip}: {exc}")
+            finally:
+                self._nodes_inflight.discard(node_ip)
+
+    async def _reap_node(self, node_ip: str):
+        ssh_command = [
+            "ssh",
+            "-i",
+            self.ssh_private_key,
+            "-o",
+            "StrictHostKeyChecking=no",
+            f"{self.ssh_user}@{node_ip}",
+            "docker stop ray_container",
+        ]
+
+        self.logger.info(f"Reaping node {node_ip} via SSH")
+        process = await asyncio.create_subprocess_exec(
+            *ssh_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await process.communicate()
+        if process.returncode != 0:
+            stdout_text = stdout.decode().strip()
+            stderr_text = stderr.decode().strip()
+            raise RuntimeError(
+                f"SSH command failed with code {process.returncode}. stdout={stdout_text} stderr={stderr_text}"
+            )
+
+    def _clear_node(self, node_ip: str):
+        to_delete = [replica for replica, data in self._unhealthy_replicas.items() if data.get("node_ip") == node_ip]
+        for replica in to_delete:
+            self._unhealthy_replicas.pop(replica, None)
+        self._nodes_marked_for_reap.pop(node_ip, None)
+
+    def _purge_stale(self):
+        if not self.retention_seconds:
+            return
+        cutoff = time.time() - self.retention_seconds
+        replica_ids = [replica_id for replica_id, data in self._unhealthy_replicas.items()
+                       if data.get("timestamp", 0) < cutoff]
+        for replica_id in replica_ids:
+            node_ip = self._unhealthy_replicas[replica_id]["node_ip"]
+            self._unhealthy_replicas.pop(replica_id, None)
+            if node_ip in self._nodes_marked_for_reap and self._nodes_marked_for_reap[node_ip] < cutoff:
+                self._nodes_marked_for_reap.pop(node_ip, None)

ray_embedding/utils.py

@@ -0,0 +1,60 @@
+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)])
+        if not nodes:
+            raise RuntimeError("Unable to locate head node for NodeReaper deployment.")
+        head_node = nodes[0]
+        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
+
+
+def get_node_reaper_handle() -> DeploymentHandle:
+    try:
+        return serve.context.get_deployment_handle(NODE_REAPER_DEPLOYMENT_NAME)
+    except Exception:
+        return serve.get_deployment(NODE_REAPER_DEPLOYMENT_NAME).get_handle(sync=False)
+
+
+def get_current_replica_tag() -> Optional[str]:
+    try:
+        context = serve.context.get_current_replica_context()
+    except Exception:
+        context = None
+    if context is None:
+        return None
+    return getattr(context, "replica_tag", None)
+
+
+def get_current_node_ip() -> Optional[str]:
+    try:
+        return get_node_ip_address()
+    except Exception:
+        return None
+
+
+def report_unhealthy_replica(error: Optional[str] = None,
+                             node_reaper: Optional[DeploymentHandle] = None) -> None:
+    replica_id = get_current_replica_tag()
+    node_ip = get_current_node_ip()
+    if not (replica_id and node_ip):
+        return
+    handle = node_reaper
+    if handle is None:
+        try:
+            handle = get_node_reaper_handle()
+        except Exception:
+            return
+    handle.report_failure.remote(replica_id, node_ip, error)

{ray_embedding-0.13.9.dist-info → ray_embedding-0.14.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.13.9
+Version: 0.14.0
 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.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/node_health.py

@@ -1,94 +0,0 @@
-import logging
-import threading
-from typing import Set, List
-
-import ray
-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=0, 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.app_name = replica_context.app_name
-        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:
-        self.logger.info(f"Checking if node {node_ip} is marked bad.")
-        with self.lock:
-            is_bad_gpu_node = node_ip in self.bad_gpu_node_ips
-            self.logger.info(f"Node {node_ip} is marked bad: {is_bad_gpu_node}")
-            return is_bad_gpu_node
-
-    async def is_bad_gpu_or_no_model_replica_on_node(self, node_ip: str):
-        self.logger.info(f"Checking if node {node_ip} is marked bad or no model replica running on the node.")
-        is_bad_gpu_node = await self.is_bad_gpu_node(node_ip)
-        is_no_model_replica_running_on_node = not await self.is_model_replica_running_on_node(node_ip)
-        return is_bad_gpu_node or is_no_model_replica_running_on_node
-
-    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}")
-
-    async def is_model_replica_running_on_node(self, node_ip: str) -> bool:
-        """
-        Return True if there is at least one replica of the self.tracked_model_deployments
-        running on the specified node_ip.
-        """
-        try:
-            self.logger.info(f"Checking if there is at least one replica of tracked_deployments={self.tracked_model_deployments} "
-                             f"running on node {node_ip}.")
-            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::{self.app_name}#{d}" for d in self.tracked_model_deployments)
-            self.logger.info(f"Checking actors with prefixes: {prefixes} in node IP {node_ip}, ID {target_node_id}")
-
-            for actor in list_actors(detail=False, filters=[("node_id", "=", target_node_id)]):
-                self.logger.info(f"Checking actor: {actor}")
-                if actor.state in ["DEPENDENCIES_UNREADY", 'PENDING_CREATION', 'ALIVE', 'RESTARTING']:
-                    for prefix in prefixes:
-                        if actor.name.startswith(prefix):
-                            self.logger.info(f"Found a replica {actor.name} of "
-                                             f"tracked_deployments={self.tracked_model_deployments} "
-                                             f"running in node IP {node_ip}, node ID {target_node_id}.")
-                            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.13.9.dist-info/RECORD

@@ -1,10 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=2R7bQ7aPc9G8H9KVoemxum6-9YxmlXQogWbhFhuslko,3762
-ray_embedding/dto.py,sha256=lk_LuVQPq3MLIMTMddqHviYXILY6V5dvbzDJuD_D_qc,1573
-ray_embedding/embedding_model.py,sha256=P2xyXCznxXmdQBK6zodOJEMvxGVRMA8Ra3O5Qi7RCh0,6013
-ray_embedding/model_router.py,sha256=fmaeXzaAJeCemzL9nUoXfdCrU-ZaCe_29fx5ayDCTC0,6845
-ray_embedding/node_health.py,sha256=bKRoFHS6cVRQBOYTcv0dRA61VDeiJjmIPT8tA0hbRIU,5350
-ray_embedding-0.13.9.dist-info/METADATA,sha256=O1ObZ9JwO7eI-6Vke5hwxmyLy4VvRWIb0IgiTD6GZzQ,1074
-ray_embedding-0.13.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.13.9.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.13.9.dist-info/RECORD,,