ray-embedding 0.10.0__py3-none-any.whl → 0.10.5__py3-none-any.whl

ray_embedding/deploy.py

@@ -19,6 +19,7 @@ def deploy_model(args: Dict[str, Any]) -> Application:
     backend: Optional[str] = args.pop("backend", "torch")
     matryoshka_dim: Optional[int] = args.pop("matryoshka_dim", None)
     trust_remote_code: Optional[bool] = args.pop("trust_remote_code", False)
+    device: Optional[str] = args.pop("device", None)
     model_kwargs: Dict[str, Any] = args.pop("model_kwargs", {})
     if "torch_dtype" in model_kwargs:
         torch_dtype = model_kwargs["torch_dtype"].strip()
@@ -33,6 +34,7 @@ def deploy_model(args: Dict[str, Any]) -> Application:
 
     deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
                                                                    backend=backend,
+                                                                   device=device,
                                                                    matryoshka_dim=matryoshka_dim,
                                                                    trust_remote_code=trust_remote_code,
                                                                    model_kwargs=model_kwargs

ray_embedding/embedding_model.py

@@ -24,24 +24,25 @@ web_api = FastAPI(title=f"Ray Embeddings - OpenAI-compatible API")
         "min_replicas": 0,
         "initial_replicas": 1,
         "max_replicas": 1,
-    },
-    user_config={
-        "max_batch_size": 8,
-        "batch_wait_timeout_s": 0.25,
     }
 )
 @serve.ingress(web_api)
 class EmbeddingModel:
-    def __init__(self, model: str, backend: Optional[str] = "torch", matryoshka_dim: Optional[int] = None,
-                 trust_remote_code: Optional[bool] = False, model_kwargs: Dict[str, Any] = None):
+    def __init__(self, model: str, backend: Optional[str] = "torch", device: Optional[str] = None,
+                 matryoshka_dim: Optional[int] = None, trust_remote_code: Optional[bool] = False,
+                 model_kwargs: Dict[str, Any] = None):
         logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(__name__)
         self.model = model
+        self.init_device = 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.torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if self.init_device is None:
+            self.torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            self.torch_device = torch.device(self.init_device)
         self.logger.info(f"Initializing embedding model: {self.model}")
         self.embedding_model = SentenceTransformer(self.model, backend=self.backend, trust_remote_code=self.trust_remote_code,
                                                    model_kwargs=self.model_kwargs)
@@ -56,12 +57,6 @@ class EmbeddingModel:
         ]
         self.logger.info(f"Successfully initialized embedding model {self.model} using device {self.torch_device}")
 
-    async def reconfigure(self, user_config: Dict):
-        assert "max_batch_size" in user_config and "batch_wait_timeout_s" in user_config, "Invalid user config"
-        self.logger.info(f"Reconfiguring dynamic batching parameters: {user_config}")
-        self.__create_embeddings_batch.set_max_batch_size(user_config["max_batch_size"])
-        self.__create_embeddings_batch.set_batch_wait_timeout_s(user_config["batch_wait_timeout_s"])
-
     @web_api.post("/v1/embeddings", response_model=EmbeddingResponse)
     async def create_embeddings(self, request: EmbeddingRequest):
         """Generate embeddings for the input text using the specified model."""
@@ -69,7 +64,6 @@ class EmbeddingModel:
             assert request.model == self.served_model_name, (
                 f"Model '{request.model}' is not supported. Use '{self.served_model_name}' instead."
             )
-
             if isinstance(request.input, str):
                 request.input = [request.input]
 
@@ -77,7 +71,7 @@ class EmbeddingModel:
 
             # Compute embeddings and convert to a PyTorch tensor on the GPU
             embeddings = self.embedding_model.encode(
-                request.input, convert_to_tensor=True, normalize_embeddings=True
+                request.input, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
             ).to(self.torch_device)
 
             if truncate_dim is not None:
@@ -99,58 +93,14 @@ class EmbeddingModel:
             self.logger.error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
