From cd1e228e55ef8cb7ca7649cc90505ab2bf76a9e6 Mon Sep 17 00:00:00 2001 From: AlexandreKempf Date: Mon, 4 Mar 2024 09:35:56 +0100 Subject: [PATCH] experiments: auto push experiments (#10323) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add auto_push for experiment on run and on save. add config.exp.auto_push and config.exp.git_remote * add func tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * separate auto_push behaviour out of commit --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Saugat Pachhai (सौगात) --- dvc/config_schema.py | 2 + dvc/repo/experiments/executor/base.py | 65 +++++++++++++++------------ tests/func/experiments/test_remote.py | 37 +++++++++++++++ 3 files changed, 76 insertions(+), 28 deletions(-) diff --git a/dvc/config_schema.py b/dvc/config_schema.py index 503e90c7f..4f712fad0 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -336,6 +336,8 @@ SCHEMA = { "params": str, "plots": str, "live": str, + "auto_push": Bool, + "git_remote": str, }, "parsing": { "bool": All(Lower, Choices("store_true", "boolean_optional")), diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index eda9abb83..855232d36 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -297,13 +297,16 @@ class BaseExecutor(ABC): exp_hash = cls.hash_exp(stages) if include_untracked: dvc.scm.add(include_untracked, force=True) # type: ignore[call-arg] - cls.commit( - dvc.scm, # type: ignore[arg-type] - exp_hash, - exp_name=info.name, - force=force, - message=message, - ) + + with cls.auto_push(dvc): + cls.commit( + dvc.scm, # type: ignore[arg-type] + exp_hash, + exp_name=info.name, + force=force, + message=message, + ) + ref: Optional[str] = dvc.scm.get_ref(EXEC_BRANCH, follow=False) exp_ref = ExpRefInfo.from_ref(ref) if ref else None untracked = dvc.scm.untracked_files() @@ -460,9 +463,6 @@ class BaseExecutor(ABC): from dvc.repo.checkout import checkout as dvc_checkout from dvc.ui import ui - auto_push = env2bool(DVC_EXP_AUTO_PUSH) - git_remote = os.getenv(DVC_EXP_GIT_REMOTE, None) - if queue is not None: queue.put((rev, os.getpid())) if log_errors and log_level is not None: @@ -483,9 +483,6 @@ class BaseExecutor(ABC): message=message, **kwargs, ) as dvc: - if auto_push: - cls._validate_remotes(dvc, git_remote) - args, kwargs = cls._repro_args(dvc) if args: targets: Optional[Union[list, str]] = args[0] @@ -519,8 +516,6 @@ class BaseExecutor(ABC): dvc, info, exp_hash, - auto_push, - git_remote, repro_force, message=message, ) @@ -550,20 +545,18 @@ class BaseExecutor(ABC): dvc, info, exp_hash, - auto_push, - git_remote, repro_force, message: Optional[str] = None, ) -> tuple[Optional[str], Optional["ExpRefInfo"], bool]: - cls.commit( - dvc.scm, - exp_hash, - exp_name=info.name, - force=repro_force, - message=message, - ) - if auto_push: - cls._auto_push(dvc, dvc.scm, git_remote) + with cls.auto_push(dvc): + cls.commit( + dvc.scm, + exp_hash, + exp_name=info.name, + force=repro_force, + message=message, + ) + ref: Optional[str] = dvc.scm.get_ref(EXEC_BRANCH, follow=False) exp_ref: Optional["ExpRefInfo"] = ExpRefInfo.from_ref(ref) if ref else None if cls.WARN_UNTRACKED: @@ -672,15 +665,30 @@ class BaseExecutor(ABC): kwargs = {} return args, kwargs + @classmethod + @contextmanager + def auto_push(cls, dvc: "Repo") -> Iterator[None]: + exp_config = dvc.config.get("exp", {}) + auto_push = env2bool(DVC_EXP_AUTO_PUSH, exp_config.get("auto_push", False)) + if not auto_push: + yield + return + + git_remote = os.getenv( + DVC_EXP_GIT_REMOTE, exp_config.get("git_remote", "origin") + ) + cls._validate_remotes(dvc, git_remote) + yield + cls._auto_push(dvc, git_remote) + @staticmethod def _auto_push( dvc: "Repo", - scm: "Git", git_remote: Optional[str], push_cache=True, run_cache=True, ): - branch = scm.get_ref(EXEC_BRANCH, follow=False) + branch = dvc.scm.get_ref(EXEC_BRANCH, follow=False) try: dvc.experiments.push( git_remote, @@ -708,6 +716,7 @@ class BaseExecutor(ABC): message: Optional[str] = None, ): """Commit stages as an experiment and return the commit SHA.""" + rev = scm.get_rev() if not scm.is_dirty(untracked_files=False): logger.debug("No changes to commit") diff --git a/tests/func/experiments/test_remote.py b/tests/func/experiments/test_remote.py index e1b14e718..62622c411 100644 --- a/tests/func/experiments/test_remote.py +++ b/tests/func/experiments/test_remote.py @@ -372,3 +372,40 @@ def test_push_pull_invalid_workspace( dvc.experiments.push(git_upstream.remote, push_cache=True) dvc.experiments.pull(git_upstream.remote, pull_cache=True) assert "failed to collect" not in caplog.text + + +@pytest.mark.parametrize( + "auto_push, expected_key", [(True, "up_to_date"), (False, "success")] +) +def test_auto_push_on_run( + tmp_dir, scm, dvc, git_upstream, local_remote, exp_stage, auto_push, expected_key +): + remote = git_upstream.remote + + with dvc.config.edit() as conf: + conf["exp"]["auto_push"] = auto_push + conf["exp"]["git_remote"] = remote + + exp_name = "foo" + dvc.experiments.run(exp_stage.addressing, params=["foo=2"], name=exp_name) + + assert first(dvc.experiments.push(name=exp_name, git_remote=remote)) == expected_key + + +@pytest.mark.parametrize( + "auto_push, expected_key", [(True, "up_to_date"), (False, "success")] +) +def test_auto_push_on_save( + tmp_dir, scm, dvc, git_upstream, local_remote, exp_stage, auto_push, expected_key +): + remote = git_upstream.remote + exp_name = "foo" + dvc.experiments.run(exp_stage.addressing, params=["foo=2"], name=exp_name) + + with dvc.config.edit() as conf: + conf["exp"]["auto_push"] = auto_push + conf["exp"]["git_remote"] = remote + + dvc.experiments.save(name=exp_name, force=True) + + assert first(dvc.experiments.push(name=exp_name, git_remote=remote)) == expected_key