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
This commit is contained in:
skshetry
2023-07-03 11:44:51 +05:45
committed by GitHub
parent 24d9a05d69
commit bfa5b1a0c3
4 changed files with 133 additions and 45 deletions
+27 -1
View File
@@ -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
+23 -44
View File
@@ -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
+14
View File
@@ -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,
]
+69
View File
@@ -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