ray-embedding 0.14.7__tar.gz

ray_embedding-0.14.7/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: ray-embedding
+Version: 0.14.7
+Summary: Deploy SentenceTransformers embedding models to a ray cluster
+Author: Crispin Almodovar
+Author-email: 
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# ray-embedding
+
+A Python library for deploying SentenceTransformers models to a ray cluster. 
+This tool encapsulates inference logic that uses SentenceTransformers 
+to load any compatible embedding model from the Hugging Face hub and 
+compute embeddings for input text.
+
+This library is meant to be used with the [embedding-models Ray cluster](https://bitbucket.org/docorto/embedding-models/src/dev/).
+
+Refer to this [Ray Serve deployment config](https://bitbucket.org/docorto/embedding-models/src/dev/serve-config/dev/serve-config.yaml) 
+to see how this library is used.
+
+### Supports the following backends
+
+- pytorch-gpu
+- pytorch-cpu
+
+### Planned:
+- onnx-gpu
+- onnx-cpu
+- openvino-cpu
+- fastembed-onnx-cpu
+
+

ray_embedding-0.14.7/README.md

@@ -0,0 +1,24 @@
+# ray-embedding
+
+A Python library for deploying SentenceTransformers models to a ray cluster. 
+This tool encapsulates inference logic that uses SentenceTransformers 
+to load any compatible embedding model from the Hugging Face hub and 
+compute embeddings for input text.
+
+This library is meant to be used with the [embedding-models Ray cluster](https://bitbucket.org/docorto/embedding-models/src/dev/).
+
+Refer to this [Ray Serve deployment config](https://bitbucket.org/docorto/embedding-models/src/dev/serve-config/dev/serve-config.yaml) 
+to see how this library is used.
+
+### Supports the following backends
+
+- pytorch-gpu
+- pytorch-cpu
+
+### Planned:
+- onnx-gpu
+- onnx-cpu
+- openvino-cpu
+- fastembed-onnx-cpu
+
+

ray_embedding-0.14.7/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"

ray_embedding-0.14.7/ray_embedding/__init__.py

@@ -0,0 +1,2 @@
+from ray_embedding.deploy import build_app
+

ray_embedding-0.14.7/ray_embedding/deploy.py

@@ -0,0 +1,86 @@
+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.model_router import ModelRouter
+from ray_embedding.node_reaper import NodeReaper, NODE_REAPER_DEPLOYMENT_NAME
+
+
+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": {"head_node": 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": {"worker_node": 1}}
+    ).bind(**model_router_kwargs)
+    
+    return router

ray_embedding-0.14.7/ray_embedding/dto.py

@@ -0,0 +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 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-0.14.7/ray_embedding/embedding_model.py

@@ -0,0 +1,122 @@
+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, report_unhealthy_replica_async
+
+
+@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
+        self.node_reaper = node_reaper
+        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 {}
+
+        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.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)
+        error_message = self._evaluate_cuda_health()
+        if error_message:
+            report_unhealthy_replica(error=error_message, node_reaper=self.node_reaper)
+            raise RuntimeError(error_message)
+
+    async def check_health(self):
+        error_message = self._evaluate_cuda_health()
+        if error_message:
+            await report_unhealthy_replica_async(error=error_message, node_reaper=self.node_reaper)
+            raise RuntimeError(error_message)
+
+    def _evaluate_cuda_health(self) -> Optional[str]:
+        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)
+            mem_info = nvmlDeviceGetMemoryInfo(handle)
+        except Exception as e:
+            return 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()
+
+        return None
+
+    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-0.14.7/ray_embedding/model_router.py

