Reduce botorch dependency creep in Ax (facebook#5161)

Summary: Pull Request resolved: facebook#5161

Reviewed By: saitcakmak

Differential Revision: D100111247

fbshipit-source-id: b74e2a913669e5a70f34f0bf24e484b3adcf36a9

Author: sdaulton

ax/analysis/healthcheck/complexity_rating.py

@@ -121,8 +121,14 @@ def compute(
         options = self.options if self.options is not None else OrchestratorOptions()
         optimization_summary = summarize_ax_optimization_complexity(
             experiment=experiment,
-            options=options,
             tier_metadata=self.tier_metadata,
+            uses_early_stopping=options.early_stopping_strategy is not None,
+            uses_global_stopping=options.global_stopping_strategy is not None,
+            tolerated_trial_failure_rate=options.tolerated_trial_failure_rate,
+            max_pending_trials=options.max_pending_trials,
+            min_failed_trials_for_failure_rate_check=(
+                options.min_failed_trials_for_failure_rate_check
+            ),
         )
 
         # Determine tier

ax/utils/common/complexity_utils.py

@@ -10,7 +10,6 @@
 from ax.adapter.parameter_utils import can_map_to_binary, is_unordered_choice
 from ax.core.experiment import Experiment
 from ax.exceptions.core import OptimizationNotConfiguredError
-from ax.orchestration.orchestrator import OrchestratorOptions
 from ax.utils.common.tier_utils import (  # noqa: F401
     check_if_in_standard,
     DEFAULT_TIER_MESSAGES,
@@ -22,8 +21,12 @@
 
 def summarize_ax_optimization_complexity(
     experiment: Experiment,
-    options: OrchestratorOptions,
     tier_metadata: dict[str, Any],
+    uses_early_stopping: bool = False,
+    uses_global_stopping: bool = False,
+    tolerated_trial_failure_rate: float | None = 0.5,
+    max_pending_trials: int | None = 10,
+    min_failed_trials_for_failure_rate_check: int | None = 5,
 ) -> OptimizationSummary:
     """Summarize the experiment's optimization complexity.
 
@@ -32,8 +35,13 @@ def summarize_ax_optimization_complexity(
 
     Args:
         experiment: The Ax Experiment.
-        options: The orchestrator options.
         tier_metadata: tier-related meta-data from the orchestrator.
+        uses_early_stopping: Whether early stopping is enabled.
+        uses_global_stopping: Whether global stopping is enabled.
+        tolerated_trial_failure_rate: The tolerated trial failure rate.
+        max_pending_trials: The maximum number of pending trials.
+        min_failed_trials_for_failure_rate_check: The minimum number of failed
+            trials before checking the failure rate.
 
     Returns:
         A dictionary summarizing the experiment.
@@ -60,8 +68,6 @@ def summarize_ax_optimization_complexity(
         else 1
     )
     num_outcome_constraints = len(optimization_config.outcome_constraints)
-    uses_early_stopping = options.early_stopping_strategy is not None
-    uses_global_stopping = options.global_stopping_strategy is not None
 
     # Check if any metrics use merge_multiple_curves
     uses_merge_multiple_curves = False
@@ -93,9 +99,9 @@ def summarize_ax_optimization_complexity(
         uses_global_stopping=uses_global_stopping,
         uses_merge_multiple_curves=uses_merge_multiple_curves,
         uses_standard_api=uses_standard_api,
-        tolerated_trial_failure_rate=options.tolerated_trial_failure_rate,
-        max_pending_trials=options.max_pending_trials,
+        tolerated_trial_failure_rate=tolerated_trial_failure_rate,
+        max_pending_trials=max_pending_trials,
         min_failed_trials_for_failure_rate_check=(
-            options.min_failed_trials_for_failure_rate_check
+            min_failed_trials_for_failure_rate_check
         ),
     )

ax/utils/common/tests/test_complexity_utils.py

@@ -8,7 +8,6 @@
 
 from ax.core.metric import Metric
 from ax.exceptions.core import OptimizationNotConfiguredError, UserInputError
-from ax.orchestration.orchestrator import OrchestratorOptions
 from ax.utils.common.complexity_utils import (
     check_if_in_standard,
     DEFAULT_TIER_MESSAGES,
@@ -29,7 +28,6 @@ class TestSummarizeAxOptimizationComplexity(TestCase):
     def setUp(self) -> None:
         super().setUp()
         self.experiment = get_experiment()
-        self.options = OrchestratorOptions()
         self.tier_metadata: dict[str, object] = {}
 
     def test_basic_experiment_summary(self) -> None:
@@ -38,7 +36,6 @@ def test_basic_experiment_summary(self) -> None:
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=self.experiment,
-            options=self.options,
             tier_metadata=self.tier_metadata,
         )
 
@@ -57,7 +54,6 @@ def test_multi_objective_experiment(self) -> None:
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=experiment,
-            options=self.options,
             tier_metadata=self.tier_metadata,
         )
 
@@ -75,7 +71,6 @@ def test_experiment_without_optimization_config_raises(self) -> None:
         ):
             summarize_ax_optimization_complexity(
                 experiment=self.experiment,
-                options=self.options,
                 tier_metadata=self.tier_metadata,
             )
 
@@ -106,7 +101,6 @@ def test_tier_metadata_extraction(self) -> None:
                 # WHEN we summarize the experiment
                 summary = summarize_ax_optimization_complexity(
                     experiment=self.experiment,
-                    options=self.options,
                     tier_metadata=tier_metadata,
                 )
 
@@ -115,21 +109,17 @@ def test_tier_metadata_extraction(self) -> None:
                 self.assertEqual(summary.uses_standard_api, expected_all_configs)
 
     def test_orchestrator_options_extraction(self) -> None:
-        # GIVEN custom orchestrator options
-        options = OrchestratorOptions(
-            tolerated_trial_failure_rate=0.25,
-            max_pending_trials=5,
-            min_failed_trials_for_failure_rate_check=10,
-        )
-
+        # GIVEN custom options
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=self.experiment,
-            options=options,
             tier_metadata=self.tier_metadata,
+            tolerated_trial_failure_rate=0.25,
+            max_pending_trials=5,
+            min_failed_trials_for_failure_rate_check=10,
         )
 
-        # THEN the summary should reflect orchestrator options
+        # THEN the summary should reflect the options
         self.assertEqual(summary.tolerated_trial_failure_rate, 0.25)
         self.assertEqual(summary.max_pending_trials, 5)
         self.assertEqual(summary.min_failed_trials_for_failure_rate_check, 10)
@@ -141,7 +131,6 @@ def test_parameter_constraints_counted(self) -> None:
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=experiment,
-            options=self.options,
             tier_metadata=self.tier_metadata,
         )
 
@@ -158,7 +147,6 @@ def test_merge_multiple_curves_detection(self) -> None:
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=self.experiment,
-            options=self.options,
             tier_metadata=self.tier_metadata,
         )