Push running checkpoint to remote

dvc/env.py

@@ -7,3 +7,4 @@
 DVCLIVE_HTML = "DVCLIVE_HTML"
 DVCLIVE_RESUME = "DVCLIVE_RESUME"
 DVC_IGNORE_ISATTY = "DVC_IGNORE_ISATTY"
+DVC_EXP_AUTO_PUSH = "DVC_EXP_AUTO_PUSH"

dvc/repo/experiments/executor/base.py

@@ -16,8 +16,10 @@
 
 from funcy import cached_property
 
+from dvc.env import DVC_EXP_AUTO_PUSH
 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,
@@ -331,6 +333,7 @@ def filter_pipeline(stages):
 
             checkpoint_func = partial(
                 cls.checkpoint_callback,
+                dvc,
                 dvc.scm,
                 name,
                 repro_force or checkpoint_reset,
@@ -361,6 +364,9 @@ def filter_pipeline(stages):
                         force=repro_force,
                         checkpoint=is_checkpoint,
                     )
+                    git_remote = os.environ.get(DVC_EXP_AUTO_PUSH, None)
+                    if git_remote:
+                        cls._auto_push(git_remote, dvc, dvc.scm)
                 except UnchangedExperimentError:
                     pass
                 ref = dvc.scm.get_ref(EXEC_BRANCH, follow=False)
@@ -393,7 +399,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 +455,27 @@ def _repro_args(cls, dvc):
             kwargs = {}
         return args, kwargs
 
+    @staticmethod
+    def _auto_push(git_remote: str, dvc: "Repo", scm: "Git"):
+        if git_remote == scm.root_dir:
+            logger.warning(
+                f"try to auto checkpoints to {git_remote} which is the "
+                "running repository dvc cache will be pushed to the "
+                "default remote while git references will not be pushed"
+            )
+
+        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 +487,10 @@ def checkpoint_callback(
             exp_rev = cls.commit(
                 scm, exp_hash, exp_name=name, force=force, checkpoint=True
             )
+
+            git_remote = os.environ.get(DVC_EXP_AUTO_PUSH, None)
+            if git_remote:
+                cls._auto_push(git_remote, 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:

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,14 @@
+import logging
+import os
+
 import pytest
 from funcy import first
 
+from dvc.env import DVC_EXP_AUTO_PUSH
 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
 
 
 @pytest.mark.parametrize("workspace", [True, False])
@@ -188,3 +193,77 @@ def test_resume_non_head_checkpoint(
     )
     new_head = first(results)
     assert orig_branch != dvc.experiments.get_branch_by_rev(new_head)
+
+
+@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
+):
+    # 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"])
+
+    # 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
+    os.environ[DVC_EXP_AUTO_PUSH] = remote
+    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
+    os.environ.pop(DVC_EXP_AUTO_PUSH)
+    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):
+    os.environ[DVC_EXP_AUTO_PUSH] = "none"
+    assert (
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+        == {}
+    )
+
+
+def test_auto_push_no_remote(dvc, scm, checkpoint_stage, git_upstream):
+    os.environ[DVC_EXP_AUTO_PUSH] = git_upstream.url
+    assert (
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+        == {}
+    )
+
+
+def test_auto_push_self_remote(
+    tmp_dir, dvc, scm, checkpoint_stage, local_remote, caplog
+):
+    root_dir = str(tmp_dir)
+    os.environ[DVC_EXP_AUTO_PUSH] = root_dir
+    assert (
+        dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"])
+        != {}
+    )
+
+    with caplog.at_level(logging.WARNING, logger="dvc"):
+        assert (
+            f"try to auto checkpoints to {root_dir} which is the "
+            "running repository dvc cache will be pushed to the "
+            "default remote while git references will not be pushed"
+            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