From 716718ca78ad3d70ff7b1ed06035f28f49372790 Mon Sep 17 00:00:00 2001
From: Daniel Shields <Daniel.Shields@twosigma.com>
Date: Fri, 9 Feb 2024 10:47:50 -0600
Subject: [PATCH 1/3] Fixed NodeError and CallError not being serializable and
 added test cases.

---
 src/uberjob/_errors.py         | 6 ++++++
 src/uberjob/_util/traceback.py | 4 ++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/uberjob/_errors.py b/src/uberjob/_errors.py
index 3ba2409..ce49693 100644
--- a/src/uberjob/_errors.py
+++ b/src/uberjob/_errors.py
@@ -27,6 +27,9 @@ def __init__(self, node: Node):
         )
         self.node = node
 
+    def __reduce__(self):
+        return NodeError, (self.node,)
+
 
 class CallError(Exception):
     """
@@ -46,6 +49,9 @@ def __init__(self, call: Call):
         )
         self.call = call
 
+    def __reduce__(self):
+        return CallError, (self.call,)
+
 
 class NotTransformedError(Exception):
     """An expected transformation was not applied."""
diff --git a/src/uberjob/_util/traceback.py b/src/uberjob/_util/traceback.py
index c342665..5e7a1f4 100644
--- a/src/uberjob/_util/traceback.py
+++ b/src/uberjob/_util/traceback.py
@@ -81,7 +81,7 @@ def recurse(frame, depth):
 def render_symbolic_traceback(stack_frame):
     stack_frames = []
     while stack_frame:
-        if stack_frame is TruncatedStackFrame:
+        if type(stack_frame) is TruncatedStackFrameType:
             stack_frames.append(stack_frame)
             break
         if "/IPython/core/" in stack_frame.path:
@@ -90,7 +90,7 @@ def render_symbolic_traceback(stack_frame):
         stack_frame = stack_frame.outer
 
     def format_stack_frame(s):
-        if s is TruncatedStackFrame:
+        if type(s) is TruncatedStackFrameType:
             return "  ... truncated"
         return f'  File "{s.path}", line {s.line}, in {s.name}'
 

From 5df87bda01b2252bf4b133b33e6ce885a3588c5c Mon Sep 17 00:00:00 2001
From: Daniel Shields <Daniel.Shields@twosigma.com>
Date: Fri, 9 Feb 2024 10:52:24 -0600
Subject: [PATCH 2/3] added test cases

---
 tests/test_plan.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tests/test_plan.py b/tests/test_plan.py
index 459f8c1..50a27d2 100644
--- a/tests/test_plan.py
+++ b/tests/test_plan.py
@@ -16,12 +16,14 @@
 import operator
 import os
 import pathlib
+import pickle
 import re
 import tempfile
 
 import networkx as nx
 
 import uberjob
+from uberjob._errors import NodeError
 from uberjob._util.traceback import get_stack_frame
 from uberjob.graph import Call
 from uberjob.progress import console_progress, default_progress, html_progress
@@ -373,3 +375,33 @@ def wrapper(*args, **kwargs):
             expected_exception_chain_traceback_summary=[["wrapper"]]
         ):
             uberjob.run(plan, output=call, retry=bad_retry2)
+
+    def test_serialize_call_error(self):
+        plan = uberjob.Plan()
+        call = plan.call(operator.truediv, 1, 0)
+        exception = None
+        try:
+            uberjob.run(plan, output=call)
+        except uberjob.CallError as e:
+            exception = e
+        self.assertIsNotNone(exception)
+        pickled_exception = pickle.dumps(exception)
+        unpickled_exception = pickle.loads(pickled_exception)
+        self.assertIsInstance(unpickled_exception, uberjob.CallError)
+        self.assertIsInstance(unpickled_exception.call, Call)
+        self.assertIs(unpickled_exception.call.fn, operator.truediv)
+
+    def test_serialize_node_error(self):
+        plan = uberjob.Plan()
+        call = plan.call(pow, 2, 2)
+        exception = None
+        try:
+            raise NodeError(call)
+        except NodeError as e:
+            exception = e
+        self.assertIsNotNone(exception)
+        pickled_exception = pickle.dumps(exception)
+        unpickled_exception = pickle.loads(pickled_exception)
+        self.assertIsInstance(unpickled_exception, NodeError)
+        self.assertIsInstance(unpickled_exception.node, Call)
+        self.assertIs(unpickled_exception.node.fn, pow)

From 1d1ba0bfc8234b56423a9a00abe1c7ccb7850acb Mon Sep 17 00:00:00 2001
From: Daniel Shields <Daniel.Shields@twosigma.com>
Date: Fri, 9 Feb 2024 11:20:26 -0600
Subject: [PATCH 3/3] update

---
 src/uberjob/_util/traceback.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/uberjob/_util/traceback.py b/src/uberjob/_util/traceback.py
index 5e7a1f4..9c7d641 100644
--- a/src/uberjob/_util/traceback.py
+++ b/src/uberjob/_util/traceback.py
@@ -48,10 +48,18 @@ def __repr__(self):
         )
 
 
+TruncatedStackFrame = None
+
+
 class TruncatedStackFrameType:
     def __repr__(self):
         return "TruncatedStackFrame"
 
+    def __new__(cls, *args, **kwargs):
+        if TruncatedStackFrame is not None:
+            return TruncatedStackFrame
+        return super().__new__(cls, *args, **kwargs)
+
 
 TruncatedStackFrame = TruncatedStackFrameType()
 
@@ -81,7 +89,7 @@ def recurse(frame, depth):
 def render_symbolic_traceback(stack_frame):
     stack_frames = []
     while stack_frame:
-        if type(stack_frame) is TruncatedStackFrameType:
+        if stack_frame is TruncatedStackFrame:
             stack_frames.append(stack_frame)
             break
         if "/IPython/core/" in stack_frame.path:
@@ -90,7 +98,7 @@ def render_symbolic_traceback(stack_frame):
         stack_frame = stack_frame.outer
 
     def format_stack_frame(s):
-        if type(s) is TruncatedStackFrameType:
+        if s is TruncatedStackFrame:
             return "  ... truncated"
         return f'  File "{s.path}", line {s.line}, in {s.name}'