Push running checkpoint to remote

dvc/env.py

@@ -7,3 +7,5 @@
 DVCLIVE_HTML = "DVCLIVE_HTML"
 DVCLIVE_RESUME = "DVCLIVE_RESUME"
 DVC_IGNORE_ISATTY = "DVC_IGNORE_ISATTY"
+DVC_EXP_GIT_REMOTE = "DVC_EXP_GIT_REMOTE"
+DVC_EXP_AUTO_PUSH = "DVC_EXP_CHECKPOINT_PUSH"

dvc/repo/experiments/executor/base.py

@@ -16,8 +16,10 @@
 
 from funcy import cached_property
 
+from dvc.env import DVC_EXP_AUTO_PUSH, DVC_EXP_GIT_REMOTE
 from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
+from dvc.repo import Repo
 from dvc.repo.experiments.base import (
     EXEC_BASELINE,
     EXEC_BRANCH,
@@ -36,7 +38,7 @@
 from dvc.stage import PipelineStage
 from dvc.stage.monitor import CheckpointKilledError
 from dvc.stage.serialize import to_lockfile
-from dvc.utils import dict_sha256
+from dvc.utils import dict_sha256, getenv_bool
 from dvc.utils.fs import remove
 
 if TYPE_CHECKING:
@@ -248,6 +250,28 @@ def on_diverged_ref(orig_ref: str, new_rev: str):
             )
         return refs
 
+    @classmethod
+    def _auto_push_check(cls, dvc: "Repo"):
+        git_remote = os.getenv(DVC_EXP_GIT_REMOTE, None)
+        if git_remote == dvc.root_dir:
+            logger.warning(
+                f"'{git_remote}' points to the current Git repo, experiment "
+                "Git refs will not be pushed. But DVC cache and run cache "
+                "will automatically be pushed to the default DVC remote "
+                "(if any) on each experiment commit."
+            )
+        try:
+            for ref in dvc.scm.iter_remote_refs(
+                git_remote, base=EXPS_NAMESPACE
+            ):
+                if ref:
+                    break
+        except BaseException as e:
+            print(e)
+            raise e
+
+        dvc.cloud.get_remote_odb()
+
     @classmethod
     def reproduce(
         cls,
@@ -292,6 +316,9 @@ def filter_pipeline(stages):
             log_errors,
             **kwargs,
         ) as dvc:
+            if getenv_bool(DVC_EXP_AUTO_PUSH):
+                cls._auto_push_check(dvc)
+
             args, kwargs = cls._repro_args(dvc)
             if args:
                 targets: Optional[Union[list, str]] = args[0]
@@ -331,6 +358,7 @@ def filter_pipeline(stages):
 
             checkpoint_func = partial(
                 cls.checkpoint_callback,
+                dvc,
                 dvc.scm,
                 name,
                 repro_force or checkpoint_reset,
@@ -361,6 +389,8 @@ def filter_pipeline(stages):
                         force=repro_force,
                         checkpoint=is_checkpoint,
                     )
+                    if getenv_bool(DVC_EXP_AUTO_PUSH):
+                        cls._auto_push(dvc, dvc.scm)
                 except UnchangedExperimentError:
                     pass
                 ref = dvc.scm.get_ref(EXEC_BRANCH, follow=False)
@@ -393,7 +423,6 @@ def _repro_dvc(
         git_url: Optional[str] = None,
         **kwargs,
     ):
-        from dvc.repo import Repo
         from dvc.utils.serialize import modify_yaml
 
         dvc = Repo(dvc_dir)
@@ -450,9 +479,21 @@ def _repro_args(cls, dvc):
             kwargs = {}
         return args, kwargs
 
