diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index b6365f606..e184f068d 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -13,6 +13,7 @@ from typing import ( Callable, Dict, Iterable, + Iterator, List, NamedTuple, Optional, @@ -259,6 +260,12 @@ class BaseExecutor(ABC): **kwargs, ) + @classmethod + def _get_top_level_paths(cls, repo: "Repo") -> Iterable["str"]: + yield from _collect_top_level_metrics(repo) + yield from _collect_top_level_params(repo) + yield from repo.index._plot_sources # pylint: disable=protected-access + @classmethod def save( cls, @@ -280,11 +287,7 @@ class BaseExecutor(ABC): os.chdir(dvc.root_dir) include_untracked = include_untracked or [] - include_untracked.extend(_collect_top_level_metrics(dvc)) - include_untracked.extend(_collect_top_level_params(dvc)) - include_untracked.extend( - dvc.index._plot_sources # pylint: disable=protected-access - ) + include_untracked.extend(cls._get_top_level_paths(dvc)) # dvc repro automatically stages dvc.lock. Running redundant `git add` # on it causes an error when exiting the detached head context. if LOCK_FILE in dvc.scm.untracked_files(): @@ -540,11 +543,19 @@ class BaseExecutor(ABC): info.name, repro_force or checkpoint_reset, ) + + def after_repro(): + paths = list(cls._get_top_level_paths(dvc)) + if paths: + logger.debug("Staging top-level files: %s", paths) + dvc.scm_context.add(paths) + stages = dvc_reproduce( dvc, *args, on_unchanged=filter_pipeline, checkpoint_func=checkpoint_func, + after_repro_callback=after_repro, **kwargs, ) @@ -614,7 +625,7 @@ class BaseExecutor(ABC): log_errors: bool = True, copy_paths: Optional[List[str]] = None, **kwargs, - ): + ) -> Iterator["Repo"]: from dvc_studio_client.post_live_metrics import post_live_metrics from dvc.repo import Repo diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index 231662a04..3ce4fadf4 100644 --- a/dvc/repo/reproduce.py +++ b/dvc/repo/reproduce.py @@ -3,8 +3,6 @@ from functools import partial from typing import TYPE_CHECKING, Iterator, List from dvc.exceptions import DvcException, ReproductionError -from dvc.repo.metrics.show import _collect_top_level_metrics -from dvc.repo.params.show import _collect_top_level_params from dvc.repo.scm_context import scm_context from dvc.stage.exceptions import CheckpointKilledError @@ -76,25 +74,15 @@ def _track_stage(stage: "Stage") -> None: return context.track_changed_files() -def _track_top_level(repo: "Repo") -> None: - context = repo.scm_context - for metric_file in _collect_top_level_metrics(repo): - context.track_file(metric_file) - for param_file in _collect_top_level_params(repo): - context.track_file(param_file) - for plot_file in repo.index._plot_sources: # pylint: disable=protected-access - context.track_file(plot_file) - context.track_changed_files() - - @locked @scm_context -def reproduce( # noqa: C901 +def reproduce( # noqa: C901, PLR0912 self: "Repo", targets=None, recursive=False, pipeline=False, all_pipelines=False, + after_repro_callback=None, **kwargs, ): from .graph import get_pipeline, get_pipelines @@ -139,12 +127,18 @@ def reproduce( # noqa: C901 ) result = _reproduce_stages(self.index.graph, list(stages), **kwargs) - _track_top_level(self) + if callable(after_repro_callback): + after_repro_callback() return result def _reproduce_stages( # noqa: C901 - graph, stages, downstream=False, single_item=False, on_unchanged=None, **kwargs + graph, + stages, + downstream=False, + single_item=False, + on_unchanged=None, + **kwargs, ): r"""Derive the evaluation of the given node for the given graph.