Decouple summarize_ax_optimization_complexity from OrchestratorOptions (#4706)

Summary:

This change decouples `summarize_ax_optimization_complexity` from requiring an `OrchestratorOptions` instance by accepting individual optional fields instead.

Differential Revision: D89778530

Author: shrutipatel31

ax/analysis/healthcheck/complexity_rating.py

@@ -138,8 +138,14 @@ def compute(
         options = none_throws(self.options)
         optimization_summary = summarize_ax_optimization_complexity(
             experiment=experiment,
-            options=options,
             tier_metadata=self.tier_metadata,
+            early_stopping_strategy=options.early_stopping_strategy,
+            global_stopping_strategy=options.global_stopping_strategy,
+            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/service/utils/orchestrator_options.py

@@ -13,6 +13,11 @@
 from ax.early_stopping.strategies import BaseEarlyStoppingStrategy
 from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
 
+# Default values for OrchestratorOptions fields
+DEFAULT_MAX_PENDING_TRIALS: int = 10
+DEFAULT_TOLERATED_TRIAL_FAILURE_RATE: float = 0.5
+DEFAULT_MIN_FAILED_TRIALS_FOR_FAILURE_RATE_CHECK: int = 5
+
 
 class TrialType(Enum):
     TRIAL = 0
@@ -125,12 +130,14 @@ class OrchestratorOptions:
             Default to False.
     """
 
-    max_pending_trials: int = 10
+    max_pending_trials: int = DEFAULT_MAX_PENDING_TRIALS
     trial_type: TrialType = TrialType.TRIAL
     batch_size: int | None = None
     total_trials: int | None = None
-    tolerated_trial_failure_rate: float = 0.5
-    min_failed_trials_for_failure_rate_check: int = 5
+    tolerated_trial_failure_rate: float = DEFAULT_TOLERATED_TRIAL_FAILURE_RATE
+    min_failed_trials_for_failure_rate_check: int = (
+        DEFAULT_MIN_FAILED_TRIALS_FOR_FAILURE_RATE_CHECK
+    )
     log_filepath: str | None = None
     logging_level: int = INFO
     ttl_seconds_for_trials: int | None = None

ax/utils/common/complexity_utils.py

@@ -12,8 +12,14 @@
 from ax.adapter.adapter_utils import can_map_to_binary, is_unordered_choice
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective
+from ax.early_stopping.strategies import BaseEarlyStoppingStrategy
 from ax.exceptions.core import OptimizationNotConfiguredError, UserInputError
-from ax.service.orchestrator import OrchestratorOptions
+from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
+from ax.service.utils.orchestrator_options import (
+    DEFAULT_MAX_PENDING_TRIALS,
+    DEFAULT_MIN_FAILED_TRIALS_FOR_FAILURE_RATE_CHECK,
+    DEFAULT_TOLERATED_TRIAL_FAILURE_RATE,
+)
 
 STANDARD_TIER_MESSAGE = """This experiment is in tier 'Standard'.
 
@@ -141,8 +147,14 @@ class OptimizationSummary:
 
 def summarize_ax_optimization_complexity(
     experiment: Experiment,
-    options: OrchestratorOptions,
     tier_metadata: dict[str, Any],
+    early_stopping_strategy: BaseEarlyStoppingStrategy | None = None,
+    global_stopping_strategy: BaseGlobalStoppingStrategy | None = None,
+    tolerated_trial_failure_rate: float | None = DEFAULT_TOLERATED_TRIAL_FAILURE_RATE,
+    max_pending_trials: int | None = DEFAULT_MAX_PENDING_TRIALS,
+    min_failed_trials_for_failure_rate_check: int | None = (
+        DEFAULT_MIN_FAILED_TRIALS_FOR_FAILURE_RATE_CHECK
+    ),
 ) -> OptimizationSummary:
     """Summarize the experiment's optimization complexity.
 
@@ -151,11 +163,25 @@ def summarize_ax_optimization_complexity(
 
     Args:
         experiment: The Ax Experiment.
-        options: The orchestrator options.
-        tier_metadata: tier-related meta-data from the orchestrator.
+        tier_metadata: Tier-related metadata from the orchestrator. Supported keys:
+            - 'user_supplied_max_trials': Maximum number of trials.
+            - 'uses_standard_api': Whether high-level configs are used (as
+                opposed to low-level Ax abstractions), ensuring the full
+                experiment configuration is known upfront.
+            - 'all_inputs_are_configs': Deprecated alias for 'uses_standard_api',
+                supported for backward compatibility.
+        early_stopping_strategy: The early stopping strategy, if any. Used to
+            determine if early stopping is enabled. Defaults to None.
+        global_stopping_strategy: The global stopping strategy, if any. Used to
+            determine if global stopping is enabled. Defaults to None.
+        tolerated_trial_failure_rate: Fraction of trials allowed to fail without
+            the whole optimization ending. Defaults to 0.5.
+        max_pending_trials: Maximum number of pending trials. Defaults to 10.
+        min_failed_trials_for_failure_rate_check: Minimum failed trials before
+            failure rate is checked. Defaults to 5.
 
     Returns:
-        A dictionary summarizing the experiment.
+        An OptimizationSummary containing experiment complexity metrics.
     """
     search_space = experiment.search_space
     optimization_config = experiment.optimization_config
@@ -179,8 +205,8 @@ 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
+    uses_early_stopping = early_stopping_strategy is not None
+    uses_global_stopping = global_stopping_strategy is not None
 
     # Check if any metrics use merge_multiple_curves
     uses_merge_multiple_curves = False
@@ -210,10 +236,10 @@ 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.service.orchestrator import OrchestratorOptions
 from ax.utils.common.complexity_utils import (
     check_if_in_standard,
     DEFAULT_TIER_MESSAGES,
@@ -30,7 +29,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:
@@ -39,7 +37,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,
         )
 
@@ -58,7 +55,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,
         )
 
@@ -76,7 +72,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,
             )
 
@@ -107,42 +102,20 @@ 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,
                 )
 
                 # THEN the summary should reflect tier metadata values
                 self.assertEqual(summary.max_trials, expected_max_trials)
                 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,
-        )
-
-        # WHEN we summarize the experiment
-        summary = summarize_ax_optimization_complexity(
-            experiment=self.experiment,
-            options=options,
-            tier_metadata=self.tier_metadata,
-        )
-
-        # THEN the summary should reflect orchestrator 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)
-
     def test_parameter_constraints_counted(self) -> None:
         # GIVEN an experiment with parameter constraints
         experiment = get_experiment(constrain_search_space=True)
 
         # WHEN we summarize the experiment
         summary = summarize_ax_optimization_complexity(
             experiment=experiment,
-            options=self.options,
             tier_metadata=self.tier_metadata,
         )
 
@@ -159,7 +132,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,
         )