[train] fix maximum recursion issue when serializing exceptions

python/ray/air/_internal/util.py

@@ -1,9 +1,10 @@
-import os
-import socket
-from contextlib import closing
+import copy
 import logging
+import os
 import queue
+import socket
 import threading
+from contextlib import closing
 from typing import Optional
 
 import numpy as np
@@ -35,7 +36,13 @@ class StartTraceback(Exception):
 
 
 def skip_exceptions(exc: Optional[Exception]) -> Exception:
-    """Skip all contained `StartTracebacks` to reduce traceback output"""
+    """Skip all contained `StartTracebacks` to reduce traceback output.
+
+    Returns a shallow copy of the exception with all `StartTracebacks` removed.
+
+    If the RAY_AIR_FULL_TRACEBACKS environment variable is set,
+    the original exception (not a copy) is returned.
+    """
     should_not_shorten = bool(int(os.environ.get("RAY_AIR_FULL_TRACEBACKS", "0")))
 
     if should_not_shorten:
@@ -45,12 +52,15 @@ def skip_exceptions(exc: Optional[Exception]) -> Exception:
         # If this is a StartTraceback, skip
         return skip_exceptions(exc.__cause__)
 
-    # Else, make sure nested exceptions are properly skipped
+    # Perform a shallow copy to prevent recursive __cause__/__context__.
+    new_exc = copy.copy(exc).with_traceback(exc.__traceback__)
+
+    # Make sure nested exceptions are properly skipped.
     cause = getattr(exc, "__cause__", None)
     if cause:
-        exc.__cause__ = skip_exceptions(cause)
+        new_exc.__cause__ = skip_exceptions(cause)
 
-    return exc
+    return new_exc
 
 
 def exception_cause(exc: Optional[Exception]) -> Optional[Exception]:

python/ray/air/tests/test_tracebacks.py

@@ -1,8 +1,10 @@
 import pytest
 
 import ray
+from ray import cloudpickle
+from tblib import pickling_support
 from ray.train import ScalingConfig
-from ray.air._internal.util import StartTraceback, skip_exceptions
+from ray.air._internal.util import StartTraceback, skip_exceptions, exception_cause
 from ray.train.data_parallel_trainer import DataParallelTrainer
 
 from ray.tune import Tuner
@@ -47,6 +49,45 @@ def test_short_traceback(levels):
     assert i == levels - start_traceback + 1
 
 
+def test_recursion():
+    """Test that the skipped exception does not point to the original exception."""
+    root_exception = None
+
+    with pytest.raises(StartTraceback) as exc_info:
+        try:
+            raise Exception("Root Exception")
+        except Exception as e:
+            root_exception = e
+            raise StartTraceback from root_exception
+
+    assert root_exception, "Root exception was not captured."
+
+    start_traceback = exc_info.value
+    skipped_exception = skip_exceptions(start_traceback)
+
+    assert (
+        root_exception != skipped_exception
+    ), "Skipped exception points to the original exception."
+
+
+def test_tblib():
+    """Test that tblib does not cause a maximum recursion error."""
+
+    with pytest.raises(Exception) as exc_info:
+        try:
+            try:
+                raise Exception("Root Exception")
+            except Exception as root_exception:
+                raise StartTraceback from root_exception
+        except Exception as start_traceback:
+            raise skip_exceptions(start_traceback) from exception_cause(start_traceback)
+
+    pickling_support.install()
+    reraised_exception = exc_info.value
+    # This should not raise a RecursionError/PicklingError.
+    cloudpickle.dumps(reraised_exception)
+
+
 def test_traceback_tuner(ray_start_2_cpus):
     """Ensure that the Tuner's stack trace is not too long."""