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

Improve parallel map

Browse source
This commit is contained in:
Andras Schmelczer 2022-06-29 09:14:35 +02:00
parent 3a8d357e3c
commit 3e393bf362
3 changed files with 30 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/great_ai/utilities/parallel_map/parallel_map.py
View file

@ -68,12 +68,13 @@ def parallel_map(
desc="Parallel map",
disable=disable_progress_bar,
total=config.input_length,
miniters=1,
)
]
ctx = mp.get_context("spawn")
start_methods = mp.get_all_start_methods()
ctx = mp.get_context("fork") if "fork" in start_methods else mp.get_context("spawn")
ctx.freeze_support()
input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
output_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
should_stop = ctx.Event()
@ -99,7 +100,6 @@ def parallel_map(
desc="Parallel map",
disable=disable_progress_bar,
total=config.input_length,
miniters=1,
)
chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length))

2
src/great_ai/utilities/parallel_map/parallel_map_configuration.py
View file

@ -14,7 +14,7 @@ class ParallelMapConfiguration(BaseModel):
input_length: Optional[int]
serialized_map_function: bytes
def pretty_print(self, prefix=" ⚙️"):
def pretty_print(self, prefix=" ⚙️ "):
logger.info(f"{prefix} concurrency: {self.concurrency}")
logger.info(f"{prefix} chunk length: {self.chunk_length}")
logger.info(