[train] Exit actor and log appropriately when poll_workers is in terminal state (#58287)

1) Add traceback to all `ControllerError`s and log it when making a
failure decision so we can see where `Worker group is not active. Call
WorkerGroup.create() to create a new worker group.` is coming from. **I
also sanity checked that this does not cause
`UserExceptionWithTraceback` to double print the traceback because this
only applies to ControllerError**
2) `_poll_workers` has the only `asyncio.sleep` in the Ray Train
controller. After waking up, it exits from the foreground asyncio task
if its state is terminal, which can happen due to the issue mentioned in
5).

---------

Signed-off-by: Timothy Seah <[email protected]>

Author: TimothySeah

python/ray/train/v2/_internal/execution/controller/controller.py

@@ -274,6 +274,12 @@ async def _poll_workers(self) -> WorkerGroupPollStatus:
                 self._health_check_interval_s - time_since_last_poll, 0
             )
             await asyncio.sleep(remaining_time)
+            if self.get_state().is_terminal():
+                logger.debug(
+                    f"Controller is unexpectedly in terminal state {self.get_state()} after "
+                    "sleeping and before polling workers. Exiting actor."
+                )
+                ray.actor.exit_actor()
 
         status = self._worker_group.poll_status(timeout=self._health_check_interval_s)
         self._latest_poll_time = time_monotonic()

python/ray/train/v2/_internal/execution/failure_handling/default.py

@@ -44,8 +44,12 @@ def _log_decision(
         logger.info(
             f"[FailurePolicy] {decision.value}\n"
             f"  Source: {error_source}\n"
-            f"  Error count: {error_count} (max allowed: {retry_limit})\n\n"
-            f"{training_failed_error}"
+            f"  Error count: {error_count} (max allowed: {retry_limit})\n\n",
+            exc_info=(
+                type(training_failed_error),
+                training_failed_error,
+                training_failed_error.__traceback__,
+            ),
         )
 
     def _is_retryable_error(self, training_failed_error: TrainingFailedError) -> bool:

python/ray/train/v2/api/exceptions.py

@@ -43,6 +43,7 @@ def __init__(self, controller_failure: Exception):
             "Training failed due to controller error:\n" + str(controller_failure)
         )
         self.controller_failure = controller_failure
+        self.with_traceback(controller_failure.__traceback__)
 
     def __reduce__(self):
         return (self.__class__, (self.controller_failure,))