Repository: treeverse/dvc (python)

Issue:
Output of `dvc dag` is unreadable if used with foreach/matrix stages

If a pipeline contains a stage that is run multiple times with different parameters (foreach/matrix), each run is treated as a separate node in the output of the `dvc dag` command. With a large number of combinations, this result becomes completely unusable. It would make more sense to treat such a stage as a single node, perhaps with a special designation.


Feature Request: Simplify dvc dag (foreach stages)

Thank you for providing DVC.

If you want to visualize the DAG using `dvc dag`  for a repository with many `foreach`  stages (e.g. many input files which need separate processing)  the dag becomes very messy, because each `foreach` entry is displayed separately. Thus, I would suggest an option to collapse `foreach` stages to a single one for display.
For example `dvc dag --collapse-foreach` could give:

``` 
+--------+                       
| inputs |                        
+--------+                
    **                   
    **           
    **                       
+---------+             
| process |     
+---------+           
``` 
Instead of `dvc dag`:

``` 
+---------------+                             +---------------+ 
| inputs@input1 |                             | inputs@input2 |      
+---------------+                             +---------------+    
               *****                     *****           
                    ****             ****                
                        ***       ***                    
                          +---------+             
                          | process |     
                          +---------+           
``` 
I already created a small script doing this from the `dot` format but I would be happy to help integrating this directly to DVC. Would that be of interest? 

[collapse_graph.py](https://github.com/user-attachments/files/24381477/collapse_graph.py)




Patch:
diff --git a/dvc/commands/dag.py b/dvc/commands/dag.py
index 221e83c638..be21abd12d 100644
--- a/dvc/commands/dag.py
+++ b/dvc/commands/dag.py
@@ -127,19 +127,80 @@ def _filter(graph, targets, full):
     return new_graph
 
 
-def _build(repo, target=None, full=False, outs=False):
+def _is_foreach_matrix_stage(node, join_string):
+    if node.endswith(".dvc"):
+        return False
+    return join_string in node
+
+
+def _collapse_foreach_matrix_get_nodes(graph):
+    from dvc.parsing import JOIN
+
+    new_nodes = set()
+    nodes_to_remove = set()
+    for _node in list(graph.nodes):
+        if not _is_foreach_matrix_stage(_node, JOIN):
+            continue
+        nodes_to_remove.add(_node)
+        new_nodes.add(_node.split(JOIN)[0])
+    return new_nodes, nodes_to_remove
+
+
+def _collapse_foreach_matrix_get_edges(graph):
+    from dvc.parsing import JOIN
+
+    new_edges = set()
+    edges_to_remove = set()
+    for _e1, _e2 in list(graph.edges):
+        _replace = False
+        _new_e1 = _e1
+        _new_e2 = _e2
+        if _is_foreach_matrix_stage(_e1, JOIN):
+            _new_e1 = _e1.split(JOIN)[0]
+            _replace = True
+        if _is_foreach_matrix_stage(_e2, JOIN):
+            _new_e2 = _e2.split(JOIN)[0]
+            _replace = True
+        if _replace:
+            edges_to_remove.add((_e1, _e2))
+            new_edges.add((_new_e1, _new_e2))
+    return new_edges, edges_to_remove
+
+
+def _collapse_foreach_matrix(graph):
+    new_nodes, nodes_to_remove = _collapse_foreach_matrix_get_nodes(graph)
+    new_edges, edges_to_remove = _collapse_foreach_matrix_get_edges(graph)
+    new_graph = graph.copy()
+    new_graph.remove_edges_from(edges_to_remove)
+    new_graph.add_nodes_from(new_nodes)
+    new_graph.add_edges_from(new_edges)
+    new_graph.remove_nodes_from(nodes_to_remove)
+    return new_graph
+
+
+def _build(repo, target=None, full=False, outs=False, collapse_foreach_matrix=False):
     targets = _collect_targets(repo, target, outs)
     graph = _transform(repo.index, outs)
-    return _filter(graph, targets, full)
+    filtered_graph = _filter(graph, targets, full)
+    if collapse_foreach_matrix:
+        filtered_graph = _collapse_foreach_matrix(filtered_graph)
+    return filtered_graph
 
 
 class CmdDAG(CmdBase):
     def run(self):
+        from dvc.exceptions import InvalidArgumentError
+
+        if self.args.outs and self.args.collapse_foreach_matrix:
+            raise InvalidArgumentError(
+                "`--outs` and `--collapse-foreach-matrix` are mutually exclusive"
+            )
         graph = _build(
             self.repo,
             target=self.args.target,
             full=self.args.full,
             outs=self.args.outs,
+            collapse_foreach_matrix=self.args.collapse_foreach_matrix,
         )
 
         if self.args.dot:
@@ -181,6 +242,12 @@ def add_parser(subparsers, parent_parser):
         dest="markdown",
         help="Print DAG with mermaid format wrapped in Markdown block.",
     )
+    dag_parser.add_argument(
+        "--collapse-foreach-matrix",
+        action="store_true",
+        default=False,
+        help="Collapse foreach/matrix stages to single stage.",
+    )
     dag_parser.add_argument(
         "--full",
         action="store_true",


Test patch:
diff --git a/tests/unit/command/test_dag.py b/tests/unit/command/test_dag.py
index d332957da2..645373ed2e 100644
--- a/tests/unit/command/test_dag.py
+++ b/tests/unit/command/test_dag.py
@@ -2,7 +2,15 @@
 import pytest
 
 from dvc.cli import main, parse_args
-from dvc.commands.dag import CmdDAG, _build, _show_ascii, _show_dot, _show_mermaid
+from dvc.commands.dag import (
+    CmdDAG,
+    _build,
+    _collapse_foreach_matrix,
+    _show_ascii,
+    _show_dot,
+    _show_mermaid,
+)
+from dvc.parsing import JOIN
 
 
 @pytest.mark.parametrize(
@@ -49,10 +57,53 @@ def repo(tmp_dir, dvc):
     return dvc
 
 
+def test_collapse_foreach_matrix(repo):
+    graph = nx.DiGraph(
+        [
+            ("2", "1"),
+            ("3@a", "2"),
+            ("3@a", "1"),
+            ("4", "1"),
+            ("3@b", "4"),
+            ("3@b", "1"),
+            ("5", "3@a"),
+            ("6", "3@b"),
+            ("7", "5"),
+            ("7", "6"),
+        ]
+    )
+    expected_graph = nx.DiGraph(
+        [
+            ("2", "1"),
+            ("3", "2"),
+            ("4", "1"),
+            ("3", "4"),
+            ("3", "1"),
+            ("5", "3"),
+            ("6", "3"),
+            ("7", "5"),
+            ("7", "6"),
+        ]
+    )
+    collapsed_graph = _collapse_foreach_matrix(graph)
+    for node in collapsed_graph.nodes:
+        assert JOIN not in node
+    for n1, n2 in collapsed_graph.edges:
+        assert JOIN not in n1
+        assert JOIN not in n2
+    assert nx.is_isomorphic(collapsed_graph, expected_graph)
+
+
 def test_build(repo):
     assert nx.is_isomorphic(_build(repo), repo.index.graph)
 
 
+def test_build_collapse(repo):
+    assert nx.is_isomorphic(
+        _build(repo, collapse_foreach_matrix=True), repo.index.graph
+    )
+
+
 def test_build_target(repo):
     graph = _build(repo, target="3")
     assert set(graph.nodes()) == {"3", "b.dvc", "a.dvc"}