Change results updating algorithm

Author: ericpai

mars/optimization/logical/core.py

@@ -135,8 +135,7 @@ def _replace_subgraph(
         self,
         graph: Optional[EntityGraph],
         nodes_to_remove: Optional[Set[EntityType]],
-        new_results: Optional[List[Entity]],
-        results_to_remove: Optional[List[Entity]],
+        new_results: Optional[List[Entity]] = None,
     ):
         """
         Replace the subgraph from the self._graph represented by a list of nodes with input graph.
@@ -149,27 +148,24 @@ def _replace_subgraph(
             The input graph. If it's none, no new node and edge will be added.
         nodes_to_remove : Set[EntityType], optional
             The nodes to be removed. All the edges connected with them are removed as well.
-        new_results : List[Entity], optional
-            The new results to be added to the graph.
-        results_to_remove : List[Entity], optional
-            The results to be removed from the graph. If a result is not in self._graph.results, it will be ignored.
+        new_results : List[Entity], optional, default None
+            The new results to be replaced to the original by their keys.
 
         Raises
         ------
         ValueError
             1. If the input key of the removed node's successor can't be found in the subgraph.
             2. Or some of the nodes of the subgraph are in removed ones.
-            3. Or the added result is not a valid output of any node in the updated graph.
+            3. Or some of the removed nodes are also in the results.
+            4. Or the key of the new result can't be found in the original results.
         """
         affected_successors = set()
-
         output_to_node = dict()
         nodes_to_remove = nodes_to_remove or set()
-        results_to_remove = results_to_remove or list()
         new_results = new_results or list()
-        final_results = set(
-            filter(lambda x: x not in results_to_remove, self._graph.results)
-        )
+        result_indices = {
+            result.key: idx for idx, result in enumerate(self._graph.results)
+        }
 
         if graph is not None:
             # Add the output key -> node of the subgraph
@@ -185,10 +181,12 @@ def _replace_subgraph(
                 for output in node.outputs:
                     output_to_node[output.key] = node
 
+        # Check if the updated result is valid
         for result in new_results:
+            if result.key not in result_indices:
+                raise ValueError(f"Unknown result {result} to replace")
             if result.key not in output_to_node:
-                raise ValueError(f"Unknown result {result} to add")
-        final_results.update(new_results)
+                raise ValueError(f"The result {result} is missing in the updated graph")
 
         for node in nodes_to_remove:
             for affected_successor in self._graph.iter_successors(node):
@@ -216,7 +214,8 @@ def _replace_subgraph(
                 pred_node = output_to_node[inp.key]
                 self._graph.add_edge(pred_node, node)
 
-        self._graph.results = list(final_results)
+        for result in new_results:
+            self._graph.results[result_indices[result.key]] = result
 
     def _add_collapsable_predecessor(self, node: EntityType, predecessor: EntityType):
         pred_original = self._records.get_original_entity(predecessor, predecessor)

mars/optimization/logical/tests/test_core.py

@@ -24,8 +24,8 @@ class _MockRule(OptimizationRule):
     def apply(self) -> bool:
         pass
 
-    def replace_subgraph(self, graph, nodes_to_remove, new_results, results_to_remove):
-        self._replace_subgraph(graph, nodes_to_remove, new_results, results_to_remove)
+    def replace_subgraph(self, graph, nodes_to_remove, new_results=None):
+        self._replace_subgraph(graph, nodes_to_remove, new_results)
 
 
 def test_replace_tileable_subgraph():
@@ -61,8 +61,9 @@ def test_replace_tileable_subgraph():
     g1 = v6.build_graph()
     v7 = v1.sub(v2)
     v8 = v7.add(v5)
+    v8._key = v6.key
+    v8.outputs[0]._key = v6.key
     g2 = v8.build_graph()
-
     # Here we use a trick way to construct the subgraph for test only
     key_to_node = dict()
     for node in g2.iter_nodes():
@@ -79,15 +80,12 @@ def test_replace_tileable_subgraph():
     c5 = g1.successors(key_to_node[s5.key])[0]
 
     new_results = [v8.outputs[0]]
-    removed_results = [
-        v6.outputs[0],
-        v8.outputs[0],  # v8.outputs[0] is not in the original results, so we ignore it.
-    ]
-    r.replace_subgraph(
-        g2, {key_to_node[op.key] for op in [v3, v4, v6]}, new_results, removed_results
-    )
+    r.replace_subgraph(g2, {key_to_node[op.key] for op in [v3, v4, v6]}, new_results)
     assert g1.results == new_results
-
+    for node in g1.iter_nodes():
+        if node.key == v8.key:
+            key_to_node[v8.key] = node
+            break
     expected_nodes = {s1, c1, v1, s2, c2, v2, s5, c5, v5, v7, v8}
     assert set(g1) == {key_to_node[n.key] for n in expected_nodes}
 
@@ -117,7 +115,7 @@ def test_replace_null_subgraph():
      s1 ---> c1 ---> v1 ---> v3(out) <--- v2 <--- c2 <--- s2
 
     Target Graph:
-      c1 ---> v1 ---> v3  <--- v2(out) <--- c2
+      c1 ---> v1 ---> v3(out)  <--- v2 <--- c2
 
     The nodes [s1, s2] will be removed.
     Subgraph is None
@@ -137,7 +135,7 @@ def test_replace_null_subgraph():
     # delete c5 s5 will fail
     with pytest.raises(ValueError):
         r.replace_subgraph(
-            None, {key_to_node[op.key] for op in [s1, s2]}, None, [v2.outputs[0]]
+            None, {key_to_node[op.key] for op in [s1, s2]}, [v2.outputs[0]]
         )
 
     assert g1.results == expected_results
@@ -161,11 +159,9 @@ def test_replace_null_subgraph():
     c2.inputs.clear()
     r.replace_subgraph(
         None,
-        {key_to_node[op.key] for op in [s1, s2]},
-        [v2.outputs[0]],
-        [v3.outputs[0]],
+        {key_to_node[op.key] for op in [s1, s2]}
     )
-    assert g1.results == [v2.outputs[0]]
+    assert g1.results == expected_results
     assert set(g1) == {key_to_node[n.key] for n in {c1, v1, c2, v2, v3}}
     expected_edges = {
         c1: [v1],
@@ -206,12 +202,12 @@ def test_replace_subgraph_without_removing_nodes():
     key_to_node = {
         node.key: node for node in itertools.chain(g1.iter_nodes(), g2.iter_nodes())
     }
-    expected_results = [v3.outputs[0], v4.outputs[0]]
+    expected_results = [v4.outputs[0]]
     c1 = g1.successors(key_to_node[s1.key])[0]
     c2 = g1.successors(key_to_node[s2.key])[0]
     c3 = g2.successors(key_to_node[s3.key])[0]
     r = _MockRule(g1, None, None)
-    r.replace_subgraph(g2, None, [v3.outputs[0]], None)
+    r.replace_subgraph(g2, None)
     assert g1.results == expected_results
     assert set(g1) == {
         key_to_node[n.key] for n in {s1, c1, v1, s2, c2, v2, s3, c3, v3, v4}