Return intermediate nodes output in pipelines

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)

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):

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