diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 30fc3835b..0b51d0444 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -4,20 +4,25 @@ import re import tempfile from collections import defaultdict from collections.abc import Mapping -from concurrent.futures import ProcessPoolExecutor, as_completed +from concurrent.futures import ( + ProcessPoolExecutor, + ThreadPoolExecutor, + as_completed, +) from contextlib import contextmanager -from functools import wraps +from functools import partial, wraps from typing import Iterable, Optional from funcy import cached_property, first -from dvc.exceptions import DvcException +from dvc.exceptions import DownloadError, DvcException 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.utils import dict_sha256, env2bool, relpath -from dvc.utils.fs import copyfile, remove +from dvc.utils.fs import remove logger = logging.getLogger(__name__) @@ -204,8 +209,7 @@ class Experiments: self.scm.add(self.args_file) def _unpack_args(self, tree=None): - args_file = os.path.join(self.exp_dvc.tmp_dir, self.PACKED_ARGS_FILE) - return ExperimentExecutor.unpack_repro_args(args_file, tree=tree) + return ExperimentExecutor.unpack_repro_args(self.args_file, tree=tree) def _update_params(self, params: dict): """Update experiment params files with the specified values.""" @@ -398,10 +402,19 @@ class Experiments: exc = future.exception() if exc is None: exp_hash = future.result() - logger.debug(f"ran exp based on {executor.baseline_rev}") self._scm_checkout(executor.baseline_rev) - self._collect_output(executor.baseline_rev, executor) - remove(self.args_file) + try: + self._collect_output(executor) + except DownloadError: + logger.error( + "Failed to collect output for experiment '%s'", + rev, + ) + continue + finally: + if os.path.exists(self.args_file): + remove(self.args_file) + try: exp_rev = self._commit(exp_hash) except UnchangedExperimentError: @@ -421,12 +434,42 @@ class Experiments: return result - def _collect_output(self, rev: str, executor: ExperimentExecutor): - logger.debug("copying tmp output from '%s'", executor.tmp_dir) - tree = self.scm.get_tree(rev) - for fname in tree.walk_files(tree.tree_root): - src = executor.path_info / relpath(fname, tree.tree_root) - copyfile(src, fname) + 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 + + from_infos = [] + to_infos = [] + names = [] + for from_info in executor.collect_output(): + 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: + func = pbar.wrap_fn(func) + # TODO: parallelize this, currently --jobs for repro applies to + # number of repro executors not download threads + with ThreadPoolExecutor(max_workers=1) as dl_executor: + fails = sum(dl_executor.map(func, from_infos, to_infos, names)) + + if fails: + raise DownloadError(fails) @scm_locked def checkout_exp(self, rev): diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 494c389f3..3788c3cf4 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -2,12 +2,15 @@ import logging import os import pickle from tempfile import TemporaryDirectory +from typing import Iterable from funcy import cached_property from dvc.path_info import PathInfo +from dvc.repo.tree import RepoTree from dvc.stage import PipelineStage from dvc.tree.base import BaseTree +from dvc.tree.local import LocalTree from dvc.utils import relpath from dvc.utils.fs import copy_fobj_to_file, makedirs @@ -32,8 +35,23 @@ class ExperimentExecutor: self.repro_args = kwargs.pop("repro_args", []) self.repro_kwargs = kwargs.pop("repro_kwargs", {}) - def run(self): - pass + @property + def tree(self) -> BaseTree: + raise NotImplementedError + + @staticmethod + def reproduce(dvc_dir, cwd=None, **kwargs): + raise NotImplementedError + + def collect_output(self) -> Iterable["PathInfo"]: + """Iterate over output pathnames for this executor. + + For DVC outs, only the .dvc file path will be yielded. DVC outs + themselves should be fetched from remote executor cache in the normal + fetch/pull way. For local executors outs will already be available via + shared local cache. + """ + raise NotImplementedError def cleanup(self): pass @@ -75,6 +93,7 @@ class LocalExecutor(ExperimentExecutor): self.tmp_dir.cleanup() raise self._config(cache_dir) + self._tree = LocalTree(self.dvc, {"url": self.dvc.root_dir}) def _config(self, cache_dir): local_config = os.path.join(self.dvc_dir, "config.local") @@ -93,6 +112,10 @@ class LocalExecutor(ExperimentExecutor): def path_info(self): return PathInfo(self.tmp_dir.name) + @property + def tree(self): + return self._tree + @staticmethod def reproduce(dvc_dir, cwd=None, **kwargs): """Run dvc repro and return the result.""" @@ -126,6 +149,14 @@ class LocalExecutor(ExperimentExecutor): # multiprocessing calls return hash_exp(stages + unchanged) + def collect_output(self) -> Iterable["PathInfo"]: + repo_tree = RepoTree(self.dvc) + 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 + ) + def cleanup(self): logger.debug("Removing tmpdir '%s'", self.tmp_dir) self.tmp_dir.cleanup()