From 9f0ee9e24d79f47bece5df92d9a3a95d5a8c857b Mon Sep 17 00:00:00 2001
From: Eric Onofrey <eonofrey@fb.com>
Date: Thu, 9 Nov 2023 12:43:43 -0800
Subject: [PATCH] Augment exp_to_df with a "reason" column (#1973)

Summary:

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

Differential Revision: D51140675
---
 ax/service/tests/test_report_utils.py | 14 ++++++++++++++
 ax/service/utils/report_utils.py      | 20 +++++++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/ax/service/tests/test_report_utils.py b/ax/service/tests/test_report_utils.py
index 792a94cbddb..17132648817 100644
--- a/ax/service/tests/test_report_utils.py
+++ b/ax/service/tests/test_report_utils.py
@@ -194,6 +194,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()
diff --git a/ax/service/utils/report_utils.py b/ax/service/utils/report_utils.py
index fc0bc6478e0..a7c3b5fec12 100644
--- a/ax/service/utils/report_utils.py
+++ b/ax/service/utils/report_utils.py
@@ -828,6 +828,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.
@@ -910,7 +928,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 [])
     )