Audit usage of trial failure reason (#4765)

Summary:
Pull Request resolved: #4765

- Updates `Experiment.to_df` to extract the column from `Trial.status_reason`. Previously, this was trying to get it from `run_metadata["fail_reason"]`, which is never populated AFAICT.
- Renames the column across analyses to `status_reason` from `fail_reason`. The motivation is that this property can represent abandoned and early stopped trial reasons, in addition to failed trials.

Reviewed By: mgarrard

Differential Revision: D90627835

fbshipit-source-id: 22b9b6aa0011ad6125f754508cd94c172130102d

Author: saitcakmak

ax/analysis/plotly/arm_effects.py

@@ -226,7 +226,7 @@ def compute(
                     "trial_index",
                     "trial_status",
                     "arm_name",
-                    "fail_reason",
+                    "status_reason",
                     "generation_node",
                     f"{self.metric_name}_mean",
                     f"{self.metric_name}_sem",

ax/analysis/plotly/tests/test_arm_effects.py

@@ -125,7 +125,7 @@ def test_compute_raw(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "foo_mean",
                 "foo_sem",
@@ -156,7 +156,7 @@ def test_compute_with_modeled(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "foo_mean",
                 "foo_sem",
@@ -343,7 +343,7 @@ def test_compute_with_relativize(self) -> None:
                         "trial_index",
                         "arm_name",
                         "trial_status",
-                        "fail_reason",
+                        "status_reason",
                         "generation_node",
                         "branin_mean",
                         "branin_sem",

ax/analysis/plotly/tests/test_scatter.py

@@ -127,7 +127,7 @@ def test_compute_raw(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -186,7 +186,7 @@ def test_compute_with_modeled(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",

ax/analysis/tests/test_summary.py

@@ -119,7 +119,7 @@ def test_compute(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "foo",
                 "bar",

ax/analysis/tests/test_utils.py

@@ -147,7 +147,7 @@ def test_prepare_arm_data_raw(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -194,7 +194,7 @@ def test_prepare_arm_data_raw(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -225,7 +225,7 @@ def test_prepare_arm_data_raw(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -334,7 +334,7 @@ def test_prepare_arm_data_use_model_predictions(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -378,7 +378,7 @@ def test_prepare_arm_data_use_model_predictions(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -411,7 +411,7 @@ def test_prepare_arm_data_use_model_predictions(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -450,7 +450,7 @@ def test_prepare_arm_data_use_model_predictions(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -493,7 +493,7 @@ def test_prepare_arm_data_use_model_predictions(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "p_feasible_mean",
                 "p_feasible_sem",
@@ -578,7 +578,7 @@ def test_prepare_arm_data_out_of_distribution_arm(self) -> None:
         self.assertFalse(np.isnan(ood_df.foo_sem.iloc[0]))
 
     def test_prepare_arm_data_includes_failure_reasons(self) -> None:
-        """Test that the fail_reason column is properly populated."""
+        """Test that the status_reason column is properly populated."""
         client = Client()
         client.configure_experiment(
             name="test_failure_reasons",
@@ -600,16 +600,16 @@ def test_prepare_arm_data_includes_failure_reasons(self) -> None:
             use_model_predictions=False,
         )
 
-        # Verify fail_reason column is populated correctly
-        self.assertIn("fail_reason", df.columns)
+        # Verify status_reason column is populated correctly
+        self.assertIn("status_reason", df.columns)
         self.assertTrue(
-            pd.isna(df[df["trial_index"] == 0]["fail_reason"].iloc[0])
+            pd.isna(df[df["trial_index"] == 0]["status_reason"].iloc[0])
         )  # Success: no reason
         self.assertEqual(
-            df[df["trial_index"] == 1]["fail_reason"].iloc[0], "Regular failure"
+            df[df["trial_index"] == 1]["status_reason"].iloc[0], "Regular failure"
         )  # Regular failure
         self.assertEqual(
-            df[df["trial_index"] == 2]["fail_reason"].iloc[0], STALE_FAIL_REASON
+            df[df["trial_index"] == 2]["status_reason"].iloc[0], STALE_FAIL_REASON
         )  # Stale failure
 
     def test_relativize_df_with_sq(self) -> None:

ax/analysis/utils.py

@@ -255,7 +255,7 @@ def prepare_arm_data(
         if trial_index != -1
         else "Additional Arm"
     )
-    df["fail_reason"] = df["trial_index"].apply(
+    df["status_reason"] = df["trial_index"].apply(
         lambda trial_index: experiment.trials[trial_index].status_reason
         if trial_index != -1
         and experiment.trials[trial_index].status_reason is not None

ax/core/experiment.py

@@ -1920,7 +1920,8 @@ def to_df(
             - trial_index: The trial index of the arm
             - arm_name: The name of the arm
             - trial_status: The status of the trial (e.g. RUNNING, SUCCEDED, FAILED)
-            - failure_reason: The reason for the failure, if applicable
+            - status_reason: The reason for the trial status (e.g., failure,
+                abandonment, early stopping), if applicable
             - generation_node: The name of the ``GenerationNode`` that generated the arm
             - **METADATA: Any metadata associated with the trial, as specified by the
                 Experiment's runner.run_metadata_report_keys field
@@ -2002,7 +2003,7 @@ def to_df(
                     "trial_index": trial.index,
                     "arm_name": arm.name,
                     "trial_status": trial.status.name,
-                    "fail_reason": trial.run_metadata.get("fail_reason", None),
+                    "status_reason": trial.status_reason,
                     "generation_node": generation_node,
                     **metadata,
                     **observed_means,

ax/core/tests/test_experiment.py

@@ -1689,7 +1689,7 @@ def test_to_df(self) -> None:
                 "trial_index",
                 "arm_name",
                 "trial_status",
-                "fail_reason",
+                "status_reason",
                 "generation_node",
                 "name",
                 "m1",