Repository: deepset-ai/haystack

Issue:
Adding the ability to use Ray Serve `async` functionality
**Is your feature request related to a problem? Please describe.**
Ray Serve uses `async` extensively. There are lots of functionalities in Ray Serve (like batch inference), which are only available if you are using `async`.
Also, ML model serving is a perfect use-case for async - especially once you use external model serving, like Ray Serve.

**Describe the solution you'd like**
Extend the `RayPipeline` functionality of Haystack to support `async` calls, deploy `async` Ray Serve deployments and call them in an `async` fashion and so also make the whole ML inference with the use of `RayPipeline` async callable. This would also allow us to use Haystack from async libraries, like FastAPI.

**Describe alternatives you've considered**
The alternative is using `sync` as-is. This mostly works, although as I have mentioned there are several important functionalities in Ray Serve which are only available if `async` is used.

**Additional context**
We have discussed this extensively with @ZanSara on Slack.
We have agreed that I will make a PR for this, although this will mean the duplication of a lot of functions in the codebas (as `await` calls can only be made from `async` methods). There is a reasonable likelihood that due the number of code duplication this would require, it will end up needing a full refactoring of the whole Haystack codebase - in which case we might need to abandon this standalone PR.
Adding the ability to use Ray Serve `async` functionality
**Is your feature request related to a problem? Please describe.**
Ray Serve uses `async` extensively. There are lots of functionalities in Ray Serve (like batch inference), which are only available if you are using `async`.
Also, ML model serving is a perfect use-case for async - especially once you use external model serving, like Ray Serve.

**Describe the solution you'd like**
Extend the `RayPipeline` functionality of Haystack to support `async` calls, deploy `async` Ray Serve deployments and call them in an `async` fashion and so also make the whole ML inference with the use of `RayPipeline` async callable. This would also allow us to use Haystack from async libraries, like FastAPI.

**Describe alternatives you've considered**
The alternative is using `sync` as-is. This mostly works, although as I have mentioned there are several important functionalities in Ray Serve which are only available if `async` is used.

**Additional context**
We have discussed this extensively with @ZanSara on Slack.
We have agreed that I will make a PR for this, although this will mean the duplication of a lot of functions in the codebas (as `await` calls can only be made from `async` methods). There is a reasonable likelihood that due the number of code duplication this would require, it will end up needing a full refactoring of the whole Haystack codebase - in which case we might need to abandon this standalone PR.


