chore: import upstream snapshot with attribution

python/ray/data/tests/test_exceptions.py

@@ -0,0 +1,129 @@
+import logging
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray.data.exceptions import SystemException, UserCodeException
+from ray.exceptions import RayTaskError
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize("log_internal_stack_trace_to_stdout", [True, False])
+def test_user_exception(
+    log_internal_stack_trace_to_stdout,
+    caplog,
+    propagate_logs,
+    restore_data_context,
+    ray_start_regular_shared,
+):
+    ctx = ray.data.DataContext.get_current()
+    ctx.log_internal_stack_trace_to_stdout = log_internal_stack_trace_to_stdout
+
+    def f(row):
+        _ = 1 / 0
+
+    with pytest.raises(UserCodeException) as exc_info:
+        ray.data.range(1).map(f).take_all()
+
+    assert issubclass(exc_info.type, RayTaskError)
+    assert issubclass(exc_info.type, UserCodeException)
+    assert ZeroDivisionError.__name__ in str(exc_info.value)
+
+    if not log_internal_stack_trace_to_stdout:
+        assert any(
+            record.levelno == logging.ERROR
+            and "Exception occurred in user code" in record.message
+            for record in caplog.records
+        ), caplog.records
+
+    assert any(
+        record.levelno == logging.ERROR
+        and "Full stack trace:" in record.message
+        and getattr(record, "hide", False) == (not log_internal_stack_trace_to_stdout)
+        for record in caplog.records
+    ), caplog.records
+
+
+def test_system_exception(caplog, propagate_logs, ray_start_regular_shared):
+    class FakeException(Exception):
+        pass
+
+    with pytest.raises(FakeException) as exc_info:
+        with patch(
+            "ray.data.dataset._ExecutionCache.get_bundle",
+            side_effect=FakeException(),
+        ):
+            ray.data.range(1).materialize()
+            assert issubclass(exc_info.type, FakeException)
+            assert issubclass(exc_info.type, SystemException)
+
+    assert any(
+        record.levelno == logging.ERROR
+        and "Exception occurred in Ray Data or Ray Core internal code."
+        in record.message
+        for record in caplog.records
+    ), caplog.records
+
+    assert any(
+        record.levelno == logging.ERROR
+        and "Full stack trace:" in record.message
+        and not getattr(record, "hide", False)
+        for record in caplog.records
+    ), caplog.records
+
+
+def test_full_traceback_logged_with_ray_debugger(
+    caplog, propagate_logs, ray_start_regular_shared, monkeypatch
+):
+    monkeypatch.setenv("RAY_DEBUG_POST_MORTEM", 1)
+
+    def f(row):
+        _ = 1 / 0
+        return row
+
+    with pytest.raises(Exception) as exc_info:
+        ray.data.range(1).map(f).take_all()
+
+    assert issubclass(exc_info.type, RayTaskError)
+    assert issubclass(exc_info.type, UserCodeException)
+    assert ZeroDivisionError.__name__ in str(exc_info.value)
+
+    assert any(
+        record.levelno == logging.ERROR
+        and "Full stack trace:" in record.message
+        and not getattr(record, "hide", False)
+        for record in caplog.records
+    ), caplog.records
+
+
+def test_raise_original_map_exception_env_var(
+    caplog, propagate_logs, restore_data_context, ray_start_regular_shared, monkeypatch
+):
+    monkeypatch.setenv("RAY_DATA_RAISE_ORIGINAL_MAP_EXCEPTION", "1")
+    ctx = ray.data.DataContext.get_current()
+    ctx.raise_original_map_exception = (
+        True  # Ensure that the context picks up the environment variable
+    )
+
+    def f(row):
+        raise ValueError("This is a test error.")
+
+    with pytest.raises(ValueError) as exc_info:
+        ray.data.range(1).map(f).take_all()
+
+    assert issubclass(exc_info.type, ValueError)
+    assert "This is a test error." in str(exc_info.value)
+
+    # Ensure that the stack trace is not cleared or replaced by UserCodeException
+    assert not any(
+        record.levelno == logging.ERROR
+        and "Exception occurred in user code" in record.message
+        for record in caplog.records
+    ), caplog.records
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))