ray-embedding 0.12.5__py3-none-any.whl → 0.13.0__py3-none-any.whl

ray_embedding/deploy.py

@@ -1,55 +1,79 @@
-import os
-
-import torch
-from ray.serve import Application
-
-from ray_embedding.dto import AppConfig, ModelDeploymentConfig, DeployedModel
-from ray_embedding.embedding_model import EmbeddingModel
-from ray_embedding.model_router import ModelRouter
-
-
-def build_model(model_config: ModelDeploymentConfig) -> 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
-                                                                   )
-    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
-
-    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)
-    return router
+import os
+
+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
+
+
+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 = {
+        "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={"scheduling_strategy": node_affinity_for_head()},
+    ).bind(**node_reaper_kwargs)
+
+    deployed_models = {model_config.served_model_name: build_model(model_config, node_reaper) for model_config in models}
+    router = ModelRouter.options(
+        name=model_router.deployment,
+        ray_actor_options={"scheduling_strategy": node_affinity_for_worker()}
+    ).bind(deployed_models, model_router.path_prefix, node_reaper)
+    
+    return router

ray_embedding/dto.py

@@ -1,51 +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
-    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: NodeReaperConfig
+    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,104 +1,112 @@
-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 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):
-        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
-        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.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}")
-
-    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:
-            raise RuntimeError(f"CUDA health check failed: {e}")
-
-        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,17 +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_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):
+    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"
 
@@ -32,6 +35,7 @@ class ModelRouter:
              "permission": []} for item in self.deployed_models.keys()
         ]
         self.logger.info(f"Successfully registered models: {self.available_models}")
+        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]
@@ -111,3 +115,20 @@ 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}
+
+    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,130 @@
+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(
+    name=NODE_REAPER_DEPLOYMENT_NAME,
+    route_prefix=None,
+    num_replicas=1,
+    ray_actor_options={"num_cpus": 0.25},
+    autoscaling_config={"min_replicas": 1, "max_replicas": 1},
+)
+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,71 @@
+from typing import Optional
+
+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() -> 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"]
+    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)
+    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.12.5.dist-info → ray_embedding-0.13.0.dist-info}/METADATA

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

ray_embedding-0.13.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
+ray_embedding/deploy.py,sha256=N3YOin0sh5_HxXoyY3laLiaOn1vA-JkBI6SjeOg3HDw,3871
+ray_embedding/dto.py,sha256=l0zVrzwjlvYqWsmY8LcoizKrW4AB3eSrGVXUNrIZp3o,1722
+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=kxxkelMFVpVo827m3JL-2Wat0HSk71Tux3cgNMivX5w,4989
+ray_embedding/utils.py,sha256=IE0uyBRIW9HVQjO1I46qvtxGQfdorPwU4FsAlYHjI9g,2186
+ray_embedding-0.13.0.dist-info/METADATA,sha256=BL9-K72yoNDNBAdtL3jmlK_gBSWNTb093ngaTvGjqdI,1094
+ray_embedding-0.13.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ray_embedding-0.13.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.13.0.dist-info/RECORD,,

ray_embedding-0.12.5.dist-info/RECORD

@@ -1,9 +0,0 @@
-ray_embedding/__init__.py,sha256=YS5LAZfRIwwVvE3C9g7hsauvjgIkqKtHyxkwMFFfAGY,46
-ray_embedding/deploy.py,sha256=1Nzb39OylxBEMDqCcfD-ByTefhANXbjLMzLo_YAkCfw,2710
-ray_embedding/dto.py,sha256=l0hxz_fdGjZtLMZS3BzQ1tLzAOiO_8NpX4i5Wdyuk6Q,1519
-ray_embedding/embedding_model.py,sha256=6iEaIg_mCpGEY-5F0uff2wTOMH1wI42u2N8DnaZE3mA,4670
-ray_embedding/model_router.py,sha256=BsOEz24ttvpDD4LZsDVg9rLhn26FxgUsDAvcjI0Feao,5917
-ray_embedding-0.12.5.dist-info/METADATA,sha256=EvWeadexmzrfUATF6dYl8c54cGCfdp5EEeW23vkni38,1094
-ray_embedding-0.12.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ray_embedding-0.12.5.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.12.5.dist-info/RECORD,,