Improve parallel map API
This commit is contained in:
parent
5dc85905b4
commit
154cc52fa6
6 changed files with 94 additions and 84 deletions
|
|
@ -103,7 +103,7 @@ def configure(
|
|||
|
||||
logger.info("Settings: configured ✅")
|
||||
for k, v in get_context().to_flat_dict().items():
|
||||
logger.info(f"🔩 {k}: {v}")
|
||||
logger.info(f" ⚙️ {k}: {v}")
|
||||
|
||||
if not is_production and not disable_se4ml_banner:
|
||||
logger.warning(
|
||||
|
|
|
|||
|
|
@ -3,13 +3,13 @@ from typing import Iterable, List, TypeVar
|
|||
T = TypeVar("T")
|
||||
|
||||
|
||||
def chunk(values: Iterable[T], chunk_length: int) -> Iterable[T]:
|
||||
assert chunk_length >= 1
|
||||
def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]:
|
||||
assert chunk_size >= 1
|
||||
|
||||
result: List[T] = []
|
||||
for v in values:
|
||||
result.append(v)
|
||||
if len(result) == chunk_length:
|
||||
if len(result) == chunk_size:
|
||||
yield result
|
||||
result = []
|
||||
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ def get_config(
|
|||
*,
|
||||
function: Callable,
|
||||
input_values: Union[Sequence, Iterable],
|
||||
chunk_length: Optional[int],
|
||||
chunk_size: Optional[int],
|
||||
concurrency: Optional[int],
|
||||
disable_logging: bool,
|
||||
) -> ParallelMapConfiguration:
|
||||
|
||||
is_input_sequence = hasattr(input_values, "__len__")
|
||||
|
|
@ -24,18 +25,18 @@ def get_config(
|
|||
concurrency = len(os.sched_getaffinity(0))
|
||||
assert concurrency >= 1, "At least one mapper process has to be created"
|
||||
|
||||
if chunk_length is None:
|
||||
if chunk_size is None:
|
||||
if is_input_sequence:
|
||||
chunk_length = max(1, ceil(len(input_values) / concurrency / 10))
|
||||
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_length`"
|
||||
"The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`"
|
||||
)
|
||||
assert chunk_length >= 1, "Chunks have to contain at least one element"
|
||||
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_length)
|
||||
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"
|
||||
|
|
@ -47,16 +48,17 @@ def get_config(
|
|||
|
||||
input_length = len(input_values) if is_input_sequence else None
|
||||
|
||||
serialized_map_function = dill.dumps(function, byref=True, recurse=True)
|
||||
|
||||
logger.info("Parallel map: configured ✅")
|
||||
config = ParallelMapConfiguration(
|
||||
concurrency=concurrency,
|
||||
chunk_count=chunk_count,
|
||||
chunk_length=chunk_length,
|
||||
chunk_size=chunk_size,
|
||||
input_length=input_length,
|
||||
serialized_map_function=serialized_map_function,
|
||||
serialized_map_function=dill.dumps(function, byref=True, recurse=True),
|
||||
function_name=function.__name__,
|
||||
)
|
||||
|
||||
config.pretty_print()
|
||||
if not disable_logging:
|
||||
logger.info("Parallel map: configured ✅")
|
||||
config.pretty_print()
|
||||
|
||||
return config
|
||||
|
|
|
|||
|
|
@ -1,17 +1,8 @@
|
|||
import multiprocessing as mp
|
||||
import queue
|
||||
from typing import (
|
||||
Callable,
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
TypeVar,
|
||||
overload,
|
||||
)
|
||||
from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload
|
||||
|
||||
from tqdm.auto import tqdm
|
||||
from tqdm.cli import tqdm
|
||||
|
||||
from ..chunk import chunk
|
||||
from .get_config import get_config
|
||||
|
|
@ -26,10 +17,11 @@ def parallel_map(
|
|||
function: Callable[[T], V],
|
||||
input_values: Sequence[T],
|
||||
*,
|
||||
chunk_length: Optional[int],
|
||||
chunk_size: Optional[int],
|
||||
concurrency: Optional[int],
|
||||
disable_progress_bar: bool,
|
||||
) -> List[V]:
|
||||
disable_logging: bool,
|
||||
unordered: Optional[bool],
|
||||
) -> Iterable[V]:
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -38,10 +30,11 @@ def parallel_map(
|
|||
function: Callable[[T], V],
|
||||
input_values: Iterable[T],
|
||||
*,
|
||||
chunk_length: int,
|
||||
chunk_size: int,
|
||||
concurrency: Optional[int],
|
||||
disable_progress_bar: bool,
|
||||
) -> List[V]:
|
||||
disable_logging: bool,
|
||||
unordered: Optional[bool],
|
||||
) -> Iterable[V]:
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -49,30 +42,32 @@ def parallel_map(
|
|||
function,
|
||||
input_values,
|
||||
*,
|
||||
chunk_length=None,
|
||||
chunk_size=None,
|
||||
concurrency=None,
|
||||
disable_progress_bar=False,
|
||||
disable_logging=False,
|
||||
unordered=False,
|
||||
):
|
||||
config = get_config(
|
||||
function=function,
|
||||
input_values=input_values,
|
||||
chunk_length=chunk_length,
|
||||
chunk_size=chunk_size,
|
||||
concurrency=concurrency,
|
||||
disable_logging=disable_logging,
|
||||
)
|
||||
|
||||
tqdm_options = dict(
|
||||
desc=f"Parallel map {config.function_name}",
|
||||
disable=disable_logging,
|
||||
total=config.input_length,
|
||||
miniters=1,
|
||||
dynamic_ncols=True,
|
||||
)
|
||||
|
||||
if config.concurrency == 1:
|
||||
return [
|
||||
function(v)
|
||||
for v in tqdm(
|
||||
input_values,
|
||||
desc="Parallel map",
|
||||
disable=disable_progress_bar,
|
||||
total=config.input_length,
|
||||
)
|
||||
]
|
||||
yield from (function(v) for v in tqdm(input_values, **tqdm_options))
|
||||
return
|
||||
|
||||
start_methods = mp.get_all_start_methods()
|
||||
ctx = mp.get_context("fork") if "fork" in start_methods else mp.get_context("spawn")
|
||||
ctx = mp.get_context("spawn")
|
||||
ctx.freeze_support()
|
||||
|
||||
input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
|
||||
|
|
@ -81,7 +76,7 @@ def parallel_map(
|
|||
|
||||
processes = [
|
||||
ctx.Process(
|
||||
name=f"parallel_map_{i}",
|
||||
name=f"parallel_map_{config.function_name}_{i}",
|
||||
target=mapper_function,
|
||||
kwargs=dict(
|
||||
input_queue=input_queue,
|
||||
|
|
@ -96,18 +91,15 @@ def parallel_map(
|
|||
for p in processes:
|
||||
p.start()
|
||||
|
||||
progress = tqdm(
|
||||
desc="Parallel map",
|
||||
disable=disable_progress_bar,
|
||||
total=config.input_length,
|
||||
)
|
||||
progress = tqdm(**tqdm_options)
|
||||
|
||||
chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length))
|
||||
indexed_results: List[Tuple[int, V]] = []
|
||||
chunks = iter(chunk(enumerate(input_values), chunk_size=config.chunk_size))
|
||||
indexed_results: Dict[int, V] = {}
|
||||
next_output_index = 0
|
||||
read_input_length = 0
|
||||
is_iteration_over = False
|
||||
try:
|
||||
while not is_iteration_over or len(indexed_results) < read_input_length:
|
||||
while not is_iteration_over or next_output_index < read_input_length:
|
||||
if not is_iteration_over:
|
||||
try:
|
||||
next_chunk = next(chunks)
|
||||
|
|
@ -118,10 +110,23 @@ def parallel_map(
|
|||
|
||||
try:
|
||||
result_chunk = output_queue.get_nowait()
|
||||
indexed_results.extend(result_chunk)
|
||||
progress.update(len(result_chunk))
|
||||
|
||||
for index, value in result_chunk:
|
||||
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
|
||||
|
||||
should_stop.set()
|
||||
except KeyboardInterrupt:
|
||||
for p in processes:
|
||||
|
|
@ -134,7 +139,3 @@ def parallel_map(
|
|||
output_queue.close()
|
||||
|
||||
progress.close()
|
||||
|
||||
results = [v for _, v in sorted(indexed_results)]
|
||||
|
||||
return results
|
||||
|
|
|
|||
|
|
@ -10,13 +10,14 @@ logger = get_logger("parallel_map")
|
|||
class ParallelMapConfiguration(BaseModel):
|
||||
concurrency: int
|
||||
chunk_count: Optional[int]
|
||||
chunk_length: int
|
||||
chunk_size: int
|
||||
input_length: Optional[int]
|
||||
serialized_map_function: bytes
|
||||
function_name: str
|
||||
|
||||
def pretty_print(self, prefix=" ⚙️ "):
|
||||
logger.info(f"{prefix} concurrency: {self.concurrency}")
|
||||
logger.info(f"{prefix} chunk length: {self.chunk_length}")
|
||||
logger.info(f"{prefix} chunk size: {self.chunk_size}")
|
||||
logger.info(
|
||||
f"{prefix} chunk count: {self.chunk_count if self.chunk_count else 'unknown'}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import unittest
|
||||
|
||||
import pytest
|
||||
|
||||
from src.great_ai.utilities import parallel_map
|
||||
|
||||
COUNT = int(1e5) + 3
|
||||
|
|
@ -10,7 +12,7 @@ class TestParallelMap(unittest.TestCase):
|
|||
inputs = range(COUNT)
|
||||
expected = [v**2 for v in range(COUNT)]
|
||||
|
||||
assert parallel_map(lambda v: v**2, inputs, concurrency=10) == expected
|
||||
assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected
|
||||
|
||||
def test_with_iterable(self) -> None:
|
||||
from time import sleep
|
||||
|
|
@ -23,38 +25,42 @@ class TestParallelMap(unittest.TestCase):
|
|||
expected = [v**3 for v in range(10)]
|
||||
|
||||
assert (
|
||||
parallel_map(lambda x: x**3, my_generator(), chunk_length=1) == expected
|
||||
list(parallel_map(lambda x: x**3, my_generator(), chunk_size=1))
|
||||
== expected
|
||||
)
|
||||
|
||||
def test_simple_case_without_progress_bar(self) -> None:
|
||||
inputs = range(COUNT)
|
||||
expected = [v**2 for v in range(COUNT)]
|
||||
|
||||
self.assertEqual(
|
||||
parallel_map(lambda v: v**2, inputs, disable_progress_bar=True), expected
|
||||
assert (
|
||||
list(parallel_map(lambda v: v**2, inputs, disable_logging=True))
|
||||
== expected
|
||||
)
|
||||
|
||||
def test_simple_case_invalid_values(self) -> None:
|
||||
inputs = range(COUNT)
|
||||
|
||||
self.assertRaises(
|
||||
AssertionError, parallel_map, lambda v: v**2, inputs, concurrency=0
|
||||
)
|
||||
self.assertRaises(
|
||||
AssertionError, parallel_map, lambda v: v**2, inputs, chunk_length=0
|
||||
)
|
||||
with pytest.raises(AssertionError):
|
||||
list(parallel_map(lambda v: v**2, inputs, concurrency=0))
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
list(parallel_map(lambda v: v**2, inputs, chunk_size=0))
|
||||
|
||||
def test_no_op(self) -> None:
|
||||
assert parallel_map(lambda v: v**2, [], disable_progress_bar=True) == []
|
||||
self.assertEqual(
|
||||
parallel_map(
|
||||
lambda v: v**2, [], disable_progress_bar=True, chunk_length=100
|
||||
),
|
||||
[],
|
||||
assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == []
|
||||
|
||||
assert (
|
||||
list(
|
||||
parallel_map(lambda v: v**2, [], disable_logging=True, chunk_size=100)
|
||||
)
|
||||
== []
|
||||
)
|
||||
self.assertEqual(
|
||||
parallel_map(
|
||||
lambda v: v**2, [], disable_progress_bar=True, concurrency=100
|
||||
),
|
||||
[],
|
||||
assert (
|
||||
list(
|
||||
parallel_map(
|
||||
lambda v: v**2, [], disable_logging=True, concurrency=100
|
||||
)
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue