Repository: facebookresearch/hydra

Issue:
[Bug] Custom OmegaConf resolvers registered at module level aren't propagated to submitit launcher processes
# 🐛 Bug
## Description
Not sure whether this is expected behavior. If it is, it should be documented somewhere. If you register a custom OmegaConf resolver at module scope, they won't be propagated to the launcher processes. Thus when you try to access a config node which uses a custom resolver, you get an exception.

## Checklist
- [x] I checked on the latest version of Hydra
- [x] I created a minimal repro

## To reproduce
```python
# foo.py
import hydra
from omegaconf import OmegaConf


def my_custom_resolver():
    return "foo"


OmegaConf.register_resolver("my_custom_resolver", my_custom_resolver)


@hydra.main(config_name="foo")
def main(cfg):
    assert cfg.x == "foo"


if __name__ == "__main__":
    main()
```
```yaml
# foo.yaml
x: ${my_custom_resolver:}
```
Regular run works fine:
```
$ python foo.py
```
Multirun with submitit throws an exception:
```
$ python foo.py -m hydra/launcher=submitit_local
[2020-08-10 19:56:14,485][HYDRA] Submitit 'local' sweep output dir : multirun/2020-08-10/19-56-14
[2020-08-10 19:56:14,487][HYDRA] 	#0 :
Traceback (most recent call last):
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra/_internal/utils.py", line 198, in run_and_report
    return func()
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra/_internal/utils.py", line 344, in <lambda>
    overrides=args.overrides,
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra/_internal/hydra.py", line 132, in multirun
    return sweeper.sweep(arguments=task_overrides)
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra/_internal/core_plugins/basic_sweeper.py", line 135, in sweep
    results = self.launcher.launch(batch, initial_job_idx=initial_job_idx)
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py", line 149, in launch
    return [j.results()[0] for j in jobs]
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py", line 149, in <listcomp>
    return [j.results()[0] for j in jobs]
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/submitit/core/core.py", line 291, in results
    raise job_exception  # pylint: disable=raising-bad-type
submitit.core.utils.FailedJobError: Job (task=0) failed during processing with trace:
----------------------
Traceback (most recent call last):
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/submitit/core/submission.py", line 47, in process_job
    result = delayed.result()
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/submitit/core/utils.py", line 123, in result
    self._result = self.function(*self.args, **self.kwargs)
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py", line 80, in __call__
    job_subdir_key="hydra.sweep.subdir",
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/hydra/core/utils.py", line 123, in run_job
    ret.return_value = task_function(task_cfg)
  File "foo.py", line 14, in main
    assert cfg.x == "foo"
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/omegaconf/dictconfig.py", line 315, in __getattr__
    self._format_and_raise(key=key, value=None, cause=e)
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/omegaconf/base.py", line 101, in _format_and_raise
    type_override=type_override,
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/omegaconf/_utils.py", line 669, in format_and_raise
    _raise(ex, cause)
  File "/private/home/calebh/miniconda3/envs/bench-detectron2/lib/python3.7/site-packages/omegaconf/_utils.py", line 583, in _raise
    raise ex  # set end OC_CAUSE=1 for full backtrace
omegaconf.errors.UnsupportedInterpolationType: Unsupported interpolation type my_custom_resolver
	full_key: x
	reference_type=Optional[Dict[Any, Any]]
	object_type=dict

----------------------
You can check full logs with 'job.stderr(0)' and 'job.stdout(0)'or at paths:
  - /private/home/calebh/scratch/hydra-bug/multirun/2020-08-10/19-56-14/.submitit/34113/34113_0_log.err
  - /private/home/calebh/scratch/hydra-bug/multirun/2020-08-10/19-56-14/.submitit/34113/34113_0_log.out
```

Moving `OmegaConf.register_resolver` inside `main` works:
```diff
--- foo.py	2020-08-10 19:51:10.401891000 -0700
+++ foo_fixed.py	2020-08-10 19:58:19.693452000 -0700
@@ -6,11 +6,9 @@
     return "foo"


-OmegaConf.register_resolver("my_custom_resolver", my_custom_resolver)
-
-
 @hydra.main(config_name="foo")
 def main(cfg):
+    OmegaConf.register_resolver("my_custom_resolver", my_custom_resolver)
     assert cfg.x == "foo"
```

## Expected Behavior
<!-- A clear and concise description of what you expected to happen. -->

## System information
- **Hydra Version**: 
    ```
    pip list | grep hydra
    hydra-core              1.0.0rc2
    hydra-submitit-launcher 1.0.0rc4
    ```
- **Python version** : 3.7.7
- **Virtual environment type and version** : conda 4.8.3
- **Operating system** : Ubuntu 18.04.3

## Additional context
I think this is an issue with how submitit is serializing the Python interpreter state to the worker processes. Might require a non-trivial change in its design



Discussion:
Thanks for reporting, will look into it.
@samuelstanton You can move the `register_resolver` inside `main` as I mentioned in the bug report
@calebho, simply moving it inside the `main` will not work perfectly - it will break normal `basic launcher` with `multirun`.
Moreover, the same behavior is happening with `joblib` - so it's not specific to `submitit` - maybe renaming an issue would help.

I had similar issue few days ago but went around it by doing this in my `main`:
```python
if OmegaConf.get_resolver("my_custom_resolver") is None:
        OmegaConf.register_resolver("my_custom_resolver", my_custom_resolver)
```
I forgot to report it - that would save some of your time - sorry.

Resolving patch:
diff --git a/hydra/_internal/hydra.py b/hydra/_internal/hydra.py
--- a/hydra/_internal/hydra.py
+++ b/hydra/_internal/hydra.py
@@ -86,6 +86,7 @@ def __init__(self, task_name: str, config_loader: ConfigLoader) -> None:
         :param config_loader: config loader
         """
         setup_globals()
+
         self.config_loader = config_loader
         JobRuntime().set("name", task_name)
 
diff --git a/hydra/core/singleton.py b/hydra/core/singleton.py
--- a/hydra/core/singleton.py
+++ b/hydra/core/singleton.py
@@ -1,6 +1,9 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from copy import deepcopy
 from typing import Any, Dict
 
+from omegaconf.basecontainer import BaseContainer
+
 
 class Singleton(type):
     _instances: Dict[type, "Singleton"] = {}
@@ -14,9 +17,13 @@ def instance(cls: Any, *args: Any, **kwargs: Any) -> Any:
         return cls(*args, **kwargs)
 
     @staticmethod
-    def get_state() -> Dict[type, "Singleton"]:
-        return Singleton._instances
+    def get_state() -> Any:
+        return {
+            "instances": Singleton._instances,
+            "omegaconf_resolvers": deepcopy(BaseContainer._resolvers),
+        }
 
     @staticmethod
-    def set_state(instances: Dict[type, "Singleton"]) -> None:
-        Singleton._instances = instances
+    def set_state(state: Any) -> None:
+        Singleton._instances = state["instances"]
+        BaseContainer._resolvers = deepcopy(state["omegaconf_resolvers"])


Test patch:
diff --git a/hydra/test_utils/configs/custom_resolver.yaml b/hydra/test_utils/configs/custom_resolver.yaml
new file mode 100644
--- /dev/null
+++ b/hydra/test_utils/configs/custom_resolver.yaml
@@ -0,0 +1 @@
+x: ${my_custom_resolver:}
diff --git a/hydra/test_utils/launcher_common_tests.py b/hydra/test_utils/launcher_common_tests.py
--- a/hydra/test_utils/launcher_common_tests.py
+++ b/hydra/test_utils/launcher_common_tests.py
@@ -208,6 +208,39 @@ def test_sweep_and_override(
                 assert job_ret.cfg == expected_conf[i]
                 verify_dir_outputs(job_ret, job_ret.overrides)
 
+    def test_sweep_with_custom_resolver(
+        self,
+        hydra_sweep_runner: TSweepRunner,
+        launcher_name: str,
+        overrides: List[str],
+        tmpdir: Path,
+    ) -> None:
+
+        overrides1 = ["hydra/launcher=" + launcher_name] + overrides
+
+        def task_func(c: DictConfig) -> Any:
+            return c.x
+
+        sweep = hydra_sweep_runner(
+            calling_file=None,
+            calling_module="hydra.test_utils.a_module",
+            task_function=task_func,
+            config_path="configs",
+            config_name="custom_resolver",
+            overrides=overrides1,
+            temp_dir=tmpdir,
+        )
+
+        def my_custom_resolver() -> Any:
+            return "foo"
+
+        OmegaConf.register_resolver("my_custom_resolver", my_custom_resolver)
+        with sweep:
+            assert sweep.returns is not None
+            job_ret = sweep.returns[0]
+            assert len(job_ret) == 1
+            assert job_ret[0].return_value == "foo"
+
 
 @pytest.mark.usefixtures("hydra_restore_singletons")
 class BatchedSweeperTestSuite: