Repository: AgnostiqHQ/covalent
Pull request #1568 (2023-04-10 19:56:29)

Reduce pion delta
<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Increment the version number in the /VERSION file. If this is a bugfix, increment the patch version (the rightmost number), for example 0.18.2 becomes 0.18.3. If it's a new feature, increment the minor version (the middle number), for example 0.18.2 becomes 0.19.0.
⚠️ Add a note to /CHANGELOG.md with your version number and the date, summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Rebase the latest changes from the develop branch. Assuming origin points to https://github.com/AgnostiqHQ/covalent, you should run git rebase origin/develop.
-->

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation and CHANGELOG accordingly.
- [ ] I have read the CONTRIBUTING document.


codecov[bot]: ## [Codecov](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1568?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) Report
> Merging [#1568](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1568?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (18ff64b) into [develop](https://codecov.io/gh/AgnostiqHQ/covalent/commit/501f82e5c81a62949cc6a588efa9d6db6fbb61e9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (501f82e) will **decrease** coverage by `0.50%`.
> The diff coverage is `89.76%`.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1568/graphs/tree.svg?width=650&height=150&src=pr&token=YGHCB3DE4P&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1568?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)

```diff
@@             Coverage Diff             @@
##           develop    #1568      +/-   ##
===========================================
- Coverage    85.87%   85.37%   -0.50%     
===========================================
  Files          157      159       +2     
  Lines         6476     6593     +117     
  Branches       145      145              
===========================================
+ Hits          5561     5629      +68     
- Misses         813      862      +49     
  Partials       102      102              
```

| Flag | Coverage Δ | |
|---|---|---|
| Dispatcher | `87.65% <99.21%> (-1.29%)` | :arrow_down: |
| Functional_Tests | `59.20% <52.80%> (+0.85%)` | :arrow_up: |
| SDK | `84.00% <77.27%> (+0.24%)` | :arrow_up: |
| UI_Frontend | `72.13% <ø> (ø)` | |

</details>

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 03eca90e..a8254f53 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Tutorial for hybrid neural network using Covalent, AWSBatch, and Qiskit Runtime.
+- Task group id in Electrons.
+- Reconstruct postprocessing method.
+
 
 ### Fixed
 
diff --git a/covalent/_workflow/electron.py b/covalent/_workflow/electron.py
index 29ae806a..e8d80de1 100644
--- a/covalent/_workflow/electron.py
+++ b/covalent/_workflow/electron.py
@@ -96,14 +96,27 @@ class Electron:
         node_id: Node id of the electron.
         metadata: Metadata to be used for the function execution.
         kwargs: Keyword arguments if any.
+        task_group_id: the group to which the task be assigned when
+        it is bound to a graph node. If unset, the group id will
+        default to node id.
+        packing_tasks: Flag to indicate whether task packing is enabled.
     """
 
-    def __init__(self, function: Callable, node_id: int = None, metadata: dict = None) -> None:
+    def __init__(
+        self,
+        function: Callable,
+        node_id: int = None,
+        metadata: dict = None,
+        task_group_id: int = None,
+        packing_tasks: bool = False,
+    ) -> None:
         if metadata is None:
             metadata = {}
         self.function = function
         self.node_id = node_id
         self.metadata = metadata
+        self.task_group_id = task_group_id
+        self.packing_tasks = packing_tasks
 
     def set_metadata(self, name: str, value: Any) -> None:
         """
@@ -242,6 +255,28 @@ class Electron:
     def __complex__(self):
         return complex()
 
+    def _get_collection_electron(self, name: str, func: Callable) -> "Electron":
+        """Get collection electron with task packing enabled.
+
+        Args:
+            name: Name of the collection node.
+            func: Function to be executed.
+
+        Returns:
+            Electron object with task packing enabled.
+
+        """
+        return (
+            Electron(function=func, metadata=self.metadata.copy())
+            if name.startswith(sublattice_prefix)
+            else Electron(
+                function=func,
+                metadata=self.metadata.copy(),
+                task_group_id=self.task_group_id,
+                packing_tasks=True,
+            )
+        )
+
     def __iter__(self):
         last_frame = inspect.currentframe().f_back
         bytecode = last_frame.f_code.co_code