@@ -0,0 +1,133 @@
+import asyncio
+import logging
+import time
+from typing import Optional, Dict, List, Tuple
+
+from fastapi import FastAPI, HTTPException
+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, node_reaper: Optional[DeploymentHandle] = None):
+        assert deployed_models, "models cannot be empty"
+        assert path_prefix, "path_prefix cannot be empty"
+
+        logging.basicConfig(level=logging.INFO)
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.deployed_models = deployed_models
+        self.path_prefix = [item.removeprefix("/").removesuffix("/") for item in path_prefix]
+        self.max_concurrency = max_concurrency
+        self.rate_limiter = asyncio.Semaphore(self.max_concurrency)
+        self.available_models = [
+            {"id": str(item),
+             "object": "model",
+             "created": int(time.time()),
+             "owned_by": "openai",
+             "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]
+        model_handle = deployed_model.deployment_handle
+        batch_size = deployed_model.batch_size
+        num_retries = deployed_model.num_retries
+
+        # Resize the inputs into batch_size items, and dispatch in parallel
+        batches = [inputs[i:i+batch_size] for i in range(0, len(inputs), batch_size)]
+        if len(inputs) > batch_size:
+            self.logger.info(f"Original input (length {len(inputs)} was resized "
+                             f"to {len(batches)} mini-batches, each with max length {batch_size}.")
+
+        # Call embedding model replicas in parallel (rate-limited)
+        tasks = [self._compute_embeddings_rate_limited(model_handle, batch, dimensions) for batch in batches]
+        all_results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Retry any failed model calls
+        for i, result in enumerate(all_results):
+            if isinstance(result, Exception):
+                self.logger.warning(f"Retrying mini-batch {i} due to exception: {result}")
+                result_retried, retries = await self._retry_failed_embedding_call(model_handle, batches[i], dimensions,
+                                                                                  num_retries)
+                if retries >= num_retries and (isinstance(result_retried, Exception) or result_retried is None):
+                    raise result_retried or ValueError(f"Failed to compute `{model}` embeddings for mini-batch {i} after {num_retries} retries.")
+
+                all_results[i] = result_retried
+
+        # Flatten the results because `all_results` is a list of lists
+        self.logger.info(f"Successfully computed embeddings from {len(batches)} mini-batches")
+        return [emb for result in all_results for emb in result]
+
+    async def _compute_embeddings_rate_limited(self, model_handle: DeploymentHandle, batch: List[str], dimensions: int):
+        async with self.rate_limiter:
+            return await model_handle.remote(batch, dimensions)
+
+    async def _retry_failed_embedding_call(self,  model_handle: DeploymentHandle, batch: List[str],
+                                           dimensions: Optional[int] = None, num_retries: Optional[int] = 2) \
+            -> Tuple[List[List[float]] | Exception, int]:
+
+        result_retried, retries = None, 0
+        while retries < num_retries:
+            try:
+                result_retried = await model_handle.remote(batch, dimensions)
+            except Exception as e:
+                result_retried = e
+                self.logger.warning(e)
+            finally:
+                retries += 1
+            if not isinstance(result_retried, Exception) and result_retried is not None:
+                break
+
+        return result_retried, retries
+
+    @web_api.post("/{path_prefix}/v1/embeddings", response_model=EmbeddingResponse)
+    async def compute_embeddings(self, path_prefix: str, request: EmbeddingRequest):
+        try:
+            assert path_prefix in self.path_prefix, f"The API path prefix specified is invalid: '{path_prefix}'"
+            assert request.model in self.deployed_models, f"The model specified is invalid: {request.model}"
+
+            inputs = request.input if isinstance(request.input, list) else [request.input]
+            self.logger.info(f"Computing embeddings for a batch of {len(inputs)} texts using model: {request.model}")
+            embeddings = await self._compute_embeddings_from_resized_batches(request.model, inputs, request.dimensions)
+            response_data = [
+                {"index": idx, "embedding": emb}
+                for idx, emb in enumerate(embeddings)
+            ]
+            return EmbeddingResponse(object="list", data=response_data, model=request.model)
+        except Exception as e:
+            status_code = 400 if isinstance(e, AssertionError) else 500
+            self.logger.error(f"Failed to create embeddings: {e}")
+            raise HTTPException(status_code=status_code, detail=str(e))
+
+    @web_api.get("/{path_prefix}/v1/models")
+    async def list_models(self, path_prefix: str):
+        """Returns the list of available models in OpenAI-compatible format."""
+        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 not self.node_reaper:
+            return
+
+        try:
+            unhealthy_node_ips = await 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-0.14.7/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-0.14.7/ray_embedding/utils.py

@@ -0,0 +1,83 @@
+from typing import Optional, Tuple
+
+from ray import serve
+from ray.serve.handle import DeploymentHandle
+from ray.util import get_node_ip_address, state
+
+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 _resolve_node_reaper_handle(node_reaper: Optional[DeploymentHandle]) -> Optional[DeploymentHandle]:
+    if node_reaper is not None:
+        return node_reaper
+    try:
+        return get_node_reaper_handle()
+    except Exception:
+        return None
+
+
+def _gather_replica_context() -> Optional[Tuple[str, str]]:
+    replica_id = get_current_replica_tag()
+    node_ip = get_current_node_ip()
+    if not (replica_id and node_ip):
+        return None
+    return replica_id, node_ip
+
+
+def report_unhealthy_replica(error: Optional[str] = None,
+                             node_reaper: Optional[DeploymentHandle] = None) -> None:
+    context = _gather_replica_context()
+    if not context:
+        return
+    handle = _resolve_node_reaper_handle(node_reaper)
+    if handle is None:
+        return
+    replica_id, node_ip = context
+    handle.report_failure.remote(replica_id, node_ip, error)
+
+
+async def report_unhealthy_replica_async(error: Optional[str] = None,
+                                         node_reaper: Optional[DeploymentHandle] = None) -> None:
+    context = _gather_replica_context()
+    if not context:
+        return
+    handle = _resolve_node_reaper_handle(node_reaper)
+    if handle is None:
+        return
+    replica_id, node_ip = context
+    await handle.report_failure.remote(replica_id, node_ip, error)

ray_embedding-0.14.7/ray_embedding.egg-info/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: ray-embedding
+Version: 0.14.7
+Summary: Deploy SentenceTransformers embedding models to a ray cluster
+Author: Crispin Almodovar
+Author-email: 
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# ray-embedding
+
+A Python library for deploying SentenceTransformers models to a ray cluster. 
+This tool encapsulates inference logic that uses SentenceTransformers 
+to load any compatible embedding model from the Hugging Face hub and 
+compute embeddings for input text.
+
+This library is meant to be used with the [embedding-models Ray cluster](https://bitbucket.org/docorto/embedding-models/src/dev/).
+
+Refer to this [Ray Serve deployment config](https://bitbucket.org/docorto/embedding-models/src/dev/serve-config/dev/serve-config.yaml) 
+to see how this library is used.
+
+### Supports the following backends
+
+- pytorch-gpu
+- pytorch-cpu
+
+### Planned:
+- onnx-gpu
+- onnx-cpu
+- openvino-cpu
+- fastembed-onnx-cpu
+
+

ray_embedding-0.14.7/ray_embedding.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+pyproject.toml
+setup.cfg
+ray_embedding/__init__.py
+ray_embedding/deploy.py
+ray_embedding/dto.py
+ray_embedding/embedding_model.py
+ray_embedding/model_router.py
+ray_embedding/node_reaper.py
+ray_embedding/utils.py
+ray_embedding.egg-info/PKG-INFO
+ray_embedding.egg-info/SOURCES.txt
+ray_embedding.egg-info/dependency_links.txt
+ray_embedding.egg-info/top_level.txt

ray_embedding-0.14.7/ray_embedding.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

ray_embedding-0.14.7/ray_embedding.egg-info/top_level.txt

@@ -0,0 +1 @@
+ray_embedding

ray_embedding-0.14.7/setup.cfg

@@ -0,0 +1,21 @@
+[metadata]
+name = ray-embedding
+version = 0.14.7
+author = Crispin Almodovar
+author_email = 
+description = Deploy SentenceTransformers embedding models to a ray cluster
+long_description = file: README.md
+long_description_content_type = text/markdown
+classifiers = 
+	Programming Language :: Python :: 3
+	License :: OSI Approved :: MIT License
+	Operating System :: OS Independent
+
+[options]
+packages = find:
+python_requires = >=3.12
+
+[egg_info]
+tag_build = 
+tag_date = 0
+