Repository: mars-project/mars

Issue:
Performance declined after `.execute()` acts differently
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Since #1201, `.execute()` acts differently, but performance declined unepxectly.

Before #1201:

```
In [1]: import mars.tensor as mt                                                

In [2]: %%time 
   ...: N = 100_000_000 
   ...: a = mt.random.uniform(-1, 1, size=(N, 2), chunk_size=10_000_000) 
   ...: print(((mt.linalg.norm(a, axis=1) < 1).sum() * 4 / N).execute()) 
   ...:  
   ...:                                                                         
3.1414312
CPU times: user 9.11 s, sys: 2.84 s, total: 11.9 s
Wall time: 3.17 s
```

After #1201:

```ipython
In [1]: import mars.tensor as mt                                                

In [2]: %%time 
   ...: N = 100_000_000 
   ...: a = mt.random.uniform(-1, 1, size=(N, 2), chunk_size=10_000_000) 
   ...: print(((mt.linalg.norm(a, axis=1) < 1).sum() * 4 / N).execute()) 
   ...:  
   ...:                                                                         
3.14149448
CPU times: user 17.3 s, sys: 5.19 s, total: 22.5 s
Wall time: 5.93 s
```



Discussion:

Resolving patch:
diff --git a/mars/executor.py b/mars/executor.py
--- a/mars/executor.py
+++ b/mars/executor.py
@@ -28,7 +28,7 @@
 import numpy as np
 
 from .operands import Fetch, ShuffleProxy
-from .graph import DirectedGraph
+from .graph import DAG, DirectedGraph
 from .config import options
 from .tiles import IterativeChunkGraphBuilder, ChunkGraphBuilder, get_tiled
 from .optimizes.runtime.optimizers.core import Optimizer
@@ -36,7 +36,7 @@
 from .graph_builder import TileableGraphBuilder
 from .context import LocalContext
 from .utils import kernel_mode, enter_build_mode, build_fetch, calc_nsplits,\
-    has_unknown_shape
+    has_unknown_shape, prune_chunk_graph
 
 try:
     from numpy.core._exceptions import UFuncTypeError
