diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 6723bdf5f..6d3f47ab0 100644 --- a/dvc/parsing/__init__.py +++ b/dvc/parsing/__init__.py @@ -232,6 +232,18 @@ class EntryDefinition: self.definition = definition self.where = where + def _resolve_wdir( + self, context: Context, name: str, wdir: Optional[str] = None + ) -> str: + if not wdir: + return self.wdir + + try: + wdir = to_str(context.resolve_str(wdir)) + except (ContextError, ParseError) as exc: + format_and_raise(exc, f"'{self.where}.{name}.wdir'", self.relpath) + return self.resolver.fs.path.join(self.wdir, wdir) + def resolve(self, **kwargs): try: return self.resolve_stage(**kwargs) @@ -249,7 +261,33 @@ class EntryDefinition: # we need to pop vars from generated/evaluated data definition = deepcopy(self.definition) + + wdir = self._resolve_wdir(context, name, definition.get(WDIR_KWD)) + vars_ = definition.pop(VARS_KWD, []) + # FIXME: Should `vars` be templatized? + check_interpolations(vars_, f"{self.where}.{name}.vars", self.relpath) + if vars_: + # Optimization: Lookahead if it has any vars, if it does not, we + # don't need to clone them. + context = Context.clone(context) + + try: + fs = self.resolver.fs + context.load_from_vars(fs, vars_, wdir, stage_name=name) + except VarsAlreadyLoaded as exc: + format_and_raise(exc, f"'{self.where}.{name}.vars'", self.relpath) + + logger.trace( # type: ignore[attr-defined] + "Context during resolution of stage %s:\n%s", name, context + ) + with context.track() as tracked_data: + # NOTE: we do not pop "wdir", and resolve it again + # this does not affect anything and is done to try to + # track the source of `wdir` interpolation. + # This works because of the side-effect that we do not + # allow overwriting and/or str interpolating complex objects. + # Fix if/when those assumptions are no longer valid. resolved = { key: self._resolve(context, value, key, skip_checks) for key, value in definition.items() diff --git a/dvc/schema.py b/dvc/schema.py index 07552232e..ca7b12bd3 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -73,12 +73,14 @@ PLOT_PSTAGE_SCHEMA = {str: Any(PLOT_PROPS_SCHEMA, [PLOT_PROPS_SCHEMA])} PARAM_PSTAGE_NON_DEFAULT_SCHEMA = {str: [str]} +VARS_SCHEMA = [str, dict] STAGE_DEFINITION = { Required(StageParams.PARAM_CMD): Any(str, list), Optional(StageParams.PARAM_WDIR): str, Optional(StageParams.PARAM_DEPS): [str], Optional(StageParams.PARAM_PARAMS): [Any(str, dict)], + Optional(VARS_KWD): VARS_SCHEMA, Optional(StageParams.PARAM_FROZEN): bool, Optional(StageParams.PARAM_META): object, Optional(StageParams.PARAM_DESC): str, @@ -120,7 +122,7 @@ SINGLE_PIPELINE_STAGE_SCHEMA = { MULTI_STAGE_SCHEMA = { PLOTS: [Any(str, SINGLE_PLOT_SCHEMA)], STAGES: SINGLE_PIPELINE_STAGE_SCHEMA, - VARS_KWD: [str, dict], + VARS_KWD: VARS_SCHEMA, StageParams.PARAM_PARAMS: [str], StageParams.PARAM_METRICS: [str], ARTIFACTS: SINGLE_ARTIFACT_SCHEMA, diff --git a/tests/func/parsing/test_errors.py b/tests/func/parsing/test_errors.py index ae86d392c..a733e4696 100644 --- a/tests/func/parsing/test_errors.py +++ b/tests/func/parsing/test_errors.py @@ -9,6 +9,7 @@ import pytest from dvc.parsing import ResolveError from dvc.parsing.context import Context from dvc.parsing.interpolate import embrace +from dvc.utils.humanize import join from . import make_entry_definition, make_foreach_def @@ -21,6 +22,18 @@ def escape_ansi(line): # Tests for the interpolated entries +@pytest.mark.parametrize("vars_", ["${file}_params.yaml", {"foo": "${foo}"}]) +def test_vars_interpolation_errors(tmp_dir, dvc, vars_): + definition = make_entry_definition(tmp_dir, "build", {"vars": [vars_]}) + with pytest.raises(ResolveError) as exc_info: + definition.resolve() + + assert ( + str(exc_info.value) == "failed to parse 'stages.build.vars' in 'dvc.yaml': " + "interpolating is not allowed" + ) + + def test_failed_to_interpolate(tmp_dir, dvc): context = Context(models={"foo": "bar"}) definition = make_entry_definition( @@ -41,6 +54,23 @@ def test_failed_to_interpolate(tmp_dir, dvc): assert definition.context == {"models": {"foo": "bar"}} +def test_local_vars_params_file_not_exist(tmp_dir, dvc): + definition = make_entry_definition( + tmp_dir, + "build", + {"vars": ["not_existing_params.yaml"], "cmd": "echo ${models.foo}"}, + ) + + with pytest.raises(ResolveError) as exc_info: + definition.resolve() + + assert ( + str(exc_info.value) == "failed to parse stage 'build' in 'dvc.yaml': " + "'not_existing_params.yaml' does not exist" + ) + assert not definition.context + + def test_specified_key_does_not_exist(tmp_dir, dvc): definition = make_entry_definition( tmp_dir, @@ -119,6 +149,25 @@ def test_interpolate_nested_iterable(tmp_dir, dvc): ) +def test_partial_vars_doesnot_exist(tmp_dir, dvc): + (tmp_dir / "test_params.yaml").dump({"sub1": "sub1", "sub2": "sub2"}) + + definition = make_entry_definition( + tmp_dir, + "build", + {"vars": ["test_params.yaml:sub3"], "cmd": "echo ${sub1} ${sub2}"}, + ) + + with pytest.raises(ResolveError) as exc_info: + definition.resolve() + + assert ( + str(exc_info.value) == "failed to parse stage 'build' in 'dvc.yaml': " + "could not find 'sub3' in 'test_params.yaml'" + ) + assert not definition.context + + # Tests foreach generated stages and their error messages @@ -235,6 +284,47 @@ def test_foreach_do_definition_item_does_not_exist(tmp_dir, dvc, key, loc): assert context == {"foo": "bar"} +@pytest.mark.parametrize( + "redefine", + [ + {"item": 5}, + {"key": 5}, + {"item": 5, "key": 10}, + {"item": {"epochs": 10}}, + ], +) +@pytest.mark.parametrize("from_file", [True, False]) +def test_item_key_in_generated_stage_vars(tmp_dir, dvc, redefine, from_file): + context = Context(foo="bar") + vars_ = [redefine] + if from_file: + (tmp_dir / "test_params.yaml").dump(redefine) + vars_ = ["test_params.yaml"] + + definition = make_foreach_def( + tmp_dir, + "build", + {"model1": {"thresh": "10"}, "model2": {"thresh": 5}}, + {"vars": vars_, "cmd": "${item}"}, + context, + ) + + with pytest.raises(ResolveError) as exc_info: + definition.resolve_all() + + message = str(exc_info.value) + assert ( + "failed to parse stage 'build@model1' in 'dvc.yaml': " + "attempted to modify reserved" in message + ) + + key_or_keys = "keys" if len(redefine) > 1 else "key" + assert f"{key_or_keys} {join(redefine)}" in message + if from_file: + assert "in 'test_params.yaml'" in message + assert context == {"foo": "bar"} + + def test_foreach_wdir_key_does_not_exist(tmp_dir, dvc): definition = make_foreach_def( tmp_dir, diff --git a/tests/func/parsing/test_foreach.py b/tests/func/parsing/test_foreach.py index 2f0b24ae7..eeda0e6f7 100644 --- a/tests/func/parsing/test_foreach.py +++ b/tests/func/parsing/test_foreach.py @@ -1,8 +1,9 @@ """Testing happy paths for the foreach.""" +import os import pytest -from dvc.parsing import DataResolver, ForeachDefinition +from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver, ForeachDefinition from dvc.parsing.context import Context @@ -269,6 +270,134 @@ def test_foreach_with_interpolated_wdir(tmp_dir, dvc): assert not any(item for item in resolver.tracked_vars.values()) +def test_foreach_with_local_vars(tmp_dir, dvc): + resolver = DataResolver(dvc, (tmp_dir / "data").fs_path, {}) + foreach_data = ["foo", "bar"] + data = { + "foreach": foreach_data, + "do": { + "vars": [{"foobar": "foobar"}], + "cmd": "echo ${item} ${foobar}", + }, + } + definition = ForeachDefinition(resolver, resolver.context, "build", data) + + assert definition.resolve_all() == { + # note that the resolver generates `wdir` relative to file's wdir + # so, this is just `foo`, not `data/foo`. + # figuring out `wdir` is the responsibility of the `load_stage`/`Stage` + "build@foo": {"cmd": "echo foo foobar"}, + "build@bar": {"cmd": "echo bar foobar"}, + } + assert not resolver.context + assert not any(item for item in resolver.tracked_vars.values()) + + +@pytest.mark.parametrize( + "local_import", + [ + "test_params.yaml", + "test_params.yaml:train", + "test_params.yaml:train,prepare", + ], +) +def test_foreach_with_imported_vars(tmp_dir, dvc, local_import): + (tmp_dir / "params.yaml").dump({"models": {"model1": {"thresh": "foo"}}}) + (tmp_dir / "test_params.yaml").dump( + {"train": {"epochs": 10}, "prepare": {"nums": 25}} + ) + resolver = DataResolver(dvc, tmp_dir.fs_path, {}) + foreach_data = ["foo", "bar"] + data = { + "foreach": foreach_data, + "do": {"vars": [local_import], "cmd": "echo ${item} ${train.epochs}"}, + } + definition = ForeachDefinition(resolver, resolver.context, "build", data) + + assert definition.resolve_all() == { + # note that the resolver generates `wdir` relative to file's wdir + # so, this is just `foo`, not `data/foo`. + # figuring out `wdir` is the responsibility of the `load_stage`/`Stage` + "build@foo": {"cmd": "echo foo 10"}, + "build@bar": {"cmd": "echo bar 10"}, + } + + assert resolver.context == {"models": {"model1": {"thresh": "foo"}}} + assert resolver.tracked_vars == { + "build@foo": {"test_params.yaml": {"train.epochs": 10}}, + "build@bar": {"test_params.yaml": {"train.epochs": 10}}, + } + + +@pytest.mark.parametrize("local_import", ["params.yaml", "params.yaml:train,prepare"]) +def test_foreach_with_interpolated_wdir_and_local_vars(tmp_dir, dvc, local_import): + (tmp_dir / "params.yaml").dump({"models": {"model1": {"thresh": "foo"}}}) + + for i in range(5): + build_dir = tmp_dir / ("model-" + str(i)) + build_dir.mkdir() + (build_dir / "params.yaml").dump( + {"train": {"epochs": 1 + i}, "prepare": {"nums": 10 * i}}, + ) + + resolver = DataResolver(dvc, tmp_dir.fs_path, {}) + data = { + "foreach": [0, 1, 2, 3, 4], + "do": { + "wdir": "model-${item}", + "vars": [local_import], + "cmd": "echo ${item} ${train.epochs} ${prepare.nums}", + }, + } + definition = ForeachDefinition(resolver, resolver.context, "build", data) + + assert definition.resolve_all() == { + # note that the resolver generates `wdir` relative to file's wdir + # so, this is just `foo`, not `data/foo`. + # figuring out `wdir` is the responsibility of the `load_stage`/`Stage` + "build@0": {"wdir": "model-0", "cmd": "echo 0 1 0"}, + "build@1": {"wdir": "model-1", "cmd": "echo 1 2 10"}, + "build@2": {"wdir": "model-2", "cmd": "echo 2 3 20"}, + "build@3": {"wdir": "model-3", "cmd": "echo 3 4 30"}, + "build@4": {"wdir": "model-4", "cmd": "echo 4 5 40"}, + } + + assert resolver.context == {"models": {"model1": {"thresh": "foo"}}} + assert resolver.tracked_vars == { + "build@0": { + os.path.join("model-0", "params.yaml"): { + "train.epochs": 1, + "prepare.nums": 0, + } + }, + "build@1": { + os.path.join("model-1", "params.yaml"): { + "train.epochs": 2, + "prepare.nums": 10, + } + }, + "build@2": { + os.path.join("model-2", "params.yaml"): { + "train.epochs": 3, + "prepare.nums": 20, + } + }, + "build@3": { + os.path.join("model-3", "params.yaml"): { + "train.epochs": 4, + "prepare.nums": 30, + } + }, + "build@4": { + os.path.join("model-4", "params.yaml"): { + "train.epochs": 5, + "prepare.nums": 40, + } + }, + } + assert resolver.context.imports == {DEFAULT_PARAMS_FILE: None} + + def test_foreach_do_syntax_is_checked_once(tmp_dir, dvc, mocker): do_def = {"cmd": "python script.py --epochs ${item}"} data = {"foreach": [0, 1, 2, 3, 4], "do": do_def} diff --git a/tests/func/parsing/test_interpolated_entry.py b/tests/func/parsing/test_interpolated_entry.py index a7306dc9b..e734629db 100644 --- a/tests/func/parsing/test_interpolated_entry.py +++ b/tests/func/parsing/test_interpolated_entry.py @@ -1,3 +1,4 @@ +import os from copy import deepcopy import pytest @@ -71,6 +72,147 @@ def test_vars_and_params_import(tmp_dir, dvc): } +def test_stage_with_wdir(tmp_dir, dvc): + """ + Test that params file from wdir are also loaded + """ + d = { + "stages": { + "stage1": { + "cmd": "echo ${dict.foo} ${dict.bar}", + "params": ["value1"], + "wdir": "data", + "vars": [DEFAULT_PARAMS_FILE], + } + } + } + + data_dir = tmp_dir / "data" + data_dir.mkdir() + (tmp_dir / DEFAULT_PARAMS_FILE).dump({"dict": {"bar": "bar"}}) + (data_dir / DEFAULT_PARAMS_FILE).dump({"dict": {"foo": "foo"}}) + resolver = DataResolver(dvc, tmp_dir.fs_path, d) + + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "stage1": { + "cmd": "echo foo bar", + "wdir": "data", + "params": ["value1"], + } + } + }, + ) + assert resolver.tracked_vars == { + "stage1": { + os.path.join("data", DEFAULT_PARAMS_FILE): {"dict.foo": "foo"}, + DEFAULT_PARAMS_FILE: {"dict.bar": "bar"}, + } + } + + +def test_with_templated_wdir(tmp_dir, dvc): + """ + Test that params from the resolved wdir are still loaded + and is used in the interpolation. + """ + d = { + "stages": { + "stage1": { + "cmd": "echo ${dict.foo} ${dict.bar}", + "params": ["value1"], + "wdir": "${dict.ws}", + "vars": [DEFAULT_PARAMS_FILE], + } + } + } + (tmp_dir / DEFAULT_PARAMS_FILE).dump({"dict": {"bar": "bar", "ws": "data"}}) + data_dir = tmp_dir / "data" + data_dir.mkdir() + (data_dir / DEFAULT_PARAMS_FILE).dump({"dict": {"foo": "foo"}}) + resolver = DataResolver(dvc, tmp_dir.fs_path, d) + + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "stage1": { + "cmd": "echo foo bar", + "wdir": "data", + "params": ["value1"], + } + } + }, + ) + assert resolver.tracked_vars == { + "stage1": { + os.path.join("data", DEFAULT_PARAMS_FILE): {"dict.foo": "foo"}, + DEFAULT_PARAMS_FILE: {"dict.bar": "bar", "dict.ws": "data"}, + } + } + assert resolver.context.imports == {"params.yaml": None} + assert resolver.context == {"dict": {"bar": "bar", "ws": "data"}} + + +def test_resolve_local_tries_to_load_globally_used_files(tmp_dir, dvc): + iterable = {"bar": "bar", "foo": "foo"} + (tmp_dir / "params.json").dump(iterable) + + d = { + "vars": ["params.json"], + "stages": { + "build": { + "cmd": "command --value ${bar}", + "params": [{"params.json": ["foo"]}], + "vars": ["params.json"], + } + }, + } + resolver = DataResolver(dvc, tmp_dir.fs_path, d) + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "build": { + "cmd": "command --value bar", + "params": [{"params.json": ["foo"]}], + } + } + }, + ) + assert resolver.tracked_vars == {"build": {"params.json": {"bar": "bar"}}} + + +def test_resolve_local_tries_to_load_globally_used_params_yaml(tmp_dir, dvc): + iterable = {"bar": "bar", "foo": "foo"} + (tmp_dir / "params.yaml").dump(iterable) + + d = { + "stages": { + "build": { + "cmd": "command --value ${bar}", + "params": [{"params.yaml": ["foo"]}], + "vars": ["params.yaml"], + } + } + } + resolver = DataResolver(dvc, tmp_dir.fs_path, d) + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "build": { + "cmd": "command --value bar", + "params": [{"params.yaml": ["foo"]}], + } + } + }, + ) + assert resolver.tracked_vars == {"build": {"params.yaml": {"bar": "bar"}}} + + def test_vars_relpath_overwrite(tmp_dir, dvc): iterable = {"bar": "bar", "foo": "foo"} (tmp_dir / "params.yaml").dump(iterable) @@ -89,6 +231,7 @@ def test_vars_relpath_overwrite(tmp_dir, dvc): assert resolver.context.imports == {"params.yaml": None} +@pytest.mark.parametrize("local", [True, False]) @pytest.mark.parametrize( "vars_", [ @@ -98,10 +241,14 @@ def test_vars_relpath_overwrite(tmp_dir, dvc): ["test_params.yaml", "test_params.yaml"], ], ) -def test_vars_load_partial(tmp_dir, dvc, vars_): +def test_vars_load_partial(tmp_dir, dvc, local, vars_): iterable = {"bar": "bar", "foo": "foo"} (tmp_dir / "test_params.yaml").dump(iterable) - d = {"vars": vars_, "stages": {"build": {"cmd": "echo ${bar}"}}} + d = {"stages": {"build": {"cmd": "echo ${bar}"}}} + if local: + d["stages"]["build"]["vars"] = vars_ + else: + d["vars"] = vars_ resolver = DataResolver(dvc, tmp_dir.fs_path, d) resolver.resolve() diff --git a/tests/func/parsing/test_resolver.py b/tests/func/parsing/test_resolver.py index 744efe5f0..5f1443437 100644 --- a/tests/func/parsing/test_resolver.py +++ b/tests/func/parsing/test_resolver.py @@ -144,6 +144,7 @@ def test_specified_params_file_not_exist(tmp_dir, dvc): ) +@pytest.mark.parametrize("local", [True, False]) @pytest.mark.parametrize( "vars_", [ @@ -152,11 +153,35 @@ def test_specified_params_file_not_exist(tmp_dir, dvc): ["test_params.yaml:sub1", "test_params.yaml:sub1,sub2"], ], ) -def test_vars_already_loaded_message(tmp_dir, dvc, vars_): - d = {"vars": vars_, "stages": {"build": {"cmd": "echo ${sub1} ${sub2}"}}} +def test_vars_already_loaded_message(tmp_dir, dvc, local, vars_): + d = {"stages": {"build": {"cmd": "echo ${sub1} ${sub2}"}}} (tmp_dir / "test_params.yaml").dump({"sub1": "sub1", "sub2": "sub2"}) + if not local: + d["vars"] = vars_ + else: + d["stages"]["build"]["vars"] = vars_ with pytest.raises(ResolveError) as exc_info: # noqa: PT012 resolver = DataResolver(dvc, tmp_dir.fs_path, d) resolver.resolve() assert "partially" in str(exc_info.value) + + +@pytest.mark.parametrize( + "vars_, loc", [(DATA, "build.vars[0]"), ("params.json", "params.json")] +) +def test_local_overwrite_error(tmp_dir, dvc, vars_, loc): + (tmp_dir / DEFAULT_PARAMS_FILE).dump(DATA) + (tmp_dir / "params.json").dump(DATA) + + d = {"stages": {"build": {"cmd": "echo ${models.foo}", "vars": [vars_]}}} + + resolver = DataResolver(dvc, tmp_dir.fs_path, d) + with pytest.raises(ResolveError) as exc_info: + resolver.resolve() + + assert ( + str(exc_info.value) == "failed to parse stage 'build' in 'dvc.yaml':\n" + f"cannot redefine 'models.bar' from '{loc}' " + "as it already exists in 'params.yaml'" + )