[WIP] repro: simplify the logic for frozen stages (#5114)

Fixes #5082

Author: efiop

dvc/repo/reproduce.py

@@ -66,28 +66,6 @@ def _track_stage(stage):
     stage.repo.scm.track_changed_files()
 
 
-def _get_active_graph(G):
-    import networkx as nx
-
-    active = G.copy()
-    for stage in G:
-        if not stage.frozen:
-            continue
-        active.remove_edges_from(G.out_edges(stage))
-        for edge in G.out_edges(stage):
-            _, to_stage = edge
-            for node in nx.dfs_preorder_nodes(G, to_stage):
-                # NOTE: `in_degree` will return InDegreeView({}) if stage
-                # no longer exists in the `active` DAG.
-                if not active.in_degree(node):
-                    # NOTE: if some edge no longer exists `remove_edges_from`
-                    # will ignore it without error.
-                    active.remove_edges_from(G.out_edges(node))
-                    active.remove_node(node)
-
-    return active
-
-
 @locked
 @scm_context
 def reproduce(
@@ -104,7 +82,7 @@ def reproduce(
     if isinstance(targets, str):
         targets = [targets]
 
-    if not all_pipelines and targets is None:
+    if not all_pipelines and not targets:
         from dvc.dvcfile import PIPELINE_FILE
 
         targets = [PIPELINE_FILE]
@@ -113,36 +91,33 @@ def reproduce(
     if not interactive:
         kwargs["interactive"] = self.config["core"].get("interactive", False)
 
-    active_graph = _get_active_graph(self.graph)
-    active_pipelines = get_pipelines(active_graph)
-
     stages = set()
     if pipeline or all_pipelines:
+        pipelines = get_pipelines(self.graph)
         if all_pipelines:
-            pipelines = active_pipelines
+            used_pipelines = pipelines
         else:
-            pipelines = []
+            used_pipelines = []
             for target in targets:
                 stage = self.stage.get_target(target)
-                pipelines.append(get_pipeline(active_pipelines, stage))
+                used_pipelines.append(get_pipeline(pipelines, stage))
 
-        for pipeline in pipelines:
-            for stage in pipeline:
-                if pipeline.in_degree(stage) == 0:
+        for pline in used_pipelines:
+            for stage in pline:
+                if pline.in_degree(stage) == 0:
                     stages.add(stage)
     else:
         for target in targets:
             stages.update(
                 self.stage.collect(
                     target,
                     recursive=recursive,
-                    graph=active_graph,
                     accept_group=accept_group,
                     glob=glob,
                 )
             )
 
-    return _reproduce_stages(active_graph, list(stages), **kwargs)
+    return _reproduce_stages(self.graph, list(stages), **kwargs)
 
 
 def _reproduce_stages(
@@ -183,15 +158,15 @@ def _reproduce_stages(
 
     The derived evaluation of _downstream_ B would be: [B, D, E]
     """
-    pipeline = _get_pipeline(G, stages, downstream, single_item)
+    steps = _get_steps(G, stages, downstream, single_item)
 
     force_downstream = kwargs.pop("force_downstream", False)
     result = []
     unchanged = []
     # `ret` is used to add a cosmetic newline.
     ret = []
     checkpoint_func = kwargs.pop("checkpoint_func", None)
-    for stage in pipeline:
+    for stage in steps:
         if ret:
             logger.info("")
 
@@ -225,9 +200,17 @@ def _reproduce_stages(
     return result
 
 
-def _get_pipeline(G, stages, downstream, single_item):
+def _get_steps(G, stages, downstream, single_item):
     import networkx as nx
 
+    active = G.copy()
+    if not single_item:
+        # NOTE: frozen stages don't matter for single_item
+        for stage in G:
+            if stage.frozen:
+                # NOTE: disconnect frozen stage from its dependencies
+                active.remove_edges_from(G.out_edges(stage))
+
     all_pipelines = []
     for stage in stages:
         if downstream:
@@ -239,20 +222,21 @@ def _get_pipeline(G, stages, downstream, single_item):
             # 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(G.copy().reverse(copy=False), stage)
+            nodes = nx.dfs_postorder_nodes(active.reverse(copy=False), stage)
             all_pipelines += reversed(list(nodes))
         else:
-            all_pipelines += nx.dfs_postorder_nodes(G, stage)
+            all_pipelines += nx.dfs_postorder_nodes(active, stage)
 
-    pipeline = []
+    steps = []
     for stage in all_pipelines:
-        if stage not in pipeline:
+        if stage not in steps:
+            # NOTE: order of steps still matters for single_item
             if single_item and stage not in stages:
                 continue
 
-            pipeline.append(stage)
+            steps.append(stage)
 
-    return pipeline
+    return steps
 
 
 def _repro_callback(experiments_callback, unchanged, stages):

tests/func/test_repro_multistage.py

@@ -170,6 +170,19 @@ def test_non_existing_stage_name(tmp_dir, dvc, run_copy):
     assert main(["freeze", ":copy-file1-file3"]) != 0
 
 
+def test_repro_frozen(tmp_dir, dvc, run_copy):
+    (data_stage,) = tmp_dir.dvc_gen("data", "foo")
+    stage0 = run_copy("data", "stage0", name="copy-data-stage0")
+    run_copy("stage0", "stage1", name="copy-data-stage1")
+    run_copy("stage1", "stage2", name="copy-data-stage2")
+
+    dvc.freeze("copy-data-stage1")
+
+    tmp_dir.gen("data", "bar")
+    stages = dvc.reproduce()
+    assert stages == [data_stage, stage0]
+
+
 def test_downstream(tmp_dir, dvc):
     # The dependency graph should look like this:
     #

tests/unit/repo/test_reproduce.py

@@ -1,37 +1,5 @@
 import mock
 
-from dvc.repo.reproduce import _get_active_graph
-
-
-def test_get_active_graph(tmp_dir, dvc):
-    (pre_foo_stage,) = tmp_dir.dvc_gen({"pre-foo": "pre-foo"})
-    foo_stage = dvc.run(
-        single_stage=True, deps=["pre-foo"], outs=["foo"], cmd="echo foo > foo"
-    )
-    bar_stage = dvc.run(
-        single_stage=True, deps=["foo"], outs=["bar"], cmd="echo bar > bar"
-    )
-    baz_stage = dvc.run(
-        single_stage=True, deps=["foo"], outs=["baz"], cmd="echo baz > baz"
-    )
-
-    dvc.freeze("bar.dvc")
-
-    graph = dvc.graph
-    active_graph = _get_active_graph(graph)
-    assert active_graph.nodes == graph.nodes
-    assert set(active_graph.edges) == {
-        (foo_stage, pre_foo_stage),
-        (baz_stage, foo_stage),
-    }
-
-    dvc.freeze("baz.dvc")
-
-    graph = dvc.graph
-    active_graph = _get_active_graph(graph)
-    assert set(active_graph.nodes) == {bar_stage, baz_stage}
-    assert not active_graph.edges
-
 
 @mock.patch("dvc.repo.reproduce._reproduce_stage", returns=[])
 def test_number_reproduces(reproduce_stage_mock, tmp_dir, dvc):