diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index f6828b8e3..eb136f09c 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -806,6 +806,10 @@ class CmdExperimentsInit(CmdBase): "plots": PLOTS, "live": DVCLIVE, } + EXP_LINK = ( + "https://dvc.org/doc" + "/user-guide/experiment-management/running-experiments" + ) def run(self): from dvc.command.stage import parse_cmd @@ -843,10 +847,34 @@ class CmdExperimentsInit(CmdBase): interactive=self.args.interactive, force=self.args.force, ) + + name = self.args.name or self.args.type + + text = ui.rich_text.assemble( + "\n" if self.args.interactive else "", + "Created ", + (name, "bright_blue"), + " stage in ", + ("dvc.yaml", "green"), + ".", + ) + if not self.args.run: + text.append_text( + ui.rich_text.assemble( + " To run, use ", + ('"dvc exp run"', "green"), + ".\nSee ", + (self.EXP_LINK, "repr.url"), + ".", + ) + ) + + ui.write(text, styled=True) if self.args.run: return self.repo.experiments.run( targets=[initialized_stage.addressing] ) + return 0 diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index 7c24b8e86..69f6e13a4 100644 --- a/dvc/repo/experiments/init.py +++ b/dvc/repo/experiments/init.py @@ -1,10 +1,18 @@ +import logging import os -from typing import TYPE_CHECKING, Dict, Iterable +from contextlib import contextmanager +from typing import TYPE_CHECKING, Dict, Iterable, cast from funcy import compact, lremove, post_processing from rich.prompt import Prompt +from rich.rule import Rule +from rich.syntax import Syntax +from dvc.exceptions import DvcException +from dvc.stage import PipelineStage +from dvc.stage.serialize import to_pipeline_file from dvc.types import OptStr +from dvc.utils.serialize import dumps_yaml if TYPE_CHECKING: from dvc.repo import Repo @@ -71,7 +79,7 @@ class SkippablePrompt(RequiredPrompt): prompt.append(_default) prompt.append(", ") - prompt.append(f"{self.skip_value} to skip", style="italic") + prompt.append(f"{self.skip_value} to omit", style="italic") prompt.append("]") prompt.append(self.prompt_suffix) return prompt @@ -92,24 +100,67 @@ def _prompts(keys: Iterable[str], defaults: Dict[str, OptStr]): yield key, value +@contextmanager +def _disable_logging(highest_level=logging.CRITICAL): + previous_level = logging.root.manager.disable + + logging.disable(highest_level) + + try: + yield + finally: + logging.disable(previous_level) + + +PIPELINE_FILE_LINK = ( + "https://dvc.org/doc/user-guide/project-structure/pipelines-files" +) + + def init_interactive( + name: str, defaults: Dict[str, str], - provided: Iterable[str], - show_heading: bool = False, + provided: Dict[str, str], + show_tree: bool = False, live: bool = False, ) -> Dict[str, str]: - primary = lremove(provided, ["cmd", "code", "data", "models", "params"]) - secondary = lremove(provided, ["live"] if live else ["metrics", "plots"]) + primary = lremove( + provided.keys(), ["cmd", "code", "data", "models", "params"] + ) + secondary = lremove( + provided.keys(), ["live"] if live else ["metrics", "plots"] + ) if not (primary or secondary): return {} - message = ( - "This command will guide you to set up your first stage in " - "[green]dvc.yaml[/green].\n" + message = ui.rich_text.assemble( + "This command will guide you to set up a ", + (name, "bright_blue"), + " stage in ", + ("dvc.yaml", "green"), + ".", ) - if show_heading: - ui.error_write(message, styled=True) + doc_link = ui.rich_text.assemble( + "See ", (PIPELINE_FILE_LINK, "repr.url"), "." + ) + ui.error_write(message, doc_link, "", sep="\n", styled=True) + + if show_tree: + from rich.tree import Tree + + tree = Tree( + "DVC assumes the following workspace structure:", + highlight=True, + ) + workspace = {**defaults, **provided} + workspace.pop("cmd", None) + if not live and "live" not in provided: + workspace.pop("live", None) + for value in sorted(workspace.values()): + tree.add(f"[green]{value}[/green]") + ui.error_write(tree, styled=True) + ui.error_write() return compact( { @@ -153,10 +204,11 @@ def init( with_live = type == "live" if interactive: defaults = init_interactive( - defaults=defaults or {}, - show_heading=not dvcfile.exists(), + name, + defaults=defaults, live=with_live, - provided=overrides.keys(), + provided=overrides, + show_tree=True, ) else: if with_live: @@ -182,7 +234,7 @@ def init( checkpoint_out = bool(context.get("live")) models = context.get("models") - return repo.stage.add( + stage = repo.stage.create( name=name, cmd=context["cmd"], deps=compact([context.get("code"), context.get("data")]), @@ -193,3 +245,19 @@ def init( force=force, **{"checkpoints" if checkpoint_out else "outs": compact([models])}, ) + + if interactive: + ui.write(Rule(style="green"), styled=True) + _yaml = dumps_yaml(to_pipeline_file(cast(PipelineStage, stage))) + syn = Syntax(_yaml, "yaml", theme="ansi_dark") + ui.error_write(syn, styled=True) + + if not interactive or ui.confirm( + "Do you want to add the above contents to dvc.yaml?" + ): + with _disable_logging(): + stage.dump(update_lock=False) + stage.ignore_outs() + else: + raise DvcException("Aborting ...") + return stage