+    @staticmethod
+    def _auto_push(dvc: "Repo", scm: "Git"):
+        git_remote = os.getenv(DVC_EXP_GIT_REMOTE)
+        branch = scm.get_ref(EXEC_BRANCH, follow=False)
+        dvc.experiments.push(
+            git_remote,
+            branch,
+            push_cache=True,
+            run_cache=True,
+        )
+
     @classmethod
     def checkpoint_callback(
         cls,
+        dvc: "Repo",
         scm: "Git",
         name: Optional[str],
         force: bool,
@@ -464,6 +505,9 @@ def checkpoint_callback(
             exp_rev = cls.commit(
                 scm, exp_hash, exp_name=name, force=force, checkpoint=True
             )
+
+            if getenv_bool(DVC_EXP_AUTO_PUSH):
+                cls._auto_push(dvc, scm)
             logger.info("Checkpoint experiment iteration '%s'.", exp_rev[:7])
         except UnchangedExperimentError:
             pass

dvc/repo/experiments/push.py

@@ -2,6 +2,7 @@
 
 from dvc.exceptions import DvcException, InvalidArgumentError
 from dvc.repo import locked
+from dvc.repo.experiments.base import ExpRefInfo
 from dvc.repo.scm_context import scm_context
 
 from .utils import exp_commits, exp_refs_by_name
@@ -12,7 +13,13 @@
 @locked
 @scm_context
 def push(
-    repo, git_remote, exp_name, *args, force=False, push_cache=False, **kwargs
+    repo,
+    git_remote,
+    exp_name: str,
+    *args,
+    force=False,
+    push_cache=False,
+    **kwargs,
 ):
     exp_ref = _get_exp_ref(repo, exp_name)
 
@@ -35,9 +42,9 @@ def on_diverged(refname: str, rev: str) -> bool:
         _push_cache(repo, exp_ref, **kwargs)
 
 
-def _get_exp_ref(repo, exp_name):
+def _get_exp_ref(repo, exp_name: str) -> ExpRefInfo:
     if exp_name.startswith("refs/"):
-        return exp_name
+        return ExpRefInfo.from_ref(exp_name)
 
     exp_refs = list(exp_refs_by_name(repo.scm, exp_name))
     if not exp_refs:

dvc/scm/git/backend/dulwich.py

@@ -353,20 +353,20 @@ def iter_remote_refs(self, url: str, base: Optional[str] = None):
         try:
             _remote, location = get_remote_repo(self.repo, url)
             client, path = get_transport_and_path(location)
+            if base:
+                yield from (
+                    os.fsdecode(ref)
+                    for ref in client.get_refs(path)
+                    if ref.startswith(os.fsencode(base))
+                )
+            else:
+                yield from (os.fsdecode(ref) for ref in client.get_refs(path))
+
         except Exception as exc:
             raise SCMError(
                 f"'{url}' is not a valid Git remote or URL"
             ) from exc
 
-        if base:
-            yield from (
-                os.fsdecode(ref)
-                for ref in client.get_refs(path)
-                if ref.startswith(os.fsencode(base))
-            )
-        else:
-            yield from (os.fsdecode(ref) for ref in client.get_refs(path))
-
     def get_refs_containing(self, rev: str, pattern: Optional[str] = None):
         raise NotImplementedError
 

dvc/utils/__init__.py

@@ -1,5 +1,6 @@
 """Helpers for other modules."""
 
+import ast
 import hashlib
 import json
 import logging
@@ -512,3 +513,11 @@ def errored_revisions(rev_data: Dict) -> List:
         if nested_contains(data, "error"):
             result.append(revision)
     return result
+
+
+def getenv_bool(name: str, default: str = "False"):
+    raw = os.getenv(name, default).title()
+    try:
+        return bool(ast.literal_eval(raw))
+    except ValueError:
+        return False

tests/func/experiments/conftest.py

@@ -88,3 +88,21 @@ def checkpoint_stage(tmp_dir, scm, dvc, mocker):
     scm.commit("init")
     stage.iterations = DEFAULT_ITERATIONS
     return stage
+
+
+@pytest.fixture
+def git_upstream(tmp_dir, erepo_dir):
+    url = "file://{}".format(erepo_dir.resolve().as_posix())
+    tmp_dir.scm.gitpython.repo.create_remote("upstream", url)
+    erepo_dir.remote = "upstream"
+    erepo_dir.url = url
+    return erepo_dir
+
+
+@pytest.fixture
+def git_downstream(tmp_dir, erepo_dir):
+    url = "file://{}".format(tmp_dir.resolve().as_posix())
+    erepo_dir.scm.gitpython.repo.create_remote("upstream", url)
+    erepo_dir.remote = "upstream"
+    erepo_dir.url = url
+    return erepo_dir

tests/func/experiments/test_checkpoints.py

@@ -1,9 +1,15 @@
+import logging
+
 import pytest
 from funcy import first
 
+from dvc.config import NoRemoteError
+from dvc.env import DVC_EXP_AUTO_PUSH, DVC_EXP_GIT_REMOTE
 from dvc.exceptions import DvcException
 from dvc.repo.experiments import MultipleBranchError
 from dvc.repo.experiments.base import EXEC_APPLY, EXEC_CHECKPOINT
+from dvc.repo.experiments.utils import exp_refs_by_rev
+from dvc.scm.base import SCMError
 
 
 @pytest.mark.parametrize("workspace", [True, False])
@@ -188,3 +194,103 @@ def test_resume_non_head_checkpoint(
     )
     new_head = first(results)
     assert orig_branch != dvc.experiments.get_branch_by_rev(new_head)
+
+
+@pytest.fixture
+def clear_env(monkeypatch):
+    yield
+    monkeypatch.delenv(DVC_EXP_GIT_REMOTE, raising=False)
+    monkeypatch.delenv(DVC_EXP_AUTO_PUSH, raising=False)
+
+
+@pytest.mark.parametrize("use_url", [True, False])
+def test_auto_push_during_iterations(
+    tmp_dir,
+    scm,
+    dvc,
+    checkpoint_stage,
+    git_upstream,
+    local_remote,
+    use_url,
+    monkeypatch,
+    clear_env,
+):
+    # set up remote repo
+    remote = git_upstream.url if use_url else git_upstream.remote
+    git_upstream.scm.fetch_refspecs(str(tmp_dir), ["master:master"])
+    monkeypatch.setenv(DVC_EXP_GIT_REMOTE, remote)
+
+    # without auto push
+    results = dvc.experiments.run(checkpoint_stage.addressing)
+    exp = first(results)
+    ref_info = first(exp_refs_by_rev(scm, exp))
+    assert git_upstream.scm.get_ref(str(ref_info)) is None
+
+    # add auto push
+    monkeypatch.setenv(DVC_EXP_AUTO_PUSH, "true")
+    results = dvc.experiments.run(checkpoint_stage.addressing)
+    assert (tmp_dir / "foo").read_text() == "4"
+    exp = first(results)
+    ref_info = first(exp_refs_by_rev(scm, exp))
+    assert git_upstream.scm.get_ref(str(ref_info)) == exp
+
+    # check the data
+    with git_upstream.dvc.config.edit() as conf:
+        conf["remote"]["local"] = local_remote.config
+        conf["core"]["remote"] = "local"
+
+    git_upstream.dvc.experiments.apply(ref_info.name)
+    git_upstream.dvc.experiments.apply(exp)
+    git_upstream.dvc.pull()
+    assert (git_upstream / "foo").read_text() == "4"
+
+    # resume the remote checkpoint
+    monkeypatch.delenv(DVC_EXP_AUTO_PUSH, raising=False)
+    with git_upstream.chdir():
+        git_upstream.dvc.experiments.run(checkpoint_stage.addressing)
+    assert (git_upstream / "foo").read_text() == "6"
+
+
+def test_auto_push_error_url(
+    dvc, scm, checkpoint_stage, local_remote, monkeypatch, clear_env
+):
+    monkeypatch.setenv(DVC_EXP_GIT_REMOTE, "true")
+    monkeypatch.setenv(DVC_EXP_AUTO_PUSH, "true")
+    with pytest.raises(SCMError):
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+
+
+def test_auto_push_no_remote(
+    dvc, scm, checkpoint_stage, git_upstream, monkeypatch, clear_env
+):
+    monkeypatch.setenv(DVC_EXP_GIT_REMOTE, git_upstream.url)
+    monkeypatch.setenv(DVC_EXP_AUTO_PUSH, "true")
+    with pytest.raises(NoRemoteError):
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+
+
+def test_auto_push_self_remote(
+    tmp_dir,
+    dvc,
+    scm,
+    checkpoint_stage,
+    local_remote,
+    caplog,
+    monkeypatch,
+    clear_env,
+):
+    root_dir = str(tmp_dir)
+    monkeypatch.setenv(DVC_EXP_GIT_REMOTE, root_dir)
+    monkeypatch.setenv(DVC_EXP_AUTO_PUSH, "true")
+    assert (
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+        != {}
+    )
+
+    with caplog.at_level(logging.WARNING, logger="dvc.repo.experiments"):
+        assert (
+            f"'{root_dir}' points to the current Git repo, experiment "
+            "Git refs will not be pushed. But DVC cache and run cache will "
+            "automatically be pushed to the default DVC remote (if any) "
+            "on each experiment commit." in caplog.text
+        )

tests/func/experiments/test_remote.py

@@ -7,24 +7,6 @@
 from dvc.repo.experiments.utils import exp_refs_by_rev
 
 
-@pytest.fixture
-def git_upstream(tmp_dir, erepo_dir):
-    url = f"file://{erepo_dir.resolve().as_posix()}"
-    tmp_dir.scm.gitpython.repo.create_remote("upstream", url)
-    erepo_dir.remote = "upstream"
-    erepo_dir.url = url
-    return erepo_dir
-
-
-@pytest.fixture
-def git_downstream(tmp_dir, erepo_dir):
-    url = f"file://{tmp_dir.resolve().as_posix()}"
-    erepo_dir.scm.gitpython.repo.create_remote("upstream", url)
-    erepo_dir.remote = "upstream"
-    erepo_dir.url = url
-    return erepo_dir
-
-
 @pytest.mark.parametrize("use_url", [True, False])
 def test_push(tmp_dir, scm, dvc, git_upstream, exp_stage, use_url):
     from dvc.exceptions import InvalidArgumentError