Add threaded parallel_map
This commit is contained in:
parent
154cc52fa6
commit
0ff0b49a79
9 changed files with 225 additions and 6 deletions
74
tests/utilities/test_threaded_parallel_map.py
Normal file
74
tests/utilities/test_threaded_parallel_map.py
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
import unittest
|
||||
|
||||
import pytest
|
||||
|
||||
from src.great_ai.utilities import threaded_parallel_map
|
||||
|
||||
COUNT = int(1e5) + 3
|
||||
|
||||
|
||||
class TestThreadedParallelMap(unittest.TestCase):
|
||||
def test_simple_case_with_progress_bar(self) -> None:
|
||||
inputs = range(COUNT)
|
||||
expected = [v**2 for v in range(COUNT)]
|
||||
|
||||
assert (
|
||||
list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=10))
|
||||
== expected
|
||||
)
|
||||
|
||||
def test_with_iterable(self) -> None:
|
||||
from time import sleep
|
||||
|
||||
def my_generator():
|
||||
for i in range(10):
|
||||
yield i
|
||||
sleep(0.1)
|
||||
|
||||
expected = [v**3 for v in range(10)]
|
||||
|
||||
assert (
|
||||
list(threaded_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)]
|
||||
|
||||
assert (
|
||||
list(threaded_parallel_map(lambda v: v**2, inputs, disable_logging=True))
|
||||
== expected
|
||||
)
|
||||
|
||||
def test_simple_case_invalid_values(self) -> None:
|
||||
inputs = range(COUNT)
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=0))
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
list(threaded_parallel_map(lambda v: v**2, inputs, chunk_size=0))
|
||||
|
||||
def test_no_op(self) -> None:
|
||||
assert (
|
||||
list(threaded_parallel_map(lambda v: v**2, [], disable_logging=True))
|
||||
== []
|
||||
)
|
||||
|
||||
assert (
|
||||
list(
|
||||
threaded_parallel_map(
|
||||
lambda v: v**2, [], disable_logging=True, chunk_size=100
|
||||
)
|
||||
)
|
||||
== []
|
||||
)
|
||||
assert (
|
||||
list(
|
||||
threaded_parallel_map(
|
||||
lambda v: v**2, [], disable_logging=True, concurrency=100
|
||||
)
|
||||
)
|
||||
== []
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue