diff --git a/dvc/commands/experiments/init.py b/dvc/commands/experiments/init.py index 04db754dd..aca762659 100644 --- a/dvc/commands/experiments/init.py +++ b/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): diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index ccee3d5c7..d30e8714f 100644 --- a/dvc/repo/experiments/init.py +++ b/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 diff --git a/tests/func/experiments/test_init.py b/tests/func/experiments/test_init.py index bba77cfa2..5f922dbc4 100644 --- a/tests/func/experiments/test_init.py +++ b/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( diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 80f38f383..0986dfd20 100644 --- a/tests/unit/command/test_experiments.py +++ b/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": {