Repository: facebookresearch/hydra

Issue:
[Ax] Fix warnings related to OmegaConf 2.0.0 compact key
Related to https://github.com/omry/omegaconf/issues/152


Discussion:

Resolving patch:
diff --git a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
--- a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
+++ b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
@@ -195,6 +195,10 @@ def sweep(self, arguments: List[str]) -> None:
 
             current_parallelism_index += 1
 
+        best_parameters = {
+            normalize_key(key, str_to_replace=".", str_to_replace_with="_",): value
+            for key, value in best_parameters.items()
+        }
         results_to_serialize = {"optimizer": "ax", "ax": best_parameters}
         OmegaConf.save(
             OmegaConf.create(results_to_serialize),
@@ -225,6 +229,7 @@ def setup_ax_client(self, arguments: List[str]) -> AxClient:
         """Method to setup the Ax Client"""
         parameters: List[Dict[str, Any]] = []
         for key, value in self.ax_params.items():
+            key = normalize_key(key, str_to_replace="_", str_to_replace_with=".")
             param = OmegaConf.to_container(value, resolve=True)
             assert isinstance(param, Dict)
             if "parameter_type" not in param:
@@ -344,3 +349,17 @@ def chunks(batch: List[Any], n: Optional[int]) -> Iterable[List[Any]]:
             raise ValueError("n must be an integer greater than 0")
         for i in range(0, len(batch), n):
             yield batch[i : i + n]
+
+
+def normalize_key(key: str, str_to_replace: str, str_to_replace_with: str) -> str:
+    """Process the key by replacing "str_to_replace" with "str_to_replace_with".
+    The "str_to_replace" escaped using r"\\" are not replaced. Finally, the r"\\" is removed.
+    """
+    str_to_escape = "\\" + str_to_replace
+    splits_to_update = key.split(str_to_escape)
+    updated_splits = [
+        current_split.replace(str_to_replace, str_to_replace_with)
+        for current_split in splits_to_update
+    ]
+    new_key = str_to_escape.join(updated_splits).replace("\\", "")
+    return new_key


Test patch:
diff --git a/plugins/hydra_ax_sweeper/tests/apps/polynomial.yaml b/plugins/hydra_ax_sweeper/tests/apps/polynomial.yaml
--- a/plugins/hydra_ax_sweeper/tests/apps/polynomial.yaml
+++ b/plugins/hydra_ax_sweeper/tests/apps/polynomial.yaml
@@ -24,10 +24,10 @@ hydra:
             max_epochs_without_improvement: 2
 
           params:
-            polynomial.x:
+            polynomial_x:
               type: range
               bounds: [-1, 1]
 
-            polynomial.y:
+            polynomial_y:
               type: choice
               values: [-1, 1]
\ No newline at end of file
diff --git a/plugins/hydra_ax_sweeper/tests/config/quadratic.yaml b/plugins/hydra_ax_sweeper/tests/config/config.yaml
similarity index 71%
rename from plugins/hydra_ax_sweeper/tests/config/quadratic.yaml
rename to plugins/hydra_ax_sweeper/tests/config/config.yaml
--- a/plugins/hydra_ax_sweeper/tests/config/quadratic.yaml
+++ b/plugins/hydra_ax_sweeper/tests/config/config.yaml
@@ -10,7 +10,10 @@ hydra:
     sweeper:
       params:
         ax_config:
-          max_trials: 10
+          client:
+            random_seed: 1
+
+          max_trials: 2
 
           experiment:
             # Defaults to minimize, set to false to maximize
@@ -22,11 +25,4 @@ hydra:
             # An Epoch is defined as a batch of trials executed in parallel
             max_epochs_without_improvement: 2
           
-          params:
-            quadratic.x:
-              type: range
-              bounds: [-1, 1]
-
-            quadratic.y:
-              type: range
-              bounds: [-1, 1]
+          
\ No newline at end of file
diff --git a/plugins/hydra_ax_sweeper/tests/config/defaults/default.yaml b/plugins/hydra_ax_sweeper/tests/config/defaults/default.yaml
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_ax_sweeper/tests/config/defaults/default.yaml
@@ -0,0 +1,2 @@
+defaults:
+  - hydra/sweeper: ax
\ No newline at end of file
diff --git a/plugins/hydra_ax_sweeper/tests/config/params/basic.yaml b/plugins/hydra_ax_sweeper/tests/config/params/basic.yaml
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_ax_sweeper/tests/config/params/basic.yaml
@@ -0,0 +1,12 @@
+hydra:
+    sweeper:
+      params:
+        ax_config:
+          params:
+            quadratic_x:
+              type: range
+              bounds: [-1, 1]
+
+            quadratic_y:
+              type: range
+              bounds: [-1, 1]
diff --git a/plugins/hydra_ax_sweeper/tests/config/params/nested_with_escape_char.yaml b/plugins/hydra_ax_sweeper/tests/config/params/nested_with_escape_char.yaml
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_ax_sweeper/tests/config/params/nested_with_escape_char.yaml
@@ -0,0 +1,11 @@
+hydra:
+    sweeper:
+      params:
+        ax_config:
+          params:
+            quadratic_x\_arg:
+              type: range
+              bounds: [-1, 1]
+            quadratic_y\_arg:
+              type: range
+              bounds: [-1, 1]
\ No newline at end of file
diff --git a/plugins/hydra_ax_sweeper/tests/config/default_quadratic.yaml b/plugins/hydra_ax_sweeper/tests/config/quadratic/basic.yaml
similarity index 59%
rename from plugins/hydra_ax_sweeper/tests/config/default_quadratic.yaml
rename to plugins/hydra_ax_sweeper/tests/config/quadratic/basic.yaml
--- a/plugins/hydra_ax_sweeper/tests/config/default_quadratic.yaml
+++ b/plugins/hydra_ax_sweeper/tests/config/quadratic/basic.yaml
@@ -1,6 +1,3 @@
-defaults:
-  - hydra/sweeper: ax
-
 quadratic:
   # To be minimized
   x: ???
diff --git a/plugins/hydra_ax_sweeper/tests/config/quadratic/nested_with_escape_char.yaml b/plugins/hydra_ax_sweeper/tests/config/quadratic/nested_with_escape_char.yaml
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_ax_sweeper/tests/config/quadratic/nested_with_escape_char.yaml
@@ -0,0 +1,3 @@
+quadratic:
+  x_arg: ???
+  y_arg: ???
\ No newline at end of file
diff --git a/plugins/hydra_ax_sweeper/tests/test_ax_sweeper_plugin.py b/plugins/hydra_ax_sweeper/tests/test_ax_sweeper_plugin.py
--- a/plugins/hydra_ax_sweeper/tests/test_ax_sweeper_plugin.py
+++ b/plugins/hydra_ax_sweeper/tests/test_ax_sweeper_plugin.py
@@ -29,12 +29,11 @@ def test_discovery() -> None:
 
 
 def quadratic(cfg: DictConfig) -> Any:
-    x = cfg.quadratic.x
-    y = cfg.quadratic.y
-    a = 100
-    b = 1
-    z = a * (x ** 2) + b * y
-    return z
+    return 100 * (cfg.quadratic.x ** 2) + 1 * cfg.quadratic.y
+
+
+def nested_quadratic_with_escape_char(cfg: DictConfig) -> Any:
+    return 100 * (cfg.quadratic.x_arg ** 2) + 1 * cfg.quadratic.y_arg
 
 
 @pytest.mark.parametrize(
@@ -53,7 +52,6 @@ def test_chunk_method_for_valid_inputs(n, expected):
     chunk_func = CoreAxSweeper.chunks
     batch = [1, 2, 3, 4, 5]
     out = list(chunk_func(batch, n))
-    print(expected)
     assert out == expected
 
 
@@ -74,13 +72,12 @@ def test_jobs_dirs(sweep_runner: TSweepRunner) -> None:
         calling_module=None,
         task_function=quadratic,
         config_path="config",
-        config_name="quadratic.yaml",
+        config_name="config.yaml",
         overrides=[
             "hydra/launcher=basic",
-            "hydra/sweeper=ax",
-            "hydra.sweeper.params.ax_config.client.random_seed=1",
             "hydra.sweeper.params.ax_config.max_trials=6",
             "hydra.sweeper.params.max_batch_size=2",
+            "params=basic",
         ],
         strict=True,
     )
@@ -100,13 +97,8 @@ def test_jobs_configured_via_config(sweep_runner: TSweepRunner) -> None:
         calling_module=None,
         task_function=quadratic,
         config_path="config",
-        config_name="quadratic.yaml",
-        overrides=[
-            "hydra/sweeper=ax",
-            "hydra/launcher=basic",
-            "hydra.sweeper.params.ax_config.client.random_seed=1",
-            "hydra.sweeper.params.ax_config.max_trials=2",
-        ],
+        config_name="config.yaml",
+        overrides=["hydra/launcher=basic", "params=basic"],
         strict=True,
     )
     with sweep:
