From 2735fbac71b6bef33a72ac1d356941a34bc0f0b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Tue, 1 Feb 2022 17:15:17 +0545 Subject: [PATCH] exp init: create workspace structure by default, show workspace tree --- dvc/repo/experiments/init.py | 65 ++++++++++++++++++++++++----- tests/func/experiments/test_init.py | 53 ++++++++++++++--------- 2 files changed, 87 insertions(+), 31 deletions(-) diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index f347ba345..7a0ea5a72 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, @@ -29,6 +30,7 @@ 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 @@ -79,18 +81,26 @@ def _disable_logging(highest_level=logging.CRITICAL): logging.disable(previous_level) -def build_workspace_tree(workspace: Dict[str, str]) -> "Tree": +def build_workspace_tree(workspace: Dict[str, str], label: str) -> "Tree": from rich.tree import Tree - tree = Tree( - "DVC assumes the following workspace structure:", - highlight=True, - ) + 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, stderr: bool = False +) -> None: + d = workspace.copy() + d.pop("cmd", None) + + if d: + ui.write(build_workspace_tree(d, label), styled=True, stderr=stderr) + ui.write(styled=True, stderr=stderr) + + PIPELINE_FILE_LINK = "https://s.dvc.org/g/pipeline-files" @@ -144,9 +154,9 @@ def init_interactive( "Enter the paths for dependencies and outputs of the command.", styled=True, ) - if workspace: - ui.error_write(build_workspace_tree(workspace), styled=True) - ui.error_write(styled=True) + + tree_label = "DVC assumes the following workspace structure:" + display_workspace_tree(workspace, tree_label, stderr=True) ret.update( compact( _prompts(prompts, defaults, validator=validator, stream=stream) @@ -192,13 +202,42 @@ def validate_prompts( raise InvalidResponse(msg_format.format(value, suffices[e.errno])) elif key in ("code", "data"): if not os.path.exists(value): - return value, ( - f"[yellow]'{value}' does not exist in the workspace. " - '"exp run" may fail.[/]' + typ = "file" if is_file(value) else "directory" + return ( + value, + f"[yellow]'{value}' does not exist, " + f"the {typ} will be created. ", ) return value +def is_file(path: str) -> bool: + _, ext = os.path.splitext(path) + return bool(ext) + + +def init_deps(deps: Iterable["Dependency"]) -> List[str]: + from dvc.dependency import ParamsDependency + from dvc.fs.local import localfs + + new_deps = [ + dep + for dep in deps + if not isinstance(dep, ParamsDependency) and not dep.exists + ] + for dep in new_deps: + fs_path = dep.fs_path + if not is_file(fs_path): + localfs.makedirs(fs_path) + continue + + localfs.makedirs(localfs.path.parent(fs_path), exist_ok=True) + with localfs.open(fs_path, "w", encoding="utf-8"): + pass + + return [dep.def_path for dep in new_deps] + + def init( repo: "Repo", name: str = "train", @@ -287,6 +326,10 @@ def init( with _disable_logging(), repo.scm_context(autostage=True, quiet=True): stage.dump(update_lock=False) stage.ignore_outs() + if not interactive: + label = "Creating experiment project structure:" + display_workspace_tree(context, label) + init_deps(stage.deps) if params: repo.scm_context.track_file(params) else: diff --git a/tests/func/experiments/test_init.py b/tests/func/experiments/test_init.py index dec48a16a..cc78508de 100644 --- a/tests/func/experiments/test_init.py +++ b/tests/func/experiments/test_init.py @@ -44,6 +44,8 @@ def test_init_simple(tmp_dir, scm, dvc, capsys): } } } + assert (tmp_dir / "data").read_text() == "data" + assert (tmp_dir / "src").is_dir() @pytest.mark.parametrize("interactive", [True, False]) @@ -142,6 +144,8 @@ def test_init_interactive_when_no_path_prompts_need_to_be_asked( } } } + assert (tmp_dir / "src").is_dir() + assert (tmp_dir / "data").is_dir() def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc): @@ -164,6 +168,8 @@ def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc): } } assert not (tmp_dir / "dvc.lock").exists() + assert (tmp_dir / "script.py").read_text() == "" + assert (tmp_dir / "data").is_dir() scm._reset() assert scm.is_tracked("dvc.yaml") assert not scm.is_tracked("params.yaml") @@ -194,6 +200,8 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys): } } } + assert (tmp_dir / "script.py").read_text() == "" + assert (tmp_dir / "data").is_dir() out, err = capsys.readouterr() assert ( @@ -209,15 +217,7 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys): def test_init_with_no_defaults_interactive(tmp_dir, dvc): - inp = io.StringIO( - "python script.py\n" - "script.py\n" - "data\n" - "model\n" - "n\n" - "metric\n" - "n\n" - ) + inp = io.StringIO("script.py\n" "data\n" "model\n" "n\n" "metric\n" "n\n") init( dvc, defaults={}, @@ -229,12 +229,14 @@ def test_init_with_no_defaults_interactive(tmp_dir, dvc): "stages": { "train": { "cmd": "python script.py", - "deps": ["python script.py", "script.py"], + "deps": ["data", "script.py"], "metrics": [{"metric": {"cache": False}}], - "outs": ["data"], + "outs": ["model"], } } } + assert (tmp_dir / "script.py").read_text() == "" + assert (tmp_dir / "data").is_dir() @pytest.mark.parametrize( @@ -258,9 +260,7 @@ def test_init_with_no_defaults_interactive(tmp_dir, dvc): ], ids=["non-interactive", "interactive"], ) -def test_init_interactive_default( - tmp_dir, scm, dvc, interactive, overrides, inp, capsys -): +def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys): (tmp_dir / "params.yaml").dump({"foo": {"bar": 1}}) init( @@ -284,6 +284,8 @@ def test_init_interactive_default( } } assert not (tmp_dir / "dvc.lock").exists() + assert (tmp_dir / "script.py").read_text() == "" + assert (tmp_dir / "data").is_dir() scm._reset() assert scm.is_tracked("dvc.yaml") assert scm.is_tracked("params.yaml") @@ -292,9 +294,11 @@ def test_init_interactive_default( out, err = capsys.readouterr() if interactive: - assert "'script.py' does not exist in the workspace." in err - assert "'data' does not exist in the workspace." in err - assert not out + 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 not out + return + assert "Creating experiment project structure: " in out @pytest.mark.timeout(5, func_only=True) @@ -365,6 +369,8 @@ def test_init_interactive_live( } } assert not (tmp_dir / "dvc.lock").exists() + assert (tmp_dir / "script.py").read_text() == "" + assert (tmp_dir / "data").is_dir() scm._reset() assert scm.is_tracked("dvc.yaml") assert scm.is_tracked("params.yaml") @@ -372,10 +378,13 @@ def test_init_interactive_live( assert scm.is_ignored("models") out, err = capsys.readouterr() + if interactive: - assert "'script.py' does not exist in the workspace." in err - assert "'data' does not exist in the workspace." in err - assert not out + 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 not out + return + assert "Creating experiment project structure: " in out @pytest.mark.parametrize( @@ -410,6 +419,8 @@ def test_init_with_type_live_and_models_plots_provided( } } } + assert (tmp_dir / "src").is_dir() + assert (tmp_dir / "data").is_dir() @pytest.mark.parametrize( @@ -443,3 +454,5 @@ def test_init_with_type_default_and_live_provided( } } } + assert (tmp_dir / "src").is_dir() + assert (tmp_dir / "data").is_dir()