Repository: iterative/dvc
Pull request #10240 (2024-01-16 21:06:56)

hydra: support plugins
Closes https://github.com/iterative/dvc/issues/9731.

We weren't supporting Hydra plugins before, which are auto-discovered from `hydra_plugins`, because the path to `hydra_plugins` was not in `sys.path`. With this PR, we look for `hydra_plugins` at the root of the DVC repo unless a different path is configured with `dvc config hydra.plugins_path`.

TODO:
- Fix test issue
- Docs PR (also include docs for https://github.com/iterative/dvc/pull/9783)


dberenbaum: The test here pass for me when running the individual test but fail for me otherwise. I'm guessing I'm missing something about how pytest manipulates `sys.path` but not sure yet how to make the tests pass consistently.

Passing for individual test:

```sh
$ pytest tests/func/utils/test_hydra.py::test_compose_and_dump_plugins
================================== test session starts ==================================
platform darwin -- Python 3.11.7, pytest-7.4.4, pluggy-1.3.0
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /Users/dave/Code/dvc
configfile: pyproject.toml
plugins: test-utils-0.0.8, timeout-2.2.0, flaky-3.7.0, cov-4.1.0, mock-3.12.0, docker-2.0.1, xdist-3.5.0, anyio-4.2.0, benchmark-4.0.0, lazy-fixture-0.6.3, hydra-core-1.3.2
collected 1 item

tests/func/utils/test_hydra.py .                                                  [100%]

=================================== 1 passed in 0.15s ===================================
```

Failing when run with other tests:

```sh
$ pytest tests/func/utils/test_hydra.py
================================== test session starts ==================================
platform darwin -- Python 3.11.7, pytest-7.4.4, pluggy-1.3.0
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /Users/dave/Code/dvc
configfile: pyproject.toml
plugins: test-utils-0.0.8, timeout-2.2.0, flaky-3.7.0, cov-4.1.0, mock-3.12.0, docker-2.0.1, xdist-3.5.0, anyio-4.2.0, benchmark-4.0.0, lazy-fixture-0.6.3, hydra-core-1.3.2
collected 64 items

tests/func/utils/test_hydra.py .s..s..s..............s........................... [ 78%]
.....F........                                                                    [100%]

======================================= FAILURES ========================================
_____________________________ test_compose_and_dump_plugins _____________________________

tmp_dir = PosixTmpDir('/private/var/folders/24/99_tf1xj3vx8k1k_jkdmnhq00000gn/T/pytest-of-dave/pytest-41/test_compose_and_dump_plugins0')

    def test_compose_and_dump_plugins(tmp_dir):
        """Ensure Hydra plugins are loaded."""
        from dvc.utils.hydra import compose_and_dump

        config = tmp_dir / "conf" / "config.yaml"
        config.parent.mkdir()
        config.write_text("foo: '${plus_10:1}'\n")
        plugin_code = "from omegaconf import OmegaConf\n"
        plugin_code += "OmegaConf.register_new_resolver('plus_10', lambda x: x + 10)"
        plugin_file = tmp_dir / "hydra_plugins" / "resolver.py"
        plugin_file.parent.mkdir()
        plugin_file.write_text(plugin_code)
        output_file = tmp_dir / "params.yaml"
>       compose_and_dump(output_file, str(config.parent), None, "config", str(tmp_dir), [])

/Users/dave/Code/dvc/tests/func/utils/test_hydra.py:279:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Users/dave/Code/dvc/dvc/utils/hydra.py:63: in compose_and_dump
    OmegaConf.resolve(cfg)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/omegaconf.py:785: in resolve
    omegaconf._impl._resolve(cfg)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/_impl.py:40: in _resolve
    _resolve_container_value(cfg, k)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/_impl.py:14: in _resolve_container_value
    resolved = node._dereference_node()
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:246: in _dereference_node
    node = self._dereference_node_impl(throw_on_resolution_failure=True)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:277: in _dereference_node_impl
    return parent._resolve_interpolation_from_parse_tree(
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:584: in _resolve_interpolation_from_parse_tree
    resolved = self.resolve_parse_tree(
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:764: in resolve_parse_tree
    return visitor.visit(parse_tree)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/antlr4/tree/Tree.py:34: in visit
    return tree.accept(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py:206: in accept
    return visitor.visitConfigValue(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar_visitor.py:101: in visitConfigValue
    return self.visit(ctx.getChild(0))
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/antlr4/tree/Tree.py:34: in visit
    return tree.accept(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py:342: in accept
    return visitor.visitText(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar_visitor.py:298: in visitText
    return self.visitInterpolation(c)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar_visitor.py:125: in visitInterpolation
    return self.visit(ctx.getChild(0))
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/antlr4/tree/Tree.py:34: in visit
    return tree.accept(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar/gen/OmegaConfGrammarParser.py:1041: in accept
    return visitor.visitInterpolationResolver(self)
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/grammar_visitor.py:179: in visitInterpolationResolver
    return self.resolver_interpolation_callback(
/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:750: in resolver_interpolation_callback
    return self._evaluate_custom_resolver(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = {'foo': '${plus_10:1}'}, key = 'foo', node = '${plus_10:1}'
inter_type = 'plus_10', inter_args = (1,), inter_args_str = ('1',)

    def _evaluate_custom_resolver(
        self,
        key: Any,
        node: Node,
        inter_type: str,
        inter_args: Tuple[Any, ...],
        inter_args_str: Tuple[str, ...],
    ) -> Any:
        from omegaconf import OmegaConf

        resolver = OmegaConf._get_resolver(inter_type)
        if resolver is not None:
            root_node = self._get_root()
            return resolver(
                root_node,
                self,
                node,
                inter_args,
                inter_args_str,
            )
        else:
>           raise UnsupportedInterpolationType(
                f"Unsupported interpolation type {inter_type}"
            )
E           omegaconf.errors.UnsupportedInterpolationType: Unsupported interpolation type plus_10

/Users/dave/micromamba/envs/dvc/lib/python3.11/site-packages/omegaconf/base.py:702: UnsupportedInterpolationType
----------------------------------- Captured log call -----------------------------------
DEBUG    hydra.core.utils:utils.py:279 Setting JobRuntime:name=app
================================ short test summary info ================================
SKIPPED [4] tests/func/utils/test_hydra.py:102: TOML dumper breaks when overriding a list/dict with other type or when handling `null` values.
FAILED tests/func/utils/test_hydra.py::test_compose_and_dump_plugins - omegaconf.errors.UnsupportedInterpolationType: Unsupported interpolation type plus_10
======================== 1 failed, 59 passed, 4 skipped in 0.81s ========================
```
skshetry: @dberenbaum, does this introduce any potential security issues for us?
efiop: @skshetry You tell us 😄 
dberenbaum: @skshetry See the discussion in https://github.com/facebookresearch/hydra/issues/2835. It seems this is left to the user if they wish to enable Hydra plugins. The user must choose to enable Hydra in DVC for this code to run. Then, it is up to them what code they want to import from the Hydra plugins path. If we want to be extra cautious, we could disable this unless the `hydra.plugins_path` is explicitly set.
dberenbaum: > Could this be a runtime setting instead of a config? Eg::
> 
> ```
> dvc exp run --hydra-plugins=<path> ...
> ```

I don't think that would be ideal since it would mean having to include that flag every time you run an experiment. What is the benefit?
dberenbaum: I don't think we are exposing any security risks beyond what Hydra is already doing. The same security risks are present if they run without DVC. If Hydra is enabled, I don't see much reason not to load the plugins since this is the expected behavior of Hydra.
skshetry: > I don't think we are exposing any security risks beyond what Hydra is already doing. The same security risks are present if they run without DVC. If Hydra is enabled, I don't see much reason not to load the plugins since this is the expected behavior of Hydra.

Reading the issue in the hydra gave me the impression that they don't load custom resolvers by default (because of the same security reason?). 
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/iterative/dvc/pull/10240?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`b46bd9c`)](https://app.codecov.io/gh/iterative/dvc/commit/b46bd9ca654f5e2dbeeb31f4eb6acaf80a7fde0a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.48% compared to head [(`c945386`)](https://app.codecov.io/gh/iterative/dvc/pull/10240?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.23%.
> Report is 6 commits behind head on main.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main   #10240      +/-   ##
==========================================
- Coverage   90.48%   90.23%   -0.26%     
==========================================
  Files         493      493              
  Lines       37596    37643      +47     
  Branches     5455     5457       +2     
==========================================
- Hits        34020    33967      -53     
- Misses       2950     3026      +76     
- Partials      626      650      +24     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iterative/dvc/pull/10240?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).


