mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
repro: move code around (#9701)
emphasizing on building the graph correctly, and preparing for being able to work with graph directly.
This commit is contained in:
skshetry
@@ -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
|
||||
|
||||
|
||||
+82
-57
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user