diff --git a/dvc/utils/hydra.py b/dvc/utils/hydra.py index 45fb759fd..5ee3180d4 100644 --- a/dvc/utils/hydra.py +++ b/dvc/utils/hydra.py @@ -1,4 +1,3 @@ -from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING, List @@ -83,32 +82,21 @@ def to_hydra_overrides(path_overrides): return parser.parse_overrides(overrides=path_overrides) +def dict_product(dicts): + import itertools + + return [dict(zip(dicts, x)) for x in itertools.product(*dicts.values())] + + def get_hydra_sweeps(path_overrides): - merged_overrides = [] + path_sweeps = {} for path, overrides in path_overrides.items(): - # `.` is reserved character in hydra syntax - # _merge_ is required to support sweeps across multiple files. - merged_overrides.extend( - [ - f"{path.replace('.', '_')}_merge_{override}" - for override in overrides - ] - ) - - hydra_overrides = to_hydra_overrides(merged_overrides) - for hydra_override in hydra_overrides: - if hydra_override.value_type == ValueType.GLOB_CHOICE_SWEEP: - raise InvalidArgumentError( - f"Glob override '{hydra_override.input_line}' " - "is not supported." - ) - - splits = BasicSweeper.split_arguments(hydra_overrides, None)[0] - sweeps = [] - for split in splits: - sweep_overrides = defaultdict(list) - for merged_override in split: - path, override = merged_override.split("_merge_") - sweep_overrides[path.replace("_", ".")].append(override) - sweeps.append(dict(sweep_overrides)) - return sweeps + overrides = to_hydra_overrides(overrides) + for override in overrides: + if override.value_type == ValueType.GLOB_CHOICE_SWEEP: + raise InvalidArgumentError( + f"Glob override '{override.input_line}' " + "is not supported." + ) + path_sweeps[path] = BasicSweeper.split_arguments(overrides, None)[0] + return dict_product(path_sweeps) diff --git a/tests/func/utils/test_hydra.py b/tests/func/utils/test_hydra.py index 8d1c0d092..8338e3cbe 100644 --- a/tests/func/utils/test_hydra.py +++ b/tests/func/utils/test_hydra.py @@ -190,6 +190,15 @@ def test_compose_and_dump(tmp_dir, suffix, overrides, expected): {"params.yaml": ["defaults/foo=2"]}, ], ), + ( + {"params.yaml": ["+foo=1,2", "~bar", "++foobar=5,6"]}, + [ + {"params.yaml": ["+foo=1", "~bar=null", "++foobar=5"]}, + {"params.yaml": ["+foo=1", "~bar=null", "++foobar=6"]}, + {"params.yaml": ["+foo=2", "~bar=null", "++foobar=5"]}, + {"params.yaml": ["+foo=2", "~bar=null", "++foobar=6"]}, + ], + ), ( {"params.yaml": ["foo=1,2", "bar=3,4"]}, [