From bfa5b1a0c3fc2ec5f6926f3bbc31cf4d22ab763a Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Mon, 3 Jul 2023 11:44:51 +0545 Subject: [PATCH] repro: make stages run in order of the graph (#9675) * repro: make stages run in order of the graph. In case of --single-item, this means it'll run in the order of arguments passed. For --downstream, this fixes running of stage out of order when multiple arguments are passed. For normal case, this enforces the same logic, so that they are never run out-of-order. * add tests * tests: repro: add test for --single-item with multiple targets * graph: support when no nodes are provided --- dvc/repo/graph.py | 28 +++++++++++++- dvc/repo/reproduce.py | 67 ++++++++++++--------------------- tests/func/repro/test_repro.py | 14 +++++++ tests/unit/repo/test_graph.py | 69 ++++++++++++++++++++++++++++++++++ 4 files changed, 133 insertions(+), 45 deletions(-) create mode 100644 tests/unit/repo/test_graph.py diff --git a/dvc/repo/graph.py b/dvc/repo/graph.py index b19677503..83f6c99b9 100644 --- a/dvc/repo/graph.py +++ b/dvc/repo/graph.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Iterator, List, Set +from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Set, TypeVar from dvc.fs import localfs from dvc.utils.fs import path_isin @@ -8,6 +8,8 @@ if TYPE_CHECKING: from dvc.stage import Stage +T = TypeVar("T") + def check_acyclic(graph: "DiGraph") -> None: import networkx as nx @@ -39,6 +41,30 @@ def get_pipelines(graph: "DiGraph"): return [graph.subgraph(c).copy() for c in nx.weakly_connected_components(graph)] +def get_subgraph_of_nodes( + graph: "DiGraph", sources: Optional[List[Any]] = None, downstream: bool = False +) -> "DiGraph": + from networkx import dfs_postorder_nodes, reverse_view + + if not sources: + return graph + + g = reverse_view(graph) if downstream else graph + nodes = [] + for source in sources: + nodes.extend(dfs_postorder_nodes(g, source)) + return graph.subgraph(nodes) + + +def get_steps( + graph: "DiGraph", sources: Optional[List[T]] = None, downstream: bool = False +) -> List[T]: + from networkx import dfs_postorder_nodes + + sub = get_subgraph_of_nodes(graph, sources, downstream=downstream) + return list(dfs_postorder_nodes(sub)) + + def collect_pipeline(stage: "Stage", graph: "DiGraph") -> Iterator["Stage"]: import networkx as nx diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index 527adcc99..673a741da 100644 --- a/dvc/repo/reproduce.py +++ b/dvc/repo/reproduce.py @@ -1,5 +1,7 @@ import logging -from typing import TYPE_CHECKING, Iterator, List +from typing import TYPE_CHECKING, Iterator, List, cast + +from funcy import ldistinct from dvc.exceptions import ReproductionError from dvc.repo.scm_context import scm_context @@ -8,6 +10,8 @@ from dvc.stage.cache import RunCacheNotSupported from . import locked if TYPE_CHECKING: + from networkx import DiGraph + from dvc.stage import Stage from . import Repo @@ -78,11 +82,12 @@ def reproduce( # noqa: C901, PLR0912 targets = [PROJECT_FILE] + targets = targets or [] interactive = kwargs.get("interactive", False) if not interactive: kwargs["interactive"] = self.config["core"].get("interactive", False) - stages = set() + stages = [] if pipeline or all_pipelines: pipelines = get_pipelines(self.index.graph) if all_pipelines: @@ -96,10 +101,10 @@ def reproduce( # noqa: C901, PLR0912 for pline in used_pipelines: for stage in pline: if pline.in_degree(stage) == 0: - stages.add(stage) + stages.append(stage) else: for target in targets: - stages.update( + stages.extend( self.stage.collect( target, recursive=recursive, @@ -114,7 +119,7 @@ def reproduce( # noqa: C901, PLR0912 except RunCacheNotSupported as e: logger.warning("Failed to pull run cache: %s", e) - return _reproduce_stages(self.index.graph, list(stages), **kwargs) + return _reproduce_stages(self.index.graph, ldistinct(stages), **kwargs) def _reproduce_stages( # noqa: C901 @@ -159,7 +164,11 @@ def _reproduce_stages( # noqa: C901 The derived evaluation of _downstream_ B would be: [B, D, E] """ - steps = _get_steps(graph, stages, downstream, single_item) + from .graph import get_steps + + if not single_item: + active = _remove_frozen_stages(graph) + stages = get_steps(active, stages, downstream=downstream) force_downstream = kwargs.pop("force_downstream", False) result: List["Stage"] = [] @@ -167,7 +176,7 @@ def _reproduce_stages( # noqa: C901 # `ret` is used to add a cosmetic newline. ret: List["Stage"] = [] - for stage in steps: + for stage in stages: if ret: logger.info("") @@ -191,40 +200,10 @@ def _reproduce_stages( # noqa: C901 return result -def _get_steps(graph, stages, downstream, single_item): - import networkx as nx - - active = graph.copy() - if not single_item: - # NOTE: frozen stages don't matter for single_item - for stage in graph: - if stage.frozen: - # NOTE: disconnect frozen stage from its dependencies - active.remove_edges_from(graph.out_edges(stage)) - - all_pipelines: List["Stage"] = [] - for stage in stages: - if downstream: - # NOTE (py3 only): - # Python's `deepcopy` defaults to pickle/unpickle the object. - # Stages are complex objects (with references to `repo`, - # `outs`, and `deps`) that cause struggles when you try - # to serialize them. We need to create a copy of the graph - # itself, and then reverse it, instead of using - # graph.reverse() directly because it calls `deepcopy` - # underneath -- unless copy=False is specified. - nodes = nx.dfs_postorder_nodes(active.reverse(copy=False), stage) - all_pipelines += reversed(list(nodes)) - else: - all_pipelines += nx.dfs_postorder_nodes(active, stage) - - steps = [] - for stage in all_pipelines: - if stage not in steps: - # NOTE: order of steps still matters for single_item - if single_item and stage not in stages: - continue - - steps.append(stage) - - return steps +def _remove_frozen_stages(graph: "DiGraph") -> "DiGraph": + g = cast("DiGraph", graph.copy()) + for stage in graph: + if stage.frozen: + # NOTE: disconnect frozen stage from its dependencies + g.remove_edges_from(graph.out_edges(stage)) + return g diff --git a/tests/func/repro/test_repro.py b/tests/func/repro/test_repro.py index 1da583db5..a2b10cb0d 100644 --- a/tests/func/repro/test_repro.py +++ b/tests/func/repro/test_repro.py @@ -1282,3 +1282,17 @@ def test_repro_rm_recursive(tmp_dir, dvc): dvc.reproduce() assert (tmp_dir / "dir").exists() assert not (tmp_dir / "dir" / "foo").exists() + + +def test_repro_single_item_with_multiple_targets(tmp_dir, dvc, copy_script): + stage1 = dvc.stage.add(cmd="echo foo > foo", outs=["foo"], name="gen-foo") + with dvc.lock: + stage1.run() + + stage2 = dvc.stage.add( + cmd="python copy.py foo bar", deps=["foo"], outs=["bar"], name="copy-foo-bar" + ) + assert dvc.reproduce(["copy-foo-bar", "gen-foo"], single_item=True) == [ + stage2, + stage1, + ] diff --git a/tests/unit/repo/test_graph.py b/tests/unit/repo/test_graph.py new file mode 100644 index 000000000..380ffa3c7 --- /dev/null +++ b/tests/unit/repo/test_graph.py @@ -0,0 +1,69 @@ +import pytest +from networkx import DiGraph +from networkx.utils import graphs_equal + +from dvc.repo.graph import get_steps, get_subgraph_of_nodes + + +@pytest.mark.parametrize( + "nodes,downstream,expected_edges", + [ + ([], False, {1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}), + ([1], False, {1: [2, 3], 2: [4, 5], 3: [6, 7]}), + ([2], False, {2: [4, 5]}), + ([3], False, {3: [6, 7]}), + ([8], False, [(8, 9)]), + ([2, 3, 8], False, {2: [4, 5], 3: [6, 7], 8: [9]}), + ([4], False, {4: []}), + ([], True, {1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}), + ([1], True, {1: []}), + ([9], True, [(8, 9)]), + ([2], True, [(1, 2)]), + ([6], True, [(1, 3), (3, 6)]), + ([2, 3, 8], True, {1: [2, 3], 8: []}), + ([4, 7], True, {1: [2, 3], 2: [4], 3: [7]}), + ], +) +def test_subgraph_of_nodes(nodes, downstream, expected_edges): + r""" + 1 + / \ + 2 3 8 + / \ / \ | + 4 5 6 7 9 + """ + graph = DiGraph({1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}) + subgraph = get_subgraph_of_nodes(graph, nodes, downstream=downstream) + expected = DiGraph(expected_edges) + assert graphs_equal(expected, subgraph) + + +@pytest.mark.parametrize( + "nodes,downstream,expected_steps", + [ + ([], False, [4, 5, 2, 6, 7, 3, 1, 9, 8]), + ([1], False, [4, 5, 2, 6, 7, 3, 1]), + ([2], False, [4, 5, 2]), + ([3], False, [6, 7, 3]), + ([8], False, [9, 8]), + ([2, 3, 8], False, [4, 5, 2, 6, 7, 3, 9, 8]), + ([4], False, [4]), + ([], True, [4, 5, 2, 6, 7, 3, 1, 9, 8]), + ([1], True, [1]), + ([9], True, [9, 8]), + ([2], True, [2, 1]), + ([6], True, [6, 3, 1]), + ([2, 3, 8], True, [8, 2, 3, 1]), + ([4, 7], True, [4, 2, 7, 3, 1]), + ], +) +def test_steps(nodes, downstream, expected_steps): + r""" + 1 + / \ + 2 3 8 + / \ / \ | + 4 5 6 7 9 + """ + graph = DiGraph({1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}) + assert get_steps(graph, nodes, downstream=downstream) == expected_steps