-    @serve.batch(max_batch_size=8, batch_wait_timeout_s=0.25)
-    async def __create_embeddings_batch(self, requests_batch: List[EmbeddingRequest]) -> List[EmbeddingResponse]:
-        self_0 = self[0] if isinstance(self, list) else self  # Ray also passes an array of self refs; just take the first one
-        embedding_model, matryoshka_dim, torch_device = self_0.embedding_model, self_0.matryoshka_dim, self_0.torch_device
-
-        inputs, truncate_dims, num_inputs_list = [], [], []
-        for request in requests_batch:
-            input_text = request.input if isinstance(request.input, list) else [request.input] # Can be a list of texts
-            inputs.extend(input_text)
-            num_inputs_list.append(len(input_text))
-            truncate_dims.append(request.dimensions or matryoshka_dim)
-
-        embeddings = embedding_model.encode(
-            inputs, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
-        ).to(torch_device)
-
-        model_name = requests_batch[0].model
-        truncate_needed = any(dim is not None for dim in truncate_dims)
-        results_batch, ix = [], 0
-
-        if truncate_needed:
-            for truncate_dim, num_inputs in zip(truncate_dims, num_inputs_list):
-                batch_embeddings = embeddings[ix: ix + num_inputs]
-                ix += num_inputs
-
-                if truncate_dim is not None:
-                    # Truncate and normalize using pytorch (faster)
-                    batch_embeddings = batch_embeddings[:, :truncate_dim]
-                    batch_embeddings = batch_embeddings / torch.norm(batch_embeddings, dim=1, keepdim=True)
-
-                batch_embeddings = batch_embeddings.cpu().tolist()
-                response_data = [
-                    {"index": emb_ix, "embedding": emb} for emb_ix, emb in enumerate(batch_embeddings)
-                ]
-                results_batch.append(EmbeddingResponse(object="list", data=response_data, model=model_name))
-        else:
-            embeddings_list = embeddings.cpu().tolist() # Move everything to CPU
-            for num_inputs in num_inputs_list:
-                batch_embeddings = embeddings_list[ix: ix + num_inputs]
-                ix += num_inputs
-
-                response_data = [
-                    {"index": emb_ix, "embedding": emb} for emb_ix, emb in enumerate(batch_embeddings)
-                ]
-                results_batch.append(EmbeddingResponse(object="list", data=response_data, model=model_name))
-
-        return results_batch
-
     @web_api.get("/v1/models")
     async def list_models(self):
         """Returns the list of available models in OpenAI-compatible format."""
         return {"object": "list", "data": self.available_models}
 
-
-
+    def check_health(self):
+        if self.init_device == "cuda" and not torch.cuda.is_available():
+            # 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")

ray_embedding-0.10.5.dist-info/METADATA

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: ray-embedding
+Version: 0.10.5
+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.10.5.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=GEKJAaV25DOlXd7Pqj7bqASM-FcP0FtFsJlI7u_U6Iw,1995
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=i0sno946g2nbnhhhcD-bZ3WoWJi5qScbm716-i6tI1I,4501
+ray_embedding-0.10.5.dist-info/METADATA,sha256=L77LrGKHM2GYVgpXGVqkRYf5LoSCkbHpar6vkwlxIp4,1094
+ray_embedding-0.10.5.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+ray_embedding-0.10.5.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.10.5.dist-info/RECORD,,

{ray_embedding-0.10.0.dist-info → ray_embedding-0.10.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.3.0)
+Generator: setuptools (80.3.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ray_embedding-0.10.0.dist-info/METADATA

@@ -1,28 +0,0 @@
-Metadata-Version: 2.4
-Name: ray-embedding
-Version: 0.10.0
-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 tool for deploying SentenceTransformers models to a ray cluster.
-
-### Supports the following backends
-
-- pytorch-gpu
-- pytorch-cpu
-
-### Planned:
-- onnx-gpu
-- onnx-cpu
-- openvino-cpu
-- fastembed-onnx-cpu
-
-

ray_embedding-0.10.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=oqOhMxS5UyZ4oGhfpL7kqvrxLO8QW41sF_FHbNSJe-w,1858
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=cyGH7CZRCAV50T51aTAxTSRCU0haCV_mX42tyhmFa6U,7007
-ray_embedding-0.10.0.dist-info/METADATA,sha256=8uU8yj2bnbxZPlL1WEaU0V6ioi0We6rcUUzTGd62yTk,605
-ray_embedding-0.10.0.dist-info/WHEEL,sha256=GHB6lJx2juba1wDgXDNlMTyM13ckjBMKf-OnwgKOCtA,91
-ray_embedding-0.10.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.10.0.dist-info/RECORD,,