[core] Improve docs for custom serialization for exceptions + add test (#56156)

Signed-off-by: sampan <sampan@anyscale.com>
Signed-off-by: Sampan S Nayak <sampansnayak2@gmail.com>
Co-authored-by: sampan <sampan@anyscale.com>

Author: sampan-s-nayak

doc/source/ray-core/objects/serialization.rst

@@ -202,6 +202,63 @@ There are at least 3 ways to define your custom serialization process:
      except TypeError:
         pass
 
+.. _custom-exception-serializer:
+
+Custom Serializers for Exceptions
+----------------------------------
+
+When Ray tasks raise exceptions that cannot be serialized with the default pickle mechanism, you can register custom serializers to handle them (Note: the serializer must be registered in the driver and all workers).
+
+.. testcode::
+
+    import ray
+    import threading
+
+    class CustomError(Exception):
+        def __init__(self, message, data):
+            self.message = message
+            self.data = data
+            self.lock = threading.Lock() # Cannot be serialized
+
+    def custom_serializer(exc):
+        return {"message": exc.message, "data": str(exc.data)}
+
+    def custom_deserializer(state):
+        return CustomError(state["message"], state["data"])
+
+    # Register in the driver
+    ray.util.register_serializer(
+        CustomError, 
+        serializer=custom_serializer, 
+        deserializer=custom_deserializer
+    )
+
+    @ray.remote
+    def task_that_registers_serializer_and_raises():
+        # Register the custom serializer in the worker
+        ray.util.register_serializer(
+            CustomError, 
+            serializer=custom_serializer, 
+            deserializer=custom_deserializer
+        )
+        
+        # Now raise the custom exception
+        raise CustomError("Something went wrong", {"complex": "data"})
+
+    # The custom exception will be properly serialized across worker boundaries
+    try:
+        ray.get(task_that_registers_serializer_and_raises.remote())
+    except ray.exceptions.RayTaskError as e:
+        print(f"Caught exception: {e.cause}")  # This will be our CustomError
+
+When a custom exception is raised in a remote task, Ray will:
+
+1. Serialize the exception using your custom serializer
+2. Wrap it in a :class:`RayTaskError <ray.exceptions.RayTaskError>`
+3. The deserialized exception will be available as ``ray_task_error.cause``
+
+Whenever serialization fails, Ray throws an :class:`UnserializableException <ray.exceptions.UnserializableException>` containing the string representation of the original stack trace.
+
 
 Troubleshooting
 ---------------

python/ray/exceptions.py

@@ -915,7 +915,7 @@ class UnserializableException(RayError):
     the original exception along with its stack trace that was captured at the
     time of serialization.
 
-    reference for more details: https://docs.ray.io/en/latest/ray-core/objects/serialization.html
+    For more details and how to handle this with custom serializers, :ref:`configuring custom exeception serializers <custom-exception-serializer>`
 
     Args:
         original_stack_trace: The string representation and stack trace of the
@@ -927,7 +927,7 @@ def __init__(self, original_stack_trace: str):
 
     def __str__(self):
         return (
-            "Failed to deserialize exception. Refer to https://docs.ray.io/en/latest/ray-core/objects/serialization.html#troubleshooting to troubleshoot.\n"
+            "Failed to deserialize exception. Refer to https://docs.ray.io/en/latest/ray-core/objects/serialization.html#custom-serializers-for-exceptions for more information.\n"
             "Original exception:\n"
             f"{self._original_stack_trace}"
         )

python/ray/tests/test_traceback.py

@@ -301,7 +301,7 @@ def __repr__(self):
 
 
 def test_unpickleable_stacktrace(shutdown_only):
-    expected_output = """Failed to deserialize exception. Refer to https://docs.ray.io/en/latest/ray-core/objects/serialization.html#troubleshooting to troubleshoot.
+    expected_output = """Failed to deserialize exception. Refer to https://docs.ray.io/en/latest/ray-core/objects/serialization.html#custom-serializers-for-exceptions for more information.
 Original exception:
 ray.exceptions.RayTaskError: ray::f() (pid=XXX, ip=YYY)
   File "FILE", line ZZ, in f
@@ -330,6 +330,43 @@ def f():
     assert clean_noqa(expected_output) == scrub_traceback(str(excinfo.value))
 
 
+def test_exception_with_registered_serializer(shutdown_only):
+    class NoPickleError(OSError):
+        def __init__(self, msg):
+            self.msg = msg
+
+        def __str__(self):
+            return f"message: {self.msg}"
+
+    def _serializer(e: NoPickleError):
+        return {"msg": e.msg}
+
+    def _deserializer(state):
+        return NoPickleError(state["msg"] + " deserialized")
+
+    @ray.remote
+    def raise_custom_exception():
+        ray.util.register_serializer(
+            NoPickleError, serializer=_serializer, deserializer=_deserializer
+        )
+        raise NoPickleError("message")
+
+    try:
+        with pytest.raises(NoPickleError) as exc_info:
+            ray.get(raise_custom_exception.remote())
+
+        # Ensure dual-typed exception and message propagation
+        assert isinstance(exc_info.value, RayTaskError)
+        # if custom serializer was not registered, this would be an instance of UnserializableException()
+        assert isinstance(exc_info.value, NoPickleError)
+        assert "message" in str(exc_info.value)
+        # modified message should not be in the exception string, only in the cause
+        assert "deserialized" not in str(exc_info.value)
+        assert "message deserialized" in str(exc_info.value.cause)
+    finally:
+        ray.util.deregister_serializer(NoPickleError)
+
+
 def test_serialization_error_message(shutdown_only):
     expected_output_ray_put = """Could not serialize the put value <unlocked _thread.lock object at ADDRESS>:\nINSPECT_SERIALIZABILITY"""  # noqa
     expected_output_task = """Could not serialize the argument <unlocked _thread.lock object at ADDRESS> for a task or actor test_traceback.test_serialization_error_message.<locals>.task_with_unserializable_arg:\nINSPECT_SERIALIZABILITY"""  # noqa