[Ax] Align optimization mode and reported SEM with Ax (#13611)

* [Ax] Align optimization mode and reported SEM with Ax

Ensure that `mode` aligns with the mode set in Ax + report SEM as None rather than as 0.0 to make use of Ax noise inference

* Account for review

* Update ax.py

* Fix lint

* Fix tests, ad additional checks

* Fix tests for python 3.6

Co-authored-by: Kai Fricke <kai@anyscale.com>

python/ray/tune/suggest/ax.py

@@ -1,7 +1,6 @@
 import copy
 from typing import Dict, List, Optional, Union
 
-from ax.service.ax_client import AxClient
 from ray.tune.result import DEFAULT_METRIC
 from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
     Quantized, Uniform
@@ -12,8 +11,17 @@ from ray.tune.utils.util import flatten_dict, unflatten_dict
 
 try:
     import ax
+    from ax.service.ax_client import AxClient
 except ImportError:
-    ax = None
+    ax = AxClient = None
+
+# This exception only exists in newer Ax releases for python 3.7
+try:
+    from ax.exceptions.generation_strategy import \
+        MaxParallelismReachedException
+except ImportError:
+    MaxParallelismReachedException = Exception
+
 import logging
 
 from ray.tune.suggest import Searcher
@@ -124,6 +132,7 @@ class AxSearch(Searcher):
         assert ax is not None, """Ax must be installed!
             You can install AxSearch with the command:
             `pip install ax-platform sqlalchemy`."""
+
         if mode:
             assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
 
@@ -151,7 +160,6 @@ class AxSearch(Searcher):
 
         self.max_concurrent = max_concurrent
 
-        self._objective_name = metric
         self._parameters = []
         self._live_trial_mapping = {}
 
@@ -179,6 +187,10 @@ class AxSearch(Searcher):
                     "`AxClient.create_experiment()`, or you should pass an "
                     "Ax search space as the `space` parameter to `AxSearch`, "
                     "or pass a `config` dict to `tune.run()`.")
+            if self._mode not in ["min", "max"]:
+                raise ValueError(
+                    "Please specify the `mode` argument when initializing "
+                    "the `AxSearch` object or pass it to `tune.run()`.")
             self._ax.create_experiment(
                 parameters=self._space,
                 objective_name=self._metric,
@@ -188,16 +200,25 @@ class AxSearch(Searcher):
         else:
             if any([
                     self._space, self._parameter_constraints,
-                    self._outcome_constraints
+                    self._outcome_constraints, self._mode, self._metric
             ]):
                 raise ValueError(
                     "If you create the Ax experiment yourself, do not pass "
                     "values for these parameters to `AxSearch`: {}.".format([
-                        "space", "parameter_constraints", "outcome_constraints"
+                        "space",
+                        "parameter_constraints",
+                        "outcome_constraints",
+                        "mode",
+                        "metric",
                     ]))
 
         exp = self._ax.experiment
-        self._objective_name = exp.optimization_config.objective.metric.name
+
+        # Update mode and metric from experiment if it has been passed
+        self._mode = "min" \
+            if exp.optimization_config.objective.minimize else "max"
+        self._metric = exp.optimization_config.objective.metric.name
+
         self._parameters = list(exp.parameters)
 
         if self._ax._enforce_sequential_optimization:
@@ -239,7 +260,10 @@ class AxSearch(Searcher):
             config = self._points_to_evaluate.pop(0)
             parameters, trial_index = self._ax.attach_trial(config)
         else:
-            parameters, trial_index = self._ax.get_next_trial()
+            try:
+                parameters, trial_index = self._ax.get_next_trial()
+            except MaxParallelismReachedException:
+                return None
 
         self._live_trial_mapping[trial_id] = trial_index
         return unflatten_dict(parameters)
@@ -255,14 +279,12 @@ class AxSearch(Searcher):
 
     def _process_result(self, trial_id, result):
         ax_trial_index = self._live_trial_mapping[trial_id]
-        metric_dict = {
-            self._objective_name: (result[self._objective_name], 0.0)
-        }
+        metric_dict = {self._metric: (result[self._metric], None)}
         outcome_names = [
             oc.metric.name for oc in
             self._ax.experiment.optimization_config.outcome_constraints
         ]
-        metric_dict.update({on: (result[on], 0.0) for on in outcome_names})
+        metric_dict.update({on: (result[on], None) for on in outcome_names})
         self._ax.complete_trial(
             trial_index=ax_trial_index, raw_data=metric_dict)
 

python/ray/tune/tests/test_sample.py

@@ -263,12 +263,14 @@ class SearchSpaceTest(unittest.TestCase):
         ]
 
         client1 = AxClient(random_seed=1234)
-        client1.create_experiment(parameters=converted_config)
-        searcher1 = AxSearch(ax_client=client1, metric="a", mode="max")
+        client1.create_experiment(
+            parameters=converted_config, objective_name="a", minimize=False)
+        searcher1 = AxSearch(ax_client=client1)
 
         client2 = AxClient(random_seed=1234)
-        client2.create_experiment(parameters=ax_config)
-        searcher2 = AxSearch(ax_client=client2, metric="a", mode="max")
+        client2.create_experiment(
+            parameters=ax_config, objective_name="a", minimize=False)
+        searcher2 = AxSearch(ax_client=client2)
 
         config1 = searcher1.suggest("0")
         config2 = searcher2.suggest("0")

python/ray/tune/tests/test_searchers.py

@@ -49,8 +49,10 @@ class InvalidValuesTest(unittest.TestCase):
         # At least one nan, inf, -inf and float
         client = AxClient(random_seed=4321)
         client.create_experiment(
-            parameters=converted_config, objective_name="_metric")
-        searcher = AxSearch(ax_client=client, metric="_metric", mode="max")
+            parameters=converted_config,
+            objective_name="_metric",
+            minimize=False)
+        searcher = AxSearch(ax_client=client)
 
         out = tune.run(
             _invalid_objective,