dvc/tests/func/experiments/test_set_params.py

import sys

import pytest

from dvc.exceptions import InvalidArgumentError

from ..utils.test_hydra import hydra_setup


@pytest.mark.parametrize(
    "changes, expected",
    [
        [["foo=baz"], "foo: baz\ngoo:\n  bag: 3.0\nlorem: false"],
        [["params.yaml:foo=baz"], "foo: baz\ngoo:\n  bag: 3.0\nlorem: false"],
    ],
)
def test_modify_params(params_repo, dvc, changes, expected):
    dvc.experiments.run(params=changes)
    # pylint: disable=unspecified-encoding
    with open("params.yaml", mode="r") as fobj:
        assert fobj.read().strip() == expected


@pytest.mark.parametrize(
    "hydra_enabled",
    [
        pytest.param(
            True,
            marks=pytest.mark.skipif(
                sys.version_info >= (3, 11), reason="unsupported on 3.11"
            ),
        ),
        False,
    ],
)
@pytest.mark.parametrize(
    "config_dir,config_name",
    [
        (None, None),
        (None, "bar"),
        ("conf", "bar"),
    ],
)
def test_hydra_compose_and_dump(
    tmp_dir, params_repo, dvc, hydra_enabled, config_dir, config_name
):
    hydra_setup(
        tmp_dir,
        config_dir=config_dir or "conf",
        config_name=config_name or "config",
    )

    dvc.experiments.run()
    assert (tmp_dir / "params.yaml").parse() == {
        "foo": [{"bar": 1}, {"baz": 2}],
        "goo": {"bag": 3.0},
        "lorem": False,
    }

    with dvc.config.edit() as conf:
        if hydra_enabled:
            conf["hydra"]["enabled"] = True
        if config_dir is not None:
            conf["hydra"]["config_dir"] = config_dir
        if config_name is not None:
            conf["hydra"]["config_name"] = config_name

    dvc.experiments.run()

    if hydra_enabled:
        assert (tmp_dir / "params.yaml").parse() == {
            "db": {"driver": "mysql", "user": "omry", "pass": "secret"},
        }

        dvc.experiments.run(params=["db=postgresql"])
        assert (tmp_dir / "params.yaml").parse() == {
            "db": {
                "driver": "postgresql",
                "user": "foo",
                "pass": "bar",
                "timeout": 10,
            }
        }
    else:
        assert (tmp_dir / "params.yaml").parse() == {
            "foo": [{"bar": 1}, {"baz": 2}],
            "goo": {"bag": 3.0},
            "lorem": False,
        }


@pytest.mark.parametrize(
    "hydra_enabled,overrides,expected",
    [
        (
            pytest.param(
                True,
                marks=pytest.mark.xfail(
                    sys.version_info >= (3, 11), reason="unsupported on 3.11"
                ),
            ),
            ["db=mysql,postgresql"],
            [
                {"params.yaml": ["db=mysql"]},
                {"params.yaml": ["db=postgresql"]},
            ],
        ),
        (
            False,
            ["foo=bar,baz"],
            [{"params.yaml": ["foo=bar"]}, {"params.yaml": ["foo=baz"]}],
        ),
    ],
)
def test_hydra_sweep(
    tmp_dir, params_repo, dvc, mocker, hydra_enabled, overrides, expected
):
    patched = mocker.patch.object(dvc.experiments, "queue_one")

    if hydra_enabled:
        hydra_setup(
            tmp_dir,
            config_dir="conf",
            config_name="config",
        )
        with dvc.config.edit() as conf:
            conf["hydra"]["enabled"] = True

    dvc.experiments.run(params=overrides, queue=True)

    assert patched.call_count == len(expected)
    for e in expected:
        patched.assert_any_call(
            mocker.ANY,
            params=e,
            reset=True,
            targets=None,
        )


def test_hydra_sweep_requires_queue(params_repo, dvc):
    with pytest.raises(InvalidArgumentError):
        dvc.experiments.run(params=["db=mysql,postgresql"])