@@ -275,8 +310,9 @@ class Electron:
                         filtered_call_before.append(elem)
                 iterable_metadata["call_before"] = filtered_call_before
 
-                get_item_electron = Electron(function=get_item, metadata=iterable_metadata)
-                yield get_item_electron(self, i)
+                # Pack with main electron unless it is a sublattice.
+                name = active_lattice.transport_graph.get_node_value(self.node_id, "name")
+                yield self._get_collection_electron(name, get_item)(self, i)
 
     def __getattr__(self, attr: str) -> "Electron":
         # This is to handle the cases where magic functions are attempted
@@ -298,8 +334,11 @@ class Electron:
                 return getattr(e, attr)
 
             get_attr.__name__ = prefix_separator + self.function.__name__ + ".__getattr__"
-            get_attr_electron = Electron(function=get_attr, metadata=self.metadata.copy())
-            return get_attr_electron(self, attr)
+
+            # Pack with main electron except for sublattices
+            name = active_lattice.transport_graph.get_node_value(self.node_id, "name")
+            bound_electron = self._get_collection_electron(name, get_attr)(self, attr)
+            return bound_electron
 
         return super().__getattr__(attr)
 
@@ -310,9 +349,8 @@ class Electron:
                 return e[key]
 
             get_item.__name__ = prefix_separator + self.function.__name__ + ".__getitem__"
-
-            get_item_electron = Electron(function=get_item, metadata=self.metadata.copy())
-            return get_item_electron(self, key)
+            name = active_lattice.transport_graph.get_node_value(self.node_id, "name")
+            return self._get_collection_electron(name, get_item)(self, key)
 
         raise StopIteration
 
@@ -383,13 +421,16 @@ class Electron:
 
             return bound_electron
 
-        # Add a node (other than sublattices) to the transport graph of the active lattice
+        # Add a node to the transport graph of the active lattice. Electrons bound to nodes will never be packed with the
+        # 'master' Electron. # Add non-sublattice node to the transport graph of the active lattice.
         self.node_id = active_lattice.transport_graph.add_node(
             name=self.function.__name__,
             function=self.function,
             metadata=self.metadata.copy(),
             function_string=get_serialized_function_str(self.function),
+            task_group_id=self.task_group_id if self.packing_tasks else None,
         )
+        self.task_group_id = self.task_group_id if self.packing_tasks else self.node_id
 
         if self.function:
             named_args, named_kwargs = get_named_params(self.function, args, kwargs)
@@ -403,7 +444,7 @@ class Electron:
                 )
 
             # For keyword arguments
-            # Filter out kwargs to be injected by call_before calldeps at execution
+            # Filter out kwargs to be injected by call_before call_deps during execution.
             call_before = self.metadata["call_before"]
             retval_keywords = {item["attributes"]["retval_keyword"]: None for item in call_before}
             for key, value in named_kwargs.items():
@@ -421,6 +462,8 @@ class Electron:
             self.function,
             metadata=self.metadata,
             node_id=self.node_id,
+            task_group_id=self.task_group_id,
+            packing_tasks=self.packing_tasks,
         )
         active_lattice._bound_electrons[self.node_id] = bound_electron
         return bound_electron
@@ -467,7 +510,12 @@ class Electron:
             def _auto_list_node(*args, **kwargs):
                 return list(args)
 
