Phase 4: Exception Groups Migration (PEP 654) (#4877)

Summary:
Pull Request resolved: #4877

Convert error collection patterns to use Python 3.11 ExceptionGroup instead
of ad-hoc exception aggregation. This provides better structured exception
handling and enables the use of except* syntax in callers.

Files changed:
- ax/core/metric.py: Update _unwrap_experiment_data(), _unwrap_trial_data_multi(),
  and _unwrap_experiment_data_multi() to raise ExceptionGroup with all collected
  errors instead of chaining to a single exception.
- ax/generation_strategy/generation_node.py: Update new_trial_limit() to collect
  all generation-blocking errors and raise as ExceptionGroup.

Reviewed By: saitcakmak

Differential Revision: D91648883

fbshipit-source-id: b9ef2a7912f26e9e6b5c1803a81dca3e008242be

Author: Balandat

ax/core/metric.py

@@ -20,7 +20,7 @@
 from ax.core.data import Data
 from ax.utils.common.base import SortableBase
 from ax.utils.common.logger import get_logger
-from ax.utils.common.result import Err, Ok, Result, UnwrapError
+from ax.utils.common.result import Err, Ok, Result
 from ax.utils.common.serialization import SerializationMixin
 
 if TYPE_CHECKING:
@@ -505,7 +505,6 @@ def _unwrap_experiment_data(cls, results: Mapping[int, MetricFetchResult]) -> Da
                 result for result in results.values() if isinstance(result, Err)
             ]
 