@@ -116,9 +108,8 @@ def test_jobs_configured_via_config(sweep_runner: TSweepRunner) -> None:
         assert returns["optimizer"] == "ax"
         assert len(returns) == 2
         best_parameters = returns["ax"]
-        assert len(best_parameters) == 2
-        assert math.isclose(best_parameters["quadratic.x"], 0.0, abs_tol=1e-4)
-        assert math.isclose(best_parameters["quadratic.y"], -1.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_x"], 0.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_y"], -1.0, abs_tol=1e-4)
 
 
 def test_jobs_configured_via_cmd(sweep_runner: TSweepRunner,) -> None:
@@ -127,14 +118,12 @@ def test_jobs_configured_via_cmd(sweep_runner: TSweepRunner,) -> None:
         calling_module=None,
         task_function=quadratic,
         config_path="config",
-        config_name="quadratic.yaml",
+        config_name="config.yaml",
         overrides=[
-            "hydra/sweeper=ax",
             "hydra/launcher=basic",
-            "hydra.sweeper.params.ax_config.client.random_seed=1",
             "quadratic.x=-5:-2",
             "quadratic.y=-2:2",
-            "hydra.sweeper.params.ax_config.max_trials=2",
+            "params=basic",
         ],
         strict=True,
     )
@@ -145,9 +134,8 @@ def test_jobs_configured_via_cmd(sweep_runner: TSweepRunner,) -> None:
         assert returns["optimizer"] == "ax"
         assert len(returns) == 2
         best_parameters = returns["ax"]
