From d6c7a065ea532a2907d6e428c1d43af7423387a8 Mon Sep 17 00:00:00 2001
From: Alex Wu <alex@anyscale.io>
Date: Wed, 21 Dec 2022 13:22:27 -0800
Subject: [PATCH] [autoscaler] Make NodeLaunchException serializable (#31256)

In order to make node launch exceptions serializable, drop the optional stack trace which isn't serializable.

Co-authored-by: Alex <alex@anyscale.com>
Signed-off-by: tmynn <hovhannes.tamoyan@gmail.com>
---
 .../ray/autoscaler/node_launch_exception.py   |  8 +++++++
 python/ray/tests/test_autoscaler_e2e.py       | 22 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/ray/autoscaler/node_launch_exception.py b/python/ray/autoscaler/node_launch_exception.py
index 8c5c2d7e9869..eb6bd25f2c61 100644
--- a/python/ray/autoscaler/node_launch_exception.py
+++ b/python/ray/autoscaler/node_launch_exception.py
@@ -27,3 +27,11 @@ def __init__(
         self.category = category
         self.description = description
         self.src_exc_info = src_exc_info
+
+    def __reduce__(self):
+        # NOTE: Since tracebacks can't be pickled, we'll drop the optional
+        # traceback if we have to serialize this object.
+        return (
+            self.__class__,
+            (self.category, self.description, None),
+        )
diff --git a/python/ray/tests/test_autoscaler_e2e.py b/python/ray/tests/test_autoscaler_e2e.py
index a1195ad403d0..6bebfb955fa6 100644
--- a/python/ray/tests/test_autoscaler_e2e.py
+++ b/python/ray/tests/test_autoscaler_e2e.py
@@ -2,11 +2,13 @@
 from ray.autoscaler._private.constants import AUTOSCALER_METRIC_PORT
 
 import ray
+import sys
 from ray._private.test_utils import (
     wait_for_condition,
     get_metric_check_condition,
 )
 from ray.cluster_utils import AutoscalingCluster
+from ray.autoscaler.node_launch_exception import NodeLaunchException
 
 
 def test_ray_status_e2e(shutdown_only):
@@ -111,8 +113,26 @@ def ping(self):
         cluster.shutdown()
 
 
+def test_node_launch_exception_serialization(shutdown_only):
+    ray.init(num_cpus=1)
+
+    exc_info = None
+    try:
+        raise Exception("Test exception.")
+    except Exception:
+        exc_info = sys.exc_info()
+    assert exc_info is not None
+
+    exc = NodeLaunchException("cat", "desc", exc_info)
+
+    after_serialization = ray.get(ray.put(exc))
+
+    assert after_serialization.category == exc.category
+    assert after_serialization.description == exc.description
+    assert after_serialization.src_exc_info is None
+
+
 if __name__ == "__main__":
-    import sys
     import os
     import pytest