(4/5) Port helpers to OSS for the new Complexity Rating Healthcheck - wheelhouse tier check (facebook#4629)

Summary: Pull Request resolved: facebook#4629

Differential Revision: D88600526

Author: shrutipatel31

ax/utils/common/tests/test_wheelhouse_utils.py

@@ -6,12 +6,15 @@
 
 # pyre-strict
 
-from ax.exceptions.core import OptimizationNotConfiguredError
+from ax.exceptions.core import OptimizationNotConfiguredError, UserInputError
 from ax.service.orchestrator import OrchestratorOptions
+
 from ax.utils.common.testutils import TestCase
 from ax.utils.common.wheelhouse_utils import (
     ADVANCED_TIER_MESSAGE,
+    check_if_in_wheelhouse,
     format_tier_message,
+    OptimizationSummary,
     summarize_ax_experiment_wheelhouse,
     UNSUPPORTED_TIER_MESSAGE,
     WHEELHOUSE_TIER_MESSAGE,
@@ -41,31 +44,13 @@ def test_basic_experiment_summary(self) -> None:
             tier_metadata=self.tier_metadata,
         )
 
-        # THEN the summary should contain all expected keys with correct values
-        expected_keys = [
-            "max_trials",
-            "num_params",
-            "num_binary",
-            "num_categorical_3_5",
-            "num_categorical_6_inf",
-            "num_parameter_constraints",
-            "num_objectives",
-            "num_outcome_constraints",
-            "uses_early_stopping",
-            "uses_global_stopping",
-            "uses_merge_multiple_curves",
-            "all_inputs_are_configs",
-            "tolerated_trial_failure_rate",
-            "max_pending_trials",
-            "min_failed_trials_for_failure_rate_check",
-        ]
-        for key in expected_keys:
-            self.assertIn(key, summary)
+        # THEN the summary should be an OptimizationSummary with correct values
+        self.assertIsInstance(summary, OptimizationSummary)
 
         # Validate specific values for single-objective experiment
-        self.assertEqual(summary["num_objectives"], 1)
-        self.assertFalse(summary["uses_early_stopping"])
-        self.assertFalse(summary["uses_global_stopping"])
+        self.assertEqual(summary.num_objectives, 1)
+        self.assertFalse(summary.uses_early_stopping)
+        self.assertFalse(summary.uses_global_stopping)
 
     def test_multi_objective_experiment(self) -> None:
         # GIVEN a multi-objective experiment
@@ -79,7 +64,7 @@ def test_multi_objective_experiment(self) -> None:
         )
 
         # THEN num_objectives should be greater than 1
-        self.assertGreater(summary["num_objectives"], 1)
+        self.assertGreater(summary.num_objectives, 1)
 
     def test_experiment_without_optimization_config_raises(self) -> None:
         # GIVEN an experiment without optimization config
@@ -128,10 +113,8 @@ def test_tier_metadata_extraction(self) -> None:
                 )
 
                 # THEN the summary should reflect tier metadata values
-                self.assertEqual(summary["max_trials"], expected_max_trials)
-                self.assertEqual(
-                    summary["all_inputs_are_configs"], expected_all_configs
-                )
+                self.assertEqual(summary.max_trials, expected_max_trials)
+                self.assertEqual(summary.all_inputs_are_configs, expected_all_configs)
 
     def test_orchestrator_options_extraction(self) -> None:
         # GIVEN custom orchestrator options
@@ -149,9 +132,9 @@ def test_orchestrator_options_extraction(self) -> None:
         )
 
         # 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)
+        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
@@ -165,7 +148,7 @@ def test_parameter_constraints_counted(self) -> None:
         )
 
         # THEN num_parameter_constraints should be greater than 0
-        self.assertGreater(summary["num_parameter_constraints"], 0)
+        self.assertGreater(summary.num_parameter_constraints, 0)
 
 
 class TestFormatTierMessage(TestCase):
@@ -241,3 +224,164 @@ def test_unknown_tier_raises_error(self) -> None:
                 why_not_is_in_wheelhouse=None,
                 why_not_supported=None,
             )