Discussion:
Hey @ZanSara ,
I think this was closed by mistake.
The mentioned [refactoring I did on the Raypipeline](https://github.com/deepset-ai/haystack/pull/2981) was NOT about adding `async`, that is still coming.
Would you mind reopening it and then I will attempt a PR on this.
Thanks
Hey @zoltan-fedor! Right, thank you for spotting this!
Hey @ZanSara ,
I think this was closed by mistake.
The mentioned [refactoring I did on the Raypipeline](https://github.com/deepset-ai/haystack/pull/2981) was NOT about adding `async`, that is still coming.
Would you mind reopening it and then I will attempt a PR on this.
Thanks
Hey @zoltan-fedor! Right, thank you for spotting this!

Resolving patch:
diff --git a/haystack/pipelines/ray.py b/haystack/pipelines/ray.py
--- a/haystack/pipelines/ray.py
+++ b/haystack/pipelines/ray.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 import inspect
-from typing import Any, Dict, List, Optional, Tuple
-
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
 from pathlib import Path
+import networkx as nx
 
 try:
     from ray import serve
@@ -11,6 +12,7 @@
     ray = None  # type: ignore
     serve = None  # type: ignore
 
+from haystack.errors import PipelineError
 from haystack.pipelines.config import (
     get_component_definitions,
     get_pipeline_definition,
@@ -19,6 +21,10 @@
 )
 from haystack.nodes.base import BaseComponent, RootNode
 from haystack.pipelines.base import Pipeline
+from haystack.schema import Document, MultiLabel
+
+
+logger = logging.getLogger(__name__)
 
 
 class RayPipeline(Pipeline):
@@ -184,7 +190,7 @@ def load_from_yaml(  # type: ignore
         :param serve_args: Optional parameters for initializing Ray Serve.
         """
         pipeline_config = read_pipeline_config_from_yaml(path)
-        return RayPipeline.load_from_config(
+        return cls.load_from_config(
             pipeline_config=pipeline_config,
             pipeline_name=pipeline_name,
             overwrite_with_env_variables=overwrite_with_env_variables,
@@ -263,9 +269,141 @@ def _add_ray_deployment_in_graph(self, handle, name: str, outgoing_edges: int, i
     def _run_node(self, node_id: str, node_input: Dict[str, Any]) -> Tuple[Dict, str]:
         return ray.get(self.graph.nodes[node_id]["component"].remote(**node_input))
 
+    async def _run_node_async(self, node_id: str, node_input: Dict[str, Any]) -> Tuple[Dict, str]:
+        # Async calling of Ray Deployments instead of using `ray.get()` as it is done
+        # in the sync version, in `_run_node()` above.
+        # See https://docs.ray.io/en/latest/ray-core/actors/async_api.html#objectrefs-as-asyncio-futures
+        return await self.graph.nodes[node_id]["component"].remote(**node_input)
+
     def _get_run_node_signature(self, node_id: str):
         return inspect.signature(self.graph.nodes[node_id]["component"].remote).parameters.keys()
 
+    # async version of the `Pipeline.run()` method
+    async def run_async(  # type: ignore
+        self,
+        query: Optional[str] = None,
+        file_paths: Optional[List[str]] = None,
+        labels: Optional[MultiLabel] = None,
+        documents: Optional[List[Document]] = None,
+        meta: Optional[Union[dict, List[dict]]] = None,
+        params: Optional[dict] = None,
+        debug: Optional[bool] = None,
+    ):
+        """
+        Runs the Pipeline, one node at a time.
+
+        :param query: The search query (for query pipelines only).
+        :param file_paths: The files to index (for indexing pipelines only).
+        :param labels: Ground-truth labels that you can use to perform an isolated evaluation of pipelines. These labels are input to nodes in the pipeline.
+        :param documents: A list of Document objects to be processed by the Pipeline Nodes.
+        :param meta: Files' metadata. Used in indexing pipelines in combination with `file_paths`.
+        :param params: A dictionary of parameters that you want to pass to the nodes.
+                       To pass a parameter to all Nodes, use: `{"top_k": 10}`.
+                       To pass a parameter to targeted Nodes, run:
+                        `{"Retriever": {"top_k": 10}, "Reader": {"top_k": 3, "debug": True}}`
+        :param debug: Specifies whether the Pipeline should instruct Nodes to collect debug information
+                      about their execution. By default, this information includes the input parameters
+                      the Nodes received and the output they generated. You can then find all debug information in the dictionary returned by this method under the key `_debug`.
+        """
+        # validate the node names
+        self._validate_node_names_in_params(params=params)
+
+        root_node = self.root_node
+        if not root_node:
+            raise PipelineError("Cannot run a pipeline with no nodes.")
+
+        node_output = None
+        queue: Dict[str, Any] = {
+            root_node: {"root_node": root_node, "params": params}
+        }  # ordered dict with "node_id" -> "input" mapping that acts as a FIFO queue
+        if query is not None:
+            queue[root_node]["query"] = query
+        if file_paths:
+            queue[root_node]["file_paths"] = file_paths
+        if labels:
+            queue[root_node]["labels"] = labels
+        if documents:
+            queue[root_node]["documents"] = documents
+        if meta:
+            queue[root_node]["meta"] = meta
+
+        i = 0  # the first item is popped off the queue unless it is a "join" node with unprocessed predecessors
+        while queue:
+            node_id = list(queue.keys())[i]
+            node_input = queue[node_id]
+            node_input["node_id"] = node_id
+
+            # Apply debug attributes to the node input params
+            # NOTE: global debug attributes will override the value specified
+            # in each node's params dictionary.
+            if debug is None and node_input:
+                if node_input.get("params", {}):
+                    debug = params.get("debug", None)  # type: ignore
+            if debug is not None:
+                if not node_input.get("params", None):
+                    node_input["params"] = {}
+                if node_id not in node_input["params"].keys():
+                    node_input["params"][node_id] = {}
+                node_input["params"][node_id]["debug"] = debug
+
+            predecessors = set(nx.ancestors(self.graph, node_id))
+            if predecessors.isdisjoint(set(queue.keys())):  # only execute if predecessor nodes are executed
+                try:
+                    logger.debug("Running node '%s` with input: %s", node_id, node_input)
+                    node_output, stream_id = await self._run_node_async(node_id, node_input)
+                except Exception as e:
+                    # The input might be a really large object with thousands of embeddings.
+                    # If you really want to see it, raise the log level.
+                    logger.debug("Exception while running node '%s' with input %s", node_id, node_input)
+                    raise Exception(
+                        f"Exception while running node '{node_id}': {e}\nEnable debug logging to see the data that was passed when the pipeline failed."
+                    ) from e
+                queue.pop(node_id)
+                #
+                if stream_id == "split":
+                    for stream_id in [key for key in node_output.keys() if key.startswith("output_")]:
+                        current_node_output = {k: v for k, v in node_output.items() if not k.startswith("output_")}
+                        current_docs = node_output.pop(stream_id)
+                        current_node_output["documents"] = current_docs
+                        next_nodes = self.get_next_nodes(node_id, stream_id)
+                        for n in next_nodes:
+                            queue[n] = current_node_output
+                else:
+                    next_nodes = self.get_next_nodes(node_id, stream_id)
+                    for n in next_nodes:  # add successor nodes with corresponding inputs to the queue
+                        if queue.get(n):  # concatenate inputs if it's a join node
+                            existing_input = queue[n]
+                            if "inputs" not in existing_input.keys():
+                                updated_input: dict = {"inputs": [existing_input, node_output], "params": params}
+                                if "_debug" in existing_input.keys() or "_debug" in node_output.keys():
+                                    updated_input["_debug"] = {
+                                        **existing_input.get("_debug", {}),
+                                        **node_output.get("_debug", {}),
+                                    }
+                                if query:
+                                    updated_input["query"] = query
+                                if file_paths:
+                                    updated_input["file_paths"] = file_paths
+                                if labels:
+                                    updated_input["labels"] = labels
+                                if documents:
+                                    updated_input["documents"] = documents
+                                if meta:
+                                    updated_input["meta"] = meta
+                            else:
+                                existing_input["inputs"].append(node_output)
+                                updated_input = existing_input
+                            queue[n] = updated_input
+                        else:
+                            queue[n] = node_output
+                i = 0
+            else:
+                i += 1  # attempt executing next node in the queue as current `node_id` has unprocessed predecessors
+
+        self.run_total += 1
+        self.send_pipeline_event_if_needed(is_indexing=file_paths is not None)
+        return node_output
+
 
 class _RayDeploymentWrapper:
     """


Test patch:
diff --git a/test/pipelines/test_ray.py b/test/pipelines/test_ray.py
--- a/test/pipelines/test_ray.py
+++ b/test/pipelines/test_ray.py
@@ -68,3 +68,31 @@ def test_load_advanced_pipeline(document_store_with_docs):
     assert prediction["query"] == "Who lives in Berlin?"
     assert prediction["answers"][0].answer == "Carla"
     assert len(prediction["answers"]) > 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+@pytest.mark.parametrize("document_store_with_docs", ["elasticsearch"], indirect=True)
+async def test_load_advanced_pipeline_async(document_store_with_docs):
+    pipeline = RayPipeline.load_from_yaml(
+        SAMPLES_PATH / "pipeline" / "ray.advanced.haystack-pipeline.yml",
+        pipeline_name="ray_query_pipeline",
+        ray_args={"num_cpus": 8},
+        serve_args={"detached": True},
+    )
+    prediction = await pipeline.run_async(
+        query="Who lives in Berlin?",
+        params={"ESRetriever1": {"top_k": 1}, "ESRetriever2": {"top_k": 2}, "Reader": {"top_k": 3}},
+    )
+
+    assert pipeline._serve_controller_client._detached is True
+    assert ray.serve.get_deployment(name="ESRetriever1").num_replicas == 2
+    assert ray.serve.get_deployment(name="ESRetriever2").num_replicas == 2
+    assert ray.serve.get_deployment(name="Reader").num_replicas == 1
+    assert ray.serve.get_deployment(name="ESRetriever1").max_concurrent_queries == 17
+    assert ray.serve.get_deployment(name="ESRetriever2").max_concurrent_queries == 15
+    assert ray.serve.get_deployment(name="ESRetriever1").ray_actor_options["num_cpus"] == 0.25
+    assert ray.serve.get_deployment(name="ESRetriever2").ray_actor_options["num_cpus"] == 0.25
+    assert prediction["query"] == "Who lives in Berlin?"
+    assert prediction["answers"][0].answer == "Carla"
+    assert len(prediction["answers"]) > 1