Move files

great_ai/utilities/parallel_map/__init__.py

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

great_ai/utilities/parallel_map/get_config.py

@@ -0,0 +1,56 @@
+import os
+from math import ceil
+from typing import Callable, Iterable, Optional, Sequence, Union
+
+from ..logger import get_logger
+from .parallel_map_configuration import ParallelMapConfiguration
+
+logger = get_logger("parallel_map")
+
+
+def get_config(
+    *,
+    function: Callable,
+    input_values: Union[Sequence, Iterable],
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+) -> ParallelMapConfiguration:
+
+    is_input_sequence = hasattr(input_values, "__len__")
+
+    if concurrency is None:
+        concurrency = len(os.sched_getaffinity(0))
+    assert concurrency >= 1, "At least one mapper process has to be created"
+
+    if chunk_size is None:
+        if is_input_sequence:
+            chunk_size = max(1, ceil(len(input_values) / concurrency / 10))
+        else:
+            raise ValueError(
+                "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`"
+            )
+    assert chunk_size >= 1, "Chunks have to contain at least one element"
+
+    chunk_count: Optional[int] = None
+    if is_input_sequence:
+        chunk_count = ceil(len(input_values) / chunk_size)
+        if chunk_count < concurrency:
+            logger.warning(
+                f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
+            )
+            concurrency = chunk_count
+
+    if concurrency == 1:
+        logger.warning("Running in series, there is no reason for parallelism")
+
+    input_length = len(input_values) if is_input_sequence else None
+
+    config = ParallelMapConfiguration(
+        concurrency=concurrency,
+        chunk_count=chunk_count,
+        chunk_size=chunk_size,
+        input_length=input_length,
+        function_name=function.__name__ if hasattr(function, "__name__") else "unknown",
+    )
+
+    return config

great_ai/utilities/parallel_map/manage_communication.py

@@ -0,0 +1,70 @@
+import multiprocessing as mp
+import queue
+import traceback
+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")
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def manage_communication(
+    *,
+    input_values: Iterable[T],
+    chunk_size: int,
+    input_queue: Union[mp.Queue, queue.Queue],
+    output_queue: Union[mp.Queue, queue.Queue],
+    unordered: bool,
+    ignore_exceptions: bool,
+) -> Iterable[V]:
+    chunks = iter(chunk(enumerate(input_values), chunk_size=chunk_size))
+    indexed_results: Dict[int, V] = {}
+    next_output_index = 0
+    read_input_length = 0
+    is_iteration_over = False
+    while not is_iteration_over or next_output_index < read_input_length:
+        if not is_iteration_over:
+            try:
+                next_chunk = next(chunks)
+                input_queue.put(next_chunk)
+                read_input_length += len(next_chunk)
+            except StopIteration:
+                is_iteration_over = True
+            except Exception as e:
+                if not ignore_exceptions:
+                    raise
+                else:
+                    logger.error(
+                        f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+                    )
+
+        try:
+            result_chunk = output_queue.get_nowait()
+
+            for index, value, exception in result_chunk:
+                if exception is not None:
+                    e, tb = exception
+                    if not ignore_exceptions:
+                        raise WorkerException from e
+                    else:
+                        logger.error(
+                            f"Exception {e} encountered in worker, traceback:\n{tb}"
+                        )
+                if unordered:
+                    yield value
+                    next_output_index += 1
+                else:
+                    indexed_results[index] = value
+
+            if not unordered:
+                while next_output_index in indexed_results:
+                    yield indexed_results[next_output_index]
+                    del indexed_results[next_output_index]
+                    next_output_index += 1
+        except queue.Empty:
+            pass

great_ai/utilities/parallel_map/manage_serial.py

@@ -0,0 +1,36 @@
+import traceback
+from typing import Callable, Iterable, TypeVar
+
+from ..logger import get_logger
+from .worker_exception import WorkerException
+
+logger = get_logger("parallel_map")
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def manage_serial(
+    *,
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    ignore_exceptions: bool,
+) -> Iterable[V]:
+    try:
+        for v in input_values:
+            try:
+                yield function(v)
+            except Exception as e:
+                if not ignore_exceptions:
+                    raise WorkerException from 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()}"
+            )

great_ai/utilities/parallel_map/mapper_function.py

@@ -0,0 +1,43 @@
+import multiprocessing as mp
+import queue
+import threading
+import traceback
+from typing import Callable, Union
+
+import dill
+
+
+def mapper_function(
+    input_queue: Union[mp.Queue, queue.Queue],
+    output_queue: Union[mp.Queue, queue.Queue],
+    should_stop: Union[mp.Event, threading.Event],
+    map_function: Union[bytes, Callable],
+) -> None:
+    try:
+        if isinstance(map_function, bytes):
+            map_function = dill.loads(map_function)
+
+        last_chunk = None
+        while not should_stop.wait(0.1):
+            if last_chunk is None:
+                try:
+                    input_chunk = input_queue.get_nowait()
+                    last_chunk = []
+                    for i, v in input_chunk:
+                        result, exception = None, None
+                        try:
+                            result = map_function(v)
+                        except Exception as e:
+                            exception = e, traceback.format_exc()
+                        last_chunk.append((i, result, exception))
+                except queue.Empty:
+                    pass
+
+            if last_chunk is not None:
+                try:
+                    output_queue.put_nowait(last_chunk)
+                    last_chunk = None
+                except queue.Full:
+                    pass
+    except (KeyboardInterrupt, BrokenPipeError):
+        pass

