exp save: simplify implementation

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 <new_name> ...\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 <new_name> ...\n"
         )
         super().__init__(msg)
         self.name = name

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)

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
+                exp_rev = _save_experiment(
+                    repo, orig_head, force, name, include_untracked
                 )
-    result = queue.collect_executor(repo.experiments, executor, save_result)
+            finally:
+                repo.scm.reset(hard=True)
 
-    exp_rev = first(result)
     return exp_rev

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"]

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):