exp init: change ui prompts structure and output

dvc/commands/experiments/init.py

@@ -1,5 +1,6 @@
 import argparse
 import logging
+from typing import TYPE_CHECKING, List
 
 from funcy import compact
 
@@ -8,6 +9,11 @@ from dvc.cli.utils import append_doc_link
 from dvc.exceptions import InvalidArgumentError
 from dvc.ui import ui
 
+if TYPE_CHECKING:
+    from dvc.dependency import Dependency
+    from dvc.stage import PipelineStage
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -29,7 +35,6 @@ class CmdExperimentsInit(CmdBase):
         "plots": PLOTS,
         "live": DVCLIVE,
     }
-    EXP_LINK = "https://s.dvc.org/g/exp/run"
 
     def run(self):
         from dvc.commands.stage import parse_cmd
@@ -58,7 +63,7 @@ class CmdExperimentsInit(CmdBase):
             }
         )
 
-        initialized_stage = init(
+        initialized_stage, initialized_deps = init(
             self.repo,
             name=self.args.name,
             type=self.args.type,
@@ -67,33 +72,43 @@ class CmdExperimentsInit(CmdBase):
             interactive=self.args.interactive,
             force=self.args.force,
         )
-
-        text = ui.rich_text.assemble(
-            "Created ",
-            (self.args.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)
+        self._post_init_display(initialized_stage, initialized_deps)
         if self.args.run:
-            return self.repo.experiments.run(
-                targets=[initialized_stage.addressing]
-            )
-
+            self.repo.experiments.run(targets=[initialized_stage.addressing])
         return 0
 
+    def _post_init_display(
+        self, stage: "PipelineStage", new_deps: List["Dependency"]
+    ) -> None:
+        from dvc.utils import humanize
+
+        path_fmt = "[green]{0}[/green]".format
+        if new_deps:
+            deps_paths = humanize.join(map(path_fmt, new_deps))
+            ui.write(f"Creating dependencies: {deps_paths}", styled=True)
+
+        ui.write(
+            f"Creating [b]{self.args.name}[/b] stage in [green]dvc.yaml[/]",
+            styled=True,
+        )
+        if stage.outs or not self.args.run:
+            # separate the above status-like messages with help/tips section
+            ui.write(styled=True)
+
+        if stage.outs:
+            outs_paths = humanize.join(map(path_fmt, stage.outs))
+            tips = f"Ensure your experiment command creates {outs_paths}."
+            ui.write(tips, styled=True)
+
+        if not self.args.run:
+            ui.write(
+                'You can now run your experiment using [b]"dvc exp run"[/].',
+                styled=True,
+            )
+        else:
+            # separate between `exp.run` output and `dvc exp init` output
+            ui.write(styled=True)
+
 
 def add_parser(experiments_subparsers, parent_parser):
 

dvc/repo/experiments/init.py

@@ -8,6 +8,7 @@ from typing import (
     Callable,
     Dict,
     Iterable,
+    List,
     Optional,
     TextIO,
     Tuple,
@@ -19,12 +20,11 @@ from funcy import compact, lremove, lsplit
 from dvc.exceptions import DvcException
 from dvc.stage import PipelineStage
 from dvc.types import OptStr
-from dvc.utils import humanize
 
 if TYPE_CHECKING:
     from dvc.repo import Repo
     from dvc.dvcfile import DVCFile
-    from rich.tree import Tree
+    from dvc.dependency import Dependency
 
 from dvc.ui import ui
 
@@ -75,70 +75,41 @@ def _disable_logging(highest_level=logging.CRITICAL):
         logging.disable(previous_level)
 
 
-def build_workspace_tree(workspace: Dict[str, str], label: str) -> "Tree":
-    from rich.tree import Tree
-
-    tree = Tree(label, highlight=True)
-    for value in sorted(workspace.values()):
-        tree.add(f"[green]{value}[/green]")
-    return tree
-
-
-def display_workspace_tree(workspace: Dict[str, str], label: str) -> None:
-    d = workspace.copy()
-    d.pop("cmd", None)
-
-    if d:
-        ui.write(build_workspace_tree(d, label), styled=True)
-    ui.write(styled=True)
-
-
-PIPELINE_FILE_LINK = "https://s.dvc.org/g/pipeline-files"
-
-
 def init_interactive(
-    name: str,
     defaults: Dict[str, str],
     provided: Dict[str, str],
     validator: Callable[[str, str], Union[str, Tuple[str, str]]] = None,
     live: bool = False,
     stream: Optional[TextIO] = None,
 ) -> Dict[str, str]:
-    command = provided.pop("cmd", None)
-    prompts = lremove(provided.keys(), ["code", "data", "models", "params"])
-    prompts.extend(
-        lremove(provided.keys(), ["live"] if live else ["metrics", "plots"])
+    command_prompts = lremove(provided.keys(), ["cmd"])
+    dependencies_prompts = lremove(provided.keys(), ["code", "data", "params"])
+    outputs_prompts = lremove(
+        provided.keys(),
+        ["models"] + (["live"] if live else ["metrics", "plots"]),
     )
 
     ret: Dict[str, str] = {}
-    if prompts or command:
-        ui.error_write(
-            f"This command will guide you to set up a [bright_blue]{name}[/]",
-            "stage in [green]dvc.yaml[/].",
-            f"\nSee [repr.url]{PIPELINE_FILE_LINK}[/].\n",
-            styled=True,
-        )
+    if "cmd" in provided:
+        ret["cmd"] = provided["cmd"]
 
-    if command:
-        ret["cmd"] = command
-    else:
-        ret.update(
-            compact(_prompts(["cmd"], allow_omission=False, stream=stream))
+    for heading, prompts, allow_omission in (
+        ("", command_prompts, False),
+        ("Enter experiment dependencies.", dependencies_prompts, True),
+        ("Enter experiment outputs.", outputs_prompts, True),
+    ):
+        if prompts and heading:
+            ui.error_write(heading, styled=True)
+        response = _prompts(
+            prompts,
+            defaults=defaults,
+            allow_omission=allow_omission,
+            validator=validator,
+            stream=stream,
         )
+        ret.update(compact(response))
         if prompts:
             ui.error_write(styled=True)
-
-    if prompts:
-        ui.error_write(
-            "Enter the paths for dependencies and outputs of the command.",
-            styled=True,
-        )
-        ret.update(
-            compact(
-                _prompts(prompts, defaults, validator=validator, stream=stream)
-            )
-        )
-        ui.error_write(styled=True)
     return ret
 
 
@@ -189,7 +160,7 @@ def is_file(path: str) -> bool:
     return bool(ext)
 
 
-def init_deps(stage: PipelineStage) -> None:
+def init_deps(stage: PipelineStage) -> List["Dependency"]:
     from funcy import rpartial
 
     from dvc.dependency import ParamsDependency
@@ -198,10 +169,6 @@ def init_deps(stage: PipelineStage) -> None:
     new_deps = [dep for dep in stage.deps if not dep.exists]
     params, deps = lsplit(rpartial(isinstance, ParamsDependency), new_deps)
 
-    if new_deps:
-        paths = map("[green]{0}[/]".format, new_deps)
-        ui.write(f"Created {humanize.join(paths)}.", styled=True)
-
     # always create a file for params, detect file/folder based on extension
     # for other dependencies
     dirs = [dep.fs_path for dep in deps if not is_file(dep.fs_path)]
@@ -213,6 +180,8 @@ def init_deps(stage: PipelineStage) -> None:
         with localfs.open(path, "w", encoding="utf-8"):
             pass
 
+    return new_deps
+
 
 def init(
     repo: "Repo",
@@ -223,7 +192,7 @@ def init(
     interactive: bool = False,
     force: bool = False,
     stream: Optional[TextIO] = None,
-) -> PipelineStage:
+) -> Tuple[PipelineStage, List["Dependency"]]:
     from dvc.dvcfile import make_dvcfile
 
     dvcfile = make_dvcfile(repo, "dvc.yaml")
@@ -236,7 +205,6 @@ def init(
 
     if interactive:
         defaults = init_interactive(
-            name,
             validator=partial(validate_prompts, repo),
             defaults=defaults,
             live=with_live,
@@ -288,10 +256,9 @@ def init(
     with _disable_logging(), repo.scm_context(autostage=True, quiet=True):
         stage.dump(update_lock=False)
         stage.ignore_outs()
-        display_workspace_tree(context, "Using experiment project structure:")
-        init_deps(stage)
+        initialized_deps = init_deps(stage)
         if params:
             repo.scm_context.track_file(params)
 
     assert isinstance(stage, PipelineStage)
-    return stage
+    return stage, initialized_deps

tests/func/experiments/test_init.py

@@ -30,7 +30,11 @@ def test_init_simple(tmp_dir, scm, dvc, capsys):
     assert main(["exp", "init", script]) == 0
     out, err = capsys.readouterr()
     assert not err
-    assert "Created train stage in dvc.yaml" in out
+    assert (
+        "Creating train stage in dvc.yaml\n\n"
+        "Ensure your experiment command creates metrics.json, plots and models"
+        '.\nYou can now run your experiment using "dvc exp run".'
+    ) in out
     assert (tmp_dir / "dvc.yaml").parse() == {
         "stages": {
             "train": {
@@ -77,7 +81,7 @@ def test_when_stage_force_if_already_exists(tmp_dir, dvc):
 
 
 @pytest.mark.parametrize("interactive", [True, False])
-def test_creates_params_file_by_default(tmp_dir, dvc, interactive, capsys):
+def test_creates_params_file_by_default(tmp_dir, dvc, interactive):
     init(
         dvc,
         interactive=interactive,
@@ -88,10 +92,6 @@ def test_creates_params_file_by_default(tmp_dir, dvc, interactive, capsys):
 
     assert (tmp_dir / "params.yaml").is_file()
     assert (tmp_dir / "params.yaml").parse() == {}
-    out, err = capsys.readouterr()
-    assert "Created src, data and params.yaml." in out
-    if interactive:
-        assert "'params.yaml' does not exist, the file will be created." in err
 
 
 def test_with_a_custom_name(tmp_dir, dvc):
@@ -152,7 +152,7 @@ def test_init_interactive_when_no_path_prompts_need_to_be_asked(
 
 def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc):
     (tmp_dir / "params.yaml").dump({"foo": 1})
-    inp = io.StringIO("python script.py\nscript.py\ndata\nmodels\nn")
+    inp = io.StringIO("python script.py\nscript.py\ndata\nn")
 
     init(
         dvc, interactive=True, stream=inp, defaults=CmdExperimentsInit.DEFAULTS
@@ -182,7 +182,7 @@ def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc):
 def test_init_interactive_params_validation(tmp_dir, dvc, capsys):
     tmp_dir.gen({"data": {"foo": "foo"}})
     (tmp_dir / "params.yaml").dump({"foo": 1})
-    inp = io.StringIO("python script.py\nscript.py\ndata\nmodels\ndata\n")
+    inp = io.StringIO("python script.py\nscript.py\ndata\ndata")
 
     init(
         dvc, stream=inp, interactive=True, defaults=CmdExperimentsInit.DEFAULTS
@@ -204,17 +204,17 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys):
     assert (tmp_dir / "data").is_dir()
 
     out, err = capsys.readouterr()
+    assert not out
     assert (
         "Path to a parameters file [params.yaml, n to omit]: "
         "'data' is a directory. "
         "Please retry with an existing parameters file.\n"
         "Path to a parameters file [params.yaml, n to omit]:"
     ) in err
-    assert "Created script.py." in out
 
 
 def test_init_with_no_defaults_interactive(tmp_dir, dvc):
-    inp = io.StringIO("script.py\n" "data\n" "model\n" "n\n" "metric\n" "n\n")
+    inp = io.StringIO("script.py\n" "data\n" "n\n" "model\n" "metric\n" "n\n")
     init(
         dvc,
         defaults={},
@@ -247,8 +247,8 @@ def test_init_with_no_defaults_interactive(tmp_dir, dvc):
                 "python script.py\n"
                 "script.py\n"
                 "data\n"
-                "models\n"
                 "params.yaml\n"
+                "models\n"
                 "metrics.json\n"
                 "plots\n"
                 "y"
@@ -290,11 +290,10 @@ def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys):
     assert scm.is_ignored("models")
     out, err = capsys.readouterr()
 
+    assert not out
     if interactive:
         assert "'script.py' does not exist, the file will be created." in err
         assert "'data' does not exist, the directory will be created." in err
-    assert "Using experiment project structure:" in out
-    assert "Created script.py and data" in out
 
 
 @pytest.mark.timeout(5, func_only=True)
@@ -309,8 +308,8 @@ def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys):
                 "python script.py\n"
                 "script.py\n"
                 "data\n"
-                "models\n"
                 "params.yaml\n"
+                "models\n"
                 "dvclive\n"
                 "y"
             ),
@@ -321,8 +320,8 @@ def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys):
             io.StringIO(
                 "script.py\n"
                 "data\n"
-                "models\n"
                 "params.yaml\n"
+                "models\n"
                 "dvclive\n"
                 "y"
             ),
@@ -375,11 +374,10 @@ def test_init_interactive_live(
 
     out, err = capsys.readouterr()
 
+    assert not out
     if interactive:
         assert "'script.py' does not exist, the file will be created." in err
         assert "'data' does not exist, the directory will be created." in err
-    assert "Using experiment project structure:" in out
-    assert "Created script.py and data" in out
 
 
 @pytest.mark.parametrize(

tests/unit/command/test_experiments.py

@@ -18,7 +18,6 @@ from dvc.commands.experiments.run import CmdExperimentsRun
 from dvc.commands.experiments.show import CmdExperimentsShow, show_experiments
 from dvc.exceptions import InvalidArgumentError
 from dvc.repo import Repo
-from dvc.stage import PipelineStage
 from tests.utils import ANY
 from tests.utils.asserts import called_once_with_subset
 
@@ -669,9 +668,9 @@ def test_show_experiments_sort_by(capsys, sort_order):
 
 @pytest.mark.parametrize("extra_args", [(), ("--run",)])
 def test_experiments_init(dvc, scm, mocker, capsys, extra_args):
+    stage = mocker.Mock(outs=[], addressing="train")
     m = mocker.patch(
-        "dvc.repo.experiments.init.init",
-        return_value=PipelineStage(dvc, path="dvc.yaml", name="stage1"),
+        "dvc.repo.experiments.init.init", return_value=(stage, [])
     )
     runner = mocker.patch("dvc.repo.experiments.run.run", return_value=0)
     cli_args = parse_args(["exp", "init", *extra_args, "cmd"])
@@ -696,22 +695,20 @@ def test_experiments_init(dvc, scm, mocker, capsys, extra_args):
         interactive=False,
         force=False,
     )
-    expected = "Created train stage in dvc.yaml."
-    if not extra_args:
-        expected += (
-            ' To run, use "dvc exp run".\n' "See https://s.dvc.org/g/exp/run."
-        )
-    assert capsys.readouterr() == (expected + "\n", "")
+
     if extra_args:
         # `parse_args` creates a new `Repo` object
-        runner.assert_called_once_with(ANY(Repo), targets=["stage1"])
+        runner.assert_called_once_with(ANY(Repo), targets=["train"])
 
 
 def test_experiments_init_config(dvc, scm, mocker):
     with dvc.config.edit() as conf:
         conf["exp"] = {"code": "new_src", "models": "new_models"}
 
-    m = mocker.patch("dvc.repo.experiments.init.init")
+    stage = mocker.Mock(outs=[])
+    m = mocker.patch(
+        "dvc.repo.experiments.init.init", return_value=(stage, [])
+    )
     cli_args = parse_args(["exp", "init", "cmd"])
     cmd = cli_args.func(cli_args)
 
@@ -738,7 +735,10 @@ def test_experiments_init_config(dvc, scm, mocker):
 
 
 def test_experiments_init_explicit(dvc, mocker):
-    m = mocker.patch("dvc.repo.experiments.init.init")
+    stage = mocker.Mock(outs=[])
+    m = mocker.patch(
+        "dvc.repo.experiments.init.init", return_value=(stage, [])
+    )
     cli_args = parse_args(["exp", "init", "--explicit", "cmd"])
     cmd = cli_args.func(cli_args)
 
@@ -769,7 +769,10 @@ def test_experiments_init_cmd_not_required_for_interactive_mode(dvc, mocker):
     cmd = cli_args.func(cli_args)
     assert isinstance(cmd, CmdExperimentsInit)
 
-    m = mocker.patch("dvc.repo.experiments.init.init")
+    stage = mocker.Mock(outs=[])
+    m = mocker.patch(
+        "dvc.repo.experiments.init.init", return_value=(stage, [])
+    )
     assert cmd.run() == 0
     assert called_once_with_subset(m, ANY(Repo), interactive=True)
 
@@ -822,7 +825,10 @@ def test_experiments_init_extra_args(extra_args, expected_kw, mocker):
     cmd = cli_args.func(cli_args)
     assert isinstance(cmd, CmdExperimentsInit)
 
-    m = mocker.patch("dvc.repo.experiments.init.init")
+    stage = mocker.Mock(outs=[])
+    m = mocker.patch(
+        "dvc.repo.experiments.init.init", return_value=(stage, [])
+    )
     assert cmd.run() == 0
     assert called_once_with_subset(m, ANY(Repo), **expected_kw)
 
@@ -832,6 +838,40 @@ def test_experiments_init_type_invalid_choice():
         parse_args(["exp", "init", "--type=invalid", "cmd"])
 
 
+@pytest.mark.parametrize("args", [[], ["--run"]])
+def test_experiments_init_displays_output_on_no_run(dvc, mocker, capsys, args):
+    stage = dvc.stage.create(
+        name="train",
+        cmd=["cmd"],
+        deps=["code", "data"],
+        params=["params.yaml"],
+        outs=["metrics.json", "plots", "models"],
+    )
+    mocker.patch(
+        "dvc.repo.experiments.init.init", return_value=(stage, stage.deps)
+    )
+    mocker.patch("dvc.repo.experiments.run.run", return_value=0)
+    cli_args = parse_args(["exp", "init", "cmd", *args])
+    cmd = cli_args.func(cli_args)
+    assert cmd.run() == 0
+
+    expected_lines = [
+        "Creating dependencies: code, data and params.yaml",
+        "Creating train stage in dvc.yaml",
+        "",
+        "Ensure your experiment command creates "
+        "metrics.json, plots and models.",
+    ]
+    if not cli_args.run:
+        expected_lines += [
+            'You can now run your experiment using "dvc exp run".',
+        ]
+
+    out, err = capsys.readouterr()
+    assert not err
+    assert out.splitlines() == expected_lines
+
+
 def test_show_experiments_pcp(tmp_dir, mocker):
     all_experiments = {
         "workspace": {