-            list_electron = Electron(function=_auto_list_node, metadata=collection_metadata)
+            list_electron = Electron(
+                function=_auto_list_node,
+                metadata=collection_metadata,
+                task_group_id=self.task_group_id,
+                packing_tasks=True,
+            )  # Group the auto-generated node with the main node.
             bound_electron = list_electron(*param_value)
             transport_graph.set_node_value(bound_electron.node_id, "name", electron_list_prefix)
             transport_graph.add_edge(
@@ -483,7 +531,12 @@ class Electron:
             def _auto_dict_node(*args, **kwargs):
                 return dict(kwargs)
 
-            dict_electron = Electron(function=_auto_dict_node, metadata=collection_metadata)
+            dict_electron = Electron(
+                function=_auto_dict_node,
+                metadata=collection_metadata,
+                task_group_id=self.task_group_id,
+                packing_tasks=True,
+            )  # Group the auto-generated node with the main node.
             bound_electron = dict_electron(**param_value)
             transport_graph.set_node_value(bound_electron.node_id, "name", electron_dict_prefix)
             transport_graph.add_edge(
diff --git a/covalent/_workflow/lattice.py b/covalent/_workflow/lattice.py
index 1a766575..e75dd8e1 100644
--- a/covalent/_workflow/lattice.py
+++ b/covalent/_workflow/lattice.py
@@ -206,8 +206,8 @@ class Lattice:
 
         Returns:
             None
-        """
 
+        """
         self.args = [TransportableObject.make_transportable(arg) for arg in args]
         self.kwargs = {k: TransportableObject.make_transportable(v) for k, v in kwargs.items()}
 
@@ -244,9 +244,12 @@ class Lattice:
                     )
                     raise
 
+        pp = Postprocessor(lattice=self)
+
         if get_config("sdk.exhaustive_postprocess") == "true":
-            pp = Postprocessor(lattice=self)
             pp.add_exhaustive_postprocess_node(self._bound_electrons.copy())
+        else:
+            pp.add_reconstruct_postprocess_node(retval, self._bound_electrons.copy())
 
         self._bound_electrons = {}  # Reset bound electrons
 
diff --git a/covalent/_workflow/postprocessing.py b/covalent/_workflow/postprocessing.py
index 1a62cbee..eb2d0608 100644
--- a/covalent/_workflow/postprocessing.py
+++ b/covalent/_workflow/postprocessing.py
@@ -22,7 +22,7 @@
 
 from builtins import list
 from dataclasses import asdict
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Set, Union
 
 from .._shared_files import logger
 from .._shared_files.context_managers import active_lattice_manager
@@ -108,7 +108,7 @@ class Postprocessor:
             self.lattice.post_processing = False
             return result
 
-    def _get_node_ids_from_retval(self, retval: Union["Electron", List, Dict]) -> List[int]:
+    def _get_node_ids_from_retval(self, retval: Union["Electron", List, Dict]) -> Set[int]:
         """Get the list of electron node ids from the return value of a lattice function.
 
         Args:
@@ -124,7 +124,7 @@ class Postprocessor:
         if isinstance(retval, Electron):
             node_ids.append(retval.node_id)
             app_log.debug(f"Preprocess: Encountered node {retval.node_id}")
-        elif isinstance(retval, list) or isinstance(retval, tuple) or isinstance(retval, set):
+        elif isinstance(retval, (list, tuple, set)):
             app_log.debug("Recursively preprocessing iterable")
             for e in retval:
                 node_ids.extend(self._get_node_ids_from_retval(e))
@@ -136,7 +136,7 @@ class Postprocessor:
             app_log.debug(f"Encountered primitive or unsupported type: {retval}")
             return []
 
-        return node_ids
+        return set(node_ids)
 
     def _get_electron_metadata(self) -> Dict:
         """Get the metadata for the postprocess electron.
@@ -154,7 +154,7 @@ class Postprocessor:
         return pp_metadata
 
     def add_exhaustive_postprocess_node(self, bound_electrons: Dict) -> None:
-        """This function adds a postprocess node to the transport graph. The postprocess node will be executed after all other nodes have been executed.
+        """This function adds a postprocess node to the transport graph based on the 'exhaustive' algorithm.
 
         Args:
             bound_electrons: Dictionary of bound electrons with node_id as key and the electron object as the value. A bound is an electron that has been assigned a node_id and is part of a transport graph.
@@ -173,3 +173,65 @@ class Postprocessor:
             )
             bound_pp = pp_electron(*filtered_ordered_electrons)
             tg.set_node_value(bound_pp.node_id, "name", postprocess_prefix)
