Repository: AgnostiqHQ/covalent
Pull request #1502 (2023-01-30 14:52:05)

Transport graph diff functionality
<!--
⚠️ 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/1502?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) Report
> Merging [#1502](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1502?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (cd95668) into [develop](https://codecov.io/gh/AgnostiqHQ/covalent/commit/3f7330e52b9136d6fa368a18152b31be9949bda3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ) (3f7330e) will **decrease** coverage by `3.16%`.
> The diff coverage is `15.10%`.

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


[![Impacted file tree graph](https://codecov.io/gh/AgnostiqHQ/covalent/pull/1502/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/1502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=AgnostiqHQ)

```diff
@@             Coverage Diff             @@
##           develop    #1502      +/-   ##
===========================================
- Coverage    84.52%   81.36%   -3.16%     
===========================================
  Files          144       70      -74     
  Lines         5743     4235    -1508     
  Branches       145        0     -145     
===========================================
- Hits          4854     3446    -1408     
- Misses         787      789       +2     
+ Partials       102        0     -102     
```

| Flag | Coverage Δ | |
|---|---|---|
| Dispatcher | `86.88% <33.33%> (-1.31%)` | :arrow_down: |
| Functional_Tests | `?` | |
| SDK | `77.85% <10.45%> (-4.31%)` | :arrow_down: |
| UI_Frontend | `?` | |

</details>

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b2cb30e8..6cec01d4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Added
+
+- TransportGraphOps class for diffing operations on transport graphs.
+- Added make derived dispatch method.
+- Apply electron updates method to _TransportGraph.
+
 ### Operations
 
 - Added job in `nightly` to trigger base executor image builds after a Covalent `pre-release`
diff --git a/covalent/_results_manager/result.py b/covalent/_results_manager/result.py
index 0f8c065c..cee43d31 100644
--- a/covalent/_results_manager/result.py
+++ b/covalent/_results_manager/result.py
@@ -235,24 +235,7 @@ Node Outputs
 
         self._num_nodes = self.lattice.transport_graph.get_internal_graph_copy().number_of_nodes()
         for node_id in range(self._num_nodes):
-
-            self.lattice.transport_graph.set_node_value(node_id, "start_time", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "end_time", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "status", Result.NEW_OBJ)
-
-            self.lattice.transport_graph.set_node_value(node_id, "output", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "error", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "sub_dispatch_id", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "sublattice_result", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "stdout", None)
-
-            self.lattice.transport_graph.set_node_value(node_id, "stderr", None)
+            self.lattice.transport_graph.reset_node(node_id)
 
     def get_node_result(self, node_id: int) -> dict:
         """Return the result of a particular node.
diff --git a/covalent/_workflow/transport.py b/covalent/_workflow/transport.py
index df2f5739..f6888754 100644
--- a/covalent/_workflow/transport.py
+++ b/covalent/_workflow/transport.py
@@ -30,6 +30,7 @@ import cloudpickle
 import networkx as nx
 
 from .._shared_files.defaults import parameter_prefix
+from .._shared_files.util_classes import RESULT_STATUS
 
 
 class TransportableObject:
@@ -299,6 +300,18 @@ class _TransportGraph:
         # IDs of nodes modified during the workflow run
         self.dirty_nodes = []
 
+        self._default_node_attrs = {
+            "start_time": None,
+            "end_time": None,
+            "status": RESULT_STATUS.NEW_OBJECT,
+            "output": None,
+            "error": None,
+            "sub_dispatch_id": None,
+            "sublattice_result": None,
+            "stdout": None,
+            "stderr": None,
+        }
+
     def add_node(self, name: str, function: Callable, metadata: Dict, **attr) -> int:
         """
         Adds a node to the graph.
@@ -439,6 +452,40 @@ class _TransportGraph:
 
         return self._graph.copy()
 
+    def reset_node(self, node_id: int) -> None:
+        """Reset node values to starting state."""
+        for node_attr, default_val in self._default_node_attrs.items():
+            self.set_node_value(node_id, node_attr, default_val)
+
+    def _replace_node(self, node_id: int, new_attrs: Dict[str, Any]) -> None:
+        """Replace node data with new attribute values and flag descendants (used in re-dispatching)."""
+        metadata = self.get_node_value(node_id, "metadata")
+        metadata.update(new_attrs["metadata"])
+
+        serialized_callable = TransportableObject.from_dict(new_attrs["function"])
+        self.set_node_value(node_id, "function", serialized_callable)
+        self.set_node_value(node_id, "function_string", new_attrs["function_string"])
+        self.set_node_value(node_id, "name", new_attrs["name"])
+        self._reset_descendants(node_id)
+
+    def _reset_descendants(self, node_id: int) -> None:
+        """Reset node and all its descendants to starting state."""
+        try:
+            if self.get_node_value(node_id, "status") == RESULT_STATUS.NEW_OBJECT:
+                return
+        except Exception:
+            return
+        self.reset_node(node_id)
+        for successor in self._graph.neighbors(node_id):
+            self._reset_descendants(successor)
+
+    def apply_electron_updates(self, electron_updates: Dict[str, Callable]) -> None:
+        """Replace transport graph node data based on the electrons that need to be updated during re-dispatching."""
+        for n in self._graph.nodes:
+            name = self.get_node_value(n, "name")
+            if name in electron_updates:
+                self._replace_node(n, electron_updates[name])
+
     def serialize(self, metadata_only: bool = False) -> bytes:
         """
         Convert transport graph object to JSON to be used in the workflow scheduler.
@@ -501,7 +548,7 @@ class _TransportGraph:
         Returns:
             str: json string representation of transport graph
 
-        Note: serialize_to_json converts metadata objects into dictionary represetations.
+        Note: serialize_to_json converts metadata objects into dictionary representations.
         """
 
         # Convert networkx.DiGraph to a format that can be converted to json .
diff --git a/covalent/_workflow/transport_graph_ops.py b/covalent/_workflow/transport_graph_ops.py
new file mode 100644
index 00000000..ea895505
--- /dev/null
+++ b/covalent/_workflow/transport_graph_ops.py
@@ -0,0 +1,204 @@
+# Copyright 2023 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+"""Module for transport graph operations."""
+
+from collections import deque
+from typing import Callable, List
+
+import networkx as nx
+
+from .._shared_files import logger
+from .transport import _TransportGraph
+
+app_log = logger.app_log
+
+
+class TransportGraphOps:
+    def __init__(self, tg):
+        self.tg = tg
+        self._status_map = {1: True, -1: False}
+
+    @staticmethod
+    def _flag_successors(A: nx.MultiDiGraph, node_statuses: dict, starting_node: int):
+        """Flag all successors of a node (including the node itself)."""
+        nodes_to_invalidate = [starting_node]
+        for node, successors in nx.bfs_successors(A, starting_node):
+            nodes_to_invalidate.extend(iter(successors))
+        for node in nodes_to_invalidate:
+            node_statuses[node] = -1
+
+    @staticmethod
+    def is_same_node(A: nx.MultiDiGraph, B: nx.MultiDiGraph, node: int) -> bool:
+        """Check if the node attributes are the same in both graphs."""
+        return A.nodes[node] == B.nodes[node]
+
+    @staticmethod
+    def is_same_edge_attributes(
+        A: nx.MultiDiGraph, B: nx.MultiDiGraph, parent: int, node: int
+    ) -> bool:
+        """Check if the edge attributes are the same in both graphs."""
+        return A.adj[parent][node] == B.adj[parent][node]
+
+    def copy_nodes_from(self, tg: _TransportGraph, nodes):
+        """Copy nodes from the transport graph in the argument."""
+        for n in nodes:
+            for k, v in tg._graph.nodes[n].items():
+                self.tg.set_node_value(n, k, v)
+
+    @staticmethod
+    def _cmp_name_and_pval(A: nx.MultiDiGraph, B: nx.MultiDiGraph, node: int) -> bool:
+        """Default node comparison function for diffing transport graphs."""
+        name_A = A.nodes[node]["name"]
+        name_B = B.nodes[node]["name"]
+
+        if name_A != name_B:
+            return False
+
+        val_A = A.nodes[node].get("value", None)
+        val_B = B.nodes[node].get("value", None)
+
+        return val_A == val_B
+
+    def _max_cbms(
+        self,
+        A: nx.MultiDiGraph,
+        B: nx.MultiDiGraph,
+        node_cmp: Callable = None,
+        edge_cmp: Callable = None,
+    ):
+        """Computes a "maximum backward-maximal common subgraph" (cbms)
+        Args:
+            A: nx.MultiDiGraph
+            B: nx.MultiDiGraph
+            node_cmp: An optional function for comparing node attributes in A and B.
+                    Defaults to testing for equality of the attribute dictionaries
+            edge_cmp: An optional function for comparing the edges between two nodes.
+                    Defaults to checking that the two sets of edges have the same attributes
+        Returns: A_node_status, B_node_status, where each is a dictionary
+            `{node: True/False}` where True means reusable.
+        Performs a modified BFS of A and B.
+        """
+        if node_cmp is None:
+            node_cmp = self.is_same_node
+        if edge_cmp is None:
+            edge_cmp = self.is_same_edge_attributes
+
+        A_node_status = {node_id: 0 for node_id in A.nodes}
+        B_node_status = {node_id: 0 for node_id in B.nodes}
+        app_log.debug(f"A node status: {A_node_status}")
+        app_log.debug(f"B node status: {B_node_status}")
+
+        virtual_root = -1
+
+        if virtual_root in A.nodes or virtual_root in B.nodes:
+            raise RuntimeError(f"Encountered forbidden node: {virtual_root}")
+
+        assert virtual_root not in B.nodes
+
+        nodes_to_visit = deque()
+        nodes_to_visit.appendleft(virtual_root)
+
+        # Add a temporary root
+        A_parentless_nodes = [node for node, deg in A.in_degree() if deg == 0]
+        B_parentless_nodes = [node for node, deg in B.in_degree() if deg == 0]
+        for node_id in A_parentless_nodes:
+            A.add_edge(virtual_root, node_id)
+
+        for node_id in B_parentless_nodes:
+            B.add_edge(virtual_root, node_id)
+
+        # Assume inductively that predecessors subgraphs are the same;
+        # this is satisfied for the root
+        while nodes_to_visit:
+            current_node = nodes_to_visit.pop()
+
+            app_log.debug(f"Visiting node {current_node}")
+            for y in A.adj[current_node]:
+                # Don't process already failed nodes
+                if A_node_status[y] == -1:
+                    continue
+
+                # Check if y is a valid child of current_node in B
+                if y not in B.adj[current_node]:
+                    app_log.debug(f"A: {y} not adjacent to node {current_node} in B")
+                    self._flag_successors(A, A_node_status, y)
+                    continue
+
+                if y in B.adj[current_node] and B_node_status[y] == -1:
+                    app_log.debug(f"A: Node {y} is marked as failed in B")
+                    self._flag_successors(A, A_node_status, y)
+                    continue
+
+                # Compare edges
+                if not edge_cmp(A, B, current_node, y):
+                    app_log.debug(f"Edges between {current_node} and {y} differ")
+                    self._flag_successors(A, A_node_status, y)
+                    self._flag_successors(B, B_node_status, y)
+                    continue
+
+                # Compare nodes
+                if not node_cmp(A, B, y):
+                    app_log.debug(f"Attributes of node {y} differ:")
+                    app_log.debug(f"A[y] = {A.nodes[y]}")
+                    app_log.debug(f"B[y] = {B.nodes[y]}")
+                    self._flag_successors(A, A_node_status, y)
+                    self._flag_successors(B, B_node_status, y)
+                    continue
+
+                # Predecessors subgraphs of y are the same in A and B, so
+                # enqueue y if it hasn't already been visited
+                assert A_node_status[y] != -1
+                if A_node_status[y] == 0:
+                    A_node_status[y] = 1
+                    B_node_status[y] = 1
+                    app_log.debug(f"Enqueueing node {y}")
+                    nodes_to_visit.appendleft(y)
+
+            # Prune children of current_node in B that aren't valid children in A
+            for y in B.adj[current_node]:
+                if B_node_status[y] == -1:
+                    continue
+                if y not in A.adj[current_node]:
+                    app_log.debug(f"B: {y} not adjacent to node {current_node} in A")
+                    self._flag_successors(B, B_node_status, y)
+                    continue
+                if y in A.adj[current_node] and B_node_status[y] == -1:
+                    app_log.debug(f"B: Node {y} is marked as failed in A")
+                    self._flag_successors(B, B_node_status, y)
+
+        A.remove_node(-1)
+        B.remove_node(-1)
+
+        app_log.debug(f"A node status: {A_node_status}")
+        app_log.debug(f"B node status: {B_node_status}")
+
+        for k, v in A_node_status.items():
+            A_node_status[k] = self._status_map[v]
+        for k, v in B_node_status.items():
+            B_node_status[k] = self._status_map[v]
+        return A_node_status, B_node_status
+
+    def get_reusable_nodes(self, tg_new: _TransportGraph) -> List[int]:
+        """Find which nodes are common between the current graph and a new graph."""
+        A = self.tg.get_internal_graph_copy()
+        B = tg_new.get_internal_graph_copy()
+        status_A, _ = self._max_cbms(A, B, node_cmp=self._cmp_name_and_pval)
+        return [k for k, v in status_A.items() if v]
diff --git a/covalent_dispatcher/_core/data_manager.py b/covalent_dispatcher/_core/data_manager.py
index b9261e17..9bd69d79 100644
--- a/covalent_dispatcher/_core/data_manager.py
+++ b/covalent_dispatcher/_core/data_manager.py
@@ -24,12 +24,14 @@ Defines the core functionality of the result service
 
 import asyncio
 import uuid
+from typing import Callable, Dict, Optional
 
 from covalent._results_manager import Result
 from covalent._shared_files import logger
 from covalent._workflow.lattice import Lattice
+from covalent._workflow.transport_graph_ops import TransportGraphOps
 
-from .._db import update, upsert
+from .._db import load, update, upsert
 from .._db.write_result_to_db import resolve_electron_id
 
 app_log = logger.app_log
@@ -144,6 +146,67 @@ def make_dispatch(
     return result_object.dispatch_id
 
 
+def _get_result_object_from_new_lattice(
+    json_lattice: str, old_result_object: Result, reuse_previous_results: bool
+) -> Result:
+    """Get new result object for re-dispatching from new lattice json."""
+    lat = Lattice.deserialize_from_json(json_lattice)
+    result_object = Result(lat, get_unique_id())
+    result_object._initialize_nodes()
+
+    if reuse_previous_results:
+        tg = result_object.lattice.transport_graph
+        tg_old = old_result_object.lattice.transport_graph
+        reusable_nodes = TransportGraphOps(tg_old).get_reusable_nodes(tg)
+        TransportGraphOps(tg_old).copy_nodes(tg, reusable_nodes)
+
+    return result_object
+
+
+def _get_result_object_from_old_result(
+    old_result_object: Result, reuse_previous_results: bool
+) -> Result:
+    """Get new result object for re-dispatching from old result object."""
+    result_object = Result(old_result_object.lattice, get_unique_id())
+    result_object._num_nodes = old_result_object._num_nodes
+
+    if not reuse_previous_results:
+        result_object._initialize_nodes()
+
+    return result_object
+
+
+def make_derived_dispatch(
+    parent_dispatch_id: str,
+    json_lattice: Optional[str] = None,
+    electron_updates: Optional[Dict[str, Callable]] = None,
+    reuse_previous_results: bool = False,
+) -> str:
+    """Make a re-dispatch from a previous dispatch."""
+    if electron_updates is None:
+        electron_updates = {}
+
+    old_result_object = load.get_result_object_from_storage(parent_dispatch_id, wait=False)
+
+    if json_lattice:
+        result_object = _get_result_object_from_new_lattice(
+            json_lattice, old_result_object, reuse_previous_results
+        )
+    else:
+        result_object = _get_result_object_from_old_result(
+            old_result_object, reuse_previous_results
+        )
+
+    result_object.lattice.transport_graph.apply_electron_updates(electron_updates)
+    result_object.lattice.transport_graph.dirty_nodes = list(
+        result_object.lattice.transport_graph._graph.nodes
+    )
+    update.persist(result_object)
+    _register_result_object(result_object)
+
+    return result_object.dispatch_id
+
+
 def get_result_object(dispatch_id: str) -> Result:
     return _registered_dispatches[dispatch_id]
 
@@ -170,9 +233,7 @@ async def persist_result(dispatch_id: str):
 
 
 async def _update_parent_electron(result_object: Result):
-    parent_eid = result_object._electron_id
-
-    if parent_eid:
+    if parent_eid := result_object._electron_id:
         dispatch_id, node_id = resolve_electron_id(parent_eid)
         status = result_object.status
         if status == Result.POSTPROCESSING_FAILED:
diff --git a/covalent_dispatcher/_db/load.py b/covalent_dispatcher/_db/load.py
new file mode 100644
index 00000000..6bf6dd72
--- /dev/null
+++ b/covalent_dispatcher/_db/load.py
@@ -0,0 +1,43 @@
+# Copyright 2023 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+"""Functions to load results from the database."""
+
+from typing import Optional
+
+from covalent._results_manager.result import Result
+from covalent._shared_files import logger
+
+from .models import Lattice
+
+app_log = logger.app_log
+log_stack_info = logger.log_stack_info
+
+
+def _result_from(lattice_record: Lattice) -> Result:
+    """Re-hydrate result object from the lattice record."""
+    pass
+
+
+def get_result_object_from_storage(
+    dispatch_id: str, wait: Optional[bool] = False, status_only: Optional[bool] = False
+):
+    """Get result object from Database."""
+    pass


Test patch:
diff --git a/tests/covalent_dispatcher_tests/_core/data_manager_test.py b/tests/covalent_dispatcher_tests/_core/data_manager_test.py
index bf71899d..54481a57 100644
--- a/tests/covalent_dispatcher_tests/_core/data_manager_test.py
+++ b/tests/covalent_dispatcher_tests/_core/data_manager_test.py
@@ -23,7 +23,7 @@ Tests for the core functionality of the dispatcher.
 """
 
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -32,6 +32,8 @@ from covalent._results_manager import Result
 from covalent._workflow.lattice import Lattice
 from covalent_dispatcher._core.data_manager import (
     _dispatch_status_queues,
+    _get_result_object_from_new_lattice,
+    _get_result_object_from_old_result,
     _register_result_object,
     _registered_dispatches,
     _update_parent_electron,
@@ -39,6 +41,7 @@ from covalent_dispatcher._core.data_manager import (
     get_result_object,
     get_status_queue,
     initialize_result_object,
+    make_derived_dispatch,
     make_dispatch,
     persist_result,
     update_node_result,
@@ -162,6 +165,138 @@ def test_make_dispatch(mocker):
     mock_register.assert_called_with(res)
 
 
+@pytest.mark.parametrize("reuse", [True, False])
+def test_get_result_object_from_new_lattice(mocker, reuse):
+    """Test the get result object from new lattice json function."""
+    lattice_mock = mocker.patch("covalent_dispatcher._core.data_manager.Lattice")
+    result_object_mock = mocker.patch("covalent_dispatcher._core.data_manager.Result")
+    transport_graph_ops_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager.TransportGraphOps"
+    )
+    old_result_mock = MagicMock()
+    res = _get_result_object_from_new_lattice(
+        json_lattice="mock-lattice",
+        old_result_object=old_result_mock,
+        reuse_previous_results=reuse,
+    )
+    assert res == result_object_mock.return_value
+    lattice_mock.deserialize_from_json.assert_called_with("mock-lattice")
+    result_object_mock()._initialize_nodes.assert_called_with()
+
+    if reuse:
+        transport_graph_ops_mock().get_reusable_nodes.assert_called_with(
+            result_object_mock().lattice.transport_graph
+        )
+        transport_graph_ops_mock().copy_nodes.assert_called_once_with(
+            result_object_mock().lattice.transport_graph,
+            transport_graph_ops_mock().get_reusable_nodes.return_value,
+        )
+
+    else:
+        transport_graph_ops_mock().get_reusable_nodes.assert_not_called()
+        transport_graph_ops_mock().copy_nodes.assert_not_called()
+
+
+@pytest.mark.parametrize("reuse", [True, False])
+def test_get_result_object_from_old_result(mocker, reuse):
+    """Test the get result object from old result function."""
+    result_object_mock = mocker.patch("covalent_dispatcher._core.data_manager.Result")
+    old_result_mock = MagicMock()
+    res = _get_result_object_from_old_result(
+        old_result_object=old_result_mock,
+        reuse_previous_results=reuse,
+    )
+    assert res == result_object_mock.return_value
+
+    if reuse:
+        result_object_mock()._initialize_nodes.assert_not_called()
+    else:
+        result_object_mock()._initialize_nodes.assert_called_with()
+
+    assert res._num_nodes == old_result_mock._num_nodes
+
+
+@pytest.mark.parametrize("reuse", [True, False])
+def test_make_derived_dispatch_from_lattice(mocker, reuse):
+    """Test the make derived dispatch function."""
+
+    def mock_func():
+        pass
+
+    mock_old_result = MagicMock()
+    mock_new_result = MagicMock()
+    mock_new_result.dispatch_id = "mock-redispatch-id"
+    mock_new_result.lattice.transport_graph._graph.nodes = ["mock-nodes"]
+    load_get_result_object_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager.load", return_value=mock_old_result
+    )
+    get_result_object_from_new_lattice_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._get_result_object_from_new_lattice",
+        return_value=mock_new_result,
+    )
+    get_result_object_from_old_result_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._get_result_object_from_old_result"
+    )
+    update_mock = mocker.patch("covalent_dispatcher._core.data_manager.update")
+    register_result_object_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._register_result_object"
+    )
+    mock_electron_updates = {"mock-electron-id": mock_func}
+    redispatch_id = make_derived_dispatch(
+        parent_dispatch_id="mock-dispatch-id",
+        json_lattice="mock-json-lattice",
+        electron_updates=mock_electron_updates,
+        reuse_previous_results=reuse,
+    )
+    load_get_result_object_mock.called_once_with("mock-dispatch-id", wait=reuse)
+    get_result_object_from_new_lattice_mock.called_once_with(
+        "mock-json-lattice", mock_old_result, reuse
+    )
+    get_result_object_from_old_result_mock.assert_not_called()
+    mock_new_result.lattice.transport_graph.apply_electron_updates.assert_called_once_with(
+        mock_electron_updates
+    )
+    update_mock().persist.called_once_with(mock_new_result)
+    register_result_object_mock.assert_called_once_with(mock_new_result)
+    assert redispatch_id == "mock-redispatch-id"
+    assert mock_new_result.lattice.transport_graph.dirty_nodes == ["mock-nodes"]
+
+
+@pytest.mark.parametrize("reuse", [True, False])
+def test_make_derived_dispatch_from_old_result(mocker, reuse):
+    """Test the make derived dispatch function."""
+    mock_old_result = MagicMock()
+    mock_new_result = MagicMock()
+    mock_new_result.dispatch_id = "mock-redispatch-id"
+    mock_new_result.lattice.transport_graph._graph.nodes = ["mock-nodes"]
+    load_get_result_object_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager.load", return_value=mock_old_result
+    )
+    get_result_object_from_new_lattice_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._get_result_object_from_new_lattice",
+    )
+    get_result_object_from_old_result_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._get_result_object_from_old_result",
+        return_value=mock_new_result,
+    )
+    update_mock = mocker.patch("covalent_dispatcher._core.data_manager.update")
+    register_result_object_mock = mocker.patch(
+        "covalent_dispatcher._core.data_manager._register_result_object"
+    )
+    redispatch_id = make_derived_dispatch(
+        parent_dispatch_id="mock-dispatch-id",
+        reuse_previous_results=reuse,
+    )
+    load_get_result_object_mock.called_once_with("mock-dispatch-id", wait=reuse)
+    get_result_object_from_new_lattice_mock.assert_not_called()
+    get_result_object_from_old_result_mock.called_once_with(mock_old_result, reuse)
+    mock_new_result.lattice.transport_graph.apply_electron_updates.assert_called_once_with({})
+    update_mock().persist.called_once_with(mock_new_result)
+    register_result_object_mock.assert_called_once_with(mock_new_result)
+    assert redispatch_id == "mock-redispatch-id"
+    assert mock_new_result.lattice.transport_graph.dirty_nodes == ["mock-nodes"]
+
+
 def test_get_result_object(mocker):
     result_object = get_mock_result()
     dispatch_id = result_object.dispatch_id
diff --git a/tests/covalent_tests/workflow/transport_graph_ops_test.py b/tests/covalent_tests/workflow/transport_graph_ops_test.py
new file mode 100644
index 00000000..20bb758e
--- /dev/null
+++ b/tests/covalent_tests/workflow/transport_graph_ops_test.py
@@ -0,0 +1,252 @@
+# Copyright 2023 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+"""Unit tests for transport graph operations module."""
+
+import pytest
+
+from covalent._workflow.transport import _TransportGraph
+from covalent._workflow.transport_graph_ops import TransportGraphOps
+
+
+def add(x, y):
+    return x + y
+
+
+def multiply(x, y):
+    return x * y
+
+
+def identity(x):
+    return x
+
+
+@pytest.fixture
+def tg():
+    """Transport graph operations fixture."""
+    tg = _TransportGraph()
+    tg.add_node(name="add", function=add, metadata={"0-mock-key": "0-mock-value"})
+    tg.add_node(name="multiply", function=multiply, metadata={"1-mock-key": "1-mock-value"})
+    tg.add_node(name="identity", function=identity, metadata={"2-mock-key": "2-mock-value"})
+    return tg
+
+
+@pytest.fixture
+def tg_2():
+    """Transport graph operations fixture - different from tg."""
+    tg_2 = _TransportGraph()
+    tg_2.add_node(name="not-add", function=add, metadata={"0- mock-key": "0-mock-value"})
+    tg_2.add_node(name="multiply", function=multiply, metadata={"1- mock-key": "1-mock-value"})
+    tg_2.add_node(name="identity", function=identity, metadata={"2- mock-key": "2-mock-value"})
+    return tg_2
+
+
+@pytest.fixture
+def tg_ops(tg):
+    """Transport graph operations fixture."""
+    return TransportGraphOps(tg)
+
+
+def test_init(tg):
+    """Test initialization of transport graph operations."""
+    tg_ops = TransportGraphOps(tg)
+    assert tg_ops.tg == tg
+    assert tg_ops._status_map == {1: True, -1: False}
+
+
+def test_flag_successors_no_successors(tg, tg_ops):
+    """Test flagging successors of a node."""
+    node_statuses = {0: 1, 1: 1, 2: 1}
+    tg_ops._flag_successors(tg._graph, node_statuses=node_statuses, starting_node=0)
+    assert node_statuses == {0: -1, 1: 1, 2: 1}
+
+
+@pytest.mark.parametrize(
+    "n_1,n_2,n_start,label,new_statuses",
+    [
+        (0, 1, 0, "01", {0: -1, 1: -1, 2: 1}),
+        (1, 2, 0, "12", {0: -1, 1: 1, 2: 1}),
+        (1, 2, 1, "12", {0: 1, 1: -1, 2: -1}),
+        (1, 2, 2, "12", {0: 1, 1: 1, 2: -1}),
+    ],
+)
+def test_flag_successors_with_one_successors(tg, tg_ops, n_1, n_2, n_start, label, new_statuses):
+    """Test flagging successors of a node."""
+    tg.add_edge(n_1, n_2, label)
+    node_statuses = {0: 1, 1: 1, 2: 1}
+    tg_ops._flag_successors(tg._graph, node_statuses=node_statuses, starting_node=n_start)
+    assert node_statuses == new_statuses
+
+
+@pytest.mark.parametrize(
+    "n_1,n_2,n_3,n_4,label_1,label_2,n_start,new_statuses",
+    [
+        (0, 1, 1, 2, "01", "12", 0, {0: -1, 1: -1, 2: -1}),
+        (0, 1, 0, 2, "01", "02", 0, {0: -1, 1: -1, 2: -1}),
+        (0, 1, 0, 2, "01", "12", 1, {0: 1, 1: -1, 2: 1}),
+    ],
+)
+def test_flag_successors_with_successors_3(
+    tg, tg_ops, n_1, n_2, n_3, n_4, label_1, n_start, label_2, new_statuses
+):
+    """Test flagging successors of a node."""
+    tg.add_edge(n_1, n_2, label_1)
+    tg.add_edge(n_3, n_4, label_2)
+    node_statuses = {0: 1, 1: 1, 2: 1}
+    tg_ops._flag_successors(tg._graph, node_statuses=node_statuses, starting_node=n_start)
+    assert node_statuses == new_statuses
+
+
+def test_is_same_node_true(tg, tg_ops):
+    """Test the is same node method."""
+    assert tg_ops.is_same_node(tg._graph, tg._graph, 0) is True
+    assert tg_ops.is_same_node(tg._graph, tg._graph, 1) is True
+
+
+def test_is_same_node_false(tg, tg_ops):
+    """Test the is same node method."""
+    tg_2 = _TransportGraph()
+    tg_2.add_node(name="multiply", function=add, metadata={"0- mock-key": "0-mock-value"})
+    assert tg_ops.is_same_node(tg._graph, tg_2._graph, 0) is False
+
+
+def test_is_same_edge_attributes_true(tg, tg_ops):
+    """Test the is same edge attributes method."""
+    tg.add_edge(0, 1, edge_name="01", kwargs={"x": 1, "y": 2})
+    assert tg_ops.is_same_edge_attributes(tg._graph, tg._graph, 0, 1) is True
+
+
+def test_is_same_edge_attributes_false(tg, tg_ops):
+    """Test the is same edge attributes method."""
+    tg.add_edge(0, 1, edge_name="01", kwargs={"x": 1, "y": 2})
+
+    tg_2 = _TransportGraph()
+    tg_2.add_node(name="add", function=add, metadata={"0- mock-key": "0-mock-value"})
+    tg_2.add_node(name="multiply", function=multiply, metadata={"1- mock-key": "1-mock-value"})
+    tg_2.add_node(name="identity", function=identity, metadata={"2- mock-key": "2-mock-value"})
+    tg_2.add_edge(0, 1, edge_name="01", kwargs={"x": 1})
+
+    assert tg_ops.is_same_edge_attributes(tg._graph, tg_2._graph, 0, 1) is False
+
+
+def test_copy_nodes(tg_ops):
+    """Test the node copying method."""
+
+    def replacement(x):
+        return x + 1
+
+    tg_new = _TransportGraph()
+    tg_new.add_node(
+        name="replacement", function=replacement, metadata={"0-mock-key": "0-mock-value"}
+    )
+    tg_new.add_node(name="multiply", function=multiply, metadata={"1-mock-key": "1-mock-value"})
+    tg_new.add_node(
+        name="replacement", function=replacement, metadata={"2-mock-key": "2-mock-value"}
+    )
+
+    tg_ops.copy_nodes_from(tg_new, [0, 2])
+    tg_ops.tg._graph.nodes(data=True)[0]["name"] == tg_ops.tg._graph.nodes(data=True)[2][
+        "name"
+    ] == "replacement"
+    tg_ops.tg._graph.nodes(data=True)[2]["name"] == "multiply"
+
+
+def test_max_cbms(tg_ops):
+    """Test method for determining a largest cbms"""
+    import networkx as nx
+
+    A = nx.MultiDiGraph()
+    B = nx.MultiDiGraph()
+    C = nx.MultiDiGraph()
+    D = nx.MultiDiGraph()
+
+    #    0     5  6
+    #  /  \
+    # 1    2
+    A.add_edge(0, 1)
+    A.add_edge(0, 2)
+    A.nodes[1]["color"] = "red"
+    A.add_node(5)
+    A.add_node(6)
+
+    #    0     5
+    #  /  \\
+    # 1    2
+    B.add_edge(0, 1)
+    B.add_edge(0, 2)
+    B.add_edge(0, 2)
+    B.nodes[1]["color"] = "black"
+    B.add_node(5)
+
+    #    0    3
+    #  /  \  /
+    # 1    2
+    C.add_edge(0, 1)
+    C.add_edge(0, 2)
+    C.add_edge(3, 2)
+
+    #    0    3
+    #  /  \  /
+    # 1    2
+    #     /
+    #   4
+    D.add_edge(0, 1)
+    D.add_edge(0, 2)
+    D.add_edge(3, 2)
+    D.add_edge(2, 4)
+
+    A_node_status, B_node_status = tg_ops._max_cbms(A, B)
+    assert A_node_status == {0: True, 1: False, 2: False, 5: True, 6: False}
+    assert B_node_status == {0: True, 1: False, 2: False, 5: True}
+
+    A_node_status, C_node_status = tg_ops._max_cbms(A, C)
+    assert A_node_status == {0: True, 1: False, 2: False, 5: False, 6: False}
+    assert C_node_status == {0: True, 1: False, 2: False, 3: False}
+
+    C_node_status, D_node_status = tg_ops._max_cbms(C, D)
+    assert C_node_status == {0: True, 1: True, 2: True, 3: True}
+    assert D_node_status == {0: True, 1: True, 2: True, 3: True, 4: False}
+
+
+def test_cmp_name_and_pval_true(tg, tg_ops):
+    """Test the name and parameter value comparison method."""
+    assert tg_ops._cmp_name_and_pval(tg._graph, tg._graph, 0) is True
+
+
+def test_cmp_name_and_pval_false(tg, tg_2, tg_ops):
+    """Test the name and parameter value comparison method."""
+    assert tg_ops._cmp_name_and_pval(tg._graph, tg_2._graph, 0) is False
+
+
+def test_get_reusable_nodes(mocker, tg, tg_2, tg_ops):
+    """Test the get reusable nodes method."""
+    max_cbms_mock = mocker.patch(
+        "covalent._workflow.transport_graph_ops.TransportGraphOps._max_cbms",
+        return_value=({"mock-key-A": "mock-value-A"}, {"mock-key-B": "mock-value-B"}),
+    )
+    reusable_nodes = tg_ops.get_reusable_nodes(tg_2)
+    assert reusable_nodes == ["mock-key-A"]
+    max_cbms_mock.assert_called_once()
+
+
+def test_get_diff_nodes_integration_test(tg_2, tg_ops):
+    """Test the get reusable nodes method."""
+    reusable_nodes = tg_ops.get_reusable_nodes(tg_2)
+    assert reusable_nodes == [1, 2]
diff --git a/tests/covalent_tests/workflow/transport_test.py b/tests/covalent_tests/workflow/transport_test.py
index dae28af8..c99f363a 100644
--- a/tests/covalent_tests/workflow/transport_test.py
+++ b/tests/covalent_tests/workflow/transport_test.py
@@ -21,6 +21,7 @@
 """Unit tests for transport graph."""
 
 import platform
+from unittest.mock import call
 
 import cloudpickle
 import networkx as nx
@@ -28,6 +29,7 @@ import pytest
 
 import covalent as ct
 from covalent._shared_files.defaults import parameter_prefix
+from covalent._shared_files.util_classes import RESULT_STATUS
 from covalent._workflow.transport import TransportableObject, _TransportGraph, encode_metadata
 from covalent.executor import LocalExecutor
 
@@ -425,3 +427,116 @@ def test_encode_metadata():
 
     # Check idempotence
     assert encode_metadata(metadata) == encode_metadata(encode_metadata(metadata))
+
+
+def test_reset_node(workflow_transport_graph, mocker):
+    """Test the node reset method."""
+    set_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.set_node_value"
+    )
+
+    node_id = 0
+    workflow_transport_graph.reset_node(node_id)
+    actual_mock_calls = set_node_value_mock.mock_calls
+    expected_mock_calls = [
+        call(node_id, node_attr, default_val)
+        for node_attr, default_val in workflow_transport_graph._default_node_attrs.items()
+    ]
+
+    for mock_call in expected_mock_calls:
+        assert mock_call in actual_mock_calls
+
+
+def test_replace_node(workflow_transport_graph, mocker):
+    """Test method that replaces node attribute values."""
+    transportable_object_from_dict_mock = mocker.patch(
+        "covalent._workflow.transport.TransportableObject.from_dict", return_value="mock-func"
+    )
+    reset_descendants_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph._reset_descendants"
+    )
+    set_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.set_node_value"
+    )
+    node_id = 0
+    new_attrs = {
+        "metadata": {"mock-key": "mock-value"},
+        "function": "mock-func",
+        "function_string": "mock-func-string",
+        "name": "mock-name",
+    }
+    workflow_transport_graph._replace_node(node_id, new_attrs)
+    transportable_object_from_dict_mock.assert_called_once_with("mock-func")
+    reset_descendants_mock.assert_called_once_with(node_id)
+    expected_set_node_value_mock_calls = [
+        call(0, "function", "mock-func"),
+        call(0, "function_string", "mock-func-string"),
+        call(0, "name", "mock-name"),
+    ]
+    actual_set_node_value_mock_calls = set_node_value_mock.mock_calls
+    for mock_call in expected_set_node_value_mock_calls:
+        assert mock_call in actual_set_node_value_mock_calls
+
+
+def test_reset_descendants_new_object(workflow_transport_graph, mocker):
+    """Test method that resets descendants of a node (including the node itself)."""
+    get_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.get_node_value",
+        return_value=RESULT_STATUS.NEW_OBJECT,
+    )
+    reset_node_mock = mocker.patch("covalent._workflow.transport._TransportGraph.reset_node")
+    res = workflow_transport_graph._reset_descendants(0)
+    reset_node_mock.assert_not_called()
+    get_node_value_mock.assert_called_once_with(0, "status")
+    assert res is None
+
+
+def test_reset_descendants_exception(workflow_transport_graph, mocker):
+    """Test method that resets descendants of a node (including the node itself) when there's an exception."""
+    get_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.get_node_value",
+        side_effect=Exception(),
+    )
+    reset_node_mock = mocker.patch("covalent._workflow.transport._TransportGraph.reset_node")
+    res = workflow_transport_graph._reset_descendants(12)
+    reset_node_mock.assert_not_called()
+    get_node_value_mock.assert_called_once_with(12, "status")
+    assert res is None
+
+
+def test_reset_descendants_resettable(workflow_transport_graph, mocker):
+    """Test method that resets descendants of a node (including the node itself)."""
+    get_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.get_node_value",
+        return_value=RESULT_STATUS.COMPLETED,
+    )
+    reset_node_mock = mocker.patch("covalent._workflow.transport._TransportGraph.reset_node")
+    res = workflow_transport_graph._reset_descendants(0)
+    reset_node_mock.assert_called_once_with(0)
+    get_node_value_mock.assert_called_once_with(0, "status")
+    assert res is None
+
+
+def test_reset_descendants_multiple(workflow_transport_graph, mocker):
+    """Test method that resets descendants of a node (including the node itself)."""
+    workflow_transport_graph._graph.add_edge(0, 1)
+    get_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.get_node_value",
+        return_value=RESULT_STATUS.COMPLETED,
+    )
+    reset_node_mock = mocker.patch("covalent._workflow.transport._TransportGraph.reset_node")
+    res = workflow_transport_graph._reset_descendants(0)
+    reset_node_mock.mock_calls = [call(0), call(1)]
+    get_node_value_mock.mock_calls = [(0, "status"), (1, "status")]
+    assert res is None
+
+
+def test_apply_electron_updates(workflow_transport_graph, mocker):
+    """Test the method that applies electron updates to the graph."""
+    get_node_value_mock = mocker.patch(
+        "covalent._workflow.transport._TransportGraph.get_node_value", return_value="mock-name"
+    )
+    replace_node_mock = mocker.patch("covalent._workflow.transport._TransportGraph._replace_node")
+    workflow_transport_graph.apply_electron_updates({"mock-name": "mock-value"})
+    get_node_value_mock.mock_calls = [call(0, "name"), call(1, "name")]
+    replace_node_mock.mock_calls = [call(0, "mock-value"), call(1, "mock-value")]