-            # TODO[mpolson64] Raise all errors in a group via PEP 654
             exceptions = [
                 (
                     err.err.exception
@@ -515,8 +514,8 @@ def _unwrap_experiment_data(cls, results: Mapping[int, MetricFetchResult]) -> Da
                 for err in errs
             ]
 
-            raise UnwrapError(errs) from (
-                exceptions[0] if len(exceptions) == 1 else Exception(exceptions)
+            raise ExceptionGroup(
+                f"Failed to fetch data for {len(errs)} trial(s)", exceptions
             )
 
         return Data.from_multiple_data(data=[ok.ok for ok in oks])
@@ -561,7 +560,6 @@ def _unwrap_trial_data_multi(
             ]
 
             if len(critical_errs) > 0:
-                # TODO[mpolson64] Raise all errors in a group via PEP 654
                 exceptions = [
                     (
                         err.err.exception
@@ -570,8 +568,9 @@ def _unwrap_trial_data_multi(
                     )
                     for err in critical_errs
                 ]
-                raise UnwrapError(critical_errs) from (
-                    exceptions[0] if len(exceptions) == 1 else Exception(exceptions)
+                raise ExceptionGroup(
+                    f"Failed to fetch data for {len(critical_errs)} critical metric(s)",
+                    exceptions,
                 )
 
         return Data.from_multiple_data(data=[ok.ok for ok in oks])
@@ -592,7 +591,6 @@ def _unwrap_experiment_data_multi(
                 result for result in flattened if isinstance(result, Err)
             ]
 
-            # TODO[mpolson64] Raise all errors in a group via PEP 654
             exceptions = [
                 (
                     err.err.exception
@@ -601,8 +599,8 @@ def _unwrap_experiment_data_multi(
                 )
                 for err in errs
             ]
-            raise UnwrapError(errs) from (
-                exceptions[0] if len(exceptions) == 1 else Exception(exceptions)
+            raise ExceptionGroup(
+                f"Failed to fetch data for {len(errs)} metric(s)", exceptions
             )
 
         return Data.from_multiple_data(data=[ok.ok for ok in oks])

ax/core/tests/test_metric.py

@@ -105,14 +105,35 @@ def test_wrap_unwrap(self) -> None:
     def test_wrap_err(self) -> None:
         err = Err(MetricFetchE(message="failed!", exception=Exception("panic!")))
 
-        with self.assertRaisesRegex(Exception, "panic"):
-            Metric._unwrap_experiment_data_multi(results={0: {"foo": err}})
-
-        with self.assertRaisesRegex(Exception, "panic"):
-            Metric._unwrap_experiment_data(results={0: err})
-
-        with self.assertRaisesRegex(Exception, "panic"):
-            Metric._unwrap_trial_data_multi(results={"foo": err})
+        # With PEP 654 ExceptionGroups, errors are now wrapped in ExceptionGroup.
+        # We verify both the top-level ExceptionGroup message and that the
+        # original "panic!" exception is preserved as a nested exception.
+        with self.subTest("experiment_data_multi"):
+            with self.assertRaisesRegex(
+                ExceptionGroup, "Failed to fetch data for 1 metric"
+            ) as cm:
+                Metric._unwrap_experiment_data_multi(results={0: {"foo": err}})
+            eg = cm.exception
+            self.assertEqual(len(eg.exceptions), 1)
+            self.assertIn("panic", str(eg.exceptions[0]))
+
+        with self.subTest("experiment_data"):
+            with self.assertRaisesRegex(
+                ExceptionGroup, "Failed to fetch data for 1 trial"
+            ) as cm:
+                Metric._unwrap_experiment_data(results={0: err})
+            eg = cm.exception
+            self.assertEqual(len(eg.exceptions), 1)
+            self.assertIn("panic", str(eg.exceptions[0]))
+
+        with self.subTest("trial_data_multi"):
+            with self.assertRaisesRegex(
+                ExceptionGroup, "Failed to fetch data for 1 critical metric"
+            ) as cm:
+                Metric._unwrap_trial_data_multi(results={"foo": err})
+            eg = cm.exception
+            self.assertEqual(len(eg.exceptions), 1)
+            self.assertIn("panic", str(eg.exceptions[0]))
 
     def test_MetricFetchE(self) -> None:
         def foo() -> bool:

ax/generation_strategy/generation_node.py

@@ -864,23 +864,27 @@ def new_trial_limit(self, raise_generation_errors: bool = False) -> int:
         ]
 
         # Raise any necessary generation errors: for any met criterion,
-        # call its `block_continued_generation_error` method The method might not
-        # raise an error, depending on its implementation on given criterion, so the
-        # error from the first met one that does block continued generation, will be
-        # raised.
+        # collect all errors from blocking criteria and raise as ExceptionGroup.
         if raise_generation_errors:
+            generation_errors: list[Exception] = []
             for criterion in trial_based_gen_blocking_criteria:
-                # TODO[mgarrard]: Raise a group of all the errors, from each gen-
-                # blocking transition criterion.
                 if criterion.is_met(
                     self.experiment,
                     curr_node=self,
                 ):
-                    criterion.block_continued_generation_error(
-                        node_name=self.name,
-                        experiment=self.experiment,
-                        trials_from_node=self.trials_from_node,
-                    )
+                    try:
+                        criterion.block_continued_generation_error(
+                            node_name=self.name,
+                            experiment=self.experiment,
+                            trials_from_node=self.trials_from_node,
+                        )
+                    except Exception as e:
+                        generation_errors.append(e)
+            if generation_errors:
+                raise ExceptionGroup(
+                    f"Generation blocked by {len(generation_errors)} criteria",
+                    generation_errors,
+                )
         if len(gen_blocking_criterion_delta_from_threshold) == 0:
             return -1
         return max(min(gen_blocking_criterion_delta_from_threshold), -1)

ax/utils/common/result.py

@@ -91,7 +91,7 @@ def unwrap(self) -> T:
         """
         Returns the contained Ok value.
 
-        Because this function may raise an UnwrapError, its use is generally
+        Because this function may raise a RuntimeError, its use is generally
         discouraged. Instead, prefer to handle the Err case explicitly, or call
         unwrap_or, unwrap_or_else, or unwrap_or_default.
         """
@@ -103,7 +103,7 @@ def unwrap_err(self) -> E:
         """
         Returns the contained Err value.
 
-        Because this function may raise an UnwrapError, its use is generally
+        Because this function may raise a RuntimeError, its use is generally
         discouraged. Instead, prefer to handle the Err case explicitly, or call
         unwrap_or, unwrap_or_else, or unwrap_or_default.
         """
@@ -178,7 +178,7 @@ def unwrap(self) -> T:
         return self._value
 
     def unwrap_err(self) -> NoReturn:
-        raise UnwrapError(f"Tried to unwrap_err {self}.")
+        raise RuntimeError(f"Tried to unwrap_err {self}.")
 
     def unwrap_or(self, default: U) -> T:
         return self._value
@@ -235,7 +235,7 @@ def map_or_else(self, default_op: Callable[[], U], op: Callable[[T], U]) -> U:
         return default_op()
 
     def unwrap(self) -> NoReturn:
-        raise UnwrapError(f"Tried to unwrap {self}.")
+        raise RuntimeError(f"Tried to unwrap {self}.")
 
     def unwrap_err(self) -> E:
         return self._value
@@ -249,17 +249,6 @@ def unwrap_or_else(self, op: Callable[[E], T]) -> T:
         return op(self._value)
 
 
-class UnwrapError(Exception):
-    """
-    Exception that indicates something has gone wrong in an unwrap call.
-
-    This should not happen in real world use and indicates a user has improperly
-    or unsafely used the Result abstraction.
-    """
-
-    pass
-
-
 class ExceptionE:
     """
     A class that holds an Exception and can be used as the E type in Result[T, E].

ax/utils/common/tests/test_result.py

@@ -5,7 +5,7 @@
 
 # pyre-strict
 
-from ax.utils.common.result import Err, Ok, Result, UnwrapError
+from ax.utils.common.result import Err, Ok, Result
 from ax.utils.common.testutils import TestCase
 
 
@@ -59,12 +59,12 @@ def h() -> int:
 
     def test_unwrap(self) -> None:
         self.assertEqual(self.ok.unwrap(), 0)
-        with self.assertRaises(UnwrapError):
+        with self.assertRaises(RuntimeError):
             self.ok.unwrap_err()
         self.assertEqual(self.ok.unwrap_or(1), 0)
         self.assertEqual(self.ok.unwrap_or_else(1), 0)
 
-        with self.assertRaises(UnwrapError):
+        with self.assertRaises(RuntimeError):
             self.err.unwrap()
         self.assertEqual(self.err.unwrap_err(), "yikes")
         self.assertEqual(self.err.unwrap_or(1), 1)