From c8d180d2b0ead19d43481933d0f76ac558e79388 Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Tue, 4 Jul 2023 21:02:53 +0545 Subject: [PATCH] repro: move code around (#9701) emphasizing on building the graph correctly, and preparing for being able to work with graph directly. --- dvc/repo/graph.py | 9 --- dvc/repo/reproduce.py | 139 ++++++++++++++++++++-------------- tests/unit/repo/test_graph.py | 33 +------- 3 files changed, 83 insertions(+), 98 deletions(-) diff --git a/dvc/repo/graph.py b/dvc/repo/graph.py index 83f6c99b9..56ea1e0ed 100644 --- a/dvc/repo/graph.py +++ b/dvc/repo/graph.py @@ -56,15 +56,6 @@ def get_subgraph_of_nodes( 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 618d83f0d..65ee8c735 100644 --- a/dvc/repo/reproduce.py +++ b/dvc/repo/reproduce.py @@ -20,19 +20,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def _reproduce_stage(stage: "Stage", **kwargs) -> Optional["Stage"]: - if stage.frozen and not stage.is_import: - logger.warning( - "%s is frozen. Its dependencies are not going to be reproduced.", - stage, - ) - - ret = stage.reproduce(**kwargs) - if ret and not kwargs.get("dry", False): - stage.dump(update_pipeline=False) - return ret - - def collect_stages( repo: "Repo", targets: Iterable[str], @@ -45,6 +32,24 @@ def collect_stages( return ldistinct(stages) +def get_subgraph( + graph: "DiGraph", + nodes: Optional[List] = None, + pipeline: bool = False, + downstream: bool = False, +) -> "DiGraph": + import networkx as nx + + from .graph import get_pipeline, get_pipelines, get_subgraph_of_nodes + + if not pipeline or not nodes: + return get_subgraph_of_nodes(graph, nodes, downstream=downstream) + + pipelines = get_pipelines(graph) + used_pipelines = [get_pipeline(pipelines, node) for node in nodes] + return nx.compose_all(used_pipelines) + + def _remove_frozen_stages(graph: "DiGraph") -> "DiGraph": g = cast("DiGraph", graph.copy()) for stage in graph: @@ -54,6 +59,17 @@ def _remove_frozen_stages(graph: "DiGraph") -> "DiGraph": return g +def get_active_graph( + graph: "DiGraph", + stages: Optional[List["Stage"]] = None, + pipeline: bool = False, + downstream: bool = False, +) -> "DiGraph": + """Return the graph to operate.""" + processed = _remove_frozen_stages(graph) + return get_subgraph(processed, stages, pipeline=pipeline, downstream=downstream) + + def plan_repro( graph: "DiGraph", stages: Optional[List["Stage"]] = None, @@ -97,53 +113,21 @@ def plan_repro( """ import networkx as nx - from .graph import get_pipeline, get_pipelines, get_steps - - active = _remove_frozen_stages(graph) - if stages and pipeline: - pipelines = get_pipelines(active) - used_pipelines = [get_pipeline(pipelines, stage) for stage in stages] - # create a disjointed union of all the pipelines - active = nx.compose_all(used_pipelines) - return get_steps(active) - return get_steps(active, stages, downstream=downstream) + active = get_active_graph(graph, stages, pipeline=pipeline, downstream=downstream) + return list(nx.dfs_postorder_nodes(active)) -@locked -@scm_context -def reproduce( - self: "Repo", - targets: Union[Iterable[str], str, None] = None, - recursive: bool = False, - pipeline: bool = False, - all_pipelines: bool = False, - downstream: bool = False, - single_item: bool = False, - glob: bool = False, - **kwargs, -): - from dvc.dvcfile import PROJECT_FILE +def _reproduce_stage(stage: "Stage", **kwargs) -> Optional["Stage"]: + if stage.frozen and not stage.is_import: + logger.warning( + "%s is frozen. Its dependencies are not going to be reproduced.", + stage, + ) - if not kwargs.get("interactive", False): - kwargs["interactive"] = self.config["core"].get("interactive", False) - - stages: List["Stage"] = [] - if not all_pipelines: - targets_list = ensure_list(targets or PROJECT_FILE) - stages = collect_stages(self, targets_list, recursive=recursive, glob=glob) - - if kwargs.get("pull", False) and kwargs.get("run_cache", True): - logger.debug("Pulling run cache") - try: - self.stage_cache.pull(None) - except RunCacheNotSupported as e: - logger.warning("Failed to pull run cache: %s", e) - - steps = stages - if pipeline or all_pipelines or not single_item: - graph = self.index.graph - steps = plan_repro(graph, stages, pipeline=pipeline, downstream=downstream) - return _reproduce_stages(steps, **kwargs) + ret = stage.reproduce(**kwargs) + if ret and not kwargs.get("dry", False): + stage.dump(update_pipeline=False) + return ret def _reproduce_stages( @@ -172,3 +156,44 @@ def _reproduce_stages( if i < len(stages) - 1: logger.info("") # add a newline return result + + +@locked +@scm_context +def reproduce( + self: "Repo", + targets: Union[Iterable[str], str, None] = None, + recursive: bool = False, + pipeline: bool = False, + all_pipelines: bool = False, + downstream: bool = False, + single_item: bool = False, + glob: bool = False, + **kwargs, +): + from dvc.dvcfile import PROJECT_FILE + + if all_pipelines or pipeline: + single_item = False + downstream = False + + if not kwargs.get("interactive", False): + kwargs["interactive"] = self.config["core"].get("interactive", False) + + stages: List["Stage"] = [] + if not all_pipelines: + targets_list = ensure_list(targets or PROJECT_FILE) + stages = collect_stages(self, targets_list, recursive=recursive, glob=glob) + + if kwargs.get("pull", False) and kwargs.get("run_cache", True): + logger.debug("Pulling run cache") + try: + self.stage_cache.pull(None) + except RunCacheNotSupported as e: + logger.warning("Failed to pull run cache: %s", e) + + steps = stages + if not single_item: + graph = self.index.graph + steps = plan_repro(graph, stages, pipeline=pipeline, downstream=downstream) + return _reproduce_stages(steps, **kwargs) diff --git a/tests/unit/repo/test_graph.py b/tests/unit/repo/test_graph.py index 380ffa3c7..5aa1988b3 100644 --- a/tests/unit/repo/test_graph.py +++ b/tests/unit/repo/test_graph.py @@ -2,7 +2,7 @@ import pytest from networkx import DiGraph from networkx.utils import graphs_equal -from dvc.repo.graph import get_steps, get_subgraph_of_nodes +from dvc.repo.graph import get_subgraph_of_nodes @pytest.mark.parametrize( @@ -36,34 +36,3 @@ def test_subgraph_of_nodes(nodes, downstream, expected_edges): 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