Augment exp_to_df with a "reason" column (#1973)

Summary:
Pull Request resolved: #1973

Augment exp_to_df with a "reason" column, which will surface the failure or abandonment reason string

Reviewed By: mpolson64

Differential Revision: D51140675

fbshipit-source-id: 8d1f5310efd1c717f4211224223230eed8c45b96

ax/service/tests/test_report_utils.py

@@ -201,6 +201,20 @@ def test_exp_to_df_trial_timing(self) -> None:
         # the last trial should have NaN for rel_time_completed
         self.assertTrue(df["time_completed"].isnull().iloc[2])
 
+    def test_exp_to_df_with_failure(self) -> None:
+        fail_reason = "test reason"
+
+        # Set up experiment with a failed trial
+        exp = get_branin_experiment(with_trial=True)
+        exp.trials[0].run()
+        exp.trials[0].mark_failed(reason=fail_reason)
+
+        df = exp_to_df(exp)
+        self.assertEqual(
+            set(EXPECTED_COLUMNS + ["reason"]) - set(df.columns), {OBJECTIVE_NAME}
+        )
+        self.assertEqual(f"{fail_reason}...", df["reason"].iloc[0])
+
     def test_exp_to_df(self) -> None:
         # MultiTypeExperiment should fail
         exp = get_multi_type_experiment()

ax/service/utils/report_utils.py

@@ -838,6 +838,24 @@ def exp_to_df(
         df=arms_df, trials_dict=trial_to_status, column_name="trial_status"
     )
 
+    # Add trial reason for failed or abandoned trials
+    trial_to_reason = {
+        index: (
+            f"{trial.failed_reason[:15]}..."
+            if trial.status.is_failed and trial.failed_reason is not None
+            else f"{trial.abandoned_reason[:15]}..."
+            if trial.status.is_abandoned and trial.abandoned_reason is not None
+            else None
+        )
+        for index, trial in trials
+    }
+
+    _merge_trials_dict_with_df(
+        df=arms_df,
+        trials_dict=trial_to_reason,
+        column_name="reason",
+    )
+
     # Add generation_method, accounting for the generic case that generator_runs is of
     # arbitrary length. Repeated methods within a trial are condensed via `set` and an
     # empty set will yield "Unknown" as the method.
@@ -920,7 +938,7 @@ def exp_to_df(
 
     exp_df = not_none(not_none(exp_df).sort_values(["trial_index"]))
     initial_column_order = (
-        ["trial_index", "arm_name", "trial_status", "generation_method"]
+        ["trial_index", "arm_name", "trial_status", "reason", "generation_method"]
         + (run_metadata_fields or [])
         + (trial_properties_fields or [])
     )