Repository: deepset-ai/haystack
Pull request: Return intermediate nodes output in pipelines
Related to #1193
---------------------------

**Proposed changes**:
These changes make nodes capable of recording some debug information during execution. This is accomplished by managing one extra key in the output dictionary, called `_debug`.

By default, the data collected includes the input, the output and the logs produced by nodes (all or some of the nodes, depending on the configuration). However, a node can choose to add its own debug information under `_debug`, and such information will be preserved. The content of each node's `_debug` entry will be in the final response (grouped by producer, see sample output below).

Note that the content of `_debug` is generally passed from node to node, but to avoid infinite recursion it is removed from the output that is stored in the`_debug` key itself (see example output below)

Details:

- Modifies  `BaseComponent.run()` to make every node accept `debug` and `debug_logs` as parameters, and if detected, saves them in the instance state as attributes. This will enable users to set these values through `pipeline.run(params={'node_name':{'debug': True}})`.

- Modifies `pipeline.run()` to accept `debug` and `debug_logs` as attributes, and to apply them to each node's parameters, overwriting whatever was set in the params (see the example below).

- Modifies `BaseComponent._dispatch_run()` to deal properly with the `_debug` key's content.

- For the logs collection, introduces an implicit decorator to `BaseComponent.run()` that, if it detects the attribute `debug`, set to `True`, in the state of the current object, will record the debug logs of the execution of a specific node and push them to their `_debug`. These logs are also printed to the console `debug_logs` is also defined and set to `True`.

Example code:

```python
from haystack.document_store import ElasticsearchDocumentStore
from haystack.retriever.sparse import ElasticsearchRetriever
from haystack.retriever.dense import DensePassageRetriever
from haystack.reader import FARMReader
from haystack.pipeline import Pipeline, JoinDocuments

def main():

    document_store_with_docs = ElasticsearchDocumentStore()
    es_retriever = ElasticsearchRetriever(document_store=document_store_with_docs)
    reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2")

    pipeline = Pipeline()
    pipeline.add_node(component=es_retriever, name="ESRetriever", inputs=["Query"])
    pipeline.add_node(component=reader, name="Reader", inputs=["ESRetriever"])

    prediction = pipeline.run(
        query="Who lives in Berlin?", 
        params={
            # New API: `debug` and `debug_logs` can be passed to single nodes as parameters
            # Note that subclasses of `BaseComponent` don't need to explicitly support them for this to work
            "ESRetriever": {"top_k": 10, "debug": True, "debug_logs": True},
            "Reader": {"top_k": 3}
        },
        # New API: the debug parameters can also be passed to `run()` directly
        # They will override any node-specific setting
        debug=True,
        debug_logs=True,
    )

    # Note: printing in JSON helps detecting circular references (`pprint` instead can deal with them)
    print("############### DEBUG LOGS #####################")
    import json
    print(json.dumps(response, default=str, indent=4),"\n")
    print("################################################")

if __name__ == "__main__":
    main()
```
Example output:
```json
{
    "answers": [],
    "_debug": {
        "ESRetriever": {
            "logs": [
                "Retriever query: {'size': '10', 'query': {'bool...",
                "POST http://localhost:9200/document/_se...",
                "> {\"size\":\"10\",\"query\":{\"bool\":{\"shou...",
                "< {\"took\":3,\"timed_out\":false,\"_shards\"...",
                "Retrieved documents with IDs: [67341323..."
            ],
            "input": {
                "root_node": "Query",
                "query": "Who lives in Berlin?",
                "top_k": 10,
                "debug": true
            },
            "output": {
                "documents": []
            }
        },
        "Reader": {
            "logs": [],
            "input": {
                "documents": [],
                "query": "Who lives in Berlin?",
                "top_k": 3,
                "debug": true
            },
            "output": {
                "answers": []
            }
        }
    },
    "documents": [],
    "root_node": "Query",
    "params": {
        "ESRetriever": {
            "top_k": 10
        },
        "Reader": {
            "top_k": 3
        },
        "Query": {
            "debug": true
        }
    },
    "query": "Who lives in Berlin?",
    "node_id": "Reader"
} 
```
------------------------

**Status (please check what you already did)**:
- [X] First draft (up for discussions & feedback)
- [x] Final code
- [x] Added tests
- [ ] Updated documentation


Problem statement:
Return Intermediate Node Output
Many users in our community have been asking to have easier ways to return the output of intermediate nodes. I can see that this could be very useful for debugging and also qualitative evaluation.

I think this feature would be very useful, though the exact design is not yet fully clear.

