Repository: kedro-org/kedro

Issue:
Use `asyncio` for async operations in Runners
## Description
We can refactor the code in `Task` that handles asynchronous loading and saving of data to use `asyncio`. 

## Context
https://github.com/kedro-org/kedro/pull/4206#pullrequestreview-2392331530




Discussion:

Patch:
diff --git a/docs/source/get_started/kedro_concepts.md b/docs/source/get_started/kedro_concepts.md
index 45a21570..9ba8f785 100644
--- a/docs/source/get_started/kedro_concepts.md
+++ b/docs/source/get_started/kedro_concepts.md
@@ -109,7 +109,6 @@ During `kedro new`, you can select which [tools to include in your project](../s
 - **Data Folder**: Adds a `data` folder structure for managing project data. The `data` folder contains multiple subfolders to store project data. We recommend you put raw data into `raw` and move processed data to other subfolders, as outlined [in this data engineering article](https://towardsdatascience.com/the-importance-of-layered-thinking-in-data-engineering-a09f685edc71).
 - **PySpark**: Adds PySpark-specific configuration files.
 
-
 ### `conf`
 
 The `conf` folder contains two subfolders for storing configuration information: `base` and `local`.
diff --git a/kedro/runner/task.py b/kedro/runner/task.py
index 437eec94..4aaa41f8 100644
--- a/kedro/runner/task.py
+++ b/kedro/runner/task.py
@@ -1,16 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 import itertools as it
+import logging
 import multiprocessing
 from collections.abc import Iterable, Iterator
-from concurrent.futures import (
-    ALL_COMPLETED,
-    Future,
-    ThreadPoolExecutor,
-    as_completed,
-    wait,
-)
 from typing import TYPE_CHECKING, Any
 
 from more_itertools import interleave
@@ -28,6 +23,8 @@ if TYPE_CHECKING:
     from kedro.io import CatalogProtocol
     from kedro.pipeline.node import Node
 
+logger = logging.getLogger(__name__)
+
 
 class TaskError(Exception):
     """``TaskError`` raised by ``Task``
@@ -78,11 +75,15 @@ class Task:
             )
             self.hook_manager = hook_manager
         if self.is_async:
-            node = self._run_node_async(
-                self.node,
-                self.catalog,
-                self.hook_manager,  # type: ignore[arg-type]
-                self.session_id,
+            import asyncio
+
+            node = asyncio.run(
+                self._run_node_async(
+                    self.node,
+                    self.catalog,
+                    self.hook_manager,  # type: ignore[arg-type]
+                    self.session_id,
+                )
             )
         else:
             node = self._run_node_sequential(
@@ -189,67 +190,99 @@ class Task:
             )
         return node
 
-    def _run_node_async(
+    async def _run_node_async(
         self,
         node: Node,
         catalog: CatalogProtocol,
         hook_manager: PluginManager,  # type: ignore[arg-type]
         session_id: str | None = None,
     ) -> Node:
-        with ThreadPoolExecutor() as pool:
-            inputs: dict[str, Future] = {}
+        input_tasks: dict[str, asyncio.Task] = {
+            name: asyncio.create_task(
+                self._async_dataset_load(name, node, catalog, hook_manager)
+            )
+            for name in node.inputs
+        }
 
-            for name in node.inputs:
-                inputs[name] = pool.submit(
-                    self._synchronous_dataset_load, name, node, catalog, hook_manager
+        # Gather results as they complete
+        inputs = {}
+        for name, input_task in input_tasks.items():
+            try:
+                inputs[name] = await input_task
+            except Exception as exc:  # pragma: no cover
+                logger.error(
+                    "Error loading input '%s' for node '%s': %s",
+                    name,
+                    node.name,
+                    exc,
+                    exc_info=True,
                 )
+                raise
 
-            wait(inputs.values(), return_when=ALL_COMPLETED)
-            inputs = {key: value.result() for key, value in inputs.items()}
-            is_async = True
-            additional_inputs = self._collect_inputs_from_hook(
-                node, catalog, inputs, is_async, hook_manager, session_id=session_id
-            )
-            inputs.update(additional_inputs)
+        is_async = True
+        additional_inputs = self._collect_inputs_from_hook(
+            node, catalog, inputs.copy(), is_async, hook_manager, session_id=session_id
+        )
+        inputs.update(additional_inputs)
 
-            outputs = self._call_node_run(
-                node, catalog, inputs, is_async, hook_manager, session_id=session_id
-            )
+        outputs = self._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
+        output_tasks = {
+            name: asyncio.create_task(
+                self._async_dataset_save(name, data, node, catalog, hook_manager)
+            )
+            for name, data in outputs.items()
+        }
+
+        for name, output_task in output_tasks.items():
+            try:
+                await output_task
+            except Exception as exc:
+                logger.error(
+                    "Error saving output '%s' for node '%s': %s",
+                    name,
+                    node.name,
+                    exc,
+                    exc_info=True,
                 )
+                raise
+
         return node
 
     @staticmethod
-    def _synchronous_dataset_load(
+    async def _async_dataset_load(
         dataset_name: str,
         node: Node,
         catalog: CatalogProtocol,
         hook_manager: PluginManager,
     ) -> Any:
-        """Minimal wrapper to ensure Hooks are run synchronously
-        within an asynchronous dataset load."""
+        """Asynchronously loads a dataset while ensuring hooks are executed in order."""
         hook_manager.hook.before_dataset_loaded(dataset_name=dataset_name, node=node)
-        return_ds = catalog.load(dataset_name)
+        return_ds = await asyncio.to_thread(catalog.load, dataset_name)
         hook_manager.hook.after_dataset_loaded(
             dataset_name=dataset_name, data=return_ds, node=node
         )
         return return_ds
 
+    @staticmethod
+    async def _async_dataset_save(
+        dataset_name: str,
+        data: Any,
+        node: Node,
+        catalog: CatalogProtocol,
+        hook_manager: PluginManager,
+    ) -> None:
+        """Asynchronously saves a dataset while ensuring hooks are executed in order."""
+        hook_manager.hook.before_dataset_saved(
+            dataset_name=dataset_name, data=data, node=node
+        )
+        await asyncio.to_thread(catalog.save, dataset_name, data)
+        hook_manager.hook.after_dataset_saved(
+            dataset_name=dataset_name, data=data, node=node
+        )
+
     @staticmethod
     def _collect_inputs_from_hook(  # noqa: PLR0913
         node: Node,
diff --git a/pyproject.toml b/pyproject.toml
index db7c1c17..3baf8ab4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -71,6 +71,7 @@ test = [
     "pytest-mock>=1.7.1, <4.0",
     "pytest-xdist[psutil]~=2.2.1",
     "pytest>=7.2,<9.0",
+    "pytest-asyncio>=0.21,<0.23",
     "s3fs>=2021.4, <2025.4",
     "requests_mock",
     # mypy related dependencies


Test patch:
diff --git a/tests/runner/test_task.py b/tests/runner/test_task.py
index a00098ca..8ae644b9 100644
--- a/tests/runner/test_task.py
+++ b/tests/runner/test_task.py
@@ -1,3 +1,5 @@
+from unittest.mock import AsyncMock, MagicMock
+
 import pytest
 
 from kedro.framework.hooks.manager import _NullPluginManager
@@ -93,3 +95,139 @@ class TestTask:
                 parallel=False,
             )
             task.execute()
+
+    @pytest.mark.asyncio
+    async def test_run_node_async(self, mocker):
+        """Test that `_run_node_async` properly loads inputs, executes the node, and saves outputs asynchronously."""
+        catalog = mocker.sentinel.catalog
+        session_id = "fake_session_id"
+        hook_manager = _NullPluginManager()
+        fake_dataset = mocker.Mock()
+        node_ = mocker.sentinel.node
+
+        # Create an async mock for dataset load/save
+        mock_async_dataset_load = mocker.patch(
+            "kedro.runner.task.Task._async_dataset_load",
+            new_callable=AsyncMock,
+            return_value=fake_dataset,
+        )
+        mock_async_dataset_save = mocker.patch(
+            "kedro.runner.task.Task._async_dataset_save", new_callable=AsyncMock
+        )
+        mock_collect_inputs_from_hook = mocker.patch(
+            "kedro.runner.task.Task._collect_inputs_from_hook",
+            return_value={"input3": fake_dataset},
+        )
+        mock_call_node_run = mocker.patch(
+            "kedro.runner.task.Task._call_node_run",
+            return_value={"output1": fake_dataset, "output2": fake_dataset},
+        )
+
+        # Mock inputs
+        node_.inputs = {"input1": fake_dataset, "input2": fake_dataset}
+
+        # Call the function
+        await Task(
+            node=node_,
+            catalog=catalog,
+            is_async=True,
+            hook_manager=hook_manager,
+            session_id=session_id,
+        )._run_node_async(node_, catalog, hook_manager, session_id)
+
+        # Assert inputs were loaded
+        assert mock_async_dataset_load.call_count == len(node_.inputs)
+        assert mock_async_dataset_load.await_count == len(node_.inputs)
+        mock_collect_inputs_from_hook.assert_called_once_with(
+            node_, catalog, node_.inputs, True, hook_manager, session_id=session_id
+        )
+
+        updated_inputs = node_.inputs.copy()
+        updated_inputs.update(mock_collect_inputs_from_hook.return_value)
+
+        mock_call_node_run.assert_called_once_with(
+            node_, catalog, updated_inputs, True, hook_manager, session_id=session_id
+        )
+        assert mock_async_dataset_save.call_count == len(
+            mock_call_node_run.return_value
+        )
+        assert mock_async_dataset_save.await_count == len(
+            mock_call_node_run.return_value
+        )
+
+    @pytest.mark.asyncio
+    async def test_async_dataset_load(self, mocker):
+        """Test that `_async_dataset_load` calls hooks and loads data in the proper order."""
+        catalog = mocker.sentinel.catalog
+        hook_manager = _NullPluginManager()
+        dataset_name = "test_dataset"
+        node_ = mocker.sentinel.node
+
+        # Create an order log list
+        call_order = []
+
+        # Set side_effects on hooks and catalog.load to log when they're called
+        def before_hook(*args, **kwargs):
+            call_order.append("before")
+
+        def after_hook(*args, **kwargs):
+            call_order.append("after")
+
+        def load_side_effect(*args, **kwargs):
+            call_order.append("load")
+            return "test_data"
+
+        catalog.load = MagicMock(side_effect=load_side_effect)
+        hook_manager.hook.before_dataset_loaded = MagicMock(side_effect=before_hook)
+        hook_manager.hook.after_dataset_loaded = MagicMock(side_effect=after_hook)
+
+        data = await Task._async_dataset_load(
+            dataset_name, node_, catalog, hook_manager
+        )
+
+        assert call_order == ["before", "load", "after"]
+        hook_manager.hook.before_dataset_loaded.assert_called_once_with(
+            dataset_name=dataset_name, node=node_
+        )
+        assert data == "test_data"
+        hook_manager.hook.after_dataset_loaded.assert_called_once_with(
+            dataset_name=dataset_name, data="test_data", node=node_
+        )
+
+    @pytest.mark.asyncio
+    async def test_async_dataset_save(self, mocker):
+        """Test that _async_dataset_save calls hooks and saves data in the proper order."""
+        # Setup test objects
+        catalog = mocker.sentinel.catalog
+        hook_manager = _NullPluginManager()
+        dataset_name = "test_dataset"
+        data = {"key": "value"}
+        node_ = mocker.sentinel.node
+
+        # Create a list to log call order
+        call_order = []
+
+        # Define side effects to log when each function is called
+        def before_hook(*args, **kwargs):
+            call_order.append("before")
+
+        def after_hook(*args, **kwargs):
+            call_order.append("after")
+
+        def save_side_effect(*args, **kwargs):
+            call_order.append("save")
+
+        catalog.save = MagicMock(side_effect=save_side_effect)
+        hook_manager.hook.before_dataset_saved = MagicMock(side_effect=before_hook)
+        hook_manager.hook.after_dataset_saved = MagicMock(side_effect=after_hook)
+
+        await Task._async_dataset_save(dataset_name, data, node_, catalog, hook_manager)
+
+        assert call_order == ["before", "save", "after"]
+        hook_manager.hook.before_dataset_saved.assert_called_once_with(
+            dataset_name=dataset_name, data=data, node=node_
+        )
+        catalog.save.assert_called_once_with(dataset_name, data)
+        hook_manager.hook.after_dataset_saved.assert_called_once_with(
+            dataset_name=dataset_name, data=data, node=node_
+        )