Fix OOM in CI by clearing chained exception tracebacks (#5776)

Author: albertvillanova

tests/conftest.py

@@ -34,6 +34,21 @@ def pytest_runtest_makereport(item, call):
     yield
     if call.when == "call" and call.excinfo is not None:
         traceback.clear_frames(call.excinfo.tb)
+        # Also clear all reachable chained exception tracebacks (both __context__ and __cause__ at
+        # every node): when OOM fires inside a try/except in the trainer, the OOM becomes __context__
+        # of the outer exception and its traceback holds frame locals (model, tensors) that prevent gc
+        # from releasing CUDA memory even after clear_frames above.
+        stack, seen = [call.excinfo.value], set()
+        while stack:
+            exc = stack.pop()
+            if exc is None or id(exc) in seen:
+                continue
+            seen.add(id(exc))
+            if exc.__traceback__ is not None:
+                traceback.clear_frames(exc.__traceback__)
+                exc.__traceback__ = None
+            stack.append(exc.__context__)
+            stack.append(exc.__cause__)
 
 
 # ============================================================================