+
+
+def get_experiment_summary(
+    max_trials: int | None = 100,
+    num_params: int = 10,
+    num_binary: int = 0,
+    num_categorical_3_5: int = 0,
+    num_categorical_6_inf: int = 0,
+    num_parameter_constraints: int = 0,
+    num_objectives: int = 1,
+    num_outcome_constraints: int = 0,
+    uses_early_stopping: bool = False,
+    uses_global_stopping: bool = False,
+    all_inputs_are_configs: bool = True,
+    tolerated_trial_failure_rate: float | None = 0.5,
+    max_pending_trials: int | None = 5,
+    min_failed_trials_for_failure_rate_check: int | None = 5,
+    non_default_advanced_options: bool | None = None,
+    uses_merge_multiple_curves: bool | None = None,
+) -> OptimizationSummary:
+    """Create an OptimizationSummary for testing."""
+    return OptimizationSummary(
+        max_trials=max_trials,
+        num_params=num_params,
+        num_binary=num_binary,
+        num_categorical_3_5=num_categorical_3_5,
+        num_categorical_6_inf=num_categorical_6_inf,
+        num_parameter_constraints=num_parameter_constraints,
+        num_objectives=num_objectives,
+        num_outcome_constraints=num_outcome_constraints,
+        uses_early_stopping=uses_early_stopping,
+        uses_global_stopping=uses_global_stopping,
+        all_inputs_are_configs=all_inputs_are_configs,
+        tolerated_trial_failure_rate=tolerated_trial_failure_rate,
+        max_pending_trials=max_pending_trials,
+        min_failed_trials_for_failure_rate_check=(
+            min_failed_trials_for_failure_rate_check
+        ),
+        non_default_advanced_options=non_default_advanced_options,
+        uses_merge_multiple_curves=uses_merge_multiple_curves,
+    )
+
+
+class TestCheckIfInWheelhouse(TestCase):
+    """Tests for check_if_in_wheelhouse."""
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.base_summary = get_experiment_summary()
+
+    def test_wheelhouse_tier_for_simple_experiment(self) -> None:
+        """Test that a simple experiment is classified as Wheelhouse tier."""
+        tier, why_not_wheelhouse, why_not_supported = check_if_in_wheelhouse(
+            self.base_summary
+        )
+
+        self.assertEqual(tier, "Wheelhouse")
+        self.assertEqual(why_not_wheelhouse, None)
+        self.assertEqual(why_not_supported, None)
+
+    def test_advanced_tier_conditions(self) -> None:
+        """Test conditions that result in Advanced tier."""
+        test_cases: list[tuple[OptimizationSummary, str]] = [
+            (get_experiment_summary(max_trials=250), "250 total trials"),
+            (get_experiment_summary(num_params=60), "60 tunable parameter(s)"),
+            (get_experiment_summary(num_binary=75), "75 binary tunable parameter(s)"),
+            (
+                get_experiment_summary(num_categorical_3_5=1),
+                "1 unordered choice parameter(s)",
+            ),
+            (
+                get_experiment_summary(num_parameter_constraints=4),
+                "4 parameter constraints",
+            ),
+            (get_experiment_summary(num_objectives=3), "3 objectives"),
+            (
+                get_experiment_summary(num_outcome_constraints=3),
+                "3 outcome constraints",
+            ),
+            (
+                get_experiment_summary(uses_early_stopping=True),
+                "Early stopping is enabled",
+            ),
+            (
+                get_experiment_summary(uses_global_stopping=True),
+                "Global stopping is enabled",
+            ),
+        ]
+
+        for summary, expected_msg in test_cases:
+            with self.subTest(expected_msg=expected_msg):
+                tier, why_not_wheelhouse, why_not_supported = check_if_in_wheelhouse(
+                    summary
+                )
+
+                self.assertEqual(tier, "Advanced")
+                self.assertIsNotNone(why_not_wheelhouse)
+                self.assertIn(expected_msg, why_not_wheelhouse[0])
+                self.assertEqual(why_not_supported, None)
+
+    def test_unsupported_tier_conditions(self) -> None:
+        """Test conditions that result in Unsupported tier."""
+        test_cases: list[tuple[OptimizationSummary, str]] = [
+            (get_experiment_summary(max_trials=510), "510 total trials"),
+            (get_experiment_summary(num_params=201), "201 tunable parameter(s)"),
+            (get_experiment_summary(num_binary=101), "101 binary tunable parameter(s)"),
+            (
+                get_experiment_summary(num_categorical_3_5=6),
+                "unordered choice parameters with more than 3 options",
+            ),
+            (
+                get_experiment_summary(num_categorical_6_inf=2),
+                "unordered choice parameters with more than 5 options",
+            ),
+            (
+                get_experiment_summary(num_parameter_constraints=6),
+                "6 parameter constraints",
+            ),
+            (get_experiment_summary(num_objectives=5), "5 objectives"),
+            (
+                get_experiment_summary(num_outcome_constraints=6),
+                "6 outcome constraints",
+            ),
+            (
+                get_experiment_summary(all_inputs_are_configs=False),
+                "Ax abstractions",
+            ),
+            (
+                get_experiment_summary(tolerated_trial_failure_rate=0.99),
+                "tolerated_trial_failure_rate=0.99",
+            ),
+            (
+                get_experiment_summary(non_default_advanced_options=True),
+                "Non-default advanced_options",
+            ),
+            (
+                get_experiment_summary(uses_merge_multiple_curves=True),
+                "merge_multiple_curves=True",
+            ),
+            (
+                get_experiment_summary(
+                    max_pending_trials=3, min_failed_trials_for_failure_rate_check=7
+                ),
+                "min_failed_trials_for_failure_rate_check=7",
+            ),
+        ]
+
+        for summary, expected_msg in test_cases:
+            with self.subTest(expected_msg=expected_msg):
+                tier, _, why_not_supported = check_if_in_wheelhouse(summary)
+
+                self.assertEqual(tier, "Unsupported")
+                self.assertIsNotNone(why_not_supported)
+                self.assertIn(expected_msg, why_not_supported[0])
+
+    def test_max_trials_none_raises(self) -> None:
+        """Test max_trials=None with all_inputs_are_configs=True raises error."""
+        summary = get_experiment_summary(all_inputs_are_configs=True, max_trials=None)
+
+        with self.assertRaisesRegex(UserInputError, "`max_trials` should not be None!"):
+            check_if_in_wheelhouse(summary)