Add NON_FAILED_STALE_ABANDONED_STATUSES to trial_statuses (facebook#4726)

Summary:
Pull Request resolved: facebook#4726

This diff introduces a reusable constant `NON_FAILED_STALE_ABANDONED_STATUSES` in `trial_status.py` to replace the repeated pattern `{*TrialStatus} - {TrialStatus.ABANDONED, TrialStatus.STALE, TrialStatus.FAILED}` used across multiple analysis files.

Changes:
- Added `NON_FAILED_STALE_ABANDONED_STATUSES` constant to `ax/core/trial_status.py`
- Updated `get_trial_statuses_with_fallback()` in `utils.py` to use the new constant
- Updated all test files to import and use the new constant

Reviewed By: bernardbeckerman

Differential Revision:
D89928534

Privacy Context Container: L1307644

fbshipit-source-id: fcd41fea15621de7a9539895d5e9677ff51a28c8

Author: shrutipatel31

ax/analysis/plotly/tests/test_arm_effects.py

@@ -13,7 +13,7 @@
 from ax.api.client import Client
 from ax.api.configs import RangeParameterConfig
 from ax.core.arm import Arm
-from ax.core.trial_status import TrialStatus
+from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
@@ -78,12 +78,10 @@ def test_trial_statuses_behavior(self) -> None:
         # When neither trial_statuses nor trial_index is provided,
         # should use default statuses (excluding ABANDONED, STALE, and FAILED)
         analysis = ArmEffectsPlot(metric_name="foo")
-        expected_statuses = {*TrialStatus} - {
-            TrialStatus.ABANDONED,
-            TrialStatus.STALE,
-            TrialStatus.FAILED,
-        }
-        self.assertEqual(set(none_throws(analysis.trial_statuses)), expected_statuses)
+        self.assertEqual(
+            set(none_throws(analysis.trial_statuses)),
+            DEFAULT_ANALYSIS_STATUSES,
+        )
 
         # When trial_statuses is explicitly provided, it should be used
         explicit_statuses = [TrialStatus.COMPLETED, TrialStatus.RUNNING]

ax/analysis/plotly/tests/test_objective_p_feasible_frontier.py

@@ -20,7 +20,7 @@
     OptimizationConfig,
 )
 from ax.core.outcome_constraint import ScalarizedOutcomeConstraint
-from ax.core.trial_status import TrialStatus
+from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus
 from ax.core.types import ComparisonOp
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
@@ -52,12 +52,10 @@ def test_trial_statuses_behavior(self) -> None:
         # When neither trial_statuses nor trial_index is provided,
         # should use default statuses (excluding ABANDONED, STALE, and FAILED)
         analysis = ObjectivePFeasibleFrontierPlot()
-        expected_statuses = {*TrialStatus} - {
-            TrialStatus.ABANDONED,
-            TrialStatus.STALE,
-            TrialStatus.FAILED,
-        }
-        self.assertEqual(set(none_throws(analysis.trial_statuses)), expected_statuses)
+        self.assertEqual(
+            set(none_throws(analysis.trial_statuses)),
+            DEFAULT_ANALYSIS_STATUSES,
+        )
 
         # When trial_statuses is explicitly provided, it should be used
         explicit_statuses = [TrialStatus.COMPLETED, TrialStatus.RUNNING]

ax/analysis/plotly/tests/test_scatter.py

@@ -14,7 +14,7 @@
 from ax.api.client import Client
 from ax.api.configs import RangeParameterConfig
 from ax.core.arm import Arm
-from ax.core.trial_status import TrialStatus
+from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
@@ -75,12 +75,10 @@ def test_trial_statuses_behavior(self) -> None:
         # When neither trial_statuses nor trial_index is provided,
         # should use default statuses (excluding ABANDONED, STALE, and FAILED)
         analysis = ScatterPlot(x_metric_name="foo", y_metric_name="bar")
-        expected_statuses = {*TrialStatus} - {
-            TrialStatus.ABANDONED,
-            TrialStatus.STALE,
-            TrialStatus.FAILED,
-        }
-        self.assertEqual(set(none_throws(analysis.trial_statuses)), expected_statuses)
+        self.assertEqual(
+            set(none_throws(analysis.trial_statuses)),
+            DEFAULT_ANALYSIS_STATUSES,
+        )
 
         # When trial_statuses is explicitly provided, it should be used
         explicit_statuses = [TrialStatus.COMPLETED, TrialStatus.RUNNING]

ax/analysis/plotly/tests/test_utils.py

@@ -11,7 +11,7 @@
     get_trial_statuses_with_fallback,
     trial_index_to_color,
 )
-from ax.core.trial_status import TrialStatus
+from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus
 from ax.utils.common.testutils import TestCase
 from pyre_extensions import none_throws
 
@@ -38,12 +38,7 @@ def test_get_trial_statuses_with_fallback_default(self) -> None:
             get_trial_statuses_with_fallback(trial_statuses=None, trial_index=None)
         )
 
-        expected_statuses = {*TrialStatus} - {
-            TrialStatus.ABANDONED,
-            TrialStatus.STALE,
-            TrialStatus.FAILED,
-        }
-        self.assertEqual(set(result), expected_statuses)
+        self.assertEqual(set(result), DEFAULT_ANALYSIS_STATUSES)
         self.assertNotIn(TrialStatus.ABANDONED, result)
         self.assertNotIn(TrialStatus.STALE, result)
         self.assertNotIn(TrialStatus.FAILED, result)

ax/analysis/plotly/utils.py

@@ -13,7 +13,7 @@
 from ax.analysis.plotly.color_constants import BOTORCH_COLOR_SCALE, LIGHT_AX_BLUE
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective, ScalarizedObjective
-from ax.core.trial_status import TrialStatus
+from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus
 from ax.exceptions.core import UnsupportedError
 from plotly import express as px
 
@@ -232,6 +232,4 @@ def get_trial_statuses_with_fallback(
         return None
     elif trial_statuses is not None:
         return [*trial_statuses]
-    return [
-        *{*TrialStatus} - {TrialStatus.ABANDONED, TrialStatus.STALE, TrialStatus.FAILED}
-    ]
+    return [*DEFAULT_ANALYSIS_STATUSES]

ax/core/trial_status.py

@@ -173,3 +173,9 @@ def __repr__(self) -> str:
     TrialStatus.STALE,
     TrialStatus.CANDIDATE,
 ]
+
+DEFAULT_ANALYSIS_STATUSES: set[TrialStatus] = set(TrialStatus) - {
+    TrialStatus.ABANDONED,
+    TrialStatus.STALE,
+    TrialStatus.FAILED,
+}