From 73b46a5789dbc0e03bebca179f09fc32de9f27ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= Date: Mon, 3 Oct 2022 16:56:22 +0200 Subject: [PATCH] exp save: initial implementation --- dvc/commands/experiments/__init__.py | 2 + dvc/commands/experiments/save.py | 83 ++++++++++++++++++++++ dvc/repo/experiments/__init__.py | 5 ++ dvc/repo/experiments/executor/local.py | 61 +++++++++++++++- dvc/repo/experiments/save.py | 32 +++++++++ tests/func/experiments/test_experiments.py | 5 +- tests/unit/command/test_experiments.py | 13 ++++ 7 files changed, 196 insertions(+), 5 deletions(-) create mode 100644 dvc/commands/experiments/save.py create mode 100644 dvc/repo/experiments/save.py diff --git a/dvc/commands/experiments/__init__.py b/dvc/commands/experiments/__init__.py index addf85b7d..adbe6b923 100644 --- a/dvc/commands/experiments/__init__.py +++ b/dvc/commands/experiments/__init__.py @@ -18,6 +18,7 @@ from dvc.commands.experiments import ( queue_worker, remove, run, + save, show, ) @@ -34,6 +35,7 @@ SUB_COMMANDS = [ queue_worker, remove, run, + save, show, ] diff --git a/dvc/commands/experiments/save.py b/dvc/commands/experiments/save.py new file mode 100644 index 000000000..10d5b0a84 --- /dev/null +++ b/dvc/commands/experiments/save.py @@ -0,0 +1,83 @@ +import argparse +import logging + +from dvc.cli.command import CmdBase +from dvc.cli.utils import append_doc_link +from dvc.exceptions import DvcException +from dvc.ui import ui + +logger = logging.getLogger(__name__) + + +class CmdExperimentsSave(CmdBase): + def run(self): + + try: + ref = self.repo.experiments.save( + name=self.args.name, force=self.args.force + ) + except DvcException: + logger.exception("failed to save experiment") + return 1 + + if self.args.json: + ui.write_json({"ref": ref}) + # fixme: add metrics + else: + name = self.repo.experiments.get_exact_name(ref) + ui.write(f"Experiment has been saved as: {name}") + ui.write( + "\nTo promote an experiment to a Git branch run:\n\n" + "\tdvc exp branch \n" + ) + if self.args.metrics: + from dvc.compare import show_metrics + + metrics = self.repo.metrics.show(revs=(ref,)) + metrics.pop("workspace", None) + show_metrics(metrics) + + return 0 + + +def add_parser(experiments_subparsers, parent_parser): + EXPERIMENTS_SAVE_HELP = "Save current workspace as a dvc experiment." + save_parser = experiments_subparsers.add_parser( + "save", + parents=[parent_parser], + description=append_doc_link(EXPERIMENTS_SAVE_HELP, "exp/save"), + help=EXPERIMENTS_SAVE_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + save_parser.add_argument( + "-f", + "--force", + action="store_true", + default=False, + help="Save even if hash value for dependencies/outputs changed.", + ) + save_parser.add_argument( + "--json", + "--show-json", + action="store_true", + default=False, + help="Show output in JSON format.", + ) + save_parser.add_argument( + "-m", + "--metrics", + action="store_true", + default=False, + help="Show metrics for the saved experiment.", + ) + save_parser.add_argument( + "-n", + "--name", + default=None, + help=( + "Human-readable experiment name. If not specified, a name will " + "be auto-generated." + ), + metavar="", + ) + save_parser.set_defaults(func=CmdExperimentsSave) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index b335c2933..27167b86c 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -479,6 +479,11 @@ class Experiments: return run(self.repo, *args, **kwargs) + def save(self, *args, **kwargs): + from dvc.repo.experiments.save import save + + return save(self.repo, *args, **kwargs) + def gc(self, *args, **kwargs): from dvc.repo.experiments.gc import gc diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index e919a1f4f..0ffe6779f 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -9,6 +9,7 @@ from scmrepo.exceptions import SCMError as _SCMError from shortuuid import uuid from dvc.lock import LockError +from dvc.exceptions import DvcException from dvc.scm import SCM, GitMergeError from dvc.utils.fs import makedirs, remove @@ -21,17 +22,18 @@ from ..refs import ( EXEC_MERGE, EXEC_NAMESPACE, EXPS_TEMP, + ExpRefInfo, ) from ..utils import EXEC_TMP_DIR, get_exp_rwlock -from .base import BaseExecutor, TaskStatus +from .base import BaseExecutor, TaskStatus, ExecutorResult if TYPE_CHECKING: from scmrepo.git import Git from dvc.repo import Repo - from ..refs import ExpRefInfo from ..stash import ExpStashEntry + from .base import ExecutorInfo logger = logging.getLogger(__name__) @@ -247,3 +249,58 @@ 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", + is_checkpoint: bool = False, + force: bool = False, + ) -> 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) + cls.commit( + dvc.scm, + exp_hash, + exp_name=info.name, + checkpoint=is_checkpoint, + force=force, + ) + ref: Optional[str] = dvc.scm.get_ref(EXEC_BRANCH, follow=False) + exp_ref = ExpRefInfo.from_ref(ref) if ref else None + # TODO: research into how untracked files should be handled + if cls.WARN_UNTRACKED: + untracked = dvc.scm.untracked_files() + if untracked: + logger.warning( + "The following untracked files were present in " + "the experiment directory after reproduction but " + "will not be included in experiment commits:\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 new file mode 100644 index 000000000..2785ca07e --- /dev/null +++ b/dvc/repo/experiments/save.py @@ -0,0 +1,32 @@ +import logging +import os +from typing import TYPE_CHECKING, Optional + +from funcy import first + +if TYPE_CHECKING: + from dvc.repo import Repo + + +logger = logging.getLogger(__name__) + + +def save( + repo: "Repo", + name: Optional[str] = None, + force: bool = False, +) -> Optional[str]: + """Save the current workspace status as an experiment. + + Returns the saved experiment's SHAs. + """ + queue = repo.experiments.workspace_queue + logger.debug("Saving workspace in %s", os.getcwd()) + + entry = repo.experiments.new(queue=queue, name=name, force=force) + executor = queue.init_executor(repo.experiments, entry) + save_result = executor.save(executor.info, force=force) + result = queue.collect_executor(repo.experiments, executor, save_result) + + exp_rev = first(result) + return exp_rev diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index d8b5606df..c1c12c6d4 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -9,10 +9,11 @@ from funcy import first from dvc.dvcfile import PIPELINE_FILE from dvc.exceptions import ReproductionError +from dvc.repo.experiments.exceptions import ExperimentExistsError from dvc.repo.experiments.queue.base import BaseStashQueue from dvc.repo.experiments.utils import exp_refs_by_rev from dvc.scm import resolve_rev -from dvc.stage.exceptions import StageFileDoesNotExistError +from dvc.stage.exceptions import StageCommitError, StageFileDoesNotExistError from dvc.utils.serialize import PythonFileCorruptedError from tests.scripts import COPY_SCRIPT @@ -44,8 +45,6 @@ def test_new_simple(tmp_dir, scm, dvc, exp_stage, mocker, name, workspace): def test_experiment_exists(tmp_dir, scm, dvc, exp_stage, mocker, workspace): - from dvc.repo.experiments.exceptions import ExperimentExistsError - dvc.experiments.run( exp_stage.addressing, name="foo", diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index c20130db8..e906411f1 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -16,6 +16,7 @@ from dvc.commands.experiments.pull import CmdExperimentsPull from dvc.commands.experiments.push import CmdExperimentsPush from dvc.commands.experiments.remove import CmdExperimentsRemove from dvc.commands.experiments.run import CmdExperimentsRun +from dvc.commands.experiments.save import CmdExperimentsSave from dvc.commands.experiments.show import CmdExperimentsShow, show_experiments from dvc.exceptions import InvalidArgumentError from dvc.repo import Repo @@ -934,3 +935,15 @@ def test_show_experiments_pcp(tmp_dir, mocker): assert kwargs["output_path"] == str(tmp_dir / "dvc_plots" / "index.html") assert kwargs["color_by"] == "Experiment" + + +def test_experiments_save(dvc, scm, mocker): + cli_args = parse_args(["exp", "save", "--name", "exp-name", "--force"]) + assert cli_args.func == CmdExperimentsSave + + cmd = cli_args.func(cli_args) + m = mocker.patch("dvc.repo.experiments.save.save", return_value="acabb") + + assert cmd.run() == 0 + + m.assert_called_once_with(cmd.repo, name="exp-name", force=True)