From 59d6d5ea98e7faebfac40908a02ae0f1d6a1cb42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= Date: Tue, 13 Dec 2022 20:13:39 +0100 Subject: [PATCH] exp save: simplify implementation --- dvc/repo/experiments/exceptions.py | 8 +-- dvc/repo/experiments/executor/local.py | 60 +------------------ dvc/repo/experiments/save.py | 67 ++++++++++++++++----- tests/func/experiments/test_save.py | 73 ++++++++++++++--------- tests/unit/repo/experiments/test_utils.py | 14 ++++- 5 files changed, 116 insertions(+), 106 deletions(-) diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py index 616b692b3..4f38f14c4 100644 --- a/dvc/repo/experiments/exceptions.py +++ b/dvc/repo/experiments/exceptions.py @@ -25,13 +25,13 @@ class BaselineMismatchError(DvcException): class ExperimentExistsError(DvcException): - def __init__(self, name: str): + def __init__(self, name: str, command: str = "run"): msg = ( "Experiment conflicts with existing experiment " f"'{name}'. To overwrite the existing experiment run:\n\n" - "\tdvc exp run -f ...\n\n" - "To run this experiment with a different name run:\n\n" - f"\tdvc exp run -n ...\n" + f"\tdvc exp {command} -f ...\n\n" + f"To {command} this experiment with a different name run:\n\n" + f"\tdvc exp {command} -n ...\n" ) super().__init__(msg) self.name = name diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index cbf725c4d..c2750c755 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -2,13 +2,12 @@ import logging import os from contextlib import ExitStack from tempfile import mkdtemp -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Optional from funcy import cached_property, retry from scmrepo.exceptions import SCMError as _SCMError from shortuuid import uuid -from dvc.exceptions import DvcException from dvc.lock import LockError from dvc.scm import SCM, GitMergeError from dvc.utils.fs import makedirs, remove @@ -25,7 +24,7 @@ from ..refs import ( ExpRefInfo, ) from ..utils import EXEC_TMP_DIR, get_exp_rwlock -from .base import BaseExecutor, ExecutorResult, TaskStatus +from .base import BaseExecutor, TaskStatus if TYPE_CHECKING: from scmrepo.git import Git @@ -33,7 +32,6 @@ if TYPE_CHECKING: from dvc.repo import Repo from ..stash import ExpStashEntry - from .base import ExecutorInfo logger = logging.getLogger(__name__) @@ -249,57 +247,3 @@ class WorkspaceExecutor(BaseLocalExecutor): checkpoint = self.scm.get_ref(EXEC_CHECKPOINT) if checkpoint and checkpoint != self._orig_checkpoint: self.scm.set_ref(EXEC_APPLY, checkpoint) - - @classmethod - def save( - cls, - info: "ExecutorInfo", - force: bool = False, - include_untracked: Optional[List[str]] = None, - ) -> ExecutorResult: - from dvc.repo import Repo - - exp_hash: Optional[str] = None - exp_ref: Optional[ExpRefInfo] = None - - dvc = Repo(os.path.join(info.root_dir, info.dvc_dir)) - old_cwd = os.getcwd() - if info.wdir: - os.chdir(os.path.join(dvc.scm.root_dir, info.wdir)) - else: - os.chdir(dvc.root_dir) - - try: - stages = dvc.commit([], force=force) - exp_hash = cls.hash_exp(stages) - if include_untracked: - dvc.scm.add(include_untracked) - cls.commit( - dvc.scm, - exp_hash, - exp_name=info.name, - force=force, - ) - 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() - if untracked: - logger.warning( - "The following untracked files were present in " - "the workspace before saving but " - "will not be included in the experiment commit:\n" - "\t%s", - ", ".join(untracked), - ) - info.result_hash = exp_hash - info.result_ref = ref - info.result_force = False - info.status = TaskStatus.SUCCESS - except DvcException: - info.status = TaskStatus.FAILED - raise - finally: - dvc.close() - os.chdir(old_cwd) - - return ExecutorResult(ref, exp_ref, info.result_force) diff --git a/dvc/repo/experiments/save.py b/dvc/repo/experiments/save.py index 260521aa8..e8114c550 100644 --- a/dvc/repo/experiments/save.py +++ b/dvc/repo/experiments/save.py @@ -2,7 +2,9 @@ import logging import os from typing import TYPE_CHECKING, List, Optional -from funcy import first +from .exceptions import ExperimentExistsError, UnchangedExperimentError +from .refs import ExpRefInfo +from .utils import check_ref_format, get_random_exp_name if TYPE_CHECKING: from dvc.repo import Repo @@ -11,6 +13,32 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _save_experiment( + repo: "Repo", + baseline_rev: str, + force: bool, + name: Optional[str], + include_untracked: Optional[List[str]], +) -> str: + repo.commit([], force=True) + + name = name or get_random_exp_name(repo.scm, baseline_rev) + ref_info = ExpRefInfo(baseline_rev, name) + check_ref_format(repo.scm.dulwich, ref_info) + ref = str(ref_info) + if repo.scm.get_ref(ref) and not force: + raise ExperimentExistsError(ref_info.name, command="save") + + repo.scm.add([], update=True) + if include_untracked: + repo.scm.add(include_untracked) + repo.scm.commit(f"dvc: commit experiment {name}", no_verify=True) + exp_rev = repo.scm.get_rev() + repo.scm.set_ref(ref, exp_rev, old_ref=None) + + return exp_rev + + def save( repo: "Repo", name: Optional[str] = None, @@ -21,24 +49,35 @@ def save( Returns the saved experiment's SHAs. """ - queue = repo.experiments.workspace_queue logger.debug("Saving workspace in %s", os.getcwd()) - staged, _, _ = repo.scm.status(untracked_files="no") - if staged: + _, _, untracked = repo.scm.status() + if include_untracked: + untracked = [ + file for file in untracked if file not in include_untracked + ] + if untracked: logger.warning( - "Your workspace contains staged Git changes which will be " - "unstaged before saving this experiment." + "The following untracked files were present in " + "the workspace before saving but " + "will not be included in the experiment commit:\n" + "\t%s", + ", ".join(untracked), ) - repo.scm.reset() - entry = repo.experiments.new(queue=queue, name=name, force=force) - executor = queue.init_executor(repo.experiments, entry) + with repo.scm.detach_head(client="dvc") as orig_head: + with repo.scm.stash_workspace() as workspace: + try: + if workspace is not None: + repo.scm.stash.apply(workspace) + else: + if not (include_untracked or force): + raise UnchangedExperimentError(orig_head) - save_result = executor.save( - executor.info, force=force, include_untracked=include_untracked - ) - result = queue.collect_executor(repo.experiments, executor, save_result) + exp_rev = _save_experiment( + repo, orig_head, force, name, include_untracked + ) + finally: + repo.scm.reset(hard=True) - exp_rev = first(result) return exp_rev diff --git a/tests/func/experiments/test_save.py b/tests/func/experiments/test_save.py index 7ad61aa2e..95f0375ed 100644 --- a/tests/func/experiments/test_save.py +++ b/tests/func/experiments/test_save.py @@ -1,16 +1,31 @@ -from contextlib import nullcontext - import pytest from funcy import first -from dvc.repo.experiments.exceptions import ExperimentExistsError +from dvc.repo.experiments.exceptions import ( + ExperimentExistsError, + InvalidArgumentError, + UnchangedExperimentError, +) from dvc.repo.experiments.utils import exp_refs_by_rev from dvc.scm import resolve_rev -from dvc.stage.exceptions import StageCommitError + + +@pytest.fixture +def modified_exp_stage(exp_stage, tmp_dir): + with open(tmp_dir / "copy.py", "a", encoding="utf-8") as fh: + fh.write("\n# dummy change") + yield + + +def test_exp_save_unchanged(tmp_dir, dvc, scm, exp_stage): + with pytest.raises(UnchangedExperimentError): + dvc.experiments.save() + + dvc.experiments.save(force=True) @pytest.mark.parametrize("name", (None, "test")) -def test_exp_save(tmp_dir, dvc, scm, exp_stage, name): +def test_exp_save(tmp_dir, dvc, scm, exp_stage, name, modified_exp_stage): baseline = scm.get_rev() exp = dvc.experiments.save(name=name) @@ -22,39 +37,39 @@ def test_exp_save(tmp_dir, dvc, scm, exp_stage, name): assert resolve_rev(scm, exp_name) == exp -@pytest.mark.parametrize( - ("force", "expected_raises"), - ( - (False, pytest.raises(StageCommitError)), - (True, nullcontext()), - ), -) -def test_exp_save_force(tmp_dir, dvc, scm, exp_stage, force, expected_raises): - with open(tmp_dir / "copy.py", "a", encoding="utf-8") as fh: - fh.write("\n# dummy change") - - with expected_raises: - dvc.experiments.save(force=force) - - -def test_exp_save_overwrite_experiment(tmp_dir, dvc, scm, exp_stage): - dvc.experiments.save(name="dummy") - - with open(tmp_dir / "copy.py", "a", encoding="utf-8") as fh: - fh.write("\n# dummy change") +def test_exp_save_overwrite_experiment( + tmp_dir, dvc, scm, exp_stage, modified_exp_stage +): + name = "dummy" + dvc.experiments.save(name=name) with pytest.raises(ExperimentExistsError): - dvc.experiments.save(name="dummy") + dvc.experiments.save(name=name) - dvc.experiments.save(name="dummy", force=True) + dvc.experiments.save(name=name, force=True) + + +@pytest.mark.parametrize( + "name", + ( + "invalid/name", + "invalid..name", + "invalid~name", + "invalid?name", + "invalidname.", + ), +) +def test_exp_save_invalid_name(tmp_dir, dvc, scm, exp_stage, name): + with pytest.raises(InvalidArgumentError): + dvc.experiments.save(name=name, force=True) def test_exp_save_after_commit(tmp_dir, dvc, scm, exp_stage): baseline = scm.get_rev() - dvc.experiments.save(name="exp-1") + dvc.experiments.save(name="exp-1", force=True) tmp_dir.scm_gen({"new_file": "new_file"}, commit="new baseline") - dvc.experiments.save(name="exp-2") + dvc.experiments.save(name="exp-2", force=True) all_exps = dvc.experiments.ls(all_commits=True) assert all_exps[baseline[:7]] == ["exp-1"] diff --git a/tests/unit/repo/experiments/test_utils.py b/tests/unit/repo/experiments/test_utils.py index 7447d6da4..ba7079511 100644 --- a/tests/unit/repo/experiments/test_utils.py +++ b/tests/unit/repo/experiments/test_utils.py @@ -32,7 +32,19 @@ def test_resolve_exp_ref(tmp_dir, scm, git_upstream, name_only, use_url): @pytest.mark.parametrize( - "name,result", [("name", True), ("group/name", False), ("na me", False)] + "name,result", + [ + ("name", True), + ("group/name", False), + ("na me", False), + ("invalid/.name", False), + ("@", pytest.param(False, marks=pytest.mark.xfail)), + (":", False), + ("^", False), + ("*", False), + ("~", False), + ("?", False), + ], ) def test_run_check_ref_format(scm, name, result):