diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index b1cc2691d..fabe357fc 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -550,18 +550,25 @@ class CmdExperimentsPush(CmdBase): def run(self): self.repo.experiments.push( - self.args.git_remote, self.args.experiment, force=self.args.force + self.args.git_remote, + self.args.experiment, + force=self.args.force, + push_cache=self.args.push_cache, + dvc_remote=self.args.dvc_remote, + jobs=self.args.jobs, + run_cache=self.args.run_cache, ) logger.info( - ( - "Pushed experiment '%s' to Git remote '%s'. " - "To push cache for this experiment to a DVC remote run:\n\n" - "\tdvc push ..." - ), + "Pushed experiment '%s' to Git remote '%s'.", self.args.experiment, self.args.git_remote, ) + if not self.args.push_cache: + logger.info( + "To push cached outputs for this experiment to DVC remote " + "storage, re-run this command without '--no-cache'." + ) return 0 @@ -570,18 +577,25 @@ class CmdExperimentsPull(CmdBase): def run(self): self.repo.experiments.pull( - self.args.git_remote, self.args.experiment, force=self.args.force + self.args.git_remote, + self.args.experiment, + force=self.args.force, + pull_cache=self.args.pull_cache, + dvc_remote=self.args.dvc_remote, + jobs=self.args.jobs, + run_cache=self.args.run_cache, ) logger.info( - ( - "Pulled experiment '%s' from Git remote '%s'. " - "To pull cache for this experiment from a DVC remote run:\n\n" - "\tdvc pull ..." - ), + "Pulled experiment '%s' from Git remote '%s'. ", self.args.experiment, self.args.git_remote, ) + if not self.args.pull_cache: + logger.info( + "To pull cached outputs for this experiment from DVC remote " + "storage, re-run this command without '--no-cache'." + ) return 0 @@ -950,7 +964,39 @@ def add_parser(subparsers, parent_parser): "-f", "--force", action="store_true", - help="Replace experiment in the remote if it already exists.", + help="Replace experiment in the Git remote if it already exists.", + ) + experiments_push_parser.add_argument( + "--no-cache", + action="store_false", + dest="push_cache", + help=( + "Do not push cached outputs for this experiment to DVC remote " + "storage." + ), + ) + experiments_push_parser.add_argument( + "-r", + "--remote", + dest="dvc_remote", + metavar="", + help="Name of the DVC remote to use when pushing cached outputs.", + ) + experiments_push_parser.add_argument( + "-j", + "--jobs", + type=int, + metavar="", + help=( + "Number of jobs to run simultaneously when pushing to DVC remote " + "storage." + ), + ) + experiments_push_parser.add_argument( + "--run-cache", + action="store_true", + default=False, + help="Push run history for all stages.", ) experiments_push_parser.add_argument( "git_remote", @@ -976,6 +1022,38 @@ def add_parser(subparsers, parent_parser): action="store_true", help="Replace local experiment already exists.", ) + experiments_pull_parser.add_argument( + "--no-cache", + action="store_false", + dest="pull_cache", + help=( + "Do not pull cached outputs for this experiment from DVC remote " + "storage." + ), + ) + experiments_pull_parser.add_argument( + "-r", + "--remote", + dest="dvc_remote", + metavar="", + help="Name of the DVC remote to use when pulling cached outputs.", + ) + experiments_pull_parser.add_argument( + "-j", + "--jobs", + type=int, + metavar="", + help=( + "Number of jobs to run simultaneously when pulling from DVC " + "remote storage." + ), + ) + experiments_pull_parser.add_argument( + "--run-cache", + action="store_true", + default=False, + help="Pull run history for all stages.", + ) experiments_pull_parser.add_argument( "git_remote", help="Git remote name or Git URL.", diff --git a/dvc/command/gc.py b/dvc/command/gc.py index 2948da26b..d418bb105 100644 --- a/dvc/command/gc.py +++ b/dvc/command/gc.py @@ -30,6 +30,8 @@ class CmdGC(CmdBase): msg += " and all git branches" elif self.args.all_tags: msg += " and all git tags" + elif self.args.all_experiments: + msg += " and all experiments" if self.args.repos: msg += " of the current and the following repos:" @@ -49,6 +51,7 @@ class CmdGC(CmdBase): all_branches=self.args.all_branches, all_tags=self.args.all_tags, all_commits=self.args.all_commits, + all_experiments=self.args.all_experiments, cloud=self.args.cloud, remote=self.args.remote, force=self.args.force, @@ -99,6 +102,12 @@ def add_parser(subparsers, parent_parser): default=False, help="Keep data files for all Git commits.", ) + gc_parser.add_argument( + "--all-experiments", + action="store_true", + default=False, + help="Keep data files for all experiments.", + ) gc_parser.add_argument( "-c", "--cloud", diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 19b1e4f50..a18cdf022 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -288,11 +288,13 @@ class Repo: with_deps=False, all_tags=False, all_commits=False, + all_experiments=False, remote=None, force=False, jobs=None, recursive=False, used_run_cache=None, + revs=None, ): """Get the stages related to the given target and collect the `info` of its outputs. @@ -301,7 +303,8 @@ class Repo: (namely, a file described as an output on a stage). The scope is, by default, the working directory, but you can use - `all_branches`/`all_tags`/`all_commits` to expand the scope. + `all_branches`/`all_tags`/`all_commits`/`all_experiments` to expand + the scope. Returns: A dictionary with Schemes (representing output's location) mapped @@ -313,9 +316,11 @@ class Repo: cache = NamedCache() for branch in self.brancher( + revs=revs, all_branches=all_branches, all_tags=all_tags, all_commits=all_commits, + all_experiments=all_experiments, ): targets = targets or [None] diff --git a/dvc/repo/brancher.py b/dvc/repo/brancher.py index 9062901c2..ed3bf9489 100644 --- a/dvc/repo/brancher.py +++ b/dvc/repo/brancher.py @@ -1,5 +1,6 @@ from funcy import group_by +from dvc.repo.experiments.utils import exp_commits from dvc.tree.local import LocalTree @@ -9,6 +10,7 @@ def brancher( # noqa: E302 all_branches=False, all_tags=False, all_commits=False, + all_experiments=False, sha_only=False, ): """Generator that iterates over specified revisions. @@ -26,7 +28,7 @@ def brancher( # noqa: E302 - empty string it there is no branches to iterate over - "workspace" if there are uncommitted changes in the SCM repo """ - if not any([revs, all_branches, all_tags, all_commits]): + if not any([revs, all_branches, all_tags, all_commits, all_experiments]): yield "" return @@ -50,6 +52,9 @@ def brancher( # noqa: E302 if all_tags: revs.extend(scm.list_tags()) + if all_experiments: + revs.extend(exp_commits(scm)) + try: if revs: for sha, names in group_by(scm.resolve_rev, revs).items(): diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 537e0b4bc..65a0ae7f0 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -31,7 +31,6 @@ from .base import ( ExpRefInfo, MultipleBranchError, ) -from .executor import BaseExecutor, LocalExecutor from .utils import exp_refs_by_rev logger = logging.getLogger(__name__) @@ -99,6 +98,8 @@ class Experiments: @cached_property def args_file(self): + from .executor import BaseExecutor + return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE) @cached_property @@ -231,6 +232,8 @@ class Experiments: def _pack_args(self, *args, **kwargs): import pickle + from .executor import BaseExecutor + if os.path.exists(self.args_file) and self.scm.is_tracked( self.args_file ): @@ -448,6 +451,8 @@ class Experiments: return result def _init_executors(self, to_run): + from .executor import LocalExecutor + executors = {} for stash_rev, item in to_run.items(): self.scm.set_ref(EXEC_HEAD, item.rev) diff --git a/dvc/repo/experiments/pull.py b/dvc/repo/experiments/pull.py index 231110bf7..487738ea8 100644 --- a/dvc/repo/experiments/pull.py +++ b/dvc/repo/experiments/pull.py @@ -4,14 +4,16 @@ from dvc.exceptions import DvcException, InvalidArgumentError from dvc.repo import locked from dvc.repo.scm_context import scm_context -from .utils import remote_exp_refs_by_name +from .utils import exp_commits, remote_exp_refs_by_name logger = logging.getLogger(__name__) @locked @scm_context -def pull(repo, git_remote, exp_name, *args, force=False, **kwargs): +def pull( + repo, git_remote, exp_name, *args, force=False, pull_cache=False, **kwargs +): exp_ref = _get_exp_ref(repo, git_remote, exp_name) def on_diverged(refname: str, rev: str) -> bool: @@ -24,11 +26,14 @@ def pull(repo, git_remote, exp_name, *args, force=False, **kwargs): ) refspec = f"{exp_ref}:{exp_ref}" - logger.debug("Pull '%s' -> '%s'", git_remote, refspec) + logger.debug("git pull experiment '%s' -> '%s'", git_remote, refspec) repo.scm.fetch_refspecs( git_remote, [refspec], force=force, on_diverged=on_diverged ) + if pull_cache: + _pull_cache(repo, exp_ref, **kwargs) + def _get_exp_ref(repo, git_remote, exp_name): if exp_name.startswith("refs/"): @@ -55,3 +60,13 @@ def _get_exp_ref(repo, git_remote, exp_name): msg.extend([f"\t{info}" for info in exp_refs]) raise InvalidArgumentError("\n".join(msg)) return exp_refs[0] + + +def _pull_cache( + repo, exp_ref, dvc_remote=None, jobs=None, run_cache=False, +): + revs = list(exp_commits(repo.scm, [exp_ref])) + logger.debug("dvc fetch experiment '%s'", exp_ref) + repo.fetch( + jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs, + ) diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py index 3e2ce56b7..6f8ae4475 100644 --- a/dvc/repo/experiments/push.py +++ b/dvc/repo/experiments/push.py @@ -4,14 +4,16 @@ from dvc.exceptions import DvcException, InvalidArgumentError from dvc.repo import locked from dvc.repo.scm_context import scm_context -from .utils import exp_refs_by_name +from .utils import exp_commits, exp_refs_by_name logger = logging.getLogger(__name__) @locked @scm_context -def push(repo, git_remote, exp_name, *args, force=False, **kwargs): +def push( + repo, git_remote, exp_name, *args, force=False, push_cache=False, **kwargs, +): exp_ref = _get_exp_ref(repo, exp_name) def on_diverged(refname: str, rev: str) -> bool: @@ -24,11 +26,14 @@ def push(repo, git_remote, exp_name, *args, force=False, **kwargs): ) refname = str(exp_ref) - logger.debug("Push '%s' -> '%s'", exp_ref, git_remote) + logger.debug("git push experiment '%s' -> '%s'", exp_ref, git_remote) repo.scm.push_refspec( git_remote, refname, refname, force=force, on_diverged=on_diverged ) + if push_cache: + _push_cache(repo, exp_ref, **kwargs) + def _get_exp_ref(repo, exp_name): if exp_name.startswith("refs/"): @@ -55,3 +60,13 @@ def _get_exp_ref(repo, exp_name): msg.extend([f"\t{info}" for info in exp_refs]) raise InvalidArgumentError("\n".join(msg)) return exp_refs[0] + + +def _push_cache( + repo, exp_ref, dvc_remote=None, jobs=None, run_cache=False, +): + revs = list(exp_commits(repo.scm, [exp_ref])) + logger.debug("dvc push experiment '%s'", exp_ref) + repo.push( + jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs, + ) diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index c743b9b0f..ebe8fcb9c 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Generator +from typing import TYPE_CHECKING, Generator, Iterable from .base import EXEC_NAMESPACE, EXPS_NAMESPACE, EXPS_STASH, ExpRefInfo @@ -71,3 +71,15 @@ def remote_exp_refs_by_baseline( if ref.startswith(EXEC_NAMESPACE) or ref == EXPS_STASH: continue yield ExpRefInfo.from_ref(ref) + + +def exp_commits( + scm: "Git", ref_infos: Iterable["ExpRefInfo"] = None +) -> Generator[str, None, None]: + """Iterate over all experiment commits.""" + shas = set() + refs = ref_infos if ref_infos else exp_refs(scm) + for ref_info in refs: + shas.update(scm.branch_revs(str(ref_info), ref_info.baseline_sha)) + shas.add(ref_info.baseline_sha) + yield from shas diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 4da58478d..a3e322131 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -22,6 +22,7 @@ def fetch( recursive=False, all_commits=False, run_cache=False, + revs=None, ): """Download data items from a cloud and imported repositories @@ -49,6 +50,7 @@ def fetch( remote=remote, jobs=jobs, recursive=recursive, + revs=revs, ) downloaded = 0 diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index 9b415b7a8..955cb3c52 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -13,7 +13,7 @@ def _raise_error_if_all_disabled(**kwargs): if not any(kwargs.values()): raise InvalidArgumentError( "Either of `-w|--workspace`, `-a|--all-branches`, `-T|--all-tags` " - "or `--all-commits` needs to be set." + "`--all-experiments` or `--all-commits` needs to be set." ) @@ -26,6 +26,7 @@ def gc( with_deps=False, all_tags=False, all_commits=False, + all_experiments=False, force=False, jobs=None, repos=None, @@ -34,12 +35,13 @@ def gc( # require `workspace` to be true to come into effect. # assume `workspace` to be enabled if any of `all_tags`, `all_commits`, - # or `all_branches` are enabled. + # `all_experiments` or `all_branches` are enabled. _raise_error_if_all_disabled( workspace=workspace, all_tags=all_tags, all_commits=all_commits, all_branches=all_branches, + all_experiments=all_experiments, ) from contextlib import ExitStack @@ -63,6 +65,7 @@ def gc( with_deps=with_deps, all_tags=all_tags, all_commits=all_commits, + all_experiments=all_experiments, remote=remote, force=force, jobs=jobs, diff --git a/dvc/repo/push.py b/dvc/repo/push.py index a950dd284..780055787 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -13,6 +13,7 @@ def push( recursive=False, all_commits=False, run_cache=False, + revs=None, ): used_run_cache = self.stage_cache.push(remote) if run_cache else [] @@ -30,6 +31,7 @@ def push( jobs=jobs, recursive=recursive, used_run_cache=used_run_cache, + revs=revs, ) return len(used_run_cache) + self.cloud.push(used, jobs, remote=remote) diff --git a/dvc/scm/git/backend/gitpython.py b/dvc/scm/git/backend/gitpython.py index 4148c7feb..2d613840d 100644 --- a/dvc/scm/git/backend/gitpython.py +++ b/dvc/scm/git/backend/gitpython.py @@ -215,7 +215,17 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method return [t.name for t in self.repo.tags] def list_all_commits(self): - return [c.hexsha for c in self.repo.iter_commits("--all")] + head = self.get_ref("HEAD") + if not head: + # Empty repo + return [] + + return [ + c.hexsha + for c in self.repo.iter_commits( + rev=head, branches=True, tags=True, remotes=True, + ) + ] def get_tree(self, rev: str, **kwargs) -> BaseTree: from dvc.tree.git import GitTree @@ -321,7 +331,12 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method ) -> Optional[str]: from git.exc import GitCommandError - if name.startswith("refs/heads/"): + if name == "HEAD": + try: + return self.repo.head.commit.hexsha + except (GitCommandError, ValueError): + return None + elif name.startswith("refs/heads/"): name = name[11:] if name in self.repo.heads: return self.repo.heads[name].commit.hexsha @@ -332,11 +347,13 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method else: if not follow: try: - return self.git.symbolic_ref(name).strip() + rev = self.git.symbolic_ref(name).strip() + return rev if rev else None except GitCommandError: pass try: - return self.git.show_ref(name, hash=True).strip() + rev = self.git.show_ref(name, hash=True).strip() + return rev if rev else None except GitCommandError: pass return None diff --git a/tests/func/experiments/test_remote.py b/tests/func/experiments/test_remote.py index 3f214f401..732f6b086 100644 --- a/tests/func/experiments/test_remote.py +++ b/tests/func/experiments/test_remote.py @@ -1,3 +1,5 @@ +import os + import pytest from funcy import first @@ -238,3 +240,33 @@ def test_pull_ambiguous_name(tmp_dir, scm, dvc, git_downstream, exp_stage): with git_downstream.scm.detach_head(ref_info_a.baseline_sha): downstream_exp.pull(remote, "foo") assert git_downstream.scm.get_ref(str(ref_info_a)) == exp_a + + +def test_push_pull_cache( + tmp_dir, scm, dvc, git_upstream, checkpoint_stage, local_remote +): + from dvc.utils.fs import remove + from tests.func.test_diff import digest + + remote = git_upstream.remote + results = dvc.experiments.run( + checkpoint_stage.addressing, params=["foo=2"] + ) + exp = first(results) + ref_info = first(exp_refs_by_rev(scm, exp)) + + dvc.experiments.push(remote, ref_info.name, push_cache=True) + for x in range(2, 6): + hash_ = digest(str(x)) + path = os.path.join(local_remote.url, hash_[:2], hash_[2:]) + assert os.path.exists(path) + assert open(path).read() == str(x) + + remove(dvc.cache.local.cache_dir) + + dvc.experiments.pull(remote, ref_info.name, pull_cache=True) + for x in range(2, 6): + hash_ = digest(str(x)) + path = os.path.join(dvc.cache.local.cache_dir, hash_[:2], hash_[2:]) + assert os.path.exists(path) + assert open(path).read() == str(x) diff --git a/tests/func/test_gc.py b/tests/func/test_gc.py index 4422f207f..269f3d278 100644 --- a/tests/func/test_gc.py +++ b/tests/func/test_gc.py @@ -271,7 +271,7 @@ def test_gc_without_workspace(tmp_dir, dvc, caplog): assert ( "Either of `-w|--workspace`, `-a|--all-branches`, `-T|--all-tags` " - "or `--all-commits` needs to be set." + "`--all-experiments` or `--all-commits` needs to be set." ) in caplog.text @@ -281,7 +281,7 @@ def test_gc_cloud_without_any_specifier(tmp_dir, dvc, caplog): assert ( "Either of `-w|--workspace`, `-a|--all-branches`, `-T|--all-tags` " - "or `--all-commits` needs to be set." + "`--all-experiments` or `--all-commits` needs to be set." ) in caplog.text @@ -390,3 +390,34 @@ def test_gc_external_output(tmp_dir, dvc, workspace): assert ( workspace / "cache" / bar_hash[:2] / bar_hash[2:] ).read_text() == "bar" + + +def test_gc_all_experiments(tmp_dir, scm, dvc): + from dvc.repo.experiments.base import ExpRefInfo + + (foo,) = tmp_dir.dvc_gen("foo", "foo", commit="foo") + foo_hash = foo.outs[0].hash_info.value + + (bar,) = tmp_dir.dvc_gen("foo", "bar", commit="bar") + bar_hash = bar.outs[0].hash_info.value + baseline = scm.get_rev() + + (baz,) = tmp_dir.dvc_gen("foo", "baz", commit="baz") + baz_hash = baz.outs[0].hash_info.value + + ref = ExpRefInfo(baseline, "exp") + scm.set_ref(str(ref), scm.get_rev()) + + dvc.gc(all_experiments=True, force=True) + + # all_experiments will include the experiment commit (baz) plus baseline + # commit (bar) + assert not ( + tmp_dir / ".dvc" / "cache" / foo_hash[:2] / foo_hash[2:] + ).exists() + assert ( + tmp_dir / ".dvc" / "cache" / bar_hash[:2] / bar_hash[2:] + ).read_text() == "bar" + assert ( + tmp_dir / ".dvc" / "cache" / baz_hash[:2] / baz_hash[2:] + ).read_text() == "baz" diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index c36a43af4..0d8434806 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -173,7 +173,19 @@ def test_experiments_list(dvc, scm, mocker): def test_experiments_push(dvc, scm, mocker): cli_args = parse_args( - ["experiments", "push", "origin", "experiment", "--force"] + [ + "experiments", + "push", + "origin", + "experiment", + "--force", + "--no-cache", + "--remote", + "my-remote", + "--jobs", + "1", + "--run-cache", + ] ) assert cli_args.func == CmdExperimentsPush @@ -182,12 +194,33 @@ def test_experiments_push(dvc, scm, mocker): assert cmd.run() == 0 - m.assert_called_once_with(cmd.repo, "origin", "experiment", force=True) + m.assert_called_once_with( + cmd.repo, + "origin", + "experiment", + force=True, + push_cache=False, + dvc_remote="my-remote", + jobs=1, + run_cache=True, + ) def test_experiments_pull(dvc, scm, mocker): cli_args = parse_args( - ["experiments", "pull", "origin", "experiment", "--force"] + [ + "experiments", + "pull", + "origin", + "experiment", + "--force", + "--no-cache", + "--remote", + "my-remote", + "--jobs", + "1", + "--run-cache", + ] ) assert cli_args.func == CmdExperimentsPull @@ -196,4 +229,13 @@ def test_experiments_pull(dvc, scm, mocker): assert cmd.run() == 0 - m.assert_called_once_with(cmd.repo, "origin", "experiment", force=True) + m.assert_called_once_with( + cmd.repo, + "origin", + "experiment", + force=True, + pull_cache=False, + dvc_remote="my-remote", + jobs=1, + run_cache=True, + ) diff --git a/tests/unit/scm/test_git.py b/tests/unit/scm/test_git.py index af94fd573..12fe77519 100644 --- a/tests/unit/scm/test_git.py +++ b/tests/unit/scm/test_git.py @@ -265,3 +265,17 @@ def test_fetch_refspecs(tmp_dir, scm, make_tmp_dir): remote_dir.scm.checkout("refs/foo/bar") assert init_rev == remote_dir.scm.get_rev() assert "0" == (remote_dir / "file").read_text() + + +def test_list_all_commits(tmp_dir, scm): + tmp_dir.scm_gen("1", "1", commit="1") + rev_a = scm.get_rev() + tmp_dir.scm_gen("2", "2", commit="2") + rev_b = scm.get_rev() + scm.tag("tag") + tmp_dir.scm_gen("3", "3", commit="3") + rev_c = scm.get_rev() + scm.gitpython.git.reset(rev_a, hard=True) + scm.set_ref("refs/foo/bar", rev_c) + + assert {rev_a, rev_b} == set(scm.list_all_commits())