diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 516c7eebd..e4c056e17 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -14,12 +14,13 @@ from typing import Iterable, Optional from funcy import cached_property, first -from dvc.exceptions import DownloadError, DvcException +from dvc.exceptions import DownloadError, DvcException, UploadError from dvc.path_info import PathInfo from dvc.progress import Tqdm from dvc.repo.experiments.executor import ExperimentExecutor, LocalExecutor from dvc.scm.git import Git from dvc.stage.serialize import to_lockfile +from dvc.tree.repo import RepoTree from dvc.utils import dict_sha256, env2bool, relpath from dvc.utils.fs import remove @@ -52,6 +53,8 @@ class UnchangedExperimentError(DvcException): class BaselineMismatchError(DvcException): def __init__(self, rev, expected): + if hasattr(rev, "hexsha"): + rev = rev.hexsha rev_str = f"{rev[:7]}" if rev is not None else "dangling commit" super().__init__( f"Experiment derived from '{rev_str}', expected '{expected[:7]}'." @@ -236,6 +239,7 @@ class Experiments: """Commit stages as an experiment and return the commit SHA.""" if not self.scm.is_dirty(): raise UnchangedExperimentError(self.scm.get_rev()) + rev = self.scm.get_rev() exp_name = f"{rev[:7]}-{exp_hash}" if branch: @@ -326,19 +330,26 @@ class Experiments: for rev in revs } - # setup executors + logger.debug( + "Reproducing experiment revs '%s'", + ", ".join((rev[:7] for rev in to_run)), + ) + + # setup executors - unstash experiment, generate executor, upload + # contents of (unstashed) exp workspace to the executor tree executors = {} for rev, baseline_rev in to_run.items(): - tree = self.scm.get_tree(rev) - repro_args, repro_kwargs = self._unpack_args(tree) + self._scm_checkout(baseline_rev) + self.scm.repo.git.stash("apply", rev) + repro_args, repro_kwargs = self._unpack_args() executor = LocalExecutor( - tree, baseline_rev, repro_args=repro_args, repro_kwargs=repro_kwargs, dvc_dir=self.dvc_dir, cache_dir=self.repo.cache.local.cache_dir, ) + self._collect_input(executor) executors[rev] = executor exec_results = self._reproduce(executors, **kwargs) @@ -421,34 +432,56 @@ class Experiments: return result + def _collect_input(self, executor: ExperimentExecutor): + """Copy (upload) input from the experiments workspace to the executor + tree. + """ + logger.debug("Collecting input for '%s'", executor.tmp_dir) + repo_tree = RepoTree(self.exp_dvc) + self._process( + executor.tree, + self.exp_dvc.tree, + executor.collect_files(self.exp_dvc.tree, repo_tree), + ) + def _collect_output(self, executor: ExperimentExecutor): """Copy (download) output from the executor tree into experiments workspace. """ - from dvc.cache.local import _log_exceptions - logger.debug("Collecting output from '%s'", executor.tmp_dir) - dest_tree = self.exp_dvc.tree - src_tree = executor.tree + self._process( + self.exp_dvc.tree, + executor.tree, + executor.collect_output(), + download=True, + ) + + @staticmethod + def _process(dest_tree, src_tree, collected_files, download=False): + from dvc.cache.local import _log_exceptions from_infos = [] to_infos = [] names = [] - for from_info in executor.collect_output(): + for from_info in collected_files: from_infos.append(from_info) fname = from_info.relative_to(src_tree.path_info) names.append(str(fname)) - to_info = dest_tree.path_info / fname to_infos.append(dest_tree.path_info / fname) - logger.debug(f"from '{from_info}' to '{to_info}'") total = len(from_infos) - func = partial( - _log_exceptions(dest_tree.download, "download"), - dir_mode=dest_tree.dir_mode, - file_mode=dest_tree.file_mode, - ) - with Tqdm(total=total, unit="file", desc="Downloading") as pbar: + if download: + func = partial( + _log_exceptions(src_tree.download, "download"), + dir_mode=dest_tree.dir_mode, + file_mode=dest_tree.file_mode, + ) + desc = "Downloading" + else: + func = partial(_log_exceptions(dest_tree.upload, "upload")) + desc = "Uploading" + + with Tqdm(total=total, unit="file", desc=desc) as pbar: func = pbar.wrap_fn(func) # TODO: parallelize this, currently --jobs for repro applies to # number of repro executors not download threads @@ -456,7 +489,9 @@ class Experiments: fails = sum(dl_executor.map(func, from_infos, to_infos, names)) if fails: - raise DownloadError(fails) + if download: + raise DownloadError(fails) + raise UploadError(fails) @scm_locked def checkout_exp(self, rev): diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 5f2fd7f46..21ce60142 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -11,8 +11,6 @@ from dvc.stage import PipelineStage from dvc.tree.base import BaseTree from dvc.tree.local import LocalTree from dvc.tree.repo import RepoTree -from dvc.utils import relpath -from dvc.utils.fs import copy_fobj_to_file, makedirs logger = logging.getLogger(__name__) @@ -21,7 +19,6 @@ class ExperimentExecutor: """Base class for executing experiments in parallel. Args: - src_tree: source tree for this experiment. baseline_rev: baseline revision that this experiment is derived from. Optional keyword args: @@ -29,8 +26,7 @@ class ExperimentExecutor: repro_kwargs: Keyword args to be passed into reproduce. """ - def __init__(self, src_tree: BaseTree, baseline_rev: str, **kwargs): - self.src_tree = src_tree + def __init__(self, baseline_rev: str, **kwargs): self.baseline_rev = baseline_rev self.repro_args = kwargs.pop("repro_args", []) self.repro_kwargs = kwargs.pop("repro_kwargs", {}) @@ -73,27 +69,29 @@ class ExperimentExecutor: class LocalExecutor(ExperimentExecutor): - """Local machine exepriment executor.""" + """Local machine experiment executor.""" + + def __init__(self, baseline_rev: str, **kwargs): + from dvc.repo import Repo - def __init__(self, src_tree: BaseTree, baseline_rev: str, **kwargs): dvc_dir = kwargs.pop("dvc_dir") cache_dir = kwargs.pop("cache_dir") - super().__init__(src_tree, baseline_rev, **kwargs) + super().__init__(baseline_rev, **kwargs) self.tmp_dir = TemporaryDirectory() - logger.debug("Init local executor in dir '%s'.", self.tmp_dir) + + # init empty DVC repo (will be overwritten when input is uploaded) + Repo.init(root_dir=self.tmp_dir.name, no_scm=True) + logger.debug( + "Init local executor in dir '%s' with baseline '%s'.", + self.tmp_dir, + baseline_rev[:7], + ) self.dvc_dir = os.path.join(self.tmp_dir.name, dvc_dir) - try: - for fname in src_tree.walk_files(src_tree.tree_root): - dest = self.path_info / relpath(fname, src_tree.tree_root) - if not os.path.exists(dest.parent): - makedirs(dest.parent) - with src_tree.open(fname, "rb") as fobj: - copy_fobj_to_file(fobj, dest) - except Exception: - self.tmp_dir.cleanup() - raise self._config(cache_dir) self._tree = LocalTree(self.dvc, {"url": self.dvc.root_dir}) + # override default CACHE_MODE since files must be writable in order + # to run repro + self._tree.CACHE_MODE = 0o644 def _config(self, cache_dir): local_config = os.path.join(self.dvc_dir, "config.local") @@ -151,11 +149,13 @@ class LocalExecutor(ExperimentExecutor): def collect_output(self) -> Iterable["PathInfo"]: repo_tree = RepoTree(self.dvc) + yield from self.collect_files(self.tree, repo_tree) + + @staticmethod + def collect_files(tree: BaseTree, repo_tree: RepoTree): for fname in repo_tree.walk_files(repo_tree.root_dir, dvcfiles=True): if not repo_tree.isdvc(fname): - yield self.tree.path_info / fname.relative_to( - repo_tree.root_dir - ) + yield tree.path_info / fname.relative_to(repo_tree.root_dir) def cleanup(self): logger.debug("Removing tmpdir '%s'", self.tmp_dir) diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 338c68a18..33d41c6d2 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -1,10 +1,6 @@ -import pytest - from tests.func.test_repro_multistage import COPY_SCRIPT -# https://github.com/iterative/dvc/issues/4401 -@pytest.mark.flaky(max_runs=3, min_passes=1) def test_new_simple(tmp_dir, scm, dvc, mocker): tmp_dir.gen("copy.py", COPY_SCRIPT) tmp_dir.gen("params.yaml", "foo: 1")