test_init.py

import io
import os

import pytest

from dvc.cli import main
from dvc.commands.experiments.init import CmdExperimentsInit
from dvc.repo.experiments.init import init
from dvc.stage.exceptions import DuplicateStageName

# the tests may hang on prompts on failure
pytestmark = pytest.mark.timeout(3, func_only=True)


@pytest.mark.timeout(5, func_only=True)
def test_init_simple(tmp_dir, scm, dvc, capsys):
    tmp_dir.gen(
        {
            CmdExperimentsInit.CODE: {"copy.py": ""},
            "data": "data",
            "params.yaml": '{"foo": 1}',
            "plots": {},
        }
    )
    code_path = os.path.join(CmdExperimentsInit.CODE, "copy.py")
    script = f"python {code_path}"

    capsys.readouterr()
    assert main(["exp", "init", script]) == 0
    out, err = capsys.readouterr()
    assert not err
    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": {
                "cmd": script,
                "deps": ["data", "src"],
                "metrics": [{"metrics.json": {"cache": False}}],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [{"plots": {"cache": False}}],
            }
        }
    }
    assert (tmp_dir / "data").read_text() == "data"
    assert (tmp_dir / "src").is_dir()


@pytest.mark.parametrize("interactive", [True, False])
def test_when_stage_already_exists_with_same_name(tmp_dir, dvc, interactive):
    (tmp_dir / "dvc.yaml").dump({"stages": {"train": {"cmd": "test"}}})
    with pytest.raises(DuplicateStageName) as exc:
        init(
            dvc,
            interactive=interactive,
            overrides={"cmd": "true"},
            defaults=CmdExperimentsInit.DEFAULTS,
        )
    assert (
        str(exc.value) == "Stage 'train' already exists in 'dvc.yaml'. "
        "Use '--force' to overwrite."
    )


def test_when_stage_force_if_already_exists(tmp_dir, dvc):
    (tmp_dir / "params.yaml").dump({"foo": 1})
    (tmp_dir / "dvc.yaml").dump({"stages": {"train": {"cmd": "test"}}})
    init(
        dvc,
        force=True,
        overrides={"cmd": "true"},
        defaults=CmdExperimentsInit.DEFAULTS,
    )
    d = (tmp_dir / "dvc.yaml").parse()
    assert d["stages"]["train"]["cmd"] == "true"


@pytest.mark.parametrize("interactive", [True, False])
def test_creates_params_file_by_default(tmp_dir, dvc, interactive):
    init(
        dvc,
        interactive=interactive,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={"cmd": "cmd"},
        stream=io.StringIO(""),
    )

    assert (tmp_dir / "params.yaml").is_file()
    assert (tmp_dir / "params.yaml").parse() == {}


def test_with_a_custom_name(tmp_dir, dvc):
    init(dvc, name="custom", overrides={"cmd": "cmd"})
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {"custom": {"cmd": "cmd"}}
    }


def test_init_with_no_defaults_non_interactive(tmp_dir, scm, dvc):
    init(dvc, defaults={}, overrides={"cmd": "python script.py"})

    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {"train": {"cmd": "python script.py"}}
    }
    scm._reset()
    assert not (tmp_dir / "dvc.lock").exists()
    assert scm.is_tracked("dvc.yaml")


@pytest.mark.parametrize(
    "extra_overrides, inp",
    [
        ({"cmd": "cmd"}, io.StringIO()),
        ({}, io.StringIO("cmd")),
    ],
)
def test_init_interactive_when_no_path_prompts_need_to_be_asked(
    tmp_dir, dvc, extra_overrides, inp
):
    """When we pass everything that's required of, it should not prompt us."""
    (tmp_dir / "params.yaml").dump({"foo": 1})
    init(
        dvc,
        interactive=True,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={**CmdExperimentsInit.DEFAULTS, **extra_overrides},
        stream=inp,
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "cmd",
                "deps": ["data", "src"],
                "metrics": [
                    {"metrics.json": {"cache": False}},
                ],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [
                    {"plots": {"cache": False}},
                ],
            }
        }
    }
    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):
    (tmp_dir / "params.yaml").dump({"foo": 1})
    inp = io.StringIO("python script.py\nscript.py\ndata\nn")

    init(
        dvc, interactive=True, stream=inp, defaults=CmdExperimentsInit.DEFAULTS
    )

    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "python script.py",
                "deps": ["data", "script.py"],
                "metrics": [{"metrics.json": {"cache": False}}],
                "outs": ["models"],
                "plots": [{"plots": {"cache": False}}],
            }
        }
    }
    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")
    assert scm.is_tracked(".gitignore")
    assert scm.is_ignored("models")


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\ndata")

    init(
        dvc, stream=inp, interactive=True, defaults=CmdExperimentsInit.DEFAULTS
    )

    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "python script.py",
                "deps": ["data", "script.py"],
                "metrics": [{"metrics.json": {"cache": False}}],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [{"plots": {"cache": False}}],
            }
        }
    }
    assert (tmp_dir / "script.py").read_text() == ""
    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


def test_init_with_no_defaults_interactive(tmp_dir, dvc):
    inp = io.StringIO("script.py\n" "data\n" "n\n" "model\n" "metric\n" "n\n")
    init(
        dvc,
        defaults={},
        overrides={"cmd": "python script.py"},
        interactive=True,
        stream=inp,
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "python script.py",
                "deps": ["data", "script.py"],
                "metrics": [{"metric": {"cache": False}}],
                "outs": ["model"],
            }
        }
    }
    assert (tmp_dir / "script.py").read_text() == ""
    assert (tmp_dir / "data").is_dir()


