ray-embedding 0.10.10__py3-none-any.whl → 0.10.12__py3-none-any.whl

ray_embedding/embedding_model.py

@@ -37,6 +37,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"
+        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
@@ -98,6 +100,11 @@ class EmbeddingModel:
         """Returns the list of available models in OpenAI-compatible format."""
         return {"object": "list", "data": self.available_models}
 
+    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" and not torch.cuda.is_available():
             # Even though CUDA was available at init time,

{ray_embedding-0.10.10.dist-info → ray_embedding-0.10.12.dist-info}/METADATA

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

ray_embedding-0.10.12.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=ZGxcG4589WcRtaM6H84YJarw0m1XqHNgfOf3PLAhM5M,1995
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=JfNt0rJYXGlbNjg5xmY14k4jQmNZKHnTQZnTi5SbNSc,4829
+ray_embedding-0.10.12.dist-info/METADATA,sha256=7EjNTy3P2zGYP6JzjrigmpbAujjaroJA-e1slz_Ra1E,1095
+ray_embedding-0.10.12.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+ray_embedding-0.10.12.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.10.12.dist-info/RECORD,,

ray_embedding-0.10.10.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=ZGxcG4589WcRtaM6H84YJarw0m1XqHNgfOf3PLAhM5M,1995
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=hIszW30di-Us0TL5Wevo8gNpD-kL-bhitU4MrNHrebc,4574
-ray_embedding-0.10.10.dist-info/METADATA,sha256=5vgf2aQCm91W6V_hXKGfJJYo8a8L2b8gqzSrYscxY88,1095
-ray_embedding-0.10.10.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-ray_embedding-0.10.10.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.10.10.dist-info/RECORD,,