Update Complexity Rating Healthcheck (facebook#4707)

Summary:

This diff refactors ComplexityRatingAnalysis to decouple it from OrchestratorOptions by accepting individual configuration parameters directly instead of the entire options object.

Differential Revision:
D89778632

Privacy Context Container: L1307644

Author: shrutipatel31

ax/analysis/healthcheck/complexity_rating.py

@@ -16,8 +16,9 @@
     HealthcheckStatus,
 )
 from ax.core.experiment import Experiment
+from ax.early_stopping.strategies import BaseEarlyStoppingStrategy
 from ax.generation_strategy.generation_strategy import GenerationStrategy
-from ax.service.orchestrator import OrchestratorOptions
+from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
 from ax.utils.common.complexity_utils import (
     check_if_in_standard,
     DEFAULT_TIER_MESSAGES,
@@ -61,16 +62,17 @@ class ComplexityRatingAnalysis(Analysis):
 
     def __init__(
         self,
-        options: OrchestratorOptions | None = None,
         tier_metadata: dict[str, Any] | None = None,
         tier_messages: TierMessages = DEFAULT_TIER_MESSAGES,
+        early_stopping_strategy: BaseEarlyStoppingStrategy | None = None,
+        global_stopping_strategy: BaseGlobalStoppingStrategy | None = None,
+        tolerated_trial_failure_rate: float | None = None,
+        max_pending_trials: int | None = None,
+        min_failed_trials_for_failure_rate_check: int | None = None,
     ) -> None:
         """Initialize the ComplexityRatingAnalysis.
 
         Args:
-            options: The orchestrator options used for the optimization.
-                Required to evaluate early stopping, global stopping, and
-                failure rate settings.
             tier_metadata: Additional tier-related metadata from the orchestrator.
                 Supported keys:
                 - 'user_supplied_max_trials': Maximum number of trials.
@@ -82,12 +84,27 @@ def __init__(
                 generic messages suitable for most users. Pass a custom TierMessages
                 instance to provide tool-specific descriptions, support SLAs,
                 links to docs, or contact information.
+            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. Default value used is 0.5.
+            max_pending_trials: Maximum number of pending trials. Default used is 10.
+            min_failed_trials_for_failure_rate_check: Minimum failed trials before
+                failure rate is checked. Default value used is 5.
         """
-        self.options = options
         self.tier_metadata: dict[str, Any] = (
             tier_metadata if tier_metadata is not None else {}
         )
         self.tier_messages = tier_messages
+        self.early_stopping_strategy = early_stopping_strategy
+        self.global_stopping_strategy = global_stopping_strategy
+        self.tolerated_trial_failure_rate = tolerated_trial_failure_rate
+        self.max_pending_trials = max_pending_trials
+        self.min_failed_trials_for_failure_rate_check = (
+            min_failed_trials_for_failure_rate_check
+        )
 
     @override
     def validate_applicable_state(
@@ -98,11 +115,6 @@ def validate_applicable_state(
     ) -> str | None:
         if experiment is None:
             return "Experiment is required for ComplexityRatingAnalysis."
-        if self.options is None:
-            return (
-                "OrchestratorOptions is required for ComplexityRatingAnalysis. "
-                "Please pass options to the constructor."
-            )
         return None
 
     @override
@@ -120,8 +132,7 @@ def compute(
 
         Note:
             This method assumes ``validate_applicable_state`` has been called
-            and returned None, ensuring ``experiment`` and ``self.options``
-            are not None.
+            and returned None, ensuring ``experiment`` is not None.
 
         Args:
             experiment: The Ax Experiment to analyze. Must not be None.
@@ -135,16 +146,15 @@ def compute(
             with key experiment metrics.
         """
         experiment = none_throws(experiment)
-        options = none_throws(self.options)
         optimization_summary = summarize_ax_optimization_complexity(
             experiment=experiment,
             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,
+            early_stopping_strategy=self.early_stopping_strategy,
+            global_stopping_strategy=self.global_stopping_strategy,
+            tolerated_trial_failure_rate=self.tolerated_trial_failure_rate,
+            max_pending_trials=self.max_pending_trials,
             min_failed_trials_for_failure_rate_check=(
-                options.min_failed_trials_for_failure_rate_check
+                self.min_failed_trials_for_failure_rate_check
             ),
         )
 

ax/analysis/healthcheck/tests/test_complexity_rating.py

@@ -20,48 +20,36 @@
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace
 from ax.core.types import ComparisonOp
-from ax.service.orchestrator import OrchestratorOptions
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import get_branin_experiment
+from ax.utils.testing.core_stubs import (
+    get_branin_experiment,
+    get_improvement_global_stopping_strategy,
+    get_percentile_early_stopping_strategy,
+)
 
 
 class TestComplexityRatingAnalysis(TestCase):
     def setUp(self) -> None:
         super().setUp()
         self.experiment = get_branin_experiment()
-        self.options = OrchestratorOptions()
         self.tier_metadata: dict[str, object] = {
             "user_supplied_max_trials": 100,
             "uses_standard_api": True,
         }
 
     def test_validate_applicable_state_requires_experiment(self) -> None:
-        healthcheck = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        )
+        healthcheck = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata)
         result = healthcheck.validate_applicable_state(experiment=None)
         self.assertIsNotNone(result)
         self.assertIn("Experiment is required", result)
 
-    def test_validate_applicable_state_requires_options(self) -> None:
-        healthcheck = ComplexityRatingAnalysis(
-            options=None, tier_metadata=self.tier_metadata
-        )
-        result = healthcheck.validate_applicable_state(experiment=self.experiment)
-        self.assertIsNotNone(result)
-        self.assertIn("OrchestratorOptions is required", result)
-
     def test_validate_applicable_state_passes_with_valid_inputs(self) -> None:
-        healthcheck = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        )
+        healthcheck = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata)
         result = healthcheck.validate_applicable_state(experiment=self.experiment)
         self.assertIsNone(result)
 
     def test_standard_configuration(self) -> None:
-        healthcheck = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        )
+        healthcheck = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata)
         card = healthcheck.compute(experiment=self.experiment)
 
         self.assertEqual(card.name, "ComplexityRatingAnalysis")
@@ -90,7 +78,7 @@ def test_parameter_counts(self) -> None:
                 ]
                 self.experiment._search_space = SearchSpace(parameters=params)
                 card = ComplexityRatingAnalysis(
-                    options=self.options, tier_metadata=self.tier_metadata
+                    tier_metadata=self.tier_metadata
                 ).compute(experiment=self.experiment)
 
                 self.assertEqual(card.get_status(), expected_status)
@@ -115,7 +103,7 @@ def test_objectives_count(self) -> None:
                     )
                 )
                 card = ComplexityRatingAnalysis(
-                    options=self.options, tier_metadata=self.tier_metadata
+                    tier_metadata=self.tier_metadata
                 ).compute(experiment=self.experiment)
 
                 self.assertEqual(card.get_status(), expected_status)
@@ -132,9 +120,9 @@ def test_constraints(self) -> None:
                     for m in metrics
                 ],
             )
-            card = ComplexityRatingAnalysis(
-                options=self.options, tier_metadata=self.tier_metadata
-            ).compute(experiment=self.experiment)
+            card = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata).compute(
+                experiment=self.experiment
+            )
 
             self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
             self.assertIn("Advanced", card.subtitle)
@@ -156,35 +144,34 @@ def test_constraints(self) -> None:
             self.experiment._search_space = SearchSpace(
                 parameters=params, parameter_constraints=parameter_constraints
             )
-            card = ComplexityRatingAnalysis(
-                options=self.options, tier_metadata=self.tier_metadata
-            ).compute(experiment=self.experiment)
+            card = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata).compute(
+                experiment=self.experiment
+            )
 
             self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
             self.assertIn("Advanced", card.subtitle)
             self.assertIn("3 parameter constraints", card.subtitle)
 
     def test_stopping_strategies(self) -> None:
-        test_cases = [
-            ("early_stopping", True, False, "Early stopping"),
-            ("global_stopping", False, True, "Global stopping"),
-        ]
+        with self.subTest(strategy="early_stopping"):
+            card = ComplexityRatingAnalysis(
+                tier_metadata=self.tier_metadata,
+                early_stopping_strategy=get_percentile_early_stopping_strategy(),
+            ).compute(experiment=self.experiment)
 
-        for name, uses_early, uses_global, expected_msg in test_cases:
-            with self.subTest(strategy=name):
-                options = OrchestratorOptions(
-                    # pyre-fixme[6]: Using a mock value for testing
-                    early_stopping_strategy="mock" if uses_early else None,
-                    # pyre-fixme[6]: Using a mock value for testing
-                    global_stopping_strategy="mock" if uses_global else None,
-                )
-                card = ComplexityRatingAnalysis(
-                    options=options, tier_metadata=self.tier_metadata
-                ).compute(experiment=self.experiment)
+            self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
+            self.assertIn("Advanced", card.subtitle)
+            self.assertIn("Early stopping", card.subtitle)
 
-                self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
-                self.assertIn("Advanced", card.subtitle)
-                self.assertIn(expected_msg, card.subtitle)
+        with self.subTest(strategy="global_stopping"):
+            card = ComplexityRatingAnalysis(
+                tier_metadata=self.tier_metadata,
+                global_stopping_strategy=get_improvement_global_stopping_strategy(),
+            ).compute(experiment=self.experiment)
+
+            self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
+            self.assertIn("Advanced", card.subtitle)
+            self.assertIn("Global stopping", card.subtitle)
 
     def test_trial_counts(self) -> None:
         test_cases = [
@@ -198,48 +185,47 @@ def test_trial_counts(self) -> None:
                     "user_supplied_max_trials": max_trials,
                     "uses_standard_api": True,
                 }
-                card = ComplexityRatingAnalysis(
-                    options=self.options, tier_metadata=tier_metadata
-                ).compute(experiment=self.experiment)
+                card = ComplexityRatingAnalysis(tier_metadata=tier_metadata).compute(
+                    experiment=self.experiment
+                )
 
                 self.assertEqual(card.get_status(), expected_status)
                 self.assertIn(expected_tier, card.subtitle)
                 self.assertIn(expected_msg, card.subtitle)
 
     def test_unsupported_configurations(self) -> None:
-        test_cases = [
-            (
-                "not_using_standard_api",
-                OrchestratorOptions(),
-                {"user_supplied_max_trials": 100, "uses_standard_api": False},
-                "uses_standard_api=False",
-            ),
-            (
-                "high_failure_rate",
-                OrchestratorOptions(tolerated_trial_failure_rate=0.95),
-                {"user_supplied_max_trials": 100, "uses_standard_api": True},
-                "0.95",
-            ),
-            (
-                "invalid_failure_rate_check",
-                OrchestratorOptions(
-                    max_pending_trials=10,
-                    min_failed_trials_for_failure_rate_check=50,
-                ),
-                {"user_supplied_max_trials": 100, "uses_standard_api": True},
-                "min_failed_trials_for_failure_rate_check",
-            ),
-        ]
+        with self.subTest(config="not_using_standard_api"):
+            tier_metadata = {
+                "user_supplied_max_trials": 100,
+                "uses_standard_api": False,
+            }
+            card = ComplexityRatingAnalysis(tier_metadata=tier_metadata).compute(
+                experiment=self.experiment
+            )
+            self.assertEqual(card.get_status(), HealthcheckStatus.FAIL)
+            self.assertIn("Unsupported", card.subtitle)
+            self.assertIn("uses_standard_api=False", card.subtitle)
 
-        for name, options, tier_metadata, expected_msg in test_cases:
-            with self.subTest(config=name):
-                card = ComplexityRatingAnalysis(
-                    options=options, tier_metadata=tier_metadata
-                ).compute(experiment=self.experiment)
+        with self.subTest(config="high_failure_rate"):
+            tier_metadata = {"user_supplied_max_trials": 100, "uses_standard_api": True}
+            card = ComplexityRatingAnalysis(
+                tier_metadata=tier_metadata,
+                tolerated_trial_failure_rate=0.95,
+            ).compute(experiment=self.experiment)
+            self.assertEqual(card.get_status(), HealthcheckStatus.FAIL)
+            self.assertIn("Unsupported", card.subtitle)
+            self.assertIn("0.95", card.subtitle)
 
-                self.assertEqual(card.get_status(), HealthcheckStatus.FAIL)
-                self.assertIn("Unsupported", card.subtitle)
-                self.assertIn(expected_msg, card.subtitle)
+        with self.subTest(config="invalid_failure_rate_check"):
+            tier_metadata = {"user_supplied_max_trials": 100, "uses_standard_api": True}
+            card = ComplexityRatingAnalysis(
+                tier_metadata=tier_metadata,
+                max_pending_trials=10,
+                min_failed_trials_for_failure_rate_check=50,
+            ).compute(experiment=self.experiment)
+            self.assertEqual(card.get_status(), HealthcheckStatus.FAIL)
+            self.assertIn("Unsupported", card.subtitle)
+            self.assertIn("min_failed_trials_for_failure_rate_check", card.subtitle)
 
     def test_unordered_choice_parameters(self) -> None:
         params = [
@@ -257,9 +243,9 @@ def test_unordered_choice_parameters(self) -> None:
 
         self.assertTrue(is_unordered_choice(params[1], min_choices=3, max_choices=5))
 
-        card = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        ).compute(experiment=self.experiment)
+        card = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata).compute(
+            experiment=self.experiment
+        )
 
         self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
         self.assertIn("Advanced", card.subtitle)
@@ -279,9 +265,9 @@ def test_binary_parameters_count(self) -> None:
         for p in params:
             self.assertTrue(can_map_to_binary(p))
 
-        card = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        ).compute(experiment=self.experiment)
+        card = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata).compute(
+            experiment=self.experiment
+        )
 
         self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
         self.assertIn("Advanced", card.subtitle)
@@ -300,10 +286,9 @@ def test_multiple_violations(self) -> None:
         experiment = self.experiment
         experiment._search_space = SearchSpace(parameters=params)
         tier_metadata = {"user_supplied_max_trials": 300, "uses_standard_api": True}
-        # pyre-ignore[6]: Using a mock value for testing
-        options = OrchestratorOptions(early_stopping_strategy="mock")
         card = ComplexityRatingAnalysis(
-            options=options, tier_metadata=tier_metadata
+            tier_metadata=tier_metadata,
+            early_stopping_strategy=get_percentile_early_stopping_strategy(),
         ).compute(experiment=experiment)
 
         self.assertEqual(card.get_status(), HealthcheckStatus.WARNING)
@@ -313,9 +298,9 @@ def test_multiple_violations(self) -> None:
         self.assertIn("Early stopping is enabled", card.subtitle)
 
     def test_dataframe_summary(self) -> None:
-        card = ComplexityRatingAnalysis(
-            options=self.options, tier_metadata=self.tier_metadata
-        ).compute(experiment=self.experiment)
+        card = ComplexityRatingAnalysis(tier_metadata=self.tier_metadata).compute(
+            experiment=self.experiment
+        )
 
         df = card.df
         self.assertIsNotNone(df)