diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 65a0ae7f0..0c95ac373 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -7,7 +7,7 @@ from concurrent.futures import CancelledError, ProcessPoolExecutor, wait from contextlib import contextmanager from functools import wraps from multiprocessing import Manager -from typing import Iterable, Mapping, Optional +from typing import Dict, Iterable, Mapping, Optional from funcy import cached_property, first @@ -160,7 +160,7 @@ class Experiments: # checkout and detach at branch (or current HEAD) if detach_rev: - head = detach_rev + head: Optional[str] = detach_rev elif branch: head = branch else: @@ -347,7 +347,7 @@ class Experiments: else: resume_rev = self.scm.resolve_rev(checkpoint_resume) allow_multiple = "params" in kwargs - branch = self.get_branch_by_rev( + branch: Optional[str] = self.get_branch_by_rev( resume_rev, allow_multiple=allow_multiple ) if not branch: @@ -392,7 +392,7 @@ class Experiments: revs: Optional[Iterable] = None, keep_stash: Optional[bool] = True, **kwargs, - ): + ) -> Mapping[str, str]: """Reproduce the specified experiments. Args: @@ -400,6 +400,10 @@ class Experiments: reproduced. keep_stash: If True, stashed experiments will be preserved if they fail to reproduce successfully. + + Returns: + dict mapping successfully reproduced experiment revs to their + hashes. """ stash_revs = self.stash_revs @@ -445,7 +449,7 @@ class Experiments: for index in to_drop: self.stash.drop(index) - result = {} + result: Dict[str, str] = {} for _, exp_result in exec_results.items(): result.update(exp_result) return result @@ -472,7 +476,7 @@ class Experiments: branch=item.branch, cache_dir=self.repo.cache.local.cache_dir, ) - executors[item.rev] = executor + executors[stash_rev] = executor for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): self.scm.remove_ref(ref) @@ -484,9 +488,11 @@ class Experiments: ) -> Mapping[str, Mapping[str, str]]: """Run dvc repro for the specified BaseExecutors in parallel. - Returns dict containing successfully executed experiments. + Returns: + dict mapping stash revs to the successfully executed experiments + for each stash rev. """ - result = defaultdict(dict) + result: Dict[str, Dict[str, str]] = defaultdict(dict) manager = Manager() pid_q = manager.Queue() @@ -608,7 +614,9 @@ class Experiments: return ref_info.baseline_sha return None - def get_branch_by_rev(self, rev: str, allow_multiple: bool = False) -> str: + def get_branch_by_rev( + self, rev: str, allow_multiple: bool = False + ) -> Optional[str]: """Returns full refname for the experiment branch containing rev.""" ref_infos = list(exp_refs_by_rev(self.scm, rev)) if not ref_infos: diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 711b50734..dae89f445 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -105,7 +105,7 @@ class BaseExecutor: return os.path.join(self.root_dir, self._dvc_dir) @staticmethod - def hash_exp(stages): + def hash_exp(stages: Iterable["PipelineStage"]): exp_data = {} for stage in stages: if isinstance(stage, PipelineStage): @@ -145,9 +145,16 @@ class BaseExecutor: self, dest_scm: "Git", force: bool = False, - on_diverged: Callable[[str], None] = None, + on_diverged: Callable[[str, bool], None] = None, ) -> Iterable[str]: - """Fetch reproduced experiments into the specified SCM.""" + """Fetch reproduced experiments into the specified SCM. + + Args: + dest_scm: Destination Git instance. + force: If True, diverged refs will be overwritten + on_diverged: Callback in the form on_diverged(ref, is_checkpoint) + to be called when an experiment ref has diverged. + """ refs = [] for ref in self.scm.iter_refs(base=EXPS_NAMESPACE): if not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH: @@ -191,13 +198,16 @@ class BaseExecutor: rev: str, cwd: Optional[str] = None, name: Optional[str] = None, - ) -> Tuple[bool, Optional[str]]: + ) -> Tuple[Optional[str], bool]: """Run dvc repro and return the result. Returns tuple of (exp_hash, force) where exp_hash is the experiment hash (or None on error) and force is a bool specifying whether or not this experiment should force overwrite any existing duplicates. """ + from dvc.repo.checkout import checkout as dvc_checkout + from dvc.repo.reproduce import reproduce as dvc_reproduce + unchanged = [] queue.put((rev, os.getpid())) @@ -207,8 +217,8 @@ class BaseExecutor: [stage for stage in stages if isinstance(stage, PipelineStage)] ) - result = None - force = False + result: Optional[str] = None + repro_force: bool = False try: dvc = Repo(dvc_dir) @@ -227,7 +237,7 @@ class BaseExecutor: args = [] kwargs = {} - force = kwargs.get("force", False) + repro_force = kwargs.get("force", False) # NOTE: for checkpoint experiments we handle persist outs slightly # differently than normal: @@ -240,10 +250,11 @@ class BaseExecutor: # be removed/does not yet exist) so that our executor workspace # is not polluted with the (persistent) out from an unrelated # experiment run - dvc.checkout(force=True, quiet=True) + dvc_checkout(dvc, force=True, quiet=True) checkpoint_func = partial(cls.checkpoint_callback, dvc.scm, name) - stages = dvc.reproduce( + stages = dvc_reproduce( + dvc, *args, on_unchanged=filter_pipeline, checkpoint_func=checkpoint_func, @@ -251,6 +262,7 @@ class BaseExecutor: ) exp_hash = cls.hash_exp(stages) + result = exp_hash exp_rev = cls.commit(dvc.scm, exp_hash, exp_name=name) if dvc.scm.get_ref(EXEC_CHECKPOINT): dvc.scm.set_ref(EXEC_CHECKPOINT, exp_rev) @@ -265,18 +277,18 @@ class BaseExecutor: # ideally we would return stages here like a normal repro() call, but # stages is not currently picklable and cannot be returned across # multiprocessing calls - return result, force + return result, repro_force @classmethod def checkpoint_callback( cls, scm: "Git", name: Optional[str], - unchanged: Iterable[PipelineStage], - stages: Iterable[PipelineStage], + unchanged: Iterable["PipelineStage"], + stages: Iterable["PipelineStage"], ): try: - exp_hash = cls.hash_exp(stages + unchanged) + exp_hash = cls.hash_exp(list(stages) + list(unchanged)) exp_rev = cls.commit(scm, exp_hash, exp_name=name) scm.set_ref(EXEC_CHECKPOINT, exp_rev) logger.info("Checkpoint experiment iteration '%s'.", exp_rev[:7]) @@ -321,7 +333,8 @@ class LocalExecutor(BaseExecutor): **kwargs, ): self._tmp_dir = TemporaryDirectory(dir=tmp_dir) - super().__init__(*args, root_dir=self._tmp_dir.name, **kwargs) + kwargs["root_dir"] = self._tmp_dir.name + super().__init__(*args, **kwargs) if cache_dir: self._config(cache_dir) logger.debug(