We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Ax-Plugin] Switch back to the usage of "." in place of "_"

> > I abandoned the idea of compact keys in OmegaConf. This is a good chance to switch back to the more intuitive usage of dot and not _ to separate the path element in keys:
> > ```yaml
> >       params:
> >         banana.x:
> >           type: range
> >           bounds: [-5, 5]
> > ```


_Originally posted by @omry in https://github.com/facebookresearch/hydra/pull/912#issuecomment-679236669_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
2 import logging
3 from dataclasses import dataclass
4 from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union
5 
6 from ax.core import types as ax_types  # type: ignore
7 from ax.service.ax_client import AxClient  # type: ignore
8 from hydra.core.config_loader import ConfigLoader
9 from hydra.core.override_parser.overrides_parser import OverridesParser
10 from hydra.core.override_parser.types import IntervalSweep, Override, Transformer
11 from hydra.core.plugins import Plugins
12 from hydra.plugins.launcher import Launcher
13 from hydra.plugins.sweeper import Sweeper
14 from hydra.types import TaskFunction
15 from omegaconf import DictConfig, OmegaConf
16 
17 from ._earlystopper import EarlyStopper
18 from .config import AxConfig, ClientConfig, ExperimentConfig
19 
20 log = logging.getLogger(__name__)
21 
22 
23 @dataclass
24 class Trial:
25     overrides: List[str]
26     trial_index: int
27 
28 
29 BatchOfTrialType = List[Trial]
30 
31 
32 def encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:
33     """Convert a list of params into a string"""
34     mandatory_keys = set(["name", "type", "bounds", "values", "value"])
35     parameter_log_string = ""
36     for parameter in parameters:
37         parameter_log_string += "\n"
38         parameter_log_string += f"{parameter['name']}: {parameter['type']}="
39         if parameter["type"] == "range":
40             parameter_log_string += f"{parameter['bounds']}"
41         elif parameter["type"] == "choice":
42             parameter_log_string += f"{parameter['values']}"
43         elif parameter["type"] == "fixed":
44             parameter_log_string += f"{parameter['value']}"
45 
46         for key, value in parameter.items():
47             if key not in mandatory_keys:
48                 parameter_log_string += f", {key} = {value}"
49     return parameter_log_string
50 
51 
52 def map_params_to_arg_list(params: Mapping[str, Union[str, float, int]]) -> List[str]:
53     """Method to map a dictionary of params to a list of string arguments"""
54     arg_list = []
55     for key in params:
56         arg_list.append(str(key) + "=" + str(params[key]))
57     return arg_list
58 
59 
60 def get_one_batch_of_trials(
61     ax_client: AxClient,
62     parallelism: Tuple[int, int],
63     num_trials_so_far: int,
64     num_max_trials_to_do: int,
65 ) -> BatchOfTrialType:
66     """Produce a batch of trials that can be run in parallel"""
67     (num_trials, max_parallelism_setting) = parallelism
68     if max_parallelism_setting == -1:
69         # Special case, we can group all the trials into one batch
70         max_parallelism_setting = num_trials - num_trials_so_far
71 
72         if num_trials == -1:
73             # This is a special case where we can run as many trials in parallel as we want.
74             # Given that num_trials is also -1, we can run all the trials in parallel.
75             max_parallelism_setting = num_max_trials_to_do
76 
77     batch_of_trials = []
78     for _ in range(max_parallelism_setting):
79         parameters, trial_index = ax_client.get_next_trial()
80         batch_of_trials.append(
81             Trial(
82                 overrides=map_params_to_arg_list(params=parameters),
83                 trial_index=trial_index,
84             )
85         )
86     return batch_of_trials
87 
88 
89 class CoreAxSweeper(Sweeper):
90     """Class to interface with the Ax Platform"""
91 
92     def __init__(self, ax_config: AxConfig, max_batch_size: Optional[int]):
93         self.config_loader: Optional[ConfigLoader] = None
94         self.config: Optional[DictConfig] = None
95         self.launcher: Optional[Launcher] = None
96 
97         self.job_results = None
98         self.experiment: ExperimentConfig = ax_config.experiment
99         self.early_stopper: EarlyStopper = EarlyStopper(
100             max_epochs_without_improvement=ax_config.early_stop.max_epochs_without_improvement,
101             epsilon=ax_config.early_stop.epsilon,
102             minimize=ax_config.early_stop.minimize,
103         )
104         self.ax_client_config: ClientConfig = ax_config.client
105         self.max_trials = ax_config.max_trials
106         self.ax_params: DictConfig = OmegaConf.create({})
107         if hasattr(ax_config, "params"):
108             self.ax_params.update(ax_config.params)
109         self.sweep_dir: str
110         self.job_idx: Optional[int] = None
111         self.max_batch_size = max_batch_size
112 
113     def setup(
114         self,
115         config: DictConfig,
116         config_loader: ConfigLoader,
117         task_function: TaskFunction,
118     ) -> None:
119         self.config = config
120         self.config_loader = config_loader
121         self.launcher = Plugins.instance().instantiate_launcher(
122             config=config, config_loader=config_loader, task_function=task_function
123         )
124         self.sweep_dir = config.hydra.sweep.dir
125 
126     def sweep(self, arguments: List[str]) -> None:
127         self.job_idx = 0
128         ax_client = self.setup_ax_client(arguments)
129 
130         num_trials_left = self.max_trials
131         max_parallelism = ax_client.get_max_parallelism()
132         current_parallelism_index = 0
133         # Index to track the parallelism value we are using right now.
134         best_parameters = {}
135         while num_trials_left > 0:
136             current_parallelism = max_parallelism[current_parallelism_index]
137             num_trials, max_parallelism_setting = current_parallelism
138             num_trials_so_far = 0
139             while (
140                 num_trials > num_trials_so_far or num_trials == -1
141             ) and num_trials_left > 0:
142                 batch_of_trials = get_one_batch_of_trials(
143                     ax_client=ax_client,
144                     parallelism=current_parallelism,
145                     num_trials_so_far=num_trials_so_far,
146                     num_max_trials_to_do=num_trials_left,
147                 )
148                 batch_of_trials_to_launch = batch_of_trials[:num_trials_left]
149 
150                 log.info(
151                     "AxSweeper is launching {} jobs".format(
152                         len(batch_of_trials_to_launch)
153                     )
154                 )
155 
156                 self.sweep_over_batches(
157                     ax_client=ax_client, batch_of_trials=batch_of_trials_to_launch
158                 )
159 
160                 num_trials_so_far += len(batch_of_trials_to_launch)
161                 num_trials_left -= len(batch_of_trials_to_launch)
162 
163                 best_parameters, predictions = ax_client.get_best_parameters()
164                 metric = predictions[0][ax_client.objective_name]
165 
166                 if self.early_stopper.should_stop(metric, best_parameters):
167                     num_trials_left = -1
168                     break
169 
170             current_parallelism_index += 1
171 
172         best_parameters = {
173             normalize_key(key, str_to_replace=".", str_to_replace_with="_"): value
174             for key, value in best_parameters.items()
175         }
176         results_to_serialize = {"optimizer": "ax", "ax": best_parameters}
177         OmegaConf.save(
178             OmegaConf.create(results_to_serialize),
179             f"{self.sweep_dir}/optimization_results.yaml",
180         )
181         log.info("Best parameters: " + str(best_parameters))
182 
183     def sweep_over_batches(
184         self, ax_client: AxClient, batch_of_trials: BatchOfTrialType
185     ) -> None:
186         assert self.launcher is not None
187         assert self.job_idx is not None
188 
189         chunked_batches = self.chunks(batch_of_trials, self.max_batch_size)
190         for batch in chunked_batches:
191             overrides = [x.overrides for x in batch]
192             self.validate_batch_is_legal(overrides)
193             rets = self.launcher.launch(
194                 job_overrides=overrides, initial_job_idx=self.job_idx
195             )
196             self.job_idx += len(rets)
197             for idx in range(len(batch)):
198                 val = rets[idx].return_value
199                 ax_client.complete_trial(
200                     trial_index=batch[idx].trial_index, raw_data=val
201                 )
202 
203     def setup_ax_client(self, arguments: List[str]) -> AxClient:
204         """Method to setup the Ax Client"""
205         parameters: List[Dict[str, Any]] = []
206         for key, value in self.ax_params.items():
207             key = normalize_key(key, str_to_replace="_", str_to_replace_with=".")
208             param = OmegaConf.to_container(value, resolve=True)
209             assert isinstance(param, Dict)
210             if param["type"] == "range":
211                 bounds = param["bounds"]
212                 if not (all(isinstance(x, int) for x in bounds)):
213                     # Type mismatch. Promote all to float
214                     param["bounds"] = [float(x) for x in bounds]
215 
216             parameters.append(param)
217             parameters[-1]["name"] = key
218         commandline_params = self.parse_commandline_args(arguments)
219         for cmd_param in commandline_params:
220             for param in parameters:
221                 if param["name"] == cmd_param["name"]:
222                     for key, value in cmd_param.items():
223                         param[key] = value
224                     break
225             else:
226                 parameters.append(cmd_param)
227 
228         log.info(
229             f"AxSweeper is optimizing the following parameters: {encoder_parameters_into_string(parameters)}"
230         )
231         ax_client = AxClient(
232             verbose_logging=self.ax_client_config.verbose_logging,
233             random_seed=self.ax_client_config.random_seed,
234         )
235         ax_client.create_experiment(parameters=parameters, **self.experiment)
236 
237         return ax_client
238 
239     def parse_commandline_args(
240         self, arguments: List[str]
241     ) -> List[Dict[str, Union[ax_types.TParamValue, List[ax_types.TParamValue]]]]:
242         """Method to parse the command line arguments and convert them into Ax parameters"""
243         parser = OverridesParser.create()
244         parsed = parser.parse_overrides(arguments)
245         parameters: List[Dict[str, Any]] = []
246         for override in parsed:
247             if override.is_sweep_override():
248                 if override.is_choice_sweep():
249                     param = create_choice_param_from_choice_override(override)
250                 elif override.is_range_sweep():
251                     param = create_choice_param_from_range_override(override)
252                 elif override.is_interval_sweep():
253                     param = create_range_param_using_interval_override(override)
254             elif not override.is_hydra_override():
255                 param = create_fixed_param_from_element_override(override)
256             parameters.append(param)
257 
258         return parameters
259 
260     @staticmethod
261     def chunks(batch: List[Any], n: Optional[int]) -> Iterable[List[Any]]:
262         """
263         Chunk the batch into chunks of upto to n items (each)
264         """
265         if n is None:
266             n = len(batch)
267         if n < 1:
268             raise ValueError("n must be an integer greater than 0")
269         for i in range(0, len(batch), n):
270             yield batch[i : i + n]
271 
272 
273 def normalize_key(key: str, str_to_replace: str, str_to_replace_with: str) -> str:
274     """Process the key by replacing "str_to_replace" with "str_to_replace_with".
275     The "str_to_replace" escaped using r"\\" are not replaced. Finally, the r"\\" is removed.
276     """
277     str_to_escape = "\\" + str_to_replace
278     splits_to_update = key.split(str_to_escape)
279     updated_splits = [
280         current_split.replace(str_to_replace, str_to_replace_with)
281         for current_split in splits_to_update
282     ]
283     new_key = str_to_escape.join(updated_splits).replace("\\", "")
284     return new_key
285 
286 
287 def create_range_param_using_interval_override(override: Override) -> Dict[str, Any]:
288     key = override.get_key_element()
289     value = override.value()
290     assert isinstance(value, IntervalSweep)
291     param = {
292         "name": key,
293         "type": "range",
294         "bounds": [value.start, value.end],
295     }
296     return param
297 
298 
299 def create_choice_param_from_choice_override(override: Override) -> Dict[str, Any]:
300     key = override.get_key_element()
301     param = {
302         "name": key,
303         "type": "choice",
304         "values": list(override.sweep_iterator(transformer=Transformer.encode)),
305     }
306     return param
307 
308 
309 def create_choice_param_from_range_override(override: Override) -> Dict[str, Any]:
310     key = override.get_key_element()
311     param = {
312         "name": key,
313         "type": "choice",
314         "values": [val for val in override.sweep_iterator()],
315     }
316 
317     return param
318 
319 
320 def create_fixed_param_from_element_override(override: Override) -> Dict[str, Any]:
321     key = override.get_key_element()
322     param = {
323         "name": key,
324         "type": "fixed",
325         "value": override.value(),
326     }
327 
328     return param
329
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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
@@ -169,10 +169,6 @@
 
             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),