@@ -596,8 +596,6 @@ def __init__(self, engine=None, storage=None, prefetch=False,
         self._chunk_result = storage if storage is not None else dict()
         self._prefetch = prefetch
 
-        # only record the executed tileable
-        self.tileable_to_tiled = weakref.WeakKeyDictionary()
         # dict structure: {tileable_key -> chunk_keys, tileable_ids}
         # dict value is a tuple object which records chunk keys and tileable id
         self.stored_tileables = dict()
@@ -764,14 +762,13 @@ def execute_tileables(self, tileables, fetch=True, n_parallel=None, n_thread=Non
         tileable_data_to_concat_keys = weakref.WeakKeyDictionary()
         tileable_data_to_chunk_keys = weakref.WeakKeyDictionary()
 
-        executed_keys = set(chunk_result)
         node_to_fetch = weakref.WeakKeyDictionary()
 
         def _generate_fetch_if_executed(nd):
             # node processor that if the node is executed
             # replace it with a fetch node
-            _keys, _to_fetch = executed_keys, node_to_fetch  # noqa: F821
-            if nd.key not in _keys:
+            _to_fetch = node_to_fetch  # noqa: F821
+            if nd.key not in chunk_result:
                 return nd
             if nd in _to_fetch:
                 return _to_fetch[nd]
@@ -813,8 +810,8 @@ def _get_tileable_graph_builder(**kwargs):
             tileable_graph_builder = _get_tileable_graph_builder()
             tileable_graph = tileable_graph_builder.build(tileables)
             chunk_graph_builder = IterativeChunkGraphBuilder(
-                graph_cls=DirectedGraph, node_processor=_generate_fetch_if_executed,
-                compose=compose, on_tile_success=_on_tile_success)
+                graph_cls=DAG, node_processor=_generate_fetch_if_executed,
+                compose=False, on_tile_success=_on_tile_success)
             intermediate_result_keys = set()
             while True:
                 # build chunk graph, tile will be done during building
@@ -829,6 +826,12 @@ def _get_tileable_graph_builder(**kwargs):
                             if inp.op not in chunk_graph_builder.interrupted_ops:
                                 for n in get_tiled(inp).chunks:
                                     temp_result_keys.add(n.key)
+                else:
+                    # if done, prune chunk graph
+                    prune_chunk_graph(chunk_graph, temp_result_keys)
+                # compose
+                if compose:
+                    chunk_graph.compose(list(temp_result_keys))
                 # execute chunk graph
                 self.execute_graph(chunk_graph, list(temp_result_keys),
                                    n_parallel=n_parallel or n_thread,
@@ -850,7 +853,6 @@ def _get_tileable_graph_builder(**kwargs):
                     delattr(chunk_graph_builder, '_prev_tileable_graph')
                     break
                 else:
-                    executed_keys = set(chunk_result)
                     intermediate_result_keys.update(temp_result_keys)
                     # add the node that failed
                     to_run_tileables = list(itertools.chain(
@@ -869,8 +871,11 @@ def _get_tileable_graph_builder(**kwargs):
                 if tileable.key in self.stored_tileables:
                     self.stored_tileables[tileable.key][0].add(tileable.id)
                 else:
-                    chunk_keys = tileable_data_to_chunk_keys[tileable_optimized.get(tileable, tileable)]
-                    self.stored_tileables[tileable.key] = tuple([{tileable.id}, set(chunk_keys)])
+                    chunk_keys = tileable_data_to_chunk_keys[
+                        tileable_optimized.get(tileable, tileable)]
+                    self.stored_tileables[tileable.key] = \
+                        tuple([{tileable.id}, set(chunk_keys)])
+
             try:
                 if fetch:
                     concat_keys = [
@@ -915,13 +920,12 @@ def fetch_tileables(self, tileables, **kw):
                 raise ValueError(
                     'Tileable object {} to fetch must be executed first'.format(tileable))
 
-        try:
-            # if chunk executed, fetch chunk mechanism will be triggered in execute_tileables
-            return self.execute_tileables(tileables, **kw)
-        finally:
-            for to_release_tileable in to_release_tileables:
-                for c in get_tiled(to_release_tileable, mapping=tileable_optimized).chunks:
-                    del self._chunk_result[c.key]
+        # if chunk executed, fetch chunk mechanism will be triggered in execute_tileables
+        result = self.execute_tileables(tileables, **kw)
+        for to_release_tileable in to_release_tileables:
+            for c in get_tiled(to_release_tileable, mapping=tileable_optimized).chunks:
+                del self._chunk_result[c.key]
+        return result
 
     def get_tileable_nsplits(self, tileable, chunk_result=None):
         chunk_idx_to_shape = OrderedDict()
diff --git a/mars/optimizes/tileable_graph/core.py b/mars/optimizes/tileable_graph/core.py
--- a/mars/optimizes/tileable_graph/core.py
+++ b/mars/optimizes/tileable_graph/core.py
@@ -60,7 +60,18 @@ def __init__(self, **kw):
     def _apply_rules(node_processor, optimizer_context):
         def inner(node):
             node = node_processor(node) if node_processor is not None else node
-            if type(node.op) in _rules:
+            if node in tileable_optimized:
+                node = tileable_optimized[node]
+            elif len(node.inputs or []) > 0 and \
+                    any(inp in tileable_optimized for inp in node.inputs):
+                new_inputs = []
+                for inp in node.inputs:
+                    if inp in tileable_optimized:
+                        new_inputs.append(tileable_optimized[inp])
+                    else:
+                        new_inputs.append(inp)
+                node.inputs = new_inputs
+            elif type(node.op) in _rules:
                 for rule in _rules[type(node.op)]:
                     ruler = rule(optimizer_context)
                     if ruler.match(node):
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -929,6 +929,19 @@ def adapt_mars_docstring(doc):
     return '\n'.join(lines)
 
 
+def prune_chunk_graph(chunk_graph, result_chunk_keys):
+    reverse_chunk_graph = chunk_graph.build_reversed()
+    marked = set()
+    for c in reverse_chunk_graph.topological_iter():
+        if c.key in result_chunk_keys:
+            marked.update(o for o in c.op.outputs)
+        elif any(inp in marked for inp in reverse_chunk_graph.iter_predecessors(c)):
+            marked.update(o for o in c.op.outputs)
+    for n in list(chunk_graph):
+        if n not in marked:
+            chunk_graph.remove_node(n)
+
+
 @functools.lru_cache(500)
 def serialize_function(function):
     return cloudpickle.dumps(function)


Test patch:
diff --git a/mars/optimizes/tileable_graph/tests/test_column_pruning.py b/mars/optimizes/tileable_graph/tests/test_column_pruning.py
--- a/mars/optimizes/tileable_graph/tests/test_column_pruning.py
+++ b/mars/optimizes/tileable_graph/tests/test_column_pruning.py
@@ -23,6 +23,8 @@
 import mars.dataframe as md
 from mars.core import ExecutableTuple
 from mars.config import option_context
+from mars.dataframe.datasource.read_csv import DataFrameReadCSV
+from mars.executor import register, Executor
 from mars.tests.core import TestBase, ExecutorForTest
 from mars.optimizes.tileable_graph.core import tileable_optimized
 
@@ -140,3 +142,28 @@ def testExecutedPruning(self):
 
         finally:
             shutil.rmtree(tempdir)
+
+    def testFetch(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            filename = os.path.join(tempdir, 'test_fetch.csv')
+            pd_df = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3],
+                                  'b': [1, 3, 4, 5, 6, 5, 4, 4, 4],
+                                  'c': list('aabaaddce'),
+                                  'd': list('abaaaddce')})
+            pd_df.to_csv(filename, index=False)
+
+            df = md.read_csv(filename)
+            df2 = df.groupby('d').agg({'b': 'min'})
+            expected = pd_df.groupby('d').agg({'b': 'min'})
+            _ = df2.execute()
+
+            def _execute_read_csv(*_):  # pragma: no cover
+                raise ValueError('cannot run read_csv again')
+
+            try:
+                register(DataFrameReadCSV, _execute_read_csv)
+
+                pd.testing.assert_frame_equal(df2.fetch(), expected)
+                pd.testing.assert_frame_equal(df2.iloc[:3].fetch(), expected.iloc[:3])
+            finally:
+                del Executor._op_runners[DataFrameReadCSV]
diff --git a/mars/tests/test_session.py b/mars/tests/test_session.py
--- a/mars/tests/test_session.py
+++ b/mars/tests/test_session.py
@@ -21,6 +21,8 @@
 
 import mars.tensor as mt
 import mars.dataframe as md
+from mars.executor import register, Executor
+from mars.tensor.datasource import ArrayDataSource
 from mars.tiles import get_tiled
 from mars.session import new_session, Session
 
@@ -281,6 +283,22 @@ def testFetch(self):
         np.testing.assert_array_equal(fetch2, raw)
         np.testing.assert_almost_equal(fetch3, raw.sum())
 
+        raw = np.random.rand(5, 10)
+        arr = mt.tensor(raw, chunk_size=5)
+        s = arr.sum()
+
+        self.assertAlmostEqual(s.execute().fetch(), raw.sum())
+
+        def _execute_ds(*_):  # pragma: no cover
+            raise ValueError('cannot run random again')
+
+        try:
+            register(ArrayDataSource, _execute_ds)
+
+            self.assertAlmostEqual(s.fetch(), raw.sum())
+        finally:
+            del Executor._op_runners[ArrayDataSource]
+
     def testDecref(self):
         sess = new_session()