+
+    def _postprocess_recursively(self, retval, **referenced_outputs):
+        from .electron import Electron
+
+        if isinstance(retval, Electron):
+            key = f"node:{retval.node_id}"
+            return referenced_outputs[key]
+        elif isinstance(retval, list):
+            return list(
+                map(lambda x: self._postprocess_recursively(x, **referenced_outputs), retval)
+            )
+        elif isinstance(retval, tuple):
+            return tuple(
+                map(lambda x: self._postprocess_recursively(x, **referenced_outputs), retval)
+            )
+        elif isinstance(retval, set):
+            return {self._postprocess_recursively(x, **referenced_outputs) for x in retval}
+        elif isinstance(retval, dict):
+            return {
+                k: self._postprocess_recursively(v, **referenced_outputs)
+                for k, v in retval.items()
+            }
+        else:
+            return retval
+
+    def add_reconstruct_postprocess_node(
+        self, retval: Union["Electron", List, Dict], bound_electrons: Dict
+    ):
+        """This function adds a postprocess node to the transport graph based on the 'eager reconstruction' algorithm.
+
+        Args:
+            retval: Return value of the lattice function.
+            bound_electrons: Dictionary of bound electrons with node_id as key and the electron object as the value. A bound is an electron that has been assigned a node_id and is part of a transport graph.
+
+        Returns:
+            None
+
+        """
+        from .electron import Electron, wait
+
+        node_id_refs = self._get_node_ids_from_retval(retval)
+        referenced_electrons = {}
+        for node_id in node_id_refs:
+            key = f"node:{node_id}"
+            referenced_electrons[key] = bound_electrons[node_id]
+
+        with active_lattice_manager.claim(self.lattice):
+            pp_metadata = self._get_electron_metadata()
+            pp_electron = Electron(function=self._postprocess_recursively, metadata=pp_metadata)
+
+            # Add pp_electron to the graph -- this will also add a
+            # parameter node in case retval is not a single electron
+            bound_pp = pp_electron(retval, **referenced_electrons)
+
+            # Edit pp electron name
+            self.lattice.transport_graph.set_node_value(
+                bound_pp.node_id, "name", f"{postprocess_prefix}reconstruct"
+            )
+
+            # Wait for non-referenced electrons
+            wait_parents = [v for k, v in bound_electrons.items() if k not in node_id_refs]
+            wait(child=bound_pp, parents=wait_parents)
diff --git a/covalent/_workflow/transport.py b/covalent/_workflow/transport.py
index f3fb7279..ecdba864 100644
--- a/covalent/_workflow/transport.py
+++ b/covalent/_workflow/transport.py
@@ -319,7 +319,9 @@ class _TransportGraph:
             "stderr": None,
         }
 
-    def add_node(self, name: str, function: Callable, metadata: Dict, **attr) -> int:
+    def add_node(
+        self, name: str, function: Callable, metadata: Dict, task_group_id: int = None, **attr
+    ) -> int:
         """
         Adds a node to the graph.
 
@@ -327,15 +329,17 @@ class _TransportGraph:
             name: The name of the node.
             function: The function to be executed.
             metadata: The metadata of the node.
+            task_group_id: The task group id of the node.
             attr: Any other attributes that need to be added to the node.
 
         Returns:
             node_key: The node id.
-        """
 
