Repository: deepset-ai/haystack
Pull request: feat: adding the ability to use Ray Serve async functionality
As I have explained in issue #2968 , a deployment on Ray Serve can be called `async`, which can make a lot of sense, as Ray Serve is typically used to serve large ML models - which means it takes a while for each request sent to a cluster of Ray Serve Deployment to return (from a few ms up to about a 100ms depending on the size of the model used).
 
As such, it makes sense to build a concurrent (eg `async) application using an async web framework like FastAPI (or recently even Flask started supporting `async`) and then send the requests to Ray Serve in an async fashion (alias without blocking), which means that while the Ray Serve Deployment is working on the request, the application is not blocked, but can work on other requests.

Unfortunately until now Haystack has only supported the calling of Ray Serve deployments in a sync fashion.
This PR is to add the ability to call the same Ray Serve deployments in an async fashion too.

This means that what is deployed into Ray Serve does NOT change, the only thing changes is how we are calling the deployment from the Ray client in Haystack.

Unfortunately making async calls require the methods containing the async call (typically `await [call]`) code to be marked as async and then that method to be called with `await`, so unfortunately some duplication of methods was required (see `pipelines/base.py::Pipeline.run()` and `pipelines/base.py::Pipeline.run_async()`) , but tried to minimize it as much as possible.

Now you can call a Ray Serve deployment async the following way.

The original (sync) call:
```
from haystack.pipelines import RayPipeline
raypipeline = RayPipeline.load_from_yaml(...)
prediction = raypipeline.run(query="Why I didn't receive my package?", params={...})
```

The new, async option was made to be very similar:
```
from haystack.pipelines import RayPipeline
raypipeline = RayPipeline.load_from_yaml(...)
prediction = await raypipeline.run_async(query="Why I didn't receive my package?", params={...})
```


### Related Issues
- fixes #2968

### Proposed Changes:
I have added the ability to call a Ray Pipeline Deployment from Haystack in an async fashion in addition to the old sync one.

### How did you test it?
I have added a dedicated test to the new async method.
I have also tested it on a live Ray Serve cluster.

### Notes for the reviewer
@ZanSara, this is what we have discussed earlier - some code repetition was unfortunately required, but managed to keep it to the minimum, see `pipelines/base.py::Pipeline.run()` and `pipelines/base.py::Pipeline.run_async()`.
The `run_async()` method is very much the same as `run()`, the only difference is the `await` call in it.

### Checklist
- [x] I have read the [contributors guidelines](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md) and the [code of conduct](https://github.com/deepset-ai/haystack/blob/main/code_of_conduct.txt)
- [x] I have updated the related issue with new insights and changes
- [x] I added tests that demonstrate the correct behavior of the change
- [x] I've used the [conventional commit convention](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title
- [x] I documented my code
- [x] I ran [pre-commit hooks](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md#installation) and fixed any issue


Problem statement:
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.

Hints from issue 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!

Commit message: Merge branch 'deepset-ai:main' into feature-ray-serve-async2

Patch to review:
diff --git a/haystack/pipelines/base.py b/haystack/pipelines/base.py
index ba8f093d08..8d3eda67aa 100644
--- a/haystack/pipelines/base.py
+++ b/haystack/pipelines/base.py
@@ -578,6 +578,132 @@ def run(  # type: ignore
         self.send_pipeline_event_if_needed(is_indexing=file_paths is not None)
         return node_output
 
+    # async version of the above `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)  # type: ignore
+                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
+
     def run_batch(  # type: ignore
         self,
         queries: Optional[List[str]] = None,
diff --git a/haystack/pipelines/ray.py b/haystack/pipelines/ray.py
index 298a4bd7bc..969865d1a3 100644
--- a/haystack/pipelines/ray.py
+++ b/haystack/pipelines/ray.py
@@ -184,7 +184,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,6 +263,14 @@ 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 in the sync version of `_run_node()`
+        # in the `RayPipeline` class in `ray.py`.
+        # See https://docs.ray.io/en/latest/ray-core/actors/async_api.html#objectrefs-as-asyncio-futures
+
+        # return ray.get(self.graph.nodes[node_id]["component"].remote(**node_input))
+        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()
 
diff --git a/pyproject.toml b/pyproject.toml
index 2e1be3c562..c0cdcdb9e1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -190,6 +190,7 @@ dev = [
   # Test
   "pytest",
   "pytest-custom_exit_code",  # used in the CI
+  "pytest-asyncio",
   "responses",
   "tox",
   "coverage",
diff --git a/test/pipelines/test_ray.py b/test/pipelines/test_ray.py
index c0471e2d74..504a7c7767 100644
--- 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


Code review comments:
diff_hunk: @@ -578,6 +578,132 @@ def run(  # type: ignore
         self.send_pipeline_event_if_needed(is_indexing=file_paths is not None)
         return node_output
 
+    # async version of the above `run()` method
+    async def run_async(  # type: ignore
path: haystack/pipelines/base.py
text: @user1:
I think for now it would make more sense to have this `run_async` in `RayPipeline`, because on regular Pipelines it will simply fail due to the lack of `_run_node_async`? Is there any concern moving it there?

@author:
No concern, that is a great idea.
I originally put `run_async` into `Pipeline` instead of `RayPipeline`, so it can be right next to the original, sync version (eg the `run` method), but you are right, it would make more sense for it to be in the `RayPipeline`.
I am moving that now and sending in a commit with that (which also makes that `type: ignore` unnecessary - see your other comment above`)

diff_hunk: @@ -578,6 +578,132 @@ def run(  # type: ignore
         self.send_pipeline_event_if_needed(is_indexing=file_paths is not None)
         return node_output
 
+    # async version of the above `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)  # type: ignore
path: haystack/pipelines/base.py
text: @user1:
Just a curiosity, what mypy is complaining about here?

@author:
The actual mypy error was `Mypy: "Pipeline" has no attribute "_run_node_async"`, as the `_run_node_async` method was placed on the `RayPipeline` and not the `Pipeline` where the  `run_async()` is.

This mypy issue will go away once I follow you other comment about moving the whole `run_async()` from the `Pipeline` to the `RayPipeline`.