Hints from issue discussion:
@oryx1729 @tholor Do either of you have strong thoughts on this topic?
One way forward here could be to add a `debug` arg to Pipeline.run() and every node.run(). If `True` the node can write information to the  `_debug` key introduced in #1321 . 
Ideally, we are not only returning this debug information at the very end but already print/log it along the way.
I have a feeling there might be a nice way to implement this via a decorator rather than manual population of a dict with input args. Happy to discuss more on this idea.

Commit message: Allow enable_debug and console_debug to be passed as arguments of run()

Patch to review:
diff --git a/haystack/__init__.py b/haystack/__init__.py
index 5dad0c9bb7..10227eb803 100644
--- a/haystack/__init__.py
+++ b/haystack/__init__.py
@@ -1,5 +1,14 @@
 import logging
 
+# Configure the root logger t0 DEBUG to allow the "debug" flag to receive the logs
+root_logger = logging.getLogger()
+root_logger.setLevel(logging.DEBUG)
+
+# Then reconfigure the StreamHandler not to display anything below WARNING as default
+stream_handler = logging.StreamHandler()
+stream_handler.setLevel(logging.WARNING)
+root_logger.addHandler(stream_handler)
+
 # Change log-levels before modules are loaded to avoid verbose log messages.
 logging.getLogger('haystack.modeling').setLevel(logging.WARNING)
 logging.getLogger('haystack.modeling.utils').setLevel(logging.INFO)
diff --git a/haystack/pipeline.py b/haystack/pipeline.py
index 01d72e7e43..a084eaa42f 100644
--- a/haystack/pipeline.py
+++ b/haystack/pipeline.py
@@ -10,6 +10,7 @@
 import pickle
 import urllib
 from functools import wraps
+from networkx.algorithms.boundary import node_boundary
 
 try:
     from ray import serve
@@ -259,9 +260,10 @@ def run(  # type: ignore
         labels: Optional[MultiLabel] = None,
         documents: Optional[List[Document]] = None,
         meta: Optional[dict] = None,
-        params: Optional[dict] = None,
+        params: Optional[dict] = None
     ):
-        node_output = None
+        node_output = {}
+        debug_output = {}
         queue = {
             self.root_node: {"root_node": self.root_node, "params": params}
         }  # ordered dict with "node_id" -> "input" mapping that acts as a FIFO queue
