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

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

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

{ray_embedding-0.10.12 → 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
 
@@ -106,8 +107,12 @@ class EmbeddingModel:
         self.check_health()
 
     def check_health(self):
-        if self.init_device == "cuda" and not torch.cuda.is_available():
+        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.12 → 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.12
+Version: 0.10.15
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

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

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