(3/5) Port helpers to OSS for the new Complexity Rating Healthcheck - format_tier_message (#4650)

Summary: Pull Request resolved: #4650

Differential Revision: D88908448

Author: shrutipatel31

ax/utils/common/complexity_utils.py

@@ -5,6 +5,7 @@
 
 # pyre-strict
 
+from collections.abc import Iterable
 from typing import Any
 
 from ax.adapter.adapter_utils import can_map_to_binary, is_unordered_choice
@@ -13,6 +14,30 @@
 from ax.exceptions.core import OptimizationNotConfiguredError
 from ax.service.orchestrator import OrchestratorOptions
 
+WHEELHOUSE_TIER_MESSAGE = """This experiment is in tier 'Wheelhouse'.
+
+Experiments belonging to this tier should not run into any problems! If an issue \
+does occur, please post to our github issues page.
+"""
+
+ADVANCED_TIER_MESSAGE = """This experiment is in tier 'Advanced'.
+
+This experiment should technically run, but uses advanced features that may not \
+be well-tested and/or may not be compatible with other advanced features. We \
+encourage users to raise issues encountered in advanced workflows just like \
+wheelhouse workflows but it is also possible that reducing the complexity of \
+your setup can solve your issue.
+"""
+
+UNSUPPORTED_TIER_MESSAGE = """This experiment is in tier 'Unsupported'.
+
+You are pushing Ax beyond its limits. Please post to our github issues page for help \
+in improving/simplifying your configuration to conform to a more \
+well-supported usage tier if possible.
+"""
+
+WIKI_TIER_MESSAGE = "https://ax.dev/docs/why-ax"
+
 
 def summarize_ax_optimization_complexity(
     experiment: Experiment,
@@ -84,3 +109,46 @@ def summarize_ax_optimization_complexity(
             options.min_failed_trials_for_failure_rate_check
         ),
     }
+
+
+def format_tier_message(
+    tier: str,
+    why_not_is_in_wheelhouse: Iterable[str] | None,
+    why_not_supported: Iterable[str] | None,
+) -> str:
+    """
+    Format the result from `check_if_in_wheelhouse` to a markdown-formatted string
+    explaining the tier.
+    """
+
+    if tier == "Wheelhouse":
+        msg = WHEELHOUSE_TIER_MESSAGE
+    else:
+        if tier == "Advanced":
+            msg = ADVANCED_TIER_MESSAGE
+        elif tier == "Unsupported":
+            msg = UNSUPPORTED_TIER_MESSAGE
+        else:
+            raise ValueError(f'Got unexpected tier "{tier}".')
+
+        # Provide user-feedback
+        if why_not_is_in_wheelhouse:
+            why_msg = "\n".join("\t- " + s for s in why_not_is_in_wheelhouse)
+            why_msg = (
+                "\n\nWhy this experiment is not in the 'Wheelhouse' tier: "
+                f"\n{why_msg}\n"
+            )
+            msg += why_msg
+        if why_not_supported:
+            why_msg = "\n".join("\t- " + s for s in why_not_supported)
+            why_msg = (
+                "\n\nWhy this experiment is not in the 'Advanced' tier: "
+                f"\n{why_msg}\n"
+            )
+            msg += why_msg
+
+    msg += (
+        "\n\nFor more information about the definition of each tier and what "
+        f"level of support you can expect: {WIKI_TIER_MESSAGE}"
+    )
+    return msg

ax/utils/common/tests/test_complexity_utils.py

@@ -8,7 +8,13 @@
 
 from ax.exceptions.core import OptimizationNotConfiguredError
 from ax.service.orchestrator import OrchestratorOptions
-from ax.utils.common.complexity_utils import summarize_ax_optimization_complexity
+from ax.utils.common.complexity_utils import (
+    ADVANCED_TIER_MESSAGE,
+    format_tier_message,
+    summarize_ax_optimization_complexity,
+    UNSUPPORTED_TIER_MESSAGE,
+    WHEELHOUSE_TIER_MESSAGE,
+)
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
     get_experiment,
@@ -160,3 +166,78 @@ def test_parameter_constraints_counted(self) -> None:
 
         # THEN num_parameter_constraints should be greater than 0
         self.assertGreater(summary["num_parameter_constraints"], 0)
+
+
+class TestFormatTierMessage(TestCase):
+    """Tests for format_tier_message."""
+
+    def test_tier_messages(self) -> None:
+        """Test formatting of tier messages for all tiers."""
+        test_cases: list[
+            tuple[
+                str,
+                list[str] | None,
+                list[str] | None,
+                str,
+                list[str],
+            ]
+        ] = [
+            (
+                "Wheelhouse",
+                None,
+                None,
+                WHEELHOUSE_TIER_MESSAGE,
+                ["tier 'Wheelhouse'"],
+            ),
+            (
+                "Advanced",
+                ["51 tunable parameters", "Early stopping is enabled"],
+                None,
+                ADVANCED_TIER_MESSAGE,
+                [
+                    "tier 'Advanced'",
+                    "Why this experiment is not in the 'Wheelhouse' tier:",
+                    "51 tunable parameters",
+                    "Early stopping is enabled",
+                ],
+            ),
+            (
+                "Unsupported",
+                ["51 tunable parameters"],
+                ["201 tunable parameters"],
+                UNSUPPORTED_TIER_MESSAGE,
+                [
+                    "tier 'Unsupported'",
+                    "Why this experiment is not in the 'Wheelhouse' tier:",
+                    "51 tunable parameters",
+                    "Why this experiment is not in the 'Advanced' tier:",
+                    "201 tunable parameters",
+                ],
+            ),
+        ]
+
+        for (
+            tier,
+            why_not_wheelhouse,
+            why_not_supported,
+            expected_message,
+            expected_contents,
+        ) in test_cases:
+            with self.subTest(tier=tier):
+                msg = format_tier_message(
+                    tier=tier,
+                    why_not_is_in_wheelhouse=why_not_wheelhouse,
+                    why_not_supported=why_not_supported,
+                )
+                self.assertIn(expected_message, msg)
+                for content in expected_contents:
+                    self.assertIn(content, msg)
+
+    def test_unknown_tier_raises_error(self) -> None:
+        """Test that unknown tier raises ValueError."""
+        with self.assertRaisesRegex(ValueError, 'Got unexpected tier "BadTier"'):
+            format_tier_message(
+                tier="BadTier",
+                why_not_is_in_wheelhouse=None,
+                why_not_supported=None,
+            )