Dump subtask graph for all backends (#3245)

Co-authored-by: 刘宝 <po.lb@antgroup.com>
mars-project · Sep 9, 2022 · a1a752f · a1a752f
1 parent 8ee68e2
commit a1a752f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 15 deletions.
diff --git a/mars/services/task/supervisor/graph_visualizer.py b/mars/services/task/supervisor/graph_visualizer.py
@@ -17,17 +17,12 @@
 from typing import Dict, List
 
 from ....core.operand import Fetch, FetchShuffle
-from ...subtask import Subtask
-from .processor import TaskProcessor
+from ...subtask import Subtask, SubtaskGraph
 
 
 class GraphVisualizer:
-    task_processor: TaskProcessor
-
-    def __init__(self, task_processor):
-        self.task_processor = task_processor
-
-    def to_dot(self):
+    @classmethod
+    def to_dot(cls, subtask_graphs: List[SubtaskGraph]):
         sio = StringIO()
         sio.write("digraph {\n")
         sio.write("splines=curved\n")
@@ -38,16 +33,14 @@ def to_dot(self):
         result_chunk_to_subtask = dict()
         line_colors = dict()
         color_iter = iter(itertools.cycle(range(1, 9)))
-        for stage_line in itertools.combinations(
-            range(len(self.task_processor.stage_processors))[::-1], 2
-        ):
+        for stage_line in itertools.combinations(range(len(subtask_graphs))[::-1], 2):
             line_colors[stage_line] = f'"/spectral9/{next(color_iter)}"'
 
-        for stage_processor in self.task_processor.stage_processors:
-            for subtask in stage_processor.subtask_graph.topological_iter():
+        for subtask_graph in subtask_graphs:
+            for subtask in subtask_graph.topological_iter():
                 current_cluster = f"cluster_{subgraph_index}"
                 sio.write(
-                    self._export_subtask_to_dot(
+                    cls._export_subtask_to_dot(
                         subtask,
                         current_cluster,
                         current_stage,

diff --git a/mars/services/task/supervisor/processor.py b/mars/services/task/supervisor/processor.py
@@ -65,6 +65,7 @@ def __init__(
             ProfilingData.init(task.task_id, task.extra_config["enable_profiling"])
 
         self._dump_subtask_graph = False
+        self._subtask_graphs = []
         if MARS_ENABLE_DUMPING_SUBTASK_GRAPH or (
             task.extra_config and task.extra_config.get("dump_subtask_graph")
         ):
@@ -221,6 +222,8 @@ async def _process_stage_chunk_graph(
                 op_to_bands=fetch_op_to_bands,
                 shuffle_fetch_type=shuffle_fetch_type,
             )
+            if self._dump_subtask_graph:
+                self._subtask_graphs.append(subtask_graph)
         stage_profiler.set(f"gen_subtask_graph({len(subtask_graph)})", timer.duration)
         logger.info(
             "Time consuming to gen a subtask graph is %ss with session id %s, task id %s, stage id %s",
@@ -419,7 +422,7 @@ def dump_subtask_graph(self):
         except ImportError:
             graphviz = None
 
-        dot = GraphVisualizer(self).to_dot()
+        dot = GraphVisualizer.to_dot(self._subtask_graphs)
         directory = tempfile.gettempdir()
         file_name = f"mars-{self.task_id}"
         logger.debug(

diff --git a/mars/services/task/supervisor/tests/test_task_manager.py b/mars/services/task/supervisor/tests/test_task_manager.py
@@ -632,6 +632,7 @@ async def test_optimization(actor_pool):
 
 
 @pytest.mark.asyncio
+@pytest.mark.ray_dag
 async def test_dump_subtask_graph(actor_pool):
     (
         execution_backend,