@pytest.mark.parametrize(
    "interactive, overrides, inp",
    [
        (False, {"cmd": "python script.py", "code": "script.py"}, None),
        (
            True,
            {},
            io.StringIO(
                "python script.py\n"
                "script.py\n"
                "data\n"
                "params.yaml\n"
                "models\n"
                "metrics.json\n"
                "plots\n"
                "y"
            ),
        ),
    ],
    ids=["non-interactive", "interactive"],
)
def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys):
    (tmp_dir / "params.yaml").dump({"foo": {"bar": 1}})

    init(
        dvc,
        interactive=interactive,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides=overrides,
        stream=inp,
    )

    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "python script.py",
                "deps": ["data", "script.py"],
                "metrics": [{"metrics.json": {"cache": False}}],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [{"plots": {"cache": False}}],
            }
        }
    }
    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")
    assert scm.is_tracked(".gitignore")
    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


@pytest.mark.timeout(5, func_only=True)
@pytest.mark.parametrize(
    "interactive, overrides, inp",
    [
        (False, {"cmd": "python script.py", "code": "script.py"}, None),
        (
            True,
            {},
            io.StringIO(
                "python script.py\n"
                "script.py\n"
                "data\n"
                "params.yaml\n"
                "models\n"
                "y"
            ),
        ),
        (
            True,
            {"cmd": "python script.py"},
            io.StringIO("script.py\n" "data\n" "params.yaml\n" "models\n" "y"),
        ),
        (
            True,
            {"cmd": "python script.py", "models": "models"},
            io.StringIO("script.py\ndata\nparams.yaml\ny"),
        ),
    ],
    ids=[
        "non-interactive",
        "interactive",
        "interactive-cmd-provided",
        "interactive-cmd-models-provided",
    ],
)
def test_init_interactive_live(
    tmp_dir, scm, dvc, interactive, overrides, inp, capsys
):
    overrides["live"] = "dvclive"

    (tmp_dir / "params.yaml").dump({"foo": {"bar": 1}})

    init(
        dvc,
        interactive=interactive,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides=overrides,
        stream=inp,
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "python script.py",
                "deps": ["data", "script.py"],
                "metrics": [{"dvclive.json": {"cache": False}}],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [
                    {os.path.join("dvclive", "scalars"): {"cache": False}}
                ],
            }
        }
    }
    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")
    assert scm.is_tracked(".gitignore")
    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


@pytest.mark.parametrize(
    "interactive, inp",
    [
        (False, None),
        (True, io.StringIO()),
    ],
)
def test_init_with_type_checkpoint_and_models_plots_provided(
    tmp_dir, dvc, interactive, inp
):
    (tmp_dir / "params.yaml").dump({"foo": 1})
    init(
        dvc,
        type="checkpoint",
        interactive=interactive,
        stream=inp,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={"cmd": "cmd", "metrics": "m", "plots": "p"},
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "cmd",
                "deps": ["data", "src"],
                "metrics": [
                    {"m": {"cache": False}},
                ],
                "outs": [{"models": {"checkpoint": True}}],
                "params": [{"params.yaml": None}],
                "plots": [
                    {"p": {"cache": False}},
                ],
            }
        }
    }
    assert (tmp_dir / "src").is_dir()
    assert (tmp_dir / "data").is_dir()


@pytest.mark.parametrize(
    "interactive, inp",
    [
        (False, None),
        (True, io.StringIO()),
    ],
)
def test_init_with_type_default_and_live_provided(
    tmp_dir, dvc, interactive, inp
):
    (tmp_dir / "params.yaml").dump({"foo": 1})
    init(
        dvc,
        interactive=interactive,
        stream=inp,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={"cmd": "cmd", "live": "live"},
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "cmd",
                "deps": ["data", "src"],
                "metrics": [
                    {"live.json": {"cache": False}},
                ],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [
                    {os.path.join("live", "scalars"): {"cache": False}},
                ],
            }
        }
    }
    assert (tmp_dir / "src").is_dir()
    assert (tmp_dir / "data").is_dir()


@pytest.mark.parametrize(
    "interactive, inp",
    [
        (False, None),
        (True, io.StringIO()),
    ],
)
def test_init_with_live_and_metrics_plots_provided(
    tmp_dir, dvc, interactive, inp
):
    (tmp_dir / "params.yaml").dump({"foo": 1})
    init(
        dvc,
        interactive=interactive,
        stream=inp,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={
            "cmd": "cmd",
            "live": "live",
            "metrics": "metrics.json",
            "plots": "plots",
        },
    )
    assert (tmp_dir / "dvc.yaml").parse() == {
        "stages": {
            "train": {
                "cmd": "cmd",
                "deps": ["data", "src"],
                "metrics": [
                    {"live.json": {"cache": False}},
                    {"metrics.json": {"cache": False}},
                ],
                "outs": ["models"],
                "params": [{"params.yaml": None}],
                "plots": [
                    {os.path.join("live", "scalars"): {"cache": False}},
                    {"plots": {"cache": False}},
                ],
            }
        }
    }
    assert (tmp_dir / "src").is_dir()
    assert (tmp_dir / "data").is_dir()


def test_gen_output_dirs(tmp_dir, dvc):
    init(
        dvc,
        defaults=CmdExperimentsInit.DEFAULTS,
        overrides={
            "cmd": "cmd",
            "models": "models/predict.h5",
            "metrics": "eval/scores.json",
            "plots": "eval/plots",
            "live": "eval/live",
        },
    )

    assert (tmp_dir / "models").is_dir()
    assert (tmp_dir / "eval").is_dir()
    assert (tmp_dir / "eval/plots").is_dir()
    assert (tmp_dir / "eval/live").is_dir()
    assert not (tmp_dir / "models/predict.h5").exists()
    assert not (tmp_dir / "eval/scores.json").exists()