Patch:
diff --git a/dvc/config_schema.py b/dvc/config_schema.py
index 9165a8e5d..206a8b652 100644
--- a/dvc/config_schema.py
+++ b/dvc/config_schema.py
@@ -346,6 +346,7 @@ SCHEMA = {
         Exclusive("config_dir", "config_source"): str,
         Exclusive("config_module", "config_source"): str,
         "config_name": str,
+        "plugins_path": str,
     },
     "studio": {
         "token": str,
diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py
index 993624662..d04056ce8 100644
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -486,11 +486,15 @@ class BaseStashQueue(ABC):
                 else:
                     config_dir = None
                 config_name = hydra_config.get("config_name", "config")
+                plugins_path = os.path.join(
+                    self.repo.root_dir, hydra_config.get("plugins_path", "")
+                )
                 compose_and_dump(
                     path,
                     config_dir,
                     config_module,
                     config_name,
+                    plugins_path,
                     overrides,
                 )
             else:
diff --git a/dvc/utils/hydra.py b/dvc/utils/hydra.py
index a6a4b936b..721b0bb17 100644
--- a/dvc/utils/hydra.py
+++ b/dvc/utils/hydra.py
@@ -13,11 +13,24 @@ if TYPE_CHECKING:
 logger = logger.getChild(__name__)
 
 
+def load_hydra_plugins(plugins_path: str):
+    import sys
+
+    from hydra.core.plugins import Plugins
+
+    sys.path.append(plugins_path)
+    try:
+        Plugins.instance()
+    finally:
+        sys.path.remove(plugins_path)
+
+
 def compose_and_dump(
     output_file: "StrPath",
     config_dir: Optional[str],
     config_module: Optional[str],
     config_name: str,
+    plugins_path: str,
     overrides: List[str],
 ) -> None:
     """Compose Hydra config and dumpt it to `output_file`.
@@ -30,6 +43,7 @@ def compose_and_dump(
             Ignored if `config_dir` is not `None`.
         config_name: Name of the config file containing defaults,
             without the .yaml extension.
+        plugins_path: Path to auto discover Hydra plugins.
         overrides: List of `Hydra Override`_ patterns.
 
     .. _Hydra Override:
@@ -47,6 +61,7 @@ def compose_and_dump(
         initialize_config_dir if config_dir else initialize_config_module
     )
 
+    load_hydra_plugins(plugins_path)
     with initialize_config(  # type: ignore[attr-defined]
         config_source, version_base=None
     ):


Test patch:
diff --git a/tests/func/utils/test_hydra.py b/tests/func/utils/test_hydra.py
index 6a448e2a9..b2c947c25 100644
--- a/tests/func/utils/test_hydra.py
+++ b/tests/func/utils/test_hydra.py
@@ -176,7 +176,9 @@ def test_compose_and_dump_overrides(tmp_dir, suffix, overrides, expected):
     output_file = tmp_dir / f"params.{suffix}"
     config_dir = hydra_setup(tmp_dir, "conf", "config")
     config_module = None
-    compose_and_dump(output_file, config_dir, config_module, config_name, overrides)
+    compose_and_dump(
+        output_file, config_dir, config_module, config_name, str(tmp_dir), overrides
+    )
     assert output_file.parse() == expected
 
 
@@ -229,7 +231,9 @@ def test_compose_and_dump_dir_module(
     )
 
     with error_context:
-        compose_and_dump(output_file, config_dir, config_module, config_name, [])
+        compose_and_dump(
+            output_file, config_dir, config_module, config_name, str(tmp_dir), []
+        )
         assert output_file.parse() == config_content
 
 
@@ -241,7 +245,7 @@ def test_compose_and_dump_yaml_handles_string(tmp_dir):
     config.parent.mkdir()
     config.write_text("foo: 'no'\n")
     output_file = tmp_dir / "params.yaml"
-    compose_and_dump(output_file, str(config.parent), None, "config", [])
+    compose_and_dump(output_file, str(config.parent), None, "config", str(tmp_dir), [])
     assert output_file.read_text() == "foo: 'no'\n"
 
 
@@ -253,12 +257,38 @@ def test_compose_and_dump_resolves_interpolation(tmp_dir):
     config.parent.mkdir()
     config.dump({"data": {"root": "path/to/root", "raw": "${.root}/raw"}})
     output_file = tmp_dir / "params.yaml"
-    compose_and_dump(output_file, str(config.parent), None, "config", [])
+    compose_and_dump(output_file, str(config.parent), None, "config", str(tmp_dir), [])
     assert output_file.parse() == {
         "data": {"root": "path/to/root", "raw": "path/to/root/raw"}
     }
 
 
+def test_compose_and_dump_plugins(tmp_dir):
+    """Ensure Hydra plugins are loaded."""
+    from hydra.core.plugins import Plugins
+
+    from dvc.utils.hydra import compose_and_dump
+
+    # clear cached plugins
+    Plugins._instances.pop(Plugins, None)
+
+    config = tmp_dir / "conf" / "config.yaml"
+    config.parent.mkdir()
+    config.write_text("foo: '${plus_10:1}'\n")
+
+    plugins = tmp_dir / "hydra_plugins"
+    plugins.mkdir()
+    (plugins / "resolver.py").write_text(
+        """\
+from omegaconf import OmegaConf
+OmegaConf.register_new_resolver('plus_10', lambda x: x + 10)"""
+    )
+
+    output_file = tmp_dir / "params.yaml"
+    compose_and_dump(output_file, str(config.parent), None, "config", str(tmp_dir), [])
+    assert output_file.read_text() == "foo: 11\n"
+
+
 @pytest.mark.parametrize(
     "overrides, expected",
     [