Refactor Runners, introduce Task class

kedro/runner/__init__.py

@@ -5,12 +5,14 @@
 from .parallel_runner import ParallelRunner
 from .runner import AbstractRunner, run_node
 from .sequential_runner import SequentialRunner
+from .task import Task
 from .thread_runner import ThreadRunner
 
 __all__ = [
     "AbstractRunner",
     "ParallelRunner",
     "SequentialRunner",
+    "Task",
     "ThreadRunner",
     "run_node",
 ]

kedro/runner/parallel_runner.py

@@ -27,7 +27,7 @@
     MemoryDataset,
     SharedMemoryDataset,
 )
-from kedro.runner.runner import AbstractRunner, run_node
+from kedro.runner.runner import AbstractRunner
 
 if TYPE_CHECKING:
     from pluggy import PluginManager
@@ -58,14 +58,10 @@ def _bootstrap_subprocess(
         configure_logging(logging_config)
 
 
-def _run_node_synchronization(  # noqa: PLR0913
-    node: Node,
-    catalog: CatalogProtocol,
-    is_async: bool = False,
-    session_id: str | None = None,
+def _run_node_synchronization(
     package_name: str | None = None,
     logging_config: dict[str, Any] | None = None,
-) -> Node:
+) -> None:
     """Run a single `Node` with inputs from and outputs to the `catalog`.
 
     A ``PluginManager`` instance is created in each subprocess because the
@@ -91,8 +87,6 @@ def _run_node_synchronization(  # noqa: PLR0913
     _register_hooks(hook_manager, settings.HOOKS)
     _register_hooks_entry_points(hook_manager, settings.DISABLE_HOOKS_FOR_PLUGINS)
 
-    return run_node(node, catalog, hook_manager, is_async, session_id)
-
 
 class ParallelRunner(AbstractRunner):
     """``ParallelRunner`` is an ``AbstractRunner`` implementation. It can
@@ -287,17 +281,20 @@ def _run(
                 ready = {n for n in todo_nodes if node_dependencies[n] <= done_nodes}
                 todo_nodes -= ready
                 for node in ready:
-                    futures.add(
-                        pool.submit(
-                            _run_node_synchronization,
-                            node,
-                            catalog,
-                            self._is_async,
-                            session_id,
-                            package_name=PACKAGE_NAME,
-                            logging_config=LOGGING,  # type: ignore[arg-type]
-                        )
+                    from kedro.runner.task import Task
+
+                    _run_node_synchronization(
+                        package_name=PACKAGE_NAME,
+                        logging_config=LOGGING,  # type: ignore[arg-type]
+                    )
+                    task = Task(
+                        node=node,
+                        catalog=catalog,
+                        hook_manager=hook_manager,
+                        is_async=self._is_async,
+                        session_id=session_id,
                     )
+                    futures.add(pool.submit(task))
                 if not futures:
                     if todo_nodes:
                         debug_data = {
@@ -319,19 +316,4 @@ def _run(
                     node = future.result()
                     done_nodes.add(node)
 
-                    # Decrement load counts, and release any datasets we
-                    # have finished with. This is particularly important
-                    # for the shared, default datasets we created above.
-                    for dataset in node.inputs:
-                        load_counts[dataset] -= 1
-                        if (
-                            load_counts[dataset] < 1
-                            and dataset not in pipeline.inputs()
-                        ):
-                            catalog.release(dataset)
-                    for dataset in node.outputs:
-                        if (
-                            load_counts[dataset] < 1
-                            and dataset not in pipeline.outputs()
-                        ):
-                            catalog.release(dataset)
+                    self._release_datasets(node, catalog, load_counts, pipeline)

kedro/runner/runner.py

@@ -5,24 +5,15 @@
 from __future__ import annotations
 
 import inspect
-import itertools as it
 import logging
 from abc import ABC, abstractmethod
 from collections import deque
-from concurrent.futures import (
-    ALL_COMPLETED,
-    Future,
-    ThreadPoolExecutor,
-    as_completed,
-    wait,
-)
-from typing import TYPE_CHECKING, Any, Collection, Iterable, Iterator
-
-from more_itertools import interleave
+from typing import TYPE_CHECKING, Any, Collection, Iterable
 
 from kedro.framework.hooks.manager import _NullPluginManager
 from kedro.io import CatalogProtocol, MemoryDataset
 from kedro.pipeline import Pipeline
+from kedro.runner.task import Task
 
 if TYPE_CHECKING:
     from pluggy import PluginManager
@@ -221,6 +212,19 @@ def _suggest_resume_scenario(
                 f"argument to your previous command:\n{postfix}"
             )
 
+    @staticmethod
+    def _release_datasets(
+        node: Node, catalog: CatalogProtocol, load_counts: dict, pipeline: Pipeline
+    ) -> None:
+        """Decrement dataset load counts and release any datasets we've finished with"""
+        for dataset in node.inputs:
+            load_counts[dataset] -= 1
+            if load_counts[dataset] < 1 and dataset not in pipeline.inputs():
+                catalog.release(dataset)
+        for dataset in node.outputs:
+            if load_counts[dataset] < 1 and dataset not in pipeline.outputs():
+                catalog.release(dataset)
+
 
 def _find_nodes_to_resume_from(
     pipeline: Pipeline, unfinished_nodes: Collection[Node], catalog: CatalogProtocol
@@ -402,6 +406,7 @@ def run_node(
         The node argument.
 
     """
+
     if is_async and inspect.isgeneratorfunction(node.func):
         raise ValueError(
             f"Async data loading and saving does not work with "
@@ -410,175 +415,6 @@ def run_node(
             f"in node {node!s}."
         )
 
-    if is_async:
-        node = _run_node_async(node, catalog, hook_manager, session_id)
-    else:
-        node = _run_node_sequential(node, catalog, hook_manager, session_id)
-
-    for name in node.confirms:
-        catalog.confirm(name)
-    return node
-
-
-def _collect_inputs_from_hook(  # noqa: PLR0913
-    node: Node,
-    catalog: CatalogProtocol,
-    inputs: dict[str, Any],
-    is_async: bool,
-    hook_manager: PluginManager,
-    session_id: str | None = None,
-) -> dict[str, Any]:
-    inputs = inputs.copy()  # shallow copy to prevent in-place modification by the hook
-    hook_response = hook_manager.hook.before_node_run(
-        node=node,
-        catalog=catalog,
-        inputs=inputs,
-        is_async=is_async,
-        session_id=session_id,
-    )
-
-    additional_inputs = {}
-    if (
-        hook_response is not None
-    ):  # all hooks on a _NullPluginManager will return None instead of a list
-        for response in hook_response:
-            if response is not None and not isinstance(response, dict):
-                response_type = type(response).__name__
-                raise TypeError(
-                    f"'before_node_run' must return either None or a dictionary mapping "
-                    f"dataset names to updated values, got '{response_type}' instead."
-                )
-            additional_inputs.update(response or {})
-
-    return additional_inputs
-
-
-def _call_node_run(  # noqa: PLR0913
-    node: Node,
-    catalog: CatalogProtocol,
-    inputs: dict[str, Any],
-    is_async: bool,
-    hook_manager: PluginManager,
-    session_id: str | None = None,
-) -> dict[str, Any]:
-    try:
-        outputs = node.run(inputs)
-    except Exception as exc:
-        hook_manager.hook.on_node_error(
-            error=exc,
-            node=node,
-            catalog=catalog,
-            inputs=inputs,
-            is_async=is_async,
-            session_id=session_id,
-        )
-        raise exc
-    hook_manager.hook.after_node_run(
-        node=node,
-        catalog=catalog,
-        inputs=inputs,
-        outputs=outputs,
-        is_async=is_async,
-        session_id=session_id,
-    )
-    return outputs
-
-
-def _run_node_sequential(
-    node: Node,
-    catalog: CatalogProtocol,
-    hook_manager: PluginManager,
-    session_id: str | None = None,
-) -> Node:
-    inputs = {}
-
-    for name in node.inputs:
-        hook_manager.hook.before_dataset_loaded(dataset_name=name, node=node)
-        inputs[name] = catalog.load(name)
-        hook_manager.hook.after_dataset_loaded(
-            dataset_name=name, data=inputs[name], node=node
-        )
-
-    is_async = False
-
-    additional_inputs = _collect_inputs_from_hook(
-        node, catalog, inputs, is_async, hook_manager, session_id=session_id
-    )
-    inputs.update(additional_inputs)
-
-    outputs = _call_node_run(
-        node, catalog, inputs, is_async, hook_manager, session_id=session_id
-    )
-
-    items: Iterable = outputs.items()
-    # if all outputs are iterators, then the node is a generator node
-    if all(isinstance(d, Iterator) for d in outputs.values()):
-        # Python dictionaries are ordered, so we are sure
-        # the keys and the chunk streams are in the same order
-        # [a, b, c]
-        keys = list(outputs.keys())
-        # [Iterator[chunk_a], Iterator[chunk_b], Iterator[chunk_c]]
-        streams = list(outputs.values())
-        # zip an endless cycle of the keys
-        # with an interleaved iterator of the streams
-        # [(a, chunk_a), (b, chunk_b), ...] until all outputs complete
-        items = zip(it.cycle(keys), interleave(*streams))
-
-    for name, data in items:
-        hook_manager.hook.before_dataset_saved(dataset_name=name, data=data, node=node)
-        catalog.save(name, data)
-        hook_manager.hook.after_dataset_saved(dataset_name=name, data=data, node=node)
-    return node
-
-
-def _run_node_async(
-    node: Node,
-    catalog: CatalogProtocol,
-    hook_manager: PluginManager,
-    session_id: str | None = None,
-) -> Node:
-    def _synchronous_dataset_load(dataset_name: str) -> Any:
-        """Minimal wrapper to ensure Hooks are run synchronously
-        within an asynchronous dataset load."""
-        hook_manager.hook.before_dataset_loaded(dataset_name=dataset_name, node=node)
-        return_ds = catalog.load(dataset_name)
-        hook_manager.hook.after_dataset_loaded(
-            dataset_name=dataset_name, data=return_ds, node=node
-        )
-        return return_ds
-
-    with ThreadPoolExecutor() as pool:
-        inputs: dict[str, Future] = {}
-
-        for name in node.inputs:
-            inputs[name] = pool.submit(_synchronous_dataset_load, name)
-
-        wait(inputs.values(), return_when=ALL_COMPLETED)
-        inputs = {key: value.result() for key, value in inputs.items()}
-        is_async = True
-        additional_inputs = _collect_inputs_from_hook(
-            node, catalog, inputs, is_async, hook_manager, session_id=session_id
-        )
-        inputs.update(additional_inputs)
-
-        outputs = _call_node_run(
-            node, catalog, inputs, is_async, hook_manager, session_id=session_id
-        )
-
-        future_dataset_mapping = {}
-        for name, data in outputs.items():
-            hook_manager.hook.before_dataset_saved(
-                dataset_name=name, data=data, node=node
-            )
-            future = pool.submit(catalog.save, name, data)
-            future_dataset_mapping[future] = (name, data)
-
-        for future in as_completed(future_dataset_mapping):
-            exception = future.exception()
-            if exception:
-                raise exception
-            name, data = future_dataset_mapping[future]
-            hook_manager.hook.after_dataset_saved(
-                dataset_name=name, data=data, node=node
-            )
+    task = Task(node, catalog, hook_manager, is_async, session_id)
+    node = task.execute()
     return node

kedro/runner/sequential_runner.py

@@ -9,7 +9,7 @@
 from itertools import chain
 from typing import TYPE_CHECKING, Any
 
-from kedro.runner.runner import AbstractRunner, run_node
+from kedro.runner.runner import AbstractRunner
 
 if TYPE_CHECKING:
     from pluggy import PluginManager
@@ -75,21 +75,22 @@ def _run(
 
         for exec_index, node in enumerate(nodes):
             try:
-                run_node(node, catalog, hook_manager, self._is_async, session_id)
+                from kedro.runner.task import Task
+
+                Task(
+                    node=node,
+                    catalog=catalog,
+                    hook_manager=hook_manager,
+                    is_async=self._is_async,
+                    session_id=session_id,
+                ).execute()
                 done_nodes.add(node)
             except Exception:
                 self._suggest_resume_scenario(pipeline, done_nodes, catalog)
                 raise
 
-            # decrement load counts and release any data sets we've finished with
-            for dataset in node.inputs:
-                load_counts[dataset] -= 1
-                if load_counts[dataset] < 1 and dataset not in pipeline.inputs():
-                    catalog.release(dataset)
-            for dataset in node.outputs:
-                if load_counts[dataset] < 1 and dataset not in pipeline.outputs():
-                    catalog.release(dataset)
+            self._release_datasets(node, catalog, load_counts, pipeline)
 
             self._logger.info(
-                "Completed %d out of %d tasks", exec_index + 1, len(nodes)
+                "Completed %d out of %d tasks", len(done_nodes), len(nodes)
             )