andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Fix typing and minor issues

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-07 14:12:44 +02:00
parent 2db2253578
commit 72ab627a34
54 changed files with 635 additions and 589 deletions
Download patch file Download diff file Expand all files Collapse all files

112
great_ai/utilities/parallel_map/parallel_map.py
View file

@ -1,11 +1,20 @@
import multiprocessing as mp
from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
from typing import (
Awaitable,
Callable,
Iterable,
Literal,
Optional,
Sequence,
TypeVar,
Union,
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
@ -15,79 +24,72 @@ V = TypeVar("V")
@overload
def parallel_map(
function: Callable[[T], V],
func: Callable[[T], Union[V, Awaitable[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,
ignore_exceptions: Literal[True],
chunk_size: Optional[int] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def parallel_map(
function: Callable[[T], V],
input_values: Iterable[T],
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: True,
ignore_exceptions: Literal[True],
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def parallel_map(
function,
input_values,
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Sequence[T],
*,
chunk_size=None,
concurrency=None,
unordered=False,
ignore_exceptions=False,
):
chunk_size: Optional[int] = ...,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
@overload
def parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
def parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: Optional[int] = None,
ignore_exceptions: bool = False,
concurrency: Optional[int] = None,
unordered: bool = False,
) -> Iterable[Optional[V]]:
config = get_config(
function=function,
function=func,
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()
@ -99,7 +101,7 @@ def parallel_map(
output_queue = manager.Queue(config.concurrency * 2)
should_stop = ctx.Event()
serialized_map_function = dill.dumps(function, byref=True, recurse=False)
serialized_map_function = dill.dumps(func, byref=True, recurse=False)
processes = [
ctx.Process(
@ -110,7 +112,7 @@ def parallel_map(
input_queue=input_queue,
output_queue=output_queue,
should_stop=should_stop,
map_function=serialized_map_function,
func=serialized_map_function,
),
)
for i in range(config.concurrency)