From 972c5386550cda88399b549727b80ac88934eae0 Mon Sep 17 00:00:00 2001 From: Tim Waterson Date: Fri, 4 Oct 2024 13:16:06 +0100 Subject: [PATCH 1/2] Introduce a new UnreadablePickledException class, to allow exceptions that fail deserialisation to be easily identified downstream. --- distributed/core.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/distributed/core.py b/distributed/core.py index a4bb031c120..a403d7944e8 100644 --- a/distributed/core.py +++ b/distributed/core.py @@ -1690,6 +1690,10 @@ def error_message(e: BaseException, status: str = "error") -> ErrorMessage: } +class UnreadablePickledException(Exception): + pass + + def clean_exception( exception: BaseException | bytes | bytearray | str | None, traceback: types.TracebackType | bytes | str | None = None, @@ -1707,7 +1711,7 @@ def clean_exception( try: exception = protocol.pickle.loads(exception) except Exception: - exception = Exception(exception) + exception = UnreadablePickledException(exception) elif isinstance(exception, str): exception = Exception(exception) From bfccb5394890b59e46676acd727ae9a29343d3bc Mon Sep 17 00:00:00 2001 From: Tim Waterson Date: Fri, 4 Oct 2024 13:44:31 +0100 Subject: [PATCH 2/2] Add a test for UnreadablePickledException behaviour. --- distributed/tests/test_core.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py index 6f4836ae08d..1549925d17d 100644 --- a/distributed/tests/test_core.py +++ b/distributed/tests/test_core.py @@ -25,6 +25,7 @@ RPCClosed, Server, Status, + UnreadablePickledException, _expects_comm, clean_exception, coerce_to_address, @@ -1357,3 +1358,13 @@ async def test_large_payload(caplog): assert response["result"] == data await comm.close() + + +@gen_test() +async def test_unreadable_pickled_exceptions_kept(): + pickled_ex = b"\x80\x04\x954\x00\x00\x00\x00\x00\x00\x00\x8c\x08__main__\x94\x8c\x14SomeUnknownException\x94\x93\x94\x8c\x08some arg\x94\x85\x94R\x94." + ex_type, ex, tb = clean_exception(pickled_ex) + assert ex_type == UnreadablePickledException + assert isinstance(ex, UnreadablePickledException) + assert ex.args[0] == pickled_ex + assert tb is None