From 57263f15fa9c23097d2e1b84a01a8a457eac6144 Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 24 Dec 2021 17:12:55 +0800 Subject: [PATCH] Pull: accept multiple as input(#7196) fix: #7196 New feature: now `dvc exp pull` can accept a list of experiment names as input. 2. add support to list input for the `exp pull` 3. add function tests for the multi exp pull 4. add unit test for the arg parse --- dvc/command/experiments/pull.py | 5 ++- dvc/repo/experiments/exceptions.py | 4 +- dvc/repo/experiments/pull.py | 56 ++++++++++++++++++-------- dvc/repo/experiments/push.py | 4 +- dvc/repo/experiments/utils.py | 3 +- dvc/types.py | 3 +- tests/func/experiments/test_remote.py | 40 +++++++++++------- tests/unit/command/test_experiments.py | 2 +- 8 files changed, 77 insertions(+), 40 deletions(-) diff --git a/dvc/command/experiments/pull.py b/dvc/command/experiments/pull.py index b84b10083..45af21a60 100644 --- a/dvc/command/experiments/pull.py +++ b/dvc/command/experiments/pull.py @@ -87,6 +87,9 @@ def add_parser(experiments_subparsers, parent_parser): metavar="", ) experiments_pull_parser.add_argument( - "experiment", help="Experiment to pull.", metavar="" + "experiment", + nargs="+", + help="Experiments to pull.", + metavar="", ) experiments_pull_parser.set_defaults(func=CmdExperimentsPull) diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py index 6befa3b65..d6a7a304e 100644 --- a/dvc/repo/experiments/exceptions.py +++ b/dvc/repo/experiments/exceptions.py @@ -1,5 +1,6 @@ +from typing import Iterable, List + from dvc.exceptions import InvalidArgumentError -from dvc.types import Iterable, List, Optional from .base import ExpRefInfo @@ -20,6 +21,7 @@ class AmbiguousExpRefInfo(InvalidArgumentError): msg.extend([f"\t{info}" for info in exp_ref_list]) super().__init__("\n".join(msg)) + class UnresolvedExpNamesError(InvalidArgumentError): def __init__( self, unresolved_list: List[str], *args, git_remote: str = None diff --git a/dvc/repo/experiments/pull.py b/dvc/repo/experiments/pull.py index 8760e523e..78b0442db 100644 --- a/dvc/repo/experiments/pull.py +++ b/dvc/repo/experiments/pull.py @@ -1,9 +1,12 @@ import logging +from typing import Iterable, Union -from dvc.exceptions import DvcException, InvalidArgumentError +from dvc.exceptions import DvcException from dvc.repo import locked from dvc.repo.scm_context import scm_context +from dvc.scm import TqdmGit +from .exceptions import UnresolvedExpNamesError from .utils import exp_commits, resolve_name logger = logging.getLogger(__name__) @@ -12,40 +15,61 @@ logger = logging.getLogger(__name__) @locked @scm_context def pull( - repo, git_remote, exp_name, *args, force=False, pull_cache=False, **kwargs + repo, + git_remote: str, + exp_names: Union[Iterable[str], str], + *args, + force: bool = False, + pull_cache: bool = False, + **kwargs, ): - exp_ref_dict = resolve_name(repo.scm, exp_name, git_remote) - exp_ref = exp_ref_dict[exp_name] - if not exp_ref: - raise InvalidArgumentError( - f"Experiment '{exp_name}' does not exist in '{git_remote}'" - ) + if isinstance(exp_names, str): + exp_names = [exp_names] + exp_ref_dict = resolve_name(repo.scm, exp_names, git_remote) + unresolved_exp_names = [ + exp_name + for exp_name, exp_ref in exp_ref_dict.items() + if exp_ref is None + ] + if unresolved_exp_names: + raise UnresolvedExpNamesError(unresolved_exp_names) + exp_ref_set = exp_ref_dict.values() + _pull(repo, git_remote, exp_ref_set, force, pull_cache, **kwargs) + + +def _pull( + repo, + git_remote: str, + exp_refs, + force: bool, + pull_cache: bool, + **kwargs, +): def on_diverged(refname: str, rev: str) -> bool: if repo.scm.get_ref(refname) == rev: return True + exp_name = refname.split("/")[-1] raise DvcException( f"Local experiment '{exp_name}' has diverged from remote " "experiment with the same name. To override the local experiment " "re-run with '--force'." ) - refspec = f"{exp_ref}:{exp_ref}" - logger.debug("git pull experiment '%s' -> '%s'", git_remote, refspec) - - from dvc.scm import TqdmGit + refspec_list = [f"{exp_ref}:{exp_ref}" for exp_ref in exp_refs] + logger.debug(f"git pull experiment '{git_remote}' -> '{refspec_list}'") with TqdmGit(desc="Fetching git refs") as pbar: repo.scm.fetch_refspecs( git_remote, - [refspec], + refspec_list, force=force, on_diverged=on_diverged, progress=pbar.update_git, ) if pull_cache: - _pull_cache(repo, exp_ref, **kwargs) + _pull_cache(repo, exp_refs, **kwargs) def _pull_cache( @@ -56,8 +80,8 @@ def _pull_cache( run_cache=False, odb=None, ): - revs = list(exp_commits(repo.scm, [exp_ref])) - logger.debug("dvc fetch experiment '%s'", exp_ref) + revs = list(exp_commits(repo.scm, exp_ref)) + logger.debug(f"dvc fetch experiment '{exp_ref}'") repo.fetch( jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs, odb=odb ) diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py index 700d86d08..6666c8c84 100644 --- a/dvc/repo/experiments/push.py +++ b/dvc/repo/experiments/push.py @@ -1,10 +1,10 @@ import logging +from typing import Iterable, Union from dvc.exceptions import DvcException from dvc.repo import locked from dvc.repo.scm_context import scm_context from dvc.scm import TqdmGit -from dvc.types import IterableStr, Union from .exceptions import UnresolvedExpNamesError from .utils import exp_commits, push_refspec, resolve_name @@ -17,7 +17,7 @@ logger = logging.getLogger(__name__) def push( repo, git_remote: str, - exp_names: Union[IterableStr, str], + exp_names: Union[Iterable[str], str], *args, force: bool = False, push_cache: bool = False, diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 239189f37..9abf301c1 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -13,6 +13,7 @@ from typing import ( from scmrepo.git import Git from dvc.exceptions import InvalidArgumentError +from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo from .base import ( EXEC_BASELINE, @@ -21,8 +22,6 @@ from .base import ( EXPS_STASH, ExpRefInfo, ) -from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo - def exp_refs(scm: "Git") -> Generator["ExpRefInfo", None, None]: diff --git a/dvc/types.py b/dvc/types.py index ee2449339..cca8e17a0 100644 --- a/dvc/types.py +++ b/dvc/types.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union if TYPE_CHECKING: from os import PathLike @@ -10,4 +10,3 @@ OptStr = Optional[str] TargetType = Union[List[str], str] DictStrAny = Dict[str, Any] DictAny = Dict[Any, Any] -IterableStr = Iterable[str] diff --git a/tests/func/experiments/test_remote.py b/tests/func/experiments/test_remote.py index e36230ac9..cc697eb96 100644 --- a/tests/func/experiments/test_remote.py +++ b/tests/func/experiments/test_remote.py @@ -151,19 +151,29 @@ def test_pull(tmp_dir, scm, dvc, git_downstream, exp_stage, use_url): remote = git_downstream.url if use_url else git_downstream.remote downstream_exp = git_downstream.tmp_dir.dvc.experiments with pytest.raises(InvalidArgumentError): - downstream_exp.pull(remote, "foo") + downstream_exp.pull(remote, ["foo"]) + results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"]) + exp1 = first(results) + ref_info1 = first(exp_refs_by_rev(scm, exp1)) results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"]) - exp = first(results) - ref_info = first(exp_refs_by_rev(scm, exp)) + exp2 = first(results) + ref_info2 = first(exp_refs_by_rev(scm, exp2)) + results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"]) + exp3 = first(results) + ref_info3 = first(exp_refs_by_rev(scm, exp3)) - downstream_exp.pull(remote, ref_info.name) - assert git_downstream.tmp_dir.scm.get_ref(str(ref_info)) == exp + downstream_exp.pull( + git_downstream.remote, [ref_info1.name, ref_info2.name], force=True + ) + assert git_downstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1 + assert git_downstream.tmp_dir.scm.get_ref(str(ref_info2)) == exp2 + assert git_downstream.tmp_dir.scm.get_ref(str(ref_info3)) is None - git_downstream.tmp_dir.scm.remove_ref(str(ref_info)) + git_downstream.tmp_dir.scm.remove_ref(str(ref_info1)) - downstream_exp.pull(remote, str(ref_info)) - assert git_downstream.tmp_dir.scm.get_ref(str(ref_info)) == exp + downstream_exp.pull(remote, [str(ref_info1)]) + assert git_downstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1 def test_pull_diverged(tmp_dir, scm, dvc, git_downstream, exp_stage): @@ -193,7 +203,7 @@ def test_pull_checkpoint(tmp_dir, scm, dvc, git_downstream, checkpoint_stage): ref_info_a = first(exp_refs_by_rev(scm, exp_a)) downstream_exp = git_downstream.tmp_dir.dvc.experiments - downstream_exp.pull(git_downstream.remote, ref_info_a.name, force=True) + downstream_exp.pull(git_downstream.remote, [ref_info_a.name], force=True) assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_a)) == exp_a results = dvc.experiments.run( @@ -202,7 +212,7 @@ def test_pull_checkpoint(tmp_dir, scm, dvc, git_downstream, checkpoint_stage): exp_b = first(results) ref_info_b = first(exp_refs_by_rev(scm, exp_b)) - downstream_exp.pull(git_downstream.remote, ref_info_b.name, force=True) + downstream_exp.pull(git_downstream.remote, [ref_info_b.name], force=True) assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_b)) == exp_b @@ -225,13 +235,13 @@ def test_pull_ambiguous_name(tmp_dir, scm, dvc, git_downstream, exp_stage): remote = git_downstream.remote downstream_exp = git_downstream.tmp_dir.dvc.experiments with pytest.raises(InvalidArgumentError): - downstream_exp.pull(remote, "foo") + downstream_exp.pull(remote, ["foo"]) - downstream_exp.pull(remote, str(ref_info_b)) + downstream_exp.pull(remote, [str(ref_info_b)]) assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_b)) == exp_b with git_downstream.tmp_dir.scm.detach_head(ref_info_a.baseline_sha): - downstream_exp.pull(remote, "foo") + downstream_exp.pull(remote, ["foo"]) assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_a)) == exp_a @@ -257,7 +267,7 @@ def test_push_pull_cache( remove(dvc.odb.local.cache_dir) - dvc.experiments.pull(remote, ref_info.name, pull_cache=True) + dvc.experiments.pull(remote, [ref_info.name], pull_cache=True) for x in range(2, checkpoint_stage.iterations + 1): hash_ = digest(str(x)) path = os.path.join(dvc.odb.local.cache_dir, hash_[:2], hash_[2:]) @@ -282,7 +292,7 @@ def test_auth_error_pull(tmp_dir, scm, dvc, http_auth_patch): GitAuthError, match=f"HTTP Git authentication is not supported: '{http_auth_patch}'", ): - dvc.experiments.pull(http_auth_patch, "foo") + dvc.experiments.pull(http_auth_patch, ["foo"]) def test_auth_error_push(tmp_dir, scm, dvc, exp_stage, http_auth_patch): diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 0802e8221..2ff01914e 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -277,7 +277,7 @@ def test_experiments_pull(dvc, scm, mocker): m.assert_called_once_with( cmd.repo, "origin", - "experiment", + ["experiment"], force=True, pull_cache=False, dvc_remote="my-remote",