great_ai/utilities/parallel_map/parallel_map.py

@@ -0,0 +1,145 @@
+import multiprocessing as mp
+from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
+
+import dill
+
+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")
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: Optional[Literal[False]],
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_size: int,
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: Optional[Literal[False]],
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: True,
+) -> Iterable[Optional[V]]:
+    ...
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_size: int,
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: True,
+) -> Iterable[Optional[V]]:
+    ...
+
+
+def parallel_map(
+    function,
+    input_values,
+    *,
+    chunk_size=None,
+    concurrency=None,
+    unordered=False,
+    ignore_exceptions=False,
+):
+    config = get_config(
+        function=function,
+        input_values=input_values,
+        chunk_size=chunk_size,
+        concurrency=concurrency,
+    )
+
+    if config.concurrency == 1:
+        yield from manage_serial(
+            function=function,
+            input_values=input_values,
+            ignore_exceptions=ignore_exceptions,
+        )
+
+    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)
+
+    should_stop = ctx.Event()
+    serialized_map_function = dill.dumps(function, byref=True, recurse=True)
+
+    processes = [
+        ctx.Process(
+            name=f"parallel_map_{config.function_name}_{i}",
+            target=mapper_function,
+            daemon=True,
+            kwargs=dict(
+                input_queue=input_queue,
+                output_queue=output_queue,
+                should_stop=should_stop,
+                map_function=serialized_map_function,
+            ),
+        )
+        for i in range(config.concurrency)
+    ]
+
+    for p in processes:
+        p.start()
+
+    try:
+        yield from manage_communication(
+            input_values=input_values,
+            chunk_size=config.chunk_size,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            unordered=unordered,
+            ignore_exceptions=ignore_exceptions,
+        )
+        should_stop.set()
+    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()
+
+        manager.shutdown()

great_ai/utilities/parallel_map/parallel_map_configuration.py

@@ -0,0 +1,15 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from ..logger import get_logger
+
+logger = get_logger("parallel_map")
+
+
+class ParallelMapConfiguration(BaseModel):
+    concurrency: int
+    chunk_count: Optional[int]
+    chunk_size: int
+    input_length: Optional[int]
+    function_name: str

great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -0,0 +1,121 @@
+import queue
+import threading
+from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
+
+from .get_config import get_config
+from .manage_communication import manage_communication
+from .manage_serial import manage_serial
+from .mapper_function import mapper_function
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+@overload
+def threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: Optional[Literal[False]],
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_size: int,
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: Optional[Literal[False]],
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: True,
+) -> Iterable[Optional[V]]:
+    ...
+
+
+@overload
+def threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_size: int,
+    concurrency: Optional[int],
+    unordered: Optional[bool],
+    ignore_exceptions: True,
+) -> Iterable[Optional[V]]:
+    ...
+
+
+def threaded_parallel_map(
+    function,
+    input_values,
+    *,
+    chunk_size=None,
+    concurrency=None,
+    unordered=False,
+    ignore_exceptions=False,
+):
+    config = get_config(
+        function=function,
+        input_values=input_values,
+        chunk_size=chunk_size,
+        concurrency=concurrency,
+    )
+
+    if config.concurrency == 1:
+        yield from manage_serial(
+            function=function,
+            input_values=input_values,
+            ignore_exceptions=ignore_exceptions,
+        )
+
+    input_queue = queue.Queue(config.concurrency * 2)
+    output_queue = queue.Queue(config.concurrency * 2)
+    should_stop = threading.Event()
+
+    threads = [
+        threading.Thread(
+            name=f"threaded_parallel_map_{config.function_name}_{i}",
+            target=mapper_function,
+            daemon=True,
+            kwargs=dict(
+                input_queue=input_queue,
+                output_queue=output_queue,
+                should_stop=should_stop,
+                map_function=function,
+            ),
+        )
+        for i in range(config.concurrency)
+    ]
+
+    for t in threads:
+        t.start()
+
+    yield from manage_communication(
+        input_values=input_values,
+        chunk_size=config.chunk_size,
+        input_queue=input_queue,
+        output_queue=output_queue,
+        unordered=unordered,
+        ignore_exceptions=ignore_exceptions,
+    )
+    should_stop.set()
+    for t in threads:
+        t.join(1)

great_ai/utilities/parallel_map/worker_exception.py

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