@@ -204,7 +200,6 @@
         """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 param["type"] == "range":
@@ -270,20 +265,6 @@
             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
-
-
 def create_range_param_using_interval_override(override: Override) -> Dict[str, Any]:
     key = override.get_key_element()
     value = override.value()

{"golden_diff": "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\n--- a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py\n+++ b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py\n@@ -169,10 +169,6 @@\n \n             current_parallelism_index += 1\n \n-        best_parameters = {\n-            normalize_key(key, str_to_replace=\".\", str_to_replace_with=\"_\"): value\n-            for key, value in best_parameters.items()\n-        }\n         results_to_serialize = {\"optimizer\": \"ax\", \"ax\": best_parameters}\n         OmegaConf.save(\n             OmegaConf.create(results_to_serialize),\n@@ -204,7 +200,6 @@\n         \"\"\"Method to setup the Ax Client\"\"\"\n         parameters: List[Dict[str, Any]] = []\n         for key, value in self.ax_params.items():\n-            key = normalize_key(key, str_to_replace=\"_\", str_to_replace_with=\".\")\n             param = OmegaConf.to_container(value, resolve=True)\n             assert isinstance(param, Dict)\n             if param[\"type\"] == \"range\":\n@@ -270,20 +265,6 @@\n             yield batch[i : i + n]\n \n \n-def normalize_key(key: str, str_to_replace: str, str_to_replace_with: str) -> str:\n-    \"\"\"Process the key by replacing \"str_to_replace\" with \"str_to_replace_with\".\n-    The \"str_to_replace\" escaped using r\"\\\\\" are not replaced. Finally, the r\"\\\\\" is removed.\n-    \"\"\"\n-    str_to_escape = \"\\\\\" + str_to_replace\n-    splits_to_update = key.split(str_to_escape)\n-    updated_splits = [\n-        current_split.replace(str_to_replace, str_to_replace_with)\n-        for current_split in splits_to_update\n-    ]\n-    new_key = str_to_escape.join(updated_splits).replace(\"\\\\\", \"\")\n-    return new_key\n-\n-\n def create_range_param_using_interval_override(override: Override) -> Dict[str, Any]:\n     key = override.get_key_element()\n     value = override.value()\n", "issue": "[Ax-Plugin] Switch back to the usage of \".\" in place of \"_\"\n\r\n> > I abandoned the idea of compact keys in OmegaConf. This is a good chance to switch back to the more intuitive usage of dot and not _ to separate the path element in keys:\r\n> > ```yaml\r\n> >       params:\r\n> >         banana.x:\r\n> >           type: range\r\n> >           bounds: [-5, 5]\r\n> > ```\r\n\r\n\r\n_Originally posted by @omry in https://github.com/facebookresearch/hydra/pull/912#issuecomment-679236669_\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport logging\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union\n\nfrom ax.core import types as ax_types  # type: ignore\nfrom ax.service.ax_client import AxClient  # type: ignore\nfrom hydra.core.config_loader import ConfigLoader\nfrom hydra.core.override_parser.overrides_parser import OverridesParser\nfrom hydra.core.override_parser.types import IntervalSweep, Override, Transformer\nfrom hydra.core.plugins import Plugins\nfrom hydra.plugins.launcher import Launcher\nfrom hydra.plugins.sweeper import Sweeper\nfrom hydra.types import TaskFunction\nfrom omegaconf import DictConfig, OmegaConf\n\nfrom ._earlystopper import EarlyStopper\nfrom .config import AxConfig, ClientConfig, ExperimentConfig\n\nlog = logging.getLogger(__name__)\n\n\n@dataclass\nclass Trial:\n    overrides: List[str]\n    trial_index: int\n\n\nBatchOfTrialType = List[Trial]\n\n\ndef encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:\n    \"\"\"Convert a list of params into a string\"\"\"\n    mandatory_keys = set([\"name\", \"type\", \"bounds\", \"values\", \"value\"])\n    parameter_log_string = \"\"\n    for parameter in parameters:\n        parameter_log_string += \"\\n\"\n        parameter_log_string += f\"{parameter['name']}: {parameter['type']}=\"\n        if parameter[\"type\"] == \"range\":\n            parameter_log_string += f\"{parameter['bounds']}\"\n        elif parameter[\"type\"] == \"choice\":\n            parameter_log_string += f\"{parameter['values']}\"\n        elif parameter[\"type\"] == \"fixed\":\n            parameter_log_string += f\"{parameter['value']}\"\n\n        for key, value in parameter.items():\n            if key not in mandatory_keys:\n                parameter_log_string += f\", {key} = {value}\"\n    return parameter_log_string\n\n\ndef map_params_to_arg_list(params: Mapping[str, Union[str, float, int]]) -> List[str]:\n    \"\"\"Method to map a dictionary of params to a list of string arguments\"\"\"\n    arg_list = []\n    for key in params:\n        arg_list.append(str(key) + \"=\" + str(params[key]))\n    return arg_list\n\n\ndef get_one_batch_of_trials(\n    ax_client: AxClient,\n    parallelism: Tuple[int, int],\n    num_trials_so_far: int,\n    num_max_trials_to_do: int,\n) -> BatchOfTrialType:\n    \"\"\"Produce a batch of trials that can be run in parallel\"\"\"\n    (num_trials, max_parallelism_setting) = parallelism\n    if max_parallelism_setting == -1:\n        # Special case, we can group all the trials into one batch\n        max_parallelism_setting = num_trials - num_trials_so_far\n\n        if num_trials == -1:\n            # This is a special case where we can run as many trials in parallel as we want.\n            # Given that num_trials is also -1, we can run all the trials in parallel.\n            max_parallelism_setting = num_max_trials_to_do\n\n    batch_of_trials = []\n    for _ in range(max_parallelism_setting):\n        parameters, trial_index = ax_client.get_next_trial()\n        batch_of_trials.append(\n            Trial(\n                overrides=map_params_to_arg_list(params=parameters),\n                trial_index=trial_index,\n            )\n        )\n    return batch_of_trials\n\n\nclass CoreAxSweeper(Sweeper):\n    \"\"\"Class to interface with the Ax Platform\"\"\"\n\n    def __init__(self, ax_config: AxConfig, max_batch_size: Optional[int]):\n        self.config_loader: Optional[ConfigLoader] = None\n        self.config: Optional[DictConfig] = None\n        self.launcher: Optional[Launcher] = None\n\n        self.job_results = None\n        self.experiment: ExperimentConfig = ax_config.experiment\n        self.early_stopper: EarlyStopper = EarlyStopper(\n            max_epochs_without_improvement=ax_config.early_stop.max_epochs_without_improvement,\n            epsilon=ax_config.early_stop.epsilon,\n            minimize=ax_config.early_stop.minimize,\n        )\n        self.ax_client_config: ClientConfig = ax_config.client\n        self.max_trials = ax_config.max_trials\n        self.ax_params: DictConfig = OmegaConf.create({})\n        if hasattr(ax_config, \"params\"):\n            self.ax_params.update(ax_config.params)\n        self.sweep_dir: str\n        self.job_idx: Optional[int] = None\n        self.max_batch_size = max_batch_size\n\n    def setup(\n        self,\n        config: DictConfig,\n        config_loader: ConfigLoader,\n        task_function: TaskFunction,\n    ) -> None:\n        self.config = config\n        self.config_loader = config_loader\n        self.launcher = Plugins.instance().instantiate_launcher(\n            config=config, config_loader=config_loader, task_function=task_function\n        )\n        self.sweep_dir = config.hydra.sweep.dir\n\n    def sweep(self, arguments: List[str]) -> None:\n        self.job_idx = 0\n        ax_client = self.setup_ax_client(arguments)\n\n        num_trials_left = self.max_trials\n        max_parallelism = ax_client.get_max_parallelism()\n        current_parallelism_index = 0\n        # Index to track the parallelism value we are using right now.\n        best_parameters = {}\n        while num_trials_left > 0:\n            current_parallelism = max_parallelism[current_parallelism_index]\n            num_trials, max_parallelism_setting = current_parallelism\n            num_trials_so_far = 0\n            while (\n                num_trials > num_trials_so_far or num_trials == -1\n            ) and num_trials_left > 0:\n                batch_of_trials = get_one_batch_of_trials(\n                    ax_client=ax_client,\n                    parallelism=current_parallelism,\n                    num_trials_so_far=num_trials_so_far,\n                    num_max_trials_to_do=num_trials_left,\n                )\n                batch_of_trials_to_launch = batch_of_trials[:num_trials_left]\n\n                log.info(\n                    \"AxSweeper is launching {} jobs\".format(\n                        len(batch_of_trials_to_launch)\n                    )\n                )\n\n                self.sweep_over_batches(\n                    ax_client=ax_client, batch_of_trials=batch_of_trials_to_launch\n                )\n\n                num_trials_so_far += len(batch_of_trials_to_launch)\n                num_trials_left -= len(batch_of_trials_to_launch)\n\n                best_parameters, predictions = ax_client.get_best_parameters()\n                metric = predictions[0][ax_client.objective_name]\n\n                if self.early_stopper.should_stop(metric, best_parameters):\n                    num_trials_left = -1\n                    break\n\n            current_parallelism_index += 1\n\n        best_parameters = {\n            normalize_key(key, str_to_replace=\".\", str_to_replace_with=\"_\"): value\n            for key, value in best_parameters.items()\n        }\n        results_to_serialize = {\"optimizer\": \"ax\", \"ax\": best_parameters}\n        OmegaConf.save(\n            OmegaConf.create(results_to_serialize),\n            f\"{self.sweep_dir}/optimization_results.yaml\",\n        )\n        log.info(\"Best parameters: \" + str(best_parameters))\n\n    def sweep_over_batches(\n        self, ax_client: AxClient, batch_of_trials: BatchOfTrialType\n    ) -> None:\n        assert self.launcher is not None\n        assert self.job_idx is not None\n\n        chunked_batches = self.chunks(batch_of_trials, self.max_batch_size)\n        for batch in chunked_batches:\n            overrides = [x.overrides for x in batch]\n            self.validate_batch_is_legal(overrides)\n            rets = self.launcher.launch(\n                job_overrides=overrides, initial_job_idx=self.job_idx\n            )\n            self.job_idx += len(rets)\n            for idx in range(len(batch)):\n                val = rets[idx].return_value\n                ax_client.complete_trial(\n                    trial_index=batch[idx].trial_index, raw_data=val\n                )\n\n    def setup_ax_client(self, arguments: List[str]) -> AxClient:\n        \"\"\"Method to setup the Ax Client\"\"\"\n        parameters: List[Dict[str, Any]] = []\n        for key, value in self.ax_params.items():\n            key = normalize_key(key, str_to_replace=\"_\", str_to_replace_with=\".\")\n            param = OmegaConf.to_container(value, resolve=True)\n            assert isinstance(param, Dict)\n            if param[\"type\"] == \"range\":\n                bounds = param[\"bounds\"]\n                if not (all(isinstance(x, int) for x in bounds)):\n                    # Type mismatch. Promote all to float\n                    param[\"bounds\"] = [float(x) for x in bounds]\n\n            parameters.append(param)\n            parameters[-1][\"name\"] = key\n        commandline_params = self.parse_commandline_args(arguments)\n        for cmd_param in commandline_params:\n            for param in parameters:\n                if param[\"name\"] == cmd_param[\"name\"]:\n                    for key, value in cmd_param.items():\n                        param[key] = value\n                    break\n            else:\n                parameters.append(cmd_param)\n\n        log.info(\n            f\"AxSweeper is optimizing the following parameters: {encoder_parameters_into_string(parameters)}\"\n        )\n        ax_client = AxClient(\n            verbose_logging=self.ax_client_config.verbose_logging,\n            random_seed=self.ax_client_config.random_seed,\n        )\n        ax_client.create_experiment(parameters=parameters, **self.experiment)\n\n        return ax_client\n\n    def parse_commandline_args(\n        self, arguments: List[str]\n    ) -> List[Dict[str, Union[ax_types.TParamValue, List[ax_types.TParamValue]]]]:\n        \"\"\"Method to parse the command line arguments and convert them into Ax parameters\"\"\"\n        parser = OverridesParser.create()\n        parsed = parser.parse_overrides(arguments)\n        parameters: List[Dict[str, Any]] = []\n        for override in parsed:\n            if override.is_sweep_override():\n                if override.is_choice_sweep():\n                    param = create_choice_param_from_choice_override(override)\n                elif override.is_range_sweep():\n                    param = create_choice_param_from_range_override(override)\n                elif override.is_interval_sweep():\n                    param = create_range_param_using_interval_override(override)\n            elif not override.is_hydra_override():\n                param = create_fixed_param_from_element_override(override)\n            parameters.append(param)\n\n        return parameters\n\n    @staticmethod\n    def chunks(batch: List[Any], n: Optional[int]) -> Iterable[List[Any]]:\n        \"\"\"\n        Chunk the batch into chunks of upto to n items (each)\n        \"\"\"\n        if n is None:\n            n = len(batch)\n        if n < 1:\n            raise ValueError(\"n must be an integer greater than 0\")\n        for i in range(0, len(batch), n):\n            yield batch[i : i + n]\n\n\ndef normalize_key(key: str, str_to_replace: str, str_to_replace_with: str) -> str:\n    \"\"\"Process the key by replacing \"str_to_replace\" with \"str_to_replace_with\".\n    The \"str_to_replace\" escaped using r\"\\\\\" are not replaced. Finally, the r\"\\\\\" is removed.\n    \"\"\"\n    str_to_escape = \"\\\\\" + str_to_replace\n    splits_to_update = key.split(str_to_escape)\n    updated_splits = [\n        current_split.replace(str_to_replace, str_to_replace_with)\n        for current_split in splits_to_update\n    ]\n    new_key = str_to_escape.join(updated_splits).replace(\"\\\\\", \"\")\n    return new_key\n\n\ndef create_range_param_using_interval_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    value = override.value()\n    assert isinstance(value, IntervalSweep)\n    param = {\n        \"name\": key,\n        \"type\": \"range\",\n        \"bounds\": [value.start, value.end],\n    }\n    return param\n\n\ndef create_choice_param_from_choice_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"choice\",\n        \"values\": list(override.sweep_iterator(transformer=Transformer.encode)),\n    }\n    return param\n\n\ndef create_choice_param_from_range_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"choice\",\n        \"values\": [val for val in override.sweep_iterator()],\n    }\n\n    return param\n\n\ndef create_fixed_param_from_element_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"fixed\",\n        \"value\": override.value(),\n    }\n\n    return param\n", "path": "plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\nimport logging\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Union\n\nfrom ax.core import types as ax_types  # type: ignore\nfrom ax.service.ax_client import AxClient  # type: ignore\nfrom hydra.core.config_loader import ConfigLoader\nfrom hydra.core.override_parser.overrides_parser import OverridesParser\nfrom hydra.core.override_parser.types import IntervalSweep, Override, Transformer\nfrom hydra.core.plugins import Plugins\nfrom hydra.plugins.launcher import Launcher\nfrom hydra.plugins.sweeper import Sweeper\nfrom hydra.types import TaskFunction\nfrom omegaconf import DictConfig, OmegaConf\n\nfrom ._earlystopper import EarlyStopper\nfrom .config import AxConfig, ClientConfig, ExperimentConfig\n\nlog = logging.getLogger(__name__)\n\n\n@dataclass\nclass Trial:\n    overrides: List[str]\n    trial_index: int\n\n\nBatchOfTrialType = List[Trial]\n\n\ndef encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:\n    \"\"\"Convert a list of params into a string\"\"\"\n    mandatory_keys = set([\"name\", \"type\", \"bounds\", \"values\", \"value\"])\n    parameter_log_string = \"\"\n    for parameter in parameters:\n        parameter_log_string += \"\\n\"\n        parameter_log_string += f\"{parameter['name']}: {parameter['type']}=\"\n        if parameter[\"type\"] == \"range\":\n            parameter_log_string += f\"{parameter['bounds']}\"\n        elif parameter[\"type\"] == \"choice\":\n            parameter_log_string += f\"{parameter['values']}\"\n        elif parameter[\"type\"] == \"fixed\":\n            parameter_log_string += f\"{parameter['value']}\"\n\n        for key, value in parameter.items():\n            if key not in mandatory_keys:\n                parameter_log_string += f\", {key} = {value}\"\n    return parameter_log_string\n\n\ndef map_params_to_arg_list(params: Mapping[str, Union[str, float, int]]) -> List[str]:\n    \"\"\"Method to map a dictionary of params to a list of string arguments\"\"\"\n    arg_list = []\n    for key in params:\n        arg_list.append(str(key) + \"=\" + str(params[key]))\n    return arg_list\n\n\ndef get_one_batch_of_trials(\n    ax_client: AxClient,\n    parallelism: Tuple[int, int],\n    num_trials_so_far: int,\n    num_max_trials_to_do: int,\n) -> BatchOfTrialType:\n    \"\"\"Produce a batch of trials that can be run in parallel\"\"\"\n    (num_trials, max_parallelism_setting) = parallelism\n    if max_parallelism_setting == -1:\n        # Special case, we can group all the trials into one batch\n        max_parallelism_setting = num_trials - num_trials_so_far\n\n        if num_trials == -1:\n            # This is a special case where we can run as many trials in parallel as we want.\n            # Given that num_trials is also -1, we can run all the trials in parallel.\n            max_parallelism_setting = num_max_trials_to_do\n\n    batch_of_trials = []\n    for _ in range(max_parallelism_setting):\n        parameters, trial_index = ax_client.get_next_trial()\n        batch_of_trials.append(\n            Trial(\n                overrides=map_params_to_arg_list(params=parameters),\n                trial_index=trial_index,\n            )\n        )\n    return batch_of_trials\n\n\nclass CoreAxSweeper(Sweeper):\n    \"\"\"Class to interface with the Ax Platform\"\"\"\n\n    def __init__(self, ax_config: AxConfig, max_batch_size: Optional[int]):\n        self.config_loader: Optional[ConfigLoader] = None\n        self.config: Optional[DictConfig] = None\n        self.launcher: Optional[Launcher] = None\n\n        self.job_results = None\n        self.experiment: ExperimentConfig = ax_config.experiment\n        self.early_stopper: EarlyStopper = EarlyStopper(\n            max_epochs_without_improvement=ax_config.early_stop.max_epochs_without_improvement,\n            epsilon=ax_config.early_stop.epsilon,\n            minimize=ax_config.early_stop.minimize,\n        )\n        self.ax_client_config: ClientConfig = ax_config.client\n        self.max_trials = ax_config.max_trials\n        self.ax_params: DictConfig = OmegaConf.create({})\n        if hasattr(ax_config, \"params\"):\n            self.ax_params.update(ax_config.params)\n        self.sweep_dir: str\n        self.job_idx: Optional[int] = None\n        self.max_batch_size = max_batch_size\n\n    def setup(\n        self,\n        config: DictConfig,\n        config_loader: ConfigLoader,\n        task_function: TaskFunction,\n    ) -> None:\n        self.config = config\n        self.config_loader = config_loader\n        self.launcher = Plugins.instance().instantiate_launcher(\n            config=config, config_loader=config_loader, task_function=task_function\n        )\n        self.sweep_dir = config.hydra.sweep.dir\n\n    def sweep(self, arguments: List[str]) -> None:\n        self.job_idx = 0\n        ax_client = self.setup_ax_client(arguments)\n\n        num_trials_left = self.max_trials\n        max_parallelism = ax_client.get_max_parallelism()\n        current_parallelism_index = 0\n        # Index to track the parallelism value we are using right now.\n        best_parameters = {}\n        while num_trials_left > 0:\n            current_parallelism = max_parallelism[current_parallelism_index]\n            num_trials, max_parallelism_setting = current_parallelism\n            num_trials_so_far = 0\n            while (\n                num_trials > num_trials_so_far or num_trials == -1\n            ) and num_trials_left > 0:\n                batch_of_trials = get_one_batch_of_trials(\n                    ax_client=ax_client,\n                    parallelism=current_parallelism,\n                    num_trials_so_far=num_trials_so_far,\n                    num_max_trials_to_do=num_trials_left,\n                )\n                batch_of_trials_to_launch = batch_of_trials[:num_trials_left]\n\n                log.info(\n                    \"AxSweeper is launching {} jobs\".format(\n                        len(batch_of_trials_to_launch)\n                    )\n                )\n\n                self.sweep_over_batches(\n                    ax_client=ax_client, batch_of_trials=batch_of_trials_to_launch\n                )\n\n                num_trials_so_far += len(batch_of_trials_to_launch)\n                num_trials_left -= len(batch_of_trials_to_launch)\n\n                best_parameters, predictions = ax_client.get_best_parameters()\n                metric = predictions[0][ax_client.objective_name]\n\n                if self.early_stopper.should_stop(metric, best_parameters):\n                    num_trials_left = -1\n                    break\n\n            current_parallelism_index += 1\n\n        results_to_serialize = {\"optimizer\": \"ax\", \"ax\": best_parameters}\n        OmegaConf.save(\n            OmegaConf.create(results_to_serialize),\n            f\"{self.sweep_dir}/optimization_results.yaml\",\n        )\n        log.info(\"Best parameters: \" + str(best_parameters))\n\n    def sweep_over_batches(\n        self, ax_client: AxClient, batch_of_trials: BatchOfTrialType\n    ) -> None:\n        assert self.launcher is not None\n        assert self.job_idx is not None\n\n        chunked_batches = self.chunks(batch_of_trials, self.max_batch_size)\n        for batch in chunked_batches:\n            overrides = [x.overrides for x in batch]\n            self.validate_batch_is_legal(overrides)\n            rets = self.launcher.launch(\n                job_overrides=overrides, initial_job_idx=self.job_idx\n            )\n            self.job_idx += len(rets)\n            for idx in range(len(batch)):\n                val = rets[idx].return_value\n                ax_client.complete_trial(\n                    trial_index=batch[idx].trial_index, raw_data=val\n                )\n\n    def setup_ax_client(self, arguments: List[str]) -> AxClient:\n        \"\"\"Method to setup the Ax Client\"\"\"\n        parameters: List[Dict[str, Any]] = []\n        for key, value in self.ax_params.items():\n            param = OmegaConf.to_container(value, resolve=True)\n            assert isinstance(param, Dict)\n            if param[\"type\"] == \"range\":\n                bounds = param[\"bounds\"]\n                if not (all(isinstance(x, int) for x in bounds)):\n                    # Type mismatch. Promote all to float\n                    param[\"bounds\"] = [float(x) for x in bounds]\n\n            parameters.append(param)\n            parameters[-1][\"name\"] = key\n        commandline_params = self.parse_commandline_args(arguments)\n        for cmd_param in commandline_params:\n            for param in parameters:\n                if param[\"name\"] == cmd_param[\"name\"]:\n                    for key, value in cmd_param.items():\n                        param[key] = value\n                    break\n            else:\n                parameters.append(cmd_param)\n\n        log.info(\n            f\"AxSweeper is optimizing the following parameters: {encoder_parameters_into_string(parameters)}\"\n        )\n        ax_client = AxClient(\n            verbose_logging=self.ax_client_config.verbose_logging,\n            random_seed=self.ax_client_config.random_seed,\n        )\n        ax_client.create_experiment(parameters=parameters, **self.experiment)\n\n        return ax_client\n\n    def parse_commandline_args(\n        self, arguments: List[str]\n    ) -> List[Dict[str, Union[ax_types.TParamValue, List[ax_types.TParamValue]]]]:\n        \"\"\"Method to parse the command line arguments and convert them into Ax parameters\"\"\"\n        parser = OverridesParser.create()\n        parsed = parser.parse_overrides(arguments)\n        parameters: List[Dict[str, Any]] = []\n        for override in parsed:\n            if override.is_sweep_override():\n                if override.is_choice_sweep():\n                    param = create_choice_param_from_choice_override(override)\n                elif override.is_range_sweep():\n                    param = create_choice_param_from_range_override(override)\n                elif override.is_interval_sweep():\n                    param = create_range_param_using_interval_override(override)\n            elif not override.is_hydra_override():\n                param = create_fixed_param_from_element_override(override)\n            parameters.append(param)\n\n        return parameters\n\n    @staticmethod\n    def chunks(batch: List[Any], n: Optional[int]) -> Iterable[List[Any]]:\n        \"\"\"\n        Chunk the batch into chunks of upto to n items (each)\n        \"\"\"\n        if n is None:\n            n = len(batch)\n        if n < 1:\n            raise ValueError(\"n must be an integer greater than 0\")\n        for i in range(0, len(batch), n):\n            yield batch[i : i + n]\n\n\ndef create_range_param_using_interval_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    value = override.value()\n    assert isinstance(value, IntervalSweep)\n    param = {\n        \"name\": key,\n        \"type\": \"range\",\n        \"bounds\": [value.start, value.end],\n    }\n    return param\n\n\ndef create_choice_param_from_choice_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"choice\",\n        \"values\": list(override.sweep_iterator(transformer=Transformer.encode)),\n    }\n    return param\n\n\ndef create_choice_param_from_range_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"choice\",\n        \"values\": [val for val in override.sweep_iterator()],\n    }\n\n    return param\n\n\ndef create_fixed_param_from_element_override(override: Override) -> Dict[str, Any]:\n    key = override.get_key_element()\n    param = {\n        \"name\": key,\n        \"type\": \"fixed\",\n        \"value\": override.value(),\n    }\n\n    return param\n", "path": "plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py"}]}