-        assert len(best_parameters) == 2
-        assert math.isclose(best_parameters["quadratic.x"], -2.0, abs_tol=1e-4)
-        assert math.isclose(best_parameters["quadratic.y"], 2.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_x"], -2.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_y"], 2.0, abs_tol=1e-4)
 
 
 def test_jobs_configured_via_cmd_and_config(sweep_runner: TSweepRunner) -> None:
@@ -156,13 +144,12 @@ def test_jobs_configured_via_cmd_and_config(sweep_runner: TSweepRunner) -> None:
         calling_module=None,
         task_function=quadratic,
         config_path="config",
-        config_name="quadratic.yaml",
+        config_name="config.yaml",
         overrides=[
-            "hydra/sweeper=ax",
             "hydra/launcher=basic",
-            "hydra.sweeper.params.ax_config.client.random_seed=1",
-            "quadratic.x=-5:-2",
             "hydra.sweeper.params.ax_config.max_trials=2",
+            "quadratic.x=-5:-2",
+            "params=basic",
         ],
         strict=True,
     )
@@ -173,9 +160,8 @@ def test_jobs_configured_via_cmd_and_config(sweep_runner: TSweepRunner) -> None:
         assert returns["optimizer"] == "ax"
         assert len(returns) == 2
         best_parameters = returns["ax"]
