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

Improve parallel map API

Browse source
This commit is contained in:
Andras Schmelczer 2022-06-30 16:54:01 +02:00
parent 5dc85905b4
commit 154cc52fa6
6 changed files with 94 additions and 84 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/great_ai/utilities/parallel_map/get_config.py
View file

@ -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