Repository: facebookresearch/hydra

Issue:
Use structured config in the Ax plugin



Discussion:
One of the parameters have the type as `Optional[Dict[str, Union[str, bool, float, int, None]]]` Specifying this type throws the following error:

```
omegaconf.errors.ValidationError: Unsupported value type : typing.Union[str, bool, float, int, NoneType]
full_key:                                                                                       
reference_type=None                                                                             
object_type=None
```

So I used the type as `Optional[Dict[str, Any]]`
Unions are not currently supported. This solution is reasonable.

Resolving patch:
diff --git a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
--- a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
+++ b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/_core.py
@@ -13,6 +13,7 @@
 from omegaconf import DictConfig, OmegaConf
 
 from ._earlystopper import EarlyStopper
+from .config import AxConfig, ClientConfig, ExperimentConfig
 
 log = logging.getLogger(__name__)
 
@@ -124,12 +125,16 @@ def get_one_batch_of_trials(
 class CoreAxSweeper:
     """Class to interface with the Ax Platform"""
 
-    def __init__(self, ax_config: DictConfig, max_batch_size: Optional[int]):
+    def __init__(self, ax_config: AxConfig, max_batch_size: Optional[int]):
         self.launcher: Optional[Launcher] = None
         self.job_results = None
-        self.experiment = ax_config.experiment
-        self.early_stopper = EarlyStopper(**ax_config.early_stop)
-        self.ax_client_config = ax_config.client
+        self.experiment: ExperimentConfig = ax_config.experiment
+        self.early_stopper: EarlyStopper = EarlyStopper(
+            max_epochs_without_improvement=ax_config.early_stop.max_epochs_without_improvement,
+            epsilon=ax_config.early_stop.epsilon,
+            minimize=ax_config.early_stop.minimize,
+        )
+        self.ax_client_config: ClientConfig = ax_config.client
         self.max_trials = ax_config.max_trials
         self.ax_params: DictConfig = OmegaConf.create({})
         if hasattr(ax_config, "params"):
diff --git a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/ax_sweeper.py b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/ax_sweeper.py
--- a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/ax_sweeper.py
+++ b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/ax_sweeper.py
@@ -8,6 +8,8 @@
 from hydra.types import TaskFunction
 from omegaconf import DictConfig
 
+from .config import AxConfig
+
 
 class AxSweeperSearchPathPlugin(SearchPathPlugin):
     """
@@ -25,7 +27,7 @@ def manipulate_search_path(self, search_path: ConfigSearchPath) -> None:
 class AxSweeper(Sweeper):
     """Class to interface with the Ax Platform"""
 
-    def __init__(self, ax_config: DictConfig, max_batch_size: Optional[int]):
+    def __init__(self, ax_config: AxConfig, max_batch_size: Optional[int]):
         from ._core import CoreAxSweeper
 
         self.sweeper = CoreAxSweeper(ax_config, max_batch_size)
diff --git a/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/config.py b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/config.py
new file mode 100644
--- /dev/null
+++ b/plugins/hydra_ax_sweeper/hydra_plugins/hydra_ax_sweeper/config.py
@@ -0,0 +1,71 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from hydra.core.config_store import ConfigStore
+from hydra.types import MISSING, ObjectConf
+
+
+@dataclass
+class EarlyStopConfig:
+    minimize: bool = MISSING
+
+    max_epochs_without_improvement: int = 10
+
+    # An improvement larger than epsilon is considered significant
+    epsilon: float = 0.00001
+
+
+@dataclass
+class ExperimentConfig:
+    # Experiment name
+    name: Optional[str] = None
+
+    # Name of metric to optimize or null if you only return one metric
+    objective_name: str = "objective"
+
+    # Defaults to minimize, set to false to maximize
+    minimize: bool = True
+
+    # For the remaining paramaters, refer the Ax documentation: https://ax.dev/api/core.html#experiment
+    parameter_constraints: Optional[List[str]] = None
+    outcome_constraints: Optional[List[str]] = None
+    status_quo: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class ClientConfig:
+
+    verbose_logging: bool = False
+
+    # set random seed here to make Ax results reproducible
+    random_seed: Optional[int] = None
+
+
+@dataclass
+class AxConfig:
+
+    # max_trials is application-specific. Tune it for your use case
+    max_trials: int = 10
+    early_stop: EarlyStopConfig = EarlyStopConfig()
+    experiment: ExperimentConfig = ExperimentConfig()
+    client: ClientConfig = ClientConfig()
+    params: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class AxSweeperParams:
+    # Maximum number of trials to run in parallel
+    max_batch_size: Optional[int] = None
+    ax_config: AxConfig = AxConfig()
+
+
+@dataclass
+class AxSweeperConf(ObjectConf):
+    cls: str = "hydra_plugins.hydra_ax_sweeper.ax_sweeper.AxSweeper"
+    params: AxSweeperParams = AxSweeperParams()
+
+
+ConfigStore.instance().store(
+    group="hydra/sweeper", name="ax", node=AxSweeperConf, provider="ax_sweeper",
+)


Test patch: