We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Ax] Fix warnings related to OmegaConf 2.0.0 compact key
Related to https://github.com/omry/omegaconf/issues/152

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

{"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@@ -195,6 +195,10 @@\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@@ -225,6 +229,7 @@\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 \"parameter_type\" not in param:\n@@ -344,3 +349,17 @@\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+\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", "issue": "[Ax] Fix warnings related to OmegaConf 2.0.0 compact key\nRelated to https://github.com/omry/omegaconf/issues/152\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 import ParameterType  # type: ignore\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.plugins import Plugins\nfrom hydra.plugins.launcher import Launcher\nfrom hydra.types import TaskFunction\nfrom omegaconf import DictConfig, OmegaConf\n\nfrom ._earlystopper import EarlyStopper\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 _is_int(string_inp: str) -> bool:\n    \"\"\"Method to check if the given string input can be parsed as integer\"\"\"\n    try:\n        int(string_inp)\n        return True\n    except ValueError:\n        return False\n\n\ndef _is_float(string_inp: str) -> bool:\n    \"\"\"Method to check if the given string input can be parsed as a float\"\"\"\n    try:\n        float(string_inp)\n        return True\n    except ValueError:\n        return False\n\n\ndef encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:\n    \"\"\"Convert a list of params into a string\"\"\"\n\n    mandatory_keys = set(\n        [\"name\", \"type\", \"bounds\", \"values\", \"value\", \"parameter_type\"]\n    )\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        parameter_log_string += f\"type = {parameter['parameter_type'].name.lower()}\"\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 infer_parameter_type(parameter: Dict[str, Any]) -> ParameterType:\n    if parameter[\"type\"] == \"range\":\n        value = parameter[\"bounds\"][0]\n    elif parameter[\"type\"] == \"choice\":\n        value = parameter[\"values\"][0]\n    elif parameter[\"type\"] == \"fixed\":\n        value = parameter[\"value\"]\n\n    if isinstance(value, int):\n        return ParameterType.INT\n    if isinstance(value, float):\n        return ParameterType.FLOAT\n    return ParameterType.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:\n    \"\"\"Class to interface with the Ax Platform\"\"\"\n\n    def __init__(self, ax_config: DictConfig, max_batch_size: Optional[int]):\n        self.launcher: Optional[Launcher] = None\n        self.job_results = None\n        self.experiment = ax_config.experiment\n        self.early_stopper = EarlyStopper(**ax_config.early_stop)\n        self.ax_client_config = 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.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        recommended_max_parallelism = ax_client.get_recommended_max_parallelism()\n        current_parallelism_index = 0\n        # Index to track the parallelism value we are using right now.\n\n        while num_trials_left > 0:\n            current_parallelism = recommended_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            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 \"parameter_type\" not in param:\n                param[\"parameter_type\"] = infer_parameter_type(param)\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\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\n        parameters = []\n        for arg in arguments:\n            key, value = arg.split(\"=\")\n            if \",\" in value:\n                # This is a Choice Parameter.\n                value_choices = [x.strip() for x in value.split(\",\")]\n                if all(_is_int(x) for x in value_choices):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": [int(x) for x in value_choices],\n                        \"parameter_type\": ParameterType.INT,\n                    }\n                elif all(_is_float(x) for x in value_choices):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": [float(x) for x in value_choices],\n                        \"parameter_type\": ParameterType.FLOAT,\n                    }\n                else:\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": value_choices,\n                        \"parameter_type\": ParameterType.STRING,\n                    }\n                parameters.append(param)\n            elif \":\" in value:\n                # This is a Range Parameter.\n                range_start, range_end = value.split(\":\")\n                if _is_int(range_start) and _is_int(range_end):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"range\",\n                        \"bounds\": [int(range_start), int(range_end)],\n                        \"parameter_type\": ParameterType.INT,\n                    }\n                elif _is_float(range_start) and _is_float(range_end):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"range\",\n                        \"bounds\": [float(range_start), float(range_end)],\n                        \"parameter_type\": ParameterType.FLOAT,\n                    }\n                else:\n                    raise ValueError(\n                        \"Input to the range parameter should be an int or a float.\"\n                    )\n                parameters.append(param)\n            else:\n                # This is a Fixed Parameter.\n                if _is_int(value):\n                    parameter_type = ParameterType.INT\n                elif _is_float(value):\n                    parameter_type = ParameterType.FLOAT\n                else:\n                    parameter_type = ParameterType.STRING\n                parameters.append(\n                    {\n                        \"name\": key,\n                        \"type\": \"fixed\",\n                        \"value\": value,\n                        \"parameter_type\": parameter_type,\n                    }\n                )\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", "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 import ParameterType  # type: ignore\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.plugins import Plugins\nfrom hydra.plugins.launcher import Launcher\nfrom hydra.types import TaskFunction\nfrom omegaconf import DictConfig, OmegaConf\n\nfrom ._earlystopper import EarlyStopper\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 _is_int(string_inp: str) -> bool:\n    \"\"\"Method to check if the given string input can be parsed as integer\"\"\"\n    try:\n        int(string_inp)\n        return True\n    except ValueError:\n        return False\n\n\ndef _is_float(string_inp: str) -> bool:\n    \"\"\"Method to check if the given string input can be parsed as a float\"\"\"\n    try:\n        float(string_inp)\n        return True\n    except ValueError:\n        return False\n\n\ndef encoder_parameters_into_string(parameters: List[Dict[str, Any]]) -> str:\n    \"\"\"Convert a list of params into a string\"\"\"\n\n    mandatory_keys = set(\n        [\"name\", \"type\", \"bounds\", \"values\", \"value\", \"parameter_type\"]\n    )\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        parameter_log_string += f\"type = {parameter['parameter_type'].name.lower()}\"\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 infer_parameter_type(parameter: Dict[str, Any]) -> ParameterType:\n    if parameter[\"type\"] == \"range\":\n        value = parameter[\"bounds\"][0]\n    elif parameter[\"type\"] == \"choice\":\n        value = parameter[\"values\"][0]\n    elif parameter[\"type\"] == \"fixed\":\n        value = parameter[\"value\"]\n\n    if isinstance(value, int):\n        return ParameterType.INT\n    if isinstance(value, float):\n        return ParameterType.FLOAT\n    return ParameterType.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:\n    \"\"\"Class to interface with the Ax Platform\"\"\"\n\n    def __init__(self, ax_config: DictConfig, max_batch_size: Optional[int]):\n        self.launcher: Optional[Launcher] = None\n        self.job_results = None\n        self.experiment = ax_config.experiment\n        self.early_stopper = EarlyStopper(**ax_config.early_stop)\n        self.ax_client_config = 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.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        recommended_max_parallelism = ax_client.get_recommended_max_parallelism()\n        current_parallelism_index = 0\n        # Index to track the parallelism value we are using right now.\n\n        while num_trials_left > 0:\n            current_parallelism = recommended_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            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 \"parameter_type\" not in param:\n                param[\"parameter_type\"] = infer_parameter_type(param)\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\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\n        parameters = []\n        for arg in arguments:\n            key, value = arg.split(\"=\")\n            if \",\" in value:\n                # This is a Choice Parameter.\n                value_choices = [x.strip() for x in value.split(\",\")]\n                if all(_is_int(x) for x in value_choices):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": [int(x) for x in value_choices],\n                        \"parameter_type\": ParameterType.INT,\n                    }\n                elif all(_is_float(x) for x in value_choices):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": [float(x) for x in value_choices],\n                        \"parameter_type\": ParameterType.FLOAT,\n                    }\n                else:\n                    param = {\n                        \"name\": key,\n                        \"type\": \"choice\",\n                        \"values\": value_choices,\n                        \"parameter_type\": ParameterType.STRING,\n                    }\n                parameters.append(param)\n            elif \":\" in value:\n                # This is a Range Parameter.\n                range_start, range_end = value.split(\":\")\n                if _is_int(range_start) and _is_int(range_end):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"range\",\n                        \"bounds\": [int(range_start), int(range_end)],\n                        \"parameter_type\": ParameterType.INT,\n                    }\n                elif _is_float(range_start) and _is_float(range_end):\n                    param = {\n                        \"name\": key,\n                        \"type\": \"range\",\n                        \"bounds\": [float(range_start), float(range_end)],\n                        \"parameter_type\": ParameterType.FLOAT,\n                    }\n                else:\n                    raise ValueError(\n                        \"Input to the range parameter should be an int or a float.\"\n                    )\n                parameters.append(param)\n            else:\n                # This is a Fixed Parameter.\n                if _is_int(value):\n                    parameter_type = ParameterType.INT\n                elif _is_float(value):\n                    parameter_type = ParameterType.FLOAT\n                else:\n                    parameter_type = ParameterType.STRING\n                parameters.append(\n                    {\n                        \"name\": key,\n                        \"type\": \"fixed\",\n                        \"value\": value,\n                        \"parameter_type\": parameter_type,\n                    }\n                )\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", "path": "plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py"}]}