From 8cd3449cffaf45226f942fc511ac36e47718759b Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sat, 16 Jul 2022 12:27:05 +0200
Subject: [PATCH] Fix exception serialisation error

---
 .../utilities/parallel_map/manage_communication.py     | 10 +++++++---
 great_ai/utilities/parallel_map/map_result.py          |  2 +-
 great_ai/utilities/parallel_map/mapper_function.py     |  4 +++-
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/great_ai/utilities/parallel_map/manage_communication.py b/great_ai/utilities/parallel_map/manage_communication.py
index d6c0740..4d53d2b 100644
--- a/great_ai/utilities/parallel_map/manage_communication.py
+++ b/great_ai/utilities/parallel_map/manage_communication.py
@@ -40,7 +40,9 @@ def manage_communication(
             except Exception as e:
                 if ignore_exceptions:
                     logger.error(
-                        f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+                        f"""Exception {e} encountered in input, traceback:\n{
+                            traceback.format_exc()
+                        }"""
                     )
                 else:
                     raise
@@ -51,10 +53,12 @@ def manage_communication(
                 if r.exception is not None:
                     if ignore_exceptions:
                         logger.error(
-                            f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}"
+                            f"""Exception {
+                                r.exception
+                            } encountered in worker, traceback:\n{r.worker_traceback}"""
                         )
                     else:
-                        raise WorkerException from r.exception
+                        raise WorkerException(r.exception)
 
                 if unordered:
 
diff --git a/great_ai/utilities/parallel_map/map_result.py b/great_ai/utilities/parallel_map/map_result.py
index 81d3dbb..481ce5d 100644
--- a/great_ai/utilities/parallel_map/map_result.py
+++ b/great_ai/utilities/parallel_map/map_result.py
@@ -4,5 +4,5 @@ from typing import Any, NamedTuple, Optional
 class MapResult(NamedTuple):
     order: int
     value: Any
-    exception: Optional[Exception] = None
+    exception: Optional[str] = None
     worker_traceback: Optional[str] = None
diff --git a/great_ai/utilities/parallel_map/mapper_function.py b/great_ai/utilities/parallel_map/mapper_function.py
index 921065a..cc4e57b 100644
--- a/great_ai/utilities/parallel_map/mapper_function.py
+++ b/great_ai/utilities/parallel_map/mapper_function.py
@@ -57,7 +57,9 @@ def mapper_function(
                                 last_chunk.append(MapResult(i, func(value)))
                             except Exception as e:
                                 last_chunk.append(
-                                    MapResult(i, None, e, traceback.format_exc())
+                                    # `e` has to be stringified in order to avoid any
+                                    # surprising serialization error when returning it
+                                    MapResult(i, None, str(e), traceback.format_exc())
                                 )
                 except queue.Empty:
                     pass