Fix hanging bug by adding manager

src/great_ai/utilities/__init__.py

@@ -6,6 +6,6 @@ from .get_sentences import get_sentences
 from .language import english_name_of_language, is_english, predict_language
 from .logger import get_logger
 from .match_names import match_names
-from .parallel_map import parallel_map, threaded_parallel_map
+from .parallel_map import WorkerException, parallel_map, threaded_parallel_map
 from .unchunk import unchunk
 from .unique import unique

src/great_ai/utilities/parallel_map/__init__.py

@@ -1,2 +1,3 @@
 from .parallel_map import parallel_map
 from .threaded_parallel_map import threaded_parallel_map
+from .worker_exception import WorkerException

src/great_ai/utilities/parallel_map/manage_communication.py

@@ -5,6 +5,7 @@ from typing import Dict, Iterable, TypeVar, Union
 
 from ..chunk import chunk
 from ..logger import get_logger
+from .worker_exception import WorkerException
 
 logger = get_logger("parallel_map")
 
@@ -36,7 +37,7 @@ def manage_communication(
                 is_iteration_over = True
             except Exception as e:
                 if not ignore_exceptions:
-                    raise e
+                    raise
                 else:
                     logger.error(
                         f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
@@ -49,7 +50,7 @@ def manage_communication(
                 if exception is not None:
                     e, tb = exception
                     if not ignore_exceptions:
-                        raise e
+                        raise WorkerException from e
                     else:
                         logger.error(
                             f"Exception {e} encountered in worker, traceback:\n{tb}"

src/great_ai/utilities/parallel_map/manage_serial.py

@@ -2,6 +2,7 @@ import traceback
 from typing import Callable, Iterable, TypeVar
 
 from ..logger import get_logger
+from .worker_exception import WorkerException
 
 logger = get_logger("parallel_map")
 
@@ -21,15 +22,15 @@ def manage_serial(
                 yield function(v)
             except Exception as e:
                 if not ignore_exceptions:
-                    raise e
+                    raise WorkerException from e
-                else:
-                    logger.error(
-                        f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
-                    )
-    except Exception as e:
-        if not ignore_exceptions:
-            raise e
                 else:
                     logger.error(
                         f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}"
                     )
+    except Exception as e:
+        if not ignore_exceptions:
+            raise
+        else:
+            logger.error(
+                f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+            )

src/great_ai/utilities/parallel_map/mapper_function.py

@@ -2,7 +2,6 @@ import multiprocessing as mp
 import queue
 import threading
 import traceback
-from time import sleep
 from typing import Callable, Union
 
 import dill
@@ -19,7 +18,7 @@ def mapper_function(
             map_function = dill.loads(map_function)
 
         last_chunk = None
-        while not should_stop.is_set():
+        while not should_stop.wait(0.1):
             if last_chunk is None:
                 try:
                     input_chunk = input_queue.get_nowait()
@@ -32,16 +31,13 @@ def mapper_function(
                             exception = e, traceback.format_exc()
                         last_chunk.append((i, result, exception))
                 except queue.Empty:
-                    sleep(
+                    pass
-                        0.1
+
-                    )  # input_queue.get(0.1) would hang in some cases with multiprocessing
+            if last_chunk is not None:
-            else:
                 try:
                     output_queue.put_nowait(last_chunk)
                     last_chunk = None
                 except queue.Full:
-                    sleep(
+                    pass
-                        0.1
-                    )  # output_queue.put(0.1) would hang in some cases with multiprocessing
     except (KeyboardInterrupt, BrokenPipeError):
         pass

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -7,6 +7,7 @@ from .get_config import get_config
 from .manage_communication import manage_communication
 from .manage_serial import manage_serial
 from .mapper_function import mapper_function
+from .worker_exception import WorkerException
 
 T = TypeVar("T")
 V = TypeVar("V")
@@ -87,13 +88,17 @@ def parallel_map(
             ignore_exceptions=ignore_exceptions,
         )
 
-    ctx = mp.get_context("spawn")
+    ctx = (
+        mp.get_context("fork")
+        if "fork" in mp.get_all_start_methods()
+        else mp.get_context("spawn")
+    )
     ctx.freeze_support()
+    manager = ctx.Manager()
+    input_queue = manager.Queue(config.concurrency * 2)
+    output_queue = manager.Queue(config.concurrency * 2)
 
-    input_queue = ctx.Queue(config.concurrency * 2)
-    output_queue = ctx.Queue(config.concurrency * 2)
     should_stop = ctx.Event()
-
     serialized_map_function = dill.dumps(function, byref=True, recurse=True)
 
     processes = [
@@ -124,13 +129,17 @@ def parallel_map(
             ignore_exceptions=ignore_exceptions,
         )
         should_stop.set()
-    except:
+    except WorkerException:
+        should_stop.set()
+        raise
+    except Exception:
         for p in processes:
             p.terminate()
+            p.kill()
         raise
     finally:
         for p in processes:
             p.join()  # terminated processes have to be joined else they remain zombies
             p.close()
-        input_queue.close()
+
-        output_queue.close()
+        manager.shutdown()

src/great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -118,4 +118,4 @@ def threaded_parallel_map(
     )
     should_stop.set()
     for t in threads:
-        t.join()
+        t.join(1)

src/great_ai/utilities/parallel_map/worker_exception.py

@@ -0,0 +1,2 @@
+class WorkerException(Exception):
+    pass

tests/utilities/test_parallel_map.py

@@ -2,7 +2,7 @@ import unittest
 
 import pytest
 
-from src.great_ai.utilities import parallel_map
+from src.great_ai.utilities import WorkerException, parallel_map
 
 COUNT = int(1e5) + 3
 
@@ -91,10 +91,10 @@ class TestParallelMap(unittest.TestCase):
         def oh_no(_):
             raise ValueError("hi")
 
-        with pytest.raises(ValueError):
+        with pytest.raises(WorkerException):
             list(parallel_map(oh_no, range(COUNT), concurrency=2))
 
-        with pytest.raises(ValueError):
+        with pytest.raises(WorkerException):
             list(parallel_map(oh_no, range(COUNT), concurrency=1))
 
     def test_ignore_worker_process_exception(self) -> None:

tests/utilities/test_threaded_parallel_map.py

@@ -2,7 +2,7 @@ import unittest
 
 import pytest
 
-from src.great_ai.utilities import threaded_parallel_map
+from src.great_ai.utilities import WorkerException, threaded_parallel_map
 
 COUNT = int(1e5) + 3
 
@@ -91,10 +91,10 @@ class TestParallelMap(unittest.TestCase):
         def oh_no(_):
             raise ValueError("hi")
 
-        with pytest.raises(ValueError):
+        with pytest.raises(WorkerException):
             list(threaded_parallel_map(oh_no, range(COUNT), concurrency=2))
 
-        with pytest.raises(ValueError):
+        with pytest.raises(WorkerException):
             list(threaded_parallel_map(oh_no, range(COUNT), concurrency=1))
 
     def test_ignore_worker_worker_exception(self) -> None: