ray-embedding 0.10.11__tar.gz → 0.10.15__tar.gz

{ray_embedding-0.10.11 → ray_embedding-0.10.15}/PKG-INFO

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

{ray_embedding-0.10.11 → ray_embedding-0.10.15}/ray_embedding/embedding_model.py

@@ -7,6 +7,7 @@ import torch
 from fastapi import FastAPI, HTTPException
 from ray import serve
 from sentence_transformers import SentenceTransformer
+from pynvml import nvmlInit, nvmlDeviceGetCount
 
 from ray_embedding.dto import EmbeddingResponse, EmbeddingRequest
 
@@ -37,7 +38,8 @@ class EmbeddingModel:
         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"
-        self.check_health()
+        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
@@ -99,9 +101,18 @@ class EmbeddingModel:
         """Returns the list of available models in OpenAI-compatible format."""
         return {"object": "list", "data": self.available_models}
 
-    def check_health(self):
+    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":
             # Even though CUDA was available at init time,
             # CUDA can become unavailable - this is a known problem in AWS EC2
             # https://github.com/ray-project/ray/issues/49594
-            raise RuntimeError("CUDA device is not available")
+            try:
+                nvmlInit()
+                assert nvmlDeviceGetCount() >= 1
+            except:
+                raise RuntimeError("CUDA device is not available")

{ray_embedding-0.10.11 → ray_embedding-0.10.15}/ray_embedding.egg-info/PKG-INFO

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

{ray_embedding-0.10.11 → ray_embedding-0.10.15}/setup.cfg

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