+        """
         node_id = len(self._graph.nodes)
         self._graph.add_node(
             node_id,
+            task_group_id=task_group_id if task_group_id is not None else node_id,
             name=name,
             function=TransportableObject(function),
             metadata=metadata,


Test patch:
diff --git a/tests/covalent_tests/workflow/electron_test.py b/tests/covalent_tests/workflow/electron_test.py
index e52bc951..03acdb5d 100644
--- a/tests/covalent_tests/workflow/electron_test.py
+++ b/tests/covalent_tests/workflow/electron_test.py
@@ -305,7 +305,7 @@ def test_autogen_dict_electrons():
     assert set(g.edges) == {(1, 0, 0), (3, 1, 0), (2, 1, 0), (0, 4, 0)}
 
 
-def test_as_transportable_dict(mocker):
+def test_as_transportable_dict():
     """Test the get transportable electron function."""
 
     @ct.electron
@@ -323,7 +323,7 @@ def test_as_transportable_dict(mocker):
     assert TransportableObject(test_func).to_dict() == transportable_electron["function"]
 
 
-def test_call_sublattice(mocker):
+def test_call_sublattice():
     """Test the sublattice logic when the __call__ method is invoked."""
 
     @ct.lattice(executor="mock")
@@ -349,3 +349,138 @@ def test_call_sublattice(mocker):
                 assert (
                     node_data["function"].get_deserialized().__name__ == "_build_sublattice_graph"
                 )
+
+
+def test_electron_auto_task_groups():
+    @ct.electron
+    def task(arr: list):
+        return sum(arr)
+
+    @ct.electron
+    @ct.lattice
+    def sublattice(x):
+        return task(x)
+
+    @ct.lattice
+    def workflow(x):
+        return sublattice(x)
+
+    workflow.build_graph([[1, 2], 3])
+    tg = workflow.transport_graph
+    assert all(tg.get_node_value(i, "task_group_id") == 0 for i in [0, 3, 4])
+    assert all(tg.get_node_value(i, "task_group_id") == i for i in [1, 2, 5, 6, 7, 8])
+
+
+def test_electron_get_attr():
+    class Point:
+        def __init__(self, x, y):
+            self.x = x
+            self.y = y
+
+    @ct.electron
+    def create_point():
+        return Point(3, 4)
+
+    @ct.electron
+    def add(a, b):
+        return a + b
+
+    @ct.lattice
+    def workflow():
+        point = create_point()
+        return add(point.x, point.y)
+
+    workflow.build_graph()
+    tg = workflow.transport_graph
+
+    # TG:
+    # 0: point
+    # 1: point.__getattr__
+    # 2: "x"
+    # 3: point.__getattr__
+    # 4: "y"
+    # 5: add
+    # 6: "postprocess"
+
+    point_electron_gid = tg.get_node_value(0, "task_group_id")
+    getitem_x_gid = tg.get_node_value(1, "task_group_id")
+    getitem_y_gid = tg.get_node_value(3, "task_group_id")
+    assert point_electron_gid == 0
+    assert getitem_x_gid == point_electron_gid
+    assert getitem_y_gid == point_electron_gid
+    assert all(tg.get_node_value(i, "task_group_id") == i for i in [2, 4, 5, 6])
+
+
+def test_electron_auto_task_groups_getitem():
+    """Test task packing with __getitem__"""
+
+    @ct.electron
+    def create_array():
+        return [3, 4]
+
+    @ct.electron
+    def add(a, b):
+        return a + b
+
+    @ct.lattice
+    def workflow():
+        arr = create_array()
+        return add(arr[0], arr[1])
+
+    workflow.build_graph()
+    tg = workflow.transport_graph
+
+    # TG:
+    # 0: arr
+    # 1: arr.__getitem__
+    # 2: 0
+    # 3: arr.__getitem__
+    # 4: 1
+    # 5: add
+    # 6: "postprocess"
+
+    arr_electron_gid = tg.get_node_value(0, "task_group_id")
+    getitem_x_gid = tg.get_node_value(1, "task_group_id")
+    getitem_y_gid = tg.get_node_value(3, "task_group_id")
+    assert arr_electron_gid == 0
+    assert getitem_x_gid == arr_electron_gid
+    assert getitem_y_gid == arr_electron_gid
+    assert all(tg.get_node_value(i, "task_group_id") == i for i in [2, 4, 5, 6])
+
+
+def test_electron_auto_task_groups_iter():
+    """Test task packing with __iter__"""
+
+    @ct.electron
+    def create_tuple():
+        return (3, 4)
+
+    @ct.electron
+    def add(a, b):
+        return a + b
+
+    @ct.lattice
+    def workflow():
+        tup = create_tuple()
+        x, y = tup
+        return add(x, y)
+
+    workflow.build_graph()
+    tg = workflow.transport_graph
+
+    # TG:
+    # 0: tup
+    # 1: tup.__getitem__
+    # 2: 0
+    # 3: tup.__getitem__
+    # 4: 1
+    # 5: add
+    # 6: "postprocess"
+
+    tup_electron_gid = tg.get_node_value(0, "task_group_id")
+    getitem_x_gid = tg.get_node_value(1, "task_group_id")
+    getitem_y_gid = tg.get_node_value(3, "task_group_id")
+    assert tup_electron_gid == 0
+    assert getitem_x_gid == tup_electron_gid
+    assert getitem_y_gid == tup_electron_gid
+    assert all(tg.get_node_value(i, "task_group_id") == i for i in [2, 4, 5, 6])
diff --git a/tests/covalent_tests/workflow/lattice_test.py b/tests/covalent_tests/workflow/lattice_test.py
index 348b5c17..d4424f0f 100644
--- a/tests/covalent_tests/workflow/lattice_test.py
+++ b/tests/covalent_tests/workflow/lattice_test.py
@@ -23,7 +23,7 @@
 from dataclasses import asdict
 
 import covalent as ct
-from covalent._shared_files.defaults import DefaultMetadataValues
+from covalent._shared_files.defaults import DefaultMetadataValues, postprocess_prefix
 
 DEFAULT_METADATA_VALUES = asdict(DefaultMetadataValues())
 
@@ -86,3 +86,24 @@ def test_lattice_executor_settings():
     assert workflow.metadata["executor"] == DEFAULT_METADATA_VALUES["executor"]
     workflow_2.build_graph(1)
     assert workflow_2.metadata["executor"] == "custom_executor"
+
+
+def test_lattice_build_graph(mocker):
+    """Test the build graph method in lattice."""
+
+    @ct.electron
+    def task(x):
+        return x
+
+    @ct.lattice
+    def workflow(x):
+        return task(x)
+
+    original_exhaustive_value = ct._shared_files.config.get_config("sdk.exhaustive_postprocess")
+    ct._shared_files.config.set_config("sdk.exhaustive_postprocess", "false")
+    workflow.build_graph(1)
+    assert workflow.transport_graph.get_node_value(2, "name") == f"{postprocess_prefix}reconstruct"
+    ct._shared_files.config.set_config("sdk.exhaustive_postprocess", "true")
+    workflow.build_graph(1)
+    assert workflow.transport_graph.get_node_value(2, "name") == postprocess_prefix
+    ct._shared_files.config.set_config("sdk.exhaustive_postprocess", original_exhaustive_value)
diff --git a/tests/covalent_tests/workflow/postprocessing_test.py b/tests/covalent_tests/workflow/postprocessing_test.py
index b0c90741..eb016638 100644
--- a/tests/covalent_tests/workflow/postprocessing_test.py
+++ b/tests/covalent_tests/workflow/postprocessing_test.py
@@ -23,6 +23,7 @@
 from unittest.mock import MagicMock, Mock
 
 import pytest
+from mock import call
 
 import covalent as ct
 from covalent._shared_files.defaults import postprocess_prefix
@@ -114,17 +115,17 @@ def test_postprocess():
 @pytest.mark.parametrize(
     "retval, node_ids",
     [
-        (Electron(function=lambda x: x, node_id=1), [1]),
+        (Electron(function=lambda x: x, node_id=1), {1}),
         (
             [Electron(function=lambda x: x, node_id=1), Electron(function=lambda x: x, node_id=2)],
-            [1, 2],
+            {1, 2},
         ),
         (
             {
                 "a": Electron(function=lambda x: x, node_id=1),
                 "b": Electron(function=lambda x: x, node_id=2),
             },
-            [1, 2],
+            {1, 2},
         ),
         ("unsupported", []),
     ],
@@ -164,3 +165,69 @@ def test_add_exhaustive_postprocess_node(postprocessor):
         postprocessor.lattice.transport_graph._graph.nodes(data=True)[0]["name"]
         == postprocess_prefix
     )
+
+
+def test_add_reconstruct_postprocess_node(postprocessor, mocker):
+    """Test method that adds eager postprocess node."""
+
+    def test_func(x):
+        return x
+
+    get_node_ids_from_retval_mock = mocker.patch(
+        "covalent._workflow.postprocessing.Postprocessor._get_node_ids_from_retval"
+    )
+    get_electron_metadata_mock = mocker.patch(
+        "covalent._workflow.postprocessing.Postprocessor._get_electron_metadata"
+    )
+
+    mock_electron = Electron(function=test_func, node_id=0)
+    mock_bound_electrons = {0: mock_electron}
+    postprocessor.add_reconstruct_postprocess_node(mock_electron, mock_bound_electrons)
+    get_electron_metadata_mock.assert_called_once_with()
+    assert get_node_ids_from_retval_mock.mock_calls == [
+        call(mock_electron),
+        call().__iter__(),
+        call().__contains__(0),
+    ]
+
+
+def test_postprocess_recursively(postprocessor, mocker):
+    """Test postprocess_recursively method."""
+
+    def test_func(x):
+        return x
+
+    mock_electron_0 = Electron(function=test_func, node_id=0)
+    res = postprocessor._postprocess_recursively(mock_electron_0, **{"node:0": "mock-output-0"})
+    assert res == "mock-output-0"
+
+    mock_electron_1 = Electron(function=test_func, node_id=1)
+    res = postprocessor._postprocess_recursively(
+        [mock_electron_0, mock_electron_1],
+        **{"node:0": "mock-output-0", "node:1": "mock-output-1"},
+    )
+    assert res == ["mock-output-0", "mock-output-1"]
+
+    res = postprocessor._postprocess_recursively(
+        {mock_electron_0, mock_electron_1},
+        **{"node:0": "mock-output-0", "node:1": "mock-output-1"},
+    )
+    assert res == {"mock-output-0", "mock-output-1"}
+
+    res = postprocessor._postprocess_recursively(
+        (mock_electron_0, mock_electron_1),
+        **{"node:0": "mock-output-0", "node:1": "mock-output-1"},
+    )
+    assert res == ("mock-output-0", "mock-output-1")
+
+    res = postprocessor._postprocess_recursively(
+        {"0": mock_electron_0, "1": mock_electron_1},
+        **{"node:0": "mock-output-0", "node:1": "mock-output-1"},
+    )
+    assert res == {"0": "mock-output-0", "1": "mock-output-1"}
+
+    res = postprocessor._postprocess_recursively(
+        "mock",
+        **{"node:0": "mock-output-0", "node:1": "mock-output-1"},
+    )
+    assert res == "mock"
diff --git a/tests/covalent_tests/workflow/transport_test.py b/tests/covalent_tests/workflow/transport_test.py
index c2c4533e..e8361eeb 100644
--- a/tests/covalent_tests/workflow/transport_test.py
+++ b/tests/covalent_tests/workflow/transport_test.py
@@ -217,16 +217,25 @@ def test_transport_graph_initialization():
     assert not tg.lattice_metadata
 
 
-def test_transport_graph_add_nodes(transport_graph):
+@pytest.mark.parametrize("task_group_id, expected_task_group_id", [(None, 0), (1984, 1984)])
+def test_transport_graph_add_nodes(transport_graph, task_group_id, expected_task_group_id):
     """Test addition of nodes (electrons) to the transport graph."""
 
     tg = transport_graph
     assert len(tg._graph.nodes) == 0
     node_id = tg.add_node(
-        name="square", kwargs={"x": 2}, function=subtask, metadata={"mock_field": "mock_value"}
+        name="square",
+        kwargs={"x": 2},
+        function=subtask,
+        metadata={"mock_field": "mock_value"},
+        task_group_id=task_group_id,
     )
     assert len(tg._graph.nodes) == 1
     assert node_id == 0
+    node_metadata = tg._graph.nodes(data=True)[0]
+    assert node_metadata["name"] == "square"
+    assert node_metadata["kwargs"] == {"x": 2}
+    assert node_metadata["task_group_id"] == expected_task_group_id
 
 
 def test_transport_graph_get_and_set_edges(workflow_transport_graph):