make constraints optional in get_experiment... functions

Summary: For get_experiment... core stubs functions that create a trial with Sobol, constraints were resulting in the logs being flooded with "Unable to round" warnings due to rejection sampling. This diff omits these constraints in test_sqa_store tests for which they are immaterial.

Differential Revision: https://internalfb.com/D46687301

fbshipit-source-id: 8113b0bf601d640d4ef469e6cc924b2c86a217c8

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -9,7 +9,7 @@
 from unittest.mock import MagicMock, Mock, patch
 
 from ax.core.arm import Arm
-from ax.core.batch_trial import LifecycleStage
+from ax.core.batch_trial import BatchTrial, LifecycleStage
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
 from ax.core.objective import Objective
@@ -71,7 +71,7 @@
 from ax.utils.common.constants import Keys
 from ax.utils.common.serialization import serialize_init_args
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import not_none
+from ax.utils.common.typeutils import checked_cast, not_none
 from ax.utils.testing.core_stubs import (
     CustomTestMetric,
     CustomTestRunner,
@@ -204,7 +204,9 @@ def testLoadExperimentTrialsInBatches(self) -> None:
     # load experiments with custom runners and metrics without a decoder.
     def testLoadExperimentSkipMetricsAndRunners(self) -> None:
         # Create a test experiment with a custom metric and runner.
-        experiment = get_experiment_with_custom_runner_and_metric()
+        experiment = get_experiment_with_custom_runner_and_metric(
+            constrain_search_space=False
+        )
 
         # Note that the experiment is created outside of the test code.
         # Confirm that it uses the custom runner and metric
@@ -287,13 +289,14 @@ def test_ExperimentSaveAndLoadReducedState(
         _mock_exp_from_sqa.reset_mock()
 
         # 2. Try case with abandoned arms.
-        exp = self.experiment
+        exp = get_experiment_with_batch_trial(constrain_search_space=False)
         save_experiment(exp)
         loaded_experiment = load_experiment(exp.name, reduced_state=True)
         # Experiments are not the same, because one has abandoned arms info.
         self.assertNotEqual(loaded_experiment, exp)
         # Remove all abandoned arms and check that all else is equal as expected.
-        exp.trials.get(0)._abandoned_arms_metadata = {}
+        t = checked_cast(BatchTrial, exp.trials[0])
+        t._abandoned_arms_metadata = {}
         self.assertEqual(loaded_experiment, exp)
         # Make sure that all relevant decoding functions were called with
         # `reduced_state=True` and correct number of times.
@@ -1608,7 +1611,7 @@ def testImmutableSearchSpaceAndOptConfigLoading(
         _mock_get_gs_sqa_imm_oc_ss,
         _mock_gr_from_sqa,
     ) -> None:
-        experiment = get_experiment_with_batch_trial()
+        experiment = get_experiment_with_batch_trial(constrain_search_space=False)
         experiment._properties = {Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True}
         save_experiment(experiment)
 

ax/utils/testing/core_stubs.py

@@ -117,10 +117,12 @@
 ##############################
 
 
-def get_experiment(with_status_quo: bool = True) -> Experiment:
+def get_experiment(
+    with_status_quo: bool = True, constrain_search_space: bool = True
+) -> Experiment:
     return Experiment(
         name="test",
-        search_space=get_search_space(),
+        search_space=get_search_space(constrain_search_space=constrain_search_space),
         optimization_config=get_optimization_config(),
         status_quo=get_status_quo() if with_status_quo else None,
         description="test description",
@@ -142,12 +144,16 @@ def get_experiment_with_map_data_type() -> Experiment:
     )
 
 
-def get_experiment_with_custom_runner_and_metric() -> Experiment:
+def get_experiment_with_custom_runner_and_metric(
+    constrain_search_space: bool = True,
+) -> Experiment:
 
     # Create experiment with custom runner and metric
     experiment = Experiment(
         name="test",
-        search_space=get_search_space(),
+        # Omit constraints to prevent Sobol rejection sampling below,
+        # which floods logs with "Unable to round" warnings.
+        search_space=get_search_space(constrain_search_space=constrain_search_space),
         optimization_config=get_optimization_config(),
         description="test description",
         tracking_metrics=[
@@ -421,8 +427,8 @@ def get_experiment_with_trial() -> Experiment:
     return trial.experiment
 
 
-def get_experiment_with_batch_trial() -> Experiment:
-    batch_trial = get_batch_trial()
+def get_experiment_with_batch_trial(constrain_search_space: bool = True) -> Experiment:
+    batch_trial = get_batch_trial(constrain_search_space=constrain_search_space)
     return batch_trial.experiment
 
 
@@ -698,20 +704,23 @@ def get_high_dimensional_branin_experiment() -> Experiment:
 ##############################
 
 
-def get_search_space() -> SearchSpace:
+def get_search_space(constrain_search_space: bool = True) -> SearchSpace:
     parameters: List[Parameter] = [
         get_range_parameter(),
         get_range_parameter2(),
         get_choice_parameter(),
         get_fixed_parameter(),
     ]
-    return SearchSpace(
-        parameters=parameters,
-        parameter_constraints=[
+    parameter_constraints = []
+    if constrain_search_space:
+        parameter_constraints = [
             get_order_constraint(),
             get_parameter_constraint(),
             get_sum_constraint1(),
-        ],
+        ]
+    return SearchSpace(
+        parameters=parameters,
+        parameter_constraints=parameter_constraints,
     )
 
 
@@ -1014,9 +1023,13 @@ def get_robust_search_space_environmental(
 
 
 def get_batch_trial(
-    abandon_arm: bool = True, experiment: Optional[Experiment] = None
+    abandon_arm: bool = True,
+    experiment: Optional[Experiment] = None,
+    constrain_search_space: bool = True,
 ) -> BatchTrial:
-    experiment = experiment or get_experiment()
+    experiment = experiment or get_experiment(
+        constrain_search_space=constrain_search_space
+    )
     batch = experiment.new_batch_trial()
     arms = get_arms_from_dict(get_arm_weights1())
     weights = get_weights_from_dict(get_arm_weights1())