Support reporting tile progress (#2954)

Author: Xuye (Chris) Qin

mars/core/__init__.py

@@ -61,5 +61,7 @@
     ChunkGraph,
     TileableGraphBuilder,
     ChunkGraphBuilder,
+    TileContext,
+    TileStatus,
 )
 from .mode import enter_mode, is_build_mode, is_eager_mode, is_kernel_mode

mars/core/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from functools import wraps
-from typing import Dict
+from typing import Dict, Tuple, Type
 
 from ..serialization.core import Placeholder, fast_id
 from ..serialization.serializables import Serializable, StringField
@@ -117,6 +117,15 @@ def key(self):
     def id(self):
         return self._id
 
+    def to_kv(self, exclude_fields: Tuple[str], accept_value_types: Tuple[Type]):
+        fields = self._FIELDS
+        field_values = self._FIELD_VALUES
+        return {
+            fields[attr_name].tag: value
+            for attr_name, value in field_values.items()
+            if attr_name not in exclude_fields and isinstance(value, accept_value_types)
+        }
+
 
 def buffered_base(func):
     @wraps(func)

mars/core/entity/utils.py

@@ -28,7 +28,12 @@ def refresh_tileable_shape(tileable):
 
 
 def tile(tileable, *tileables: TileableType):
-    from ..graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
+    from ..graph import (
+        TileableGraph,
+        TileableGraphBuilder,
+        ChunkGraphBuilder,
+        TileContext,
+    )
 
     raw_tileables = target_tileables = [tileable] + list(tileables)
     target_tileables = [t.data if hasattr(t, "data") else t for t in target_tileables]
@@ -38,7 +43,7 @@ def tile(tileable, *tileables: TileableType):
     next(tileable_graph_builder.build())
 
     # tile
-    tile_context = dict()
+    tile_context = TileContext()
     chunk_graph_builder = ChunkGraphBuilder(
         tileable_graph, fuse_enabled=False, tile_context=tile_context
     )

mars/core/graph/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .builder import TileableGraphBuilder, ChunkGraphBuilder
+from .builder import TileableGraphBuilder, ChunkGraphBuilder, TileContext, TileStatus
 from .core import DirectedGraph, DAG, GraphContainsCycleError
 from .entity import TileableGraph, ChunkGraph, EntityGraph

mars/core/graph/builder/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .chunk import ChunkGraphBuilder
+from .chunk import ChunkGraphBuilder, TileContext, TileStatus
 from .tileable import TileableGraphBuilder

mars/core/graph/builder/chunk.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import dataclasses
+import functools
 from typing import (
     Callable,
     Dict,
@@ -35,6 +36,7 @@
 
 
 tile_gen_type = Generator[List[ChunkType], List[ChunkType], List[TileableType]]
+DEFAULT_UPDATED_PROGRESS = 0.4
 
 
 @dataclasses.dataclass
@@ -44,14 +46,84 @@ class _TileableHandler:
     last_need_processes: List[EntityType] = None
 
 
+@dataclasses.dataclass
+class _TileableTileInfo:
+    curr_iter: int
+    # incremental progress for this iteration
+    tile_progress: float
+    # newly generated chunks by a tileable in this iteration
+    generated_chunks: List[ChunkType] = dataclasses.field(default_factory=list)
+
+
+class TileContext(Dict[TileableType, TileableType]):
+    _tileables = Set[TileableType]
+    _tileable_to_progress: Dict[TileableType, float]
+    _tileable_to_tile_infos: Dict[TileableType, List[_TileableTileInfo]]
+
+    def __init__(self, *args, **kw):
+        super().__init__(*args, **kw)
+        self._tileables = None
+        self._tileable_to_progress = dict()
+        self._tileable_to_tile_infos = dict()
+
+    def set_tileables(self, tileables: Set[TileableType]):
+        self._tileables = tileables
+
+    def __setitem__(self, key, value):
+        self._tileable_to_progress.pop(key, None)
+        return super().__setitem__(key, value)
+
+    def set_progress(self, tileable: TileableType, progress: float):
+        assert 0.0 <= progress <= 1.0
+        last_progress = self._tileable_to_progress.get(tileable, 0.0)
+        self._tileable_to_progress[tileable] = max(progress, last_progress)
+
+    def get_progress(self, tileable: TileableType) -> float:
+        if tileable in self:
+            return 1.0
+        else:
+            return self._tileable_to_progress.get(tileable, 0.0)
+
+    def get_all_progress(self) -> float:
+        return sum(self.get_progress(t) for t in self._tileables) / len(self._tileables)
+
+    def record_tileable_tile_info(
+        self, tileable: TileableType, curr_iter: int, generated_chunks: List[ChunkType]
+    ):
+        if tileable not in self._tileable_to_tile_infos:
+            self._tileable_to_tile_infos[tileable] = []
+        prev_progress = sum(
+            info.tile_progress for info in self._tileable_to_tile_infos[tileable]
+        )
+        curr_progress = self.get_progress(tileable)
+        infos = self._tileable_to_tile_infos[tileable]
+        infos.append(
+            _TileableTileInfo(
+                curr_iter=curr_iter,
+                tile_progress=curr_progress - prev_progress,
+                generated_chunks=generated_chunks,
+            )
+        )
+
+    def get_tileable_tile_infos(self) -> Dict[TileableType, List[_TileableTileInfo]]:
+        return {t: self._tileable_to_tile_infos.get(t, list()) for t in self._tileables}
+
+
+@dataclasses.dataclass
+class TileStatus:
+    entities: List[EntityType] = None
+    progress: float = None
+
+
 class Tiler:
+    _cur_iter: int
     _cur_chunk_graph: Optional[ChunkGraph]
     _tileable_handlers: Iterable[_TileableHandler]
 
     def __init__(
         self,
         tileable_graph: TileableGraph,
-        tile_context: Dict[TileableType, TileableType],
+        tile_context: TileContext,
         processed_chunks: Set[ChunkType],
         chunk_to_fetch: Dict[ChunkType, ChunkType],
         add_nodes: Callable,
@@ -60,13 +132,31 @@ def __init__(
         self._tile_context = tile_context
         self._processed_chunks = processed_chunks
         self._chunk_to_fetch = chunk_to_fetch
-        self._add_nodes = add_nodes
+        self._add_nodes = self._wrap_add_nodes(add_nodes)
+        self._curr_iter = 0
         self._cur_chunk_graph = None
         self._tileable_handlers = (
             _TileableHandler(tileable, self._tile_handler(tileable))
             for tileable in tileable_graph.topological_iter()
         )
 
+    def _wrap_add_nodes(self, add_nodes: Callable):
+        @functools.wraps(add_nodes)
+        def inner(
+            chunk_graph: ChunkGraph,
+            chunks: List[ChunkType],
+            visited: Set[ChunkType],
+            tileable: TileableType,
+        ):
+            prev_chunks = set(chunk_graph)
+            add_nodes(chunk_graph, chunks, visited)
+            new_chunks = set(chunk_graph)
+            self._tile_context.record_tileable_tile_info(
+                tileable, self._curr_iter, list(new_chunks - prev_chunks)
+            )
+
+        return inner
+
     @staticmethod
     def _get_data(entity: EntityType):
         return entity.data if hasattr(entity, "data") else entity
@@ -119,6 +209,17 @@ def _tile(
     ):
         try:
             need_process = next(tile_handler)
+
+            if isinstance(need_process, TileStatus):
+                # process tile that returns progress
+                self._tile_context.set_progress(tileable, need_process.progress)
+                need_process = need_process.entities
+            else:
+                # if progress not specified, we just update 0.4 * rest progress
+                progress = self._tile_context.get_progress(tileable)
+                new_progress = progress + (1.0 - progress) * DEFAULT_UPDATED_PROGRESS
+                self._tile_context.set_progress(tileable, new_progress)
+
             chunks = []
             if need_process is not None:
                 for t in need_process:
@@ -127,7 +228,7 @@ def _tile(
                     elif isinstance(t, TILEABLE_TYPE):
                         to_update_tileables.append(self._get_data(t))
             # not finished yet
-            self._add_nodes(chunk_graph, chunks.copy(), visited)
+            self._add_nodes(chunk_graph, chunks.copy(), visited, tileable)
             next_tileable_handlers.append(
                 _TileableHandler(tileable, tile_handler, need_process)
             )
@@ -145,8 +246,8 @@ def _tile(
                 if chunks is None:  # pragma: no cover
                     raise ValueError(f"tileable({out}) is still coarse after tile")
                 chunks = [self._get_data(c) for c in chunks]
-                self._add_nodes(chunk_graph, chunks, visited)
                 self._tile_context[out] = tiled_tileable
+                self._add_nodes(chunk_graph, chunks, visited, tileable)
 
     def _gen_result_chunks(
         self,
@@ -227,6 +328,8 @@ def _iter(self):
         # prune unused chunks
         prune_chunk_graph(chunk_graph)
 
+        self._curr_iter += 1
+
         return to_update_tileables
 
     def __iter__(self):
@@ -278,12 +381,13 @@ def __init__(
         self,
         graph: TileableGraph,
         fuse_enabled: bool = True,
-        tile_context: Dict[TileableType, TileableType] = None,
+        tile_context: TileContext = None,
         tiler_cls: Union[Type[Tiler], Callable] = None,
     ):
         super().__init__(graph)
         self.fuse_enabled = fuse_enabled
-        self.tile_context = dict() if tile_context is None else tile_context
+        self.tile_context = TileContext() if tile_context is None else tile_context
+        self.tile_context.set_tileables(set(graph))
 
         self._processed_chunks: Set[ChunkType] = set()
         self._chunk_to_fetch: Dict[ChunkType, ChunkType] = dict()

mars/dataframe/merge/merge.py

@@ -21,7 +21,7 @@
 import pandas as pd
 
 from ... import opcodes as OperandDef
-from ...core import OutputType, recursive_tile
+from ...core import OutputType, recursive_tile, TileStatus
 from ...core.context import get_context
 from ...core.operand import OperandStage, MapReduceOperand
 from ...serialization.serializables import (
@@ -609,7 +609,7 @@ def tile(cls, op: "DataFrameMerge"):
             auto_merge_before
             and len(left.chunks) + len(right.chunks) > auto_merge_threshold
         ):
-            yield [left, right] + left.chunks + right.chunks
+            yield TileStatus([left, right] + left.chunks + right.chunks, progress=0.2)
             left = auto_merge_chunks(ctx, left)
             right = auto_merge_chunks(ctx, right)
 
@@ -626,7 +626,7 @@ def tile(cls, op: "DataFrameMerge"):
             right_on = _prepare_shuffle_on(op.right_index, op.right_on, op.on)
             if op.how == "inner" and op.bloom_filter:
                 if has_unknown_shape(left, right):
-                    yield left.chunks + right.chunks
+                    yield TileStatus(left.chunks + right.chunks, progress=0.3)
                 small_one = right if len(left.chunks) > len(right.chunks) else left
                 logger.debug(
                     "Apply bloom filter for operand %s, use DataFrame %s to build bloom filter.",
@@ -637,7 +637,9 @@ def tile(cls, op: "DataFrameMerge"):
                     *cls._apply_bloom_filter(left, right, left_on, right_on, op)
                 )
                 # auto merge after bloom filter
-                yield [left, right] + left.chunks + right.chunks
+                yield TileStatus(
+                    [left, right] + left.chunks + right.chunks, progress=0.5
+                )
                 left = auto_merge_chunks(ctx, left)
                 right = auto_merge_chunks(ctx, right)
 
@@ -660,7 +662,9 @@ def tile(cls, op: "DataFrameMerge"):
         ):
             # if how=="inner", output data size will reduce greatly with high probability，
             # use auto_merge_chunks to combine small chunks.
-            yield ret[0].chunks  # trigger execution for chunks
+            yield TileStatus(
+                ret[0].chunks, progress=0.8
+            )  # trigger execution for chunks
             return [auto_merge_chunks(get_context(), ret[0])]
         else:
             return ret

mars/dataframe/merge/tests/test_merge_execution.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 
 from ....core.graph.builder.utils import build_graph
 from ...datasource.dataframe import from_pandas
@@ -597,7 +598,8 @@ def test_merge_with_bloom_filter(setup):
     )
 
 
-def test_merge_on_duplicate_columns(setup):
+@pytest.mark.parametrize("auto_merge", ["none", "both", "before", "after"])
+def test_merge_on_duplicate_columns(setup, auto_merge):
     raw1 = pd.DataFrame(
         [["foo", 1, "bar"], ["bar", 2, "foo"], ["baz", 3, "foo"]],
         columns=["lkey", "value", "value"],
@@ -611,7 +613,7 @@ def test_merge_on_duplicate_columns(setup):
     df1 = from_pandas(raw1, chunk_size=2)
     df2 = from_pandas(raw2, chunk_size=3)
 
-    r = df1.merge(df2, left_on="lkey", right_on="rkey", auto_merge="none")
+    r = df1.merge(df2, left_on="lkey", right_on="rkey", auto_merge=auto_merge)
     result = r.execute().fetch()
     expected = raw1.merge(raw2, left_on="lkey", right_on="rkey")
     pd.testing.assert_frame_equal(expected, result)

mars/optimization/logical/chunk/tests/test_column_pruning.py

@@ -19,7 +19,13 @@
 import pytest
 
 from ..... import dataframe as md
-from .....core import enter_mode, TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
+from .....core import (
+    enter_mode,
+    TileableGraph,
+    TileableGraphBuilder,
+    ChunkGraphBuilder,
+    TileContext,
+)
 from .. import optimize
 
 
@@ -47,7 +53,7 @@ def test_groupby_read_csv(gen_data1):
     df2 = df1[["a", "b"]]
     graph = TileableGraph([df2.data])
     next(TileableGraphBuilder(graph).build())
-    context = dict()
+    context = TileContext()
     chunk_graph_builder = ChunkGraphBuilder(
         graph, fuse_enabled=False, tile_context=context
     )

mars/optimization/logical/chunk/tests/test_head.py

@@ -19,7 +19,13 @@
 import pytest
 
 from ..... import dataframe as md
-from .....core import enter_mode, TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
+from .....core import (
+    enter_mode,
+    TileableGraph,
+    TileableGraphBuilder,
+    ChunkGraphBuilder,
+    TileContext,
+)
 from .. import optimize
 
 
@@ -47,7 +53,7 @@ def test_read_csv_head(gen_data1):
     df2 = df1.head(5)
     graph = TileableGraph([df2.data])
     next(TileableGraphBuilder(graph).build())
-    context = dict()
+    context = TileContext()
     chunk_graph_builder = ChunkGraphBuilder(
         graph, fuse_enabled=False, tile_context=context
     )