@@ -286,6 +288,7 @@ def run(  # type: ignore
                 try:
                     logger.debug(f"Running node `{node_id}` with input `{node_input}`")
                     node_output, stream_id = self.graph.nodes[node_id]["component"]._dispatch_run(**node_input)
+                    debug_output.update(node_output.get("_debug", {}))
                 except Exception as e:
                     tb = traceback.format_exc()
                     raise Exception(f"Exception while running node `{node_id}` with input `{node_input}`: {e}, full stack trace: {tb}")
@@ -315,6 +318,8 @@ def run(  # type: ignore
                 i = 0
             else:
                 i += 1  # attempt executing next node in the queue as current `node_id` has unprocessed predecessors
+        if debug_output:
+            node_output["_debug"] = debug_output
         return node_output
 
     def get_next_nodes(self, node_id: str, stream_id: str):
diff --git a/haystack/schema.py b/haystack/schema.py
index 191905f80e..4be3b08425 100644
--- a/haystack/schema.py
+++ b/haystack/schema.py
@@ -1,10 +1,16 @@
-from typing import Any, Optional, Dict, List
+from typing import Any, Optional, Dict, List, Callable, Tuple, Optional
+
 from uuid import uuid4
 from copy import deepcopy
 import mmh3
 import numpy as np
 from abc import abstractmethod
 import inspect
+import logging
+import io
+from functools import wraps
+import types
+
 
 class Document:
     def __init__(
@@ -249,6 +255,66 @@ def __str__(self):
         return str(self.to_dict())
 
 
+class InMemoryLogger(io.TextIOBase):
+    """
+    Implementation of a logger that keeps track
+    of the log lines in a list called `logs`,
+    from where they can be accessed freely.
+    """
+
+    def __init__(self, *args):
+        io.TextIOBase.__init__(self, *args)
+        self.logs = []
+    
+    def write(self, x):
+        self.logs.append(x[:100])
+
+
+def record_debug_logs(func: Callable, to_console_too: Optional[bool] = False) -> Callable:
+    """
+    Captures the debug logs of the wrapped function and 
+    saves them in the `_debug` key of the output dictionary. 
+    If `to_console_too` is True, dumps the same logs to the console as well.
+
+    Used in `BaseComponent.__getattribute__()` to wrap `run()` functions.
+    This makes sure that every implementation of `run()` by a subclass will
+    be automagically decorated with this method when requested.
+
+    :param func: the function to decorate (must be an implementation of 
+                 `BaseComponent.run()`).
+    :param to_console_too: whether the captured logs should also be displayed
+                           in the console during the execution of the pipeline.
+    """
+    @wraps(func)
+    def inner(*args, **kwargs) -> Tuple[Dict[str, Any], str]:
+
+        with InMemoryLogger() as logs_container:
+            logger = logging.getLogger()
+
+            # Adds a handler that stores the logs in a variable
+            handler = logging.StreamHandler(logs_container)
+            handler.setLevel(logging.DEBUG)
+            logger.addHandler(handler)    
+
+            # Add a handler that prints DEBUG messages in the console
+            if to_console_too:
+                handler_console = logging.StreamHandler()
+                handler_console.setLevel(logging.DEBUG)
+                logger.addHandler(handler_console)
+
+            output, stream = func(*args, **kwargs)
+            output["_debug"] = logs_container.logs
+
+            # Remove both handlers            
+            logger.removeHandler(handler)
+            if to_console_too:
+                logger.removeHandler(handler_console)
+
+            return output, stream
+
+    return inner
+
+
 class BaseComponent:
     """
     A base class for implementing nodes in a Pipeline.
@@ -266,6 +332,30 @@ def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         cls.subclasses[cls.__name__] = cls
 
+    def __getattribute__(self, name):
+        """
+        This modified `__getattribute__` method automagically decorates
+        every `BaseComponent.run()` implementation with the 
+        `record_debug_logs` decorator defined above. 
+        
+        This decorator makes the function collect its debug logs into a 
+        `_debug` key of the output dictionary.
+
+        The logs collection is not always performed. Before applying the decorator,
+        it checks for an instance attribute called `enable_debug` to know
+        whether it should or not. The decorator is applied if the attribute is 
+        defined and True.
+
+        In addition, the value of the instance attribute `console_debug` is
+        passed to the decorator. If it's defined and True, the same logs collected in 
+        `_debug` are also printed in the console during the execution.
+        """
+        if name == "run":
+            func = getattr(type(self), "run")
+            if getattr(self, 'enable_debug', False):
+                return types.MethodType(record_debug_logs(func, getattr(self, 'console_debug', False)), self)
+        return object.__getattribute__(self, name)
+
     @classmethod
     def get_subclass(cls, component_type: str):
         if component_type not in cls.subclasses.keys():
@@ -317,7 +407,7 @@ def run(
         documents: Optional[List[Document]] = None,
         meta: Optional[dict] = None,
         params: Optional[dict] = None,
-    ):
+    ) -> Tuple[Dict, str]:
         """
         Method that will be executed when the node in the graph is called.
 
@@ -330,12 +420,13 @@ def run(
         """
         pass
 
-    def _dispatch_run(self, **kwargs):
+    def _dispatch_run(self, **kwargs) -> Tuple[Dict, str]:
         """
         The Pipelines call this method which in turn executes the run() method of Component.
 
         It takes care of the following:
           - inspect run() signature to validate if all necessary arguments are available
+          - pop `enable_debug` and `console_debug` and sets them on the instance to control debug output
           - call run() with the corresponding arguments and gather output
           - collate _debug information if present
           - merge component output with the preceding output and pass it on to the subsequent Component in the Pipeline
@@ -345,13 +436,21 @@ def _dispatch_run(self, **kwargs):
 
         run_signature_args = inspect.signature(self.run).parameters.keys()
 
-        run_params = {}
+        run_params: Dict[str, Any] = {}
         for key, value in params.items():
             if key == self.name:  # targeted params for this node
                 if isinstance(value, dict):
+                    
+                    # Debug attributes
+                    if "enable_debug" in value.keys():
+                        self.enable_debug = value.pop("enable_debug")
+                    if "console_debug" in value.keys():
+                        self.console_debug = value.pop("console_debug")
+
                     for _k, _v in value.items():
                         if _k not in run_signature_args:
                             raise Exception(f"Invalid parameter '{_k}' for the node '{self.name}'.")
+
                 run_params.update(**value)
             elif key in run_signature_args:  # global params
                 run_params[key] = value


Code review comments:
diff_hunk: @@ -10,6 +10,7 @@
 import pickle
 import urllib
 from functools import wraps
+from networkx.algorithms.boundary import node_boundary
path: haystack/pipeline.py
text: @user1:
This import doesn't seem to be used in this module