-        assert len(best_parameters) == 2
-        assert math.isclose(best_parameters["quadratic.x"], -2.0, abs_tol=1e-4)
-        assert math.isclose(best_parameters["quadratic.y"], 1.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_x"], -2.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_y"], 1.0, abs_tol=1e-4)
 
 
 def test_configuration_set_via_cmd_and_default_config(
@@ -186,13 +172,12 @@ def test_configuration_set_via_cmd_and_default_config(
         calling_module=None,
         task_function=quadratic,
         config_path="config",
-        config_name="default_quadratic.yaml",
+        config_name="config.yaml",
         overrides=[
-            "hydra/sweeper=ax",
             "hydra/launcher=basic",
-            "hydra.sweeper.params.ax_config.client.random_seed=1",
             "hydra.sweeper.params.ax_config.max_trials=2",
             "hydra.sweeper.params.ax_config.early_stop.max_epochs_without_improvement=2",
+            "quadratic=basic",
             "quadratic.x=-5:-2",
             "quadratic.y=-1:1",
         ],
@@ -206,8 +191,8 @@ def test_configuration_set_via_cmd_and_default_config(
         returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
         assert isinstance(returns, DictConfig)
         best_parameters = returns["ax"]
-        assert "quadratic.x" in best_parameters
-        assert "quadratic.y" in best_parameters
+        assert "quadratic_x" in best_parameters
+        assert "quadratic_y" in best_parameters
 
 
 def test_ax_logging(tmpdir: Path) -> None:
@@ -251,3 +236,52 @@ def test_example_app(tmpdir: Path) -> None:
 # class TestAxSweeper(LauncherTestSuite):
 #     def task_function(self, cfg):
 #         return 100 * (cfg.quadratic.x ** 2) + 1 * cfg.quadratic.y
+
+
+@pytest.mark.parametrize(
+    "overrides", [[], ["quadratic.x_arg=-1:1"]],
+)
+def test_jobs_configured_via_nested_config(
+    sweep_runner: TSweepRunner, overrides: list,
+) -> None:
+    sweep = sweep_runner(
+        calling_file="tests/test_ax_sweeper_plugin.py",
+        calling_module=None,
+        task_function=nested_quadratic_with_escape_char,
+        config_path="config",
+        config_name="config.yaml",
+        overrides=[
+            "hydra/launcher=basic",
+            "quadratic=nested_with_escape_char",
+            "params=nested_with_escape_char",
+        ]
+        + overrides,
+        strict=True,
+    )
+    with sweep:
+        assert sweep.returns is None
+        returns = OmegaConf.load(f"{sweep.temp_dir}/optimization_results.yaml")
+        assert isinstance(returns, DictConfig)
+        assert returns["optimizer"] == "ax"
+        assert len(returns) == 2
+        best_parameters = returns["ax"]
+        assert math.isclose(best_parameters["quadratic_x_arg"], 0.0, abs_tol=1e-4)
+        assert math.isclose(best_parameters["quadratic_y_arg"], -1.0, abs_tol=1e-4)
+
+
+@pytest.mark.parametrize(
+    "inp, str_to_replace, str_to_replace_with, expected",
+    [
+        ("apple", ".", "_", "apple"),
+        ("a.pple", ".", "_", "a_pple"),
+        ("a.p.ple", ".", "_", "a_p_ple"),
+        (r"a\.pple", ".", "_", "a.pple"),
+        (r"a.p\.pl.e", ".", "_", "a_p.pl_e"),
+    ],
+)
+def test_process_key_method(
+    inp: str, str_to_replace: str, str_to_replace_with: str, expected: str
+):
+    from hydra_plugins.hydra_ax_sweeper._core import normalize_key
+
+    assert normalize_key(inp, str_to_replace, str_to_replace_with) == expected