experiments: support experiments cache collection for exp push/pull/gc (#5064)

* experiments: add utils.exp_commits for iterating over experiment commits

* repo: add support for all_experiments in brancher and used_cache

* git: make all_commits return commits in standard refs instead of --all

* gc: add --all-experiments to preserve cache for experiments

* update gc tests

* repo: support pushing and fetching specific revs internally

* exp push: push cache for experiment by default

* exp pull: pull cache for experiment by default

* fix empty git repo bug
This commit is contained in:
Peter Rowlands (변기호)
2020-12-09 22:35:06 +09:00
committed by GitHub
parent a7d665157c
commit 7bf9a0e805
16 changed files with 322 additions and 35 deletions
+91 -13
View File
@@ -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="<name>",
help="Name of the DVC remote to use when pushing cached outputs.",
)
experiments_push_parser.add_argument(
"-j",
"--jobs",
type=int,
metavar="<number>",
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="<name>",
help="Name of the DVC remote to use when pulling cached outputs.",
)
experiments_pull_parser.add_argument(
"-j",
"--jobs",
type=int,
metavar="<number>",
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.",
+9
View File
@@ -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",
+6 -1
View File
@@ -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]
+6 -1
View File
@@ -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():
+6 -1
View File
@@ -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)
+18 -3
View File
@@ -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,
)
+18 -3
View File
@@ -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,
)
+13 -1
View File
@@ -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
+2
View File
@@ -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
+5 -2
View File
@@ -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,
+2
View File
@@ -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)
+21 -4
View File
@@ -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
+32
View File
@@ -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)
+33 -2
View File
@@ -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"
+46 -4
View File
@@ -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,
)
+14
View File
@@ -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())