Merge pull request #46 from python-thread/perf/memory-optimisation

Perf: memory optimisation for ParallelProcessing

Author: caffeine-addictt

src/thread/thread.py

@@ -32,7 +32,7 @@ class ParallelProcessing: ...
   HookFunction,
 )
 from typing_extensions import Generic, ParamSpec
-from typing import List, Callable, Optional, Union, Mapping, Sequence, Tuple
+from typing import List, Callable, Optional, Union, Mapping, Sequence, Tuple, Generator
 
 
 Threads: set['Thread'] = set()
@@ -387,15 +387,19 @@ def _wrap_function(self, function: TargetFunction) -> TargetFunction:
     @wraps(function)
     def wrapper(
       index: int,
-      data_chunk: Sequence[_Dataset_T],
+      length: int,
+      data_chunk: Generator[_Dataset_T, None, None],
       *args: _Target_P.args,
       **kwargs: _Target_P.kwargs,
       ) -> List[_Target_T]:
       computed: List[Data_Out] = []
-      for i, data_entry in enumerate(data_chunk):
+
+      i = 0
+      for data_entry in data_chunk:
         v = function(data_entry, *args, **kwargs)
         computed.append(v)
-        self._threads[index].progress = round((i + 1) / len(data_chunk), 5)
+        self._threads[index].progress = round((i + 1) / length, 5)
+        i += 1
 
       self._completed += 1
       if self._completed == len(self._threads):
@@ -507,15 +511,23 @@ def start(self) -> None:
       i: v for i, v in self.overflow_kwargs.items() if i != 'name' and i != 'args'
     }
 
-    for i, data_chunk in enumerate(chunk_split(self.dataset, max_threads)):
+    i = 0
+    for chunkStart, chunkEnd in chunk_split(len(self.dataset), max_threads):
       chunk_thread = Thread(
         target=self.function,
-        args=[i, data_chunk, *parsed_args, *self.overflow_args],
+        args=[
+          i,
+          chunkEnd - chunkStart,
+          (self.dataset[x] for x in range(chunkStart, chunkEnd)),
+          *parsed_args,
+          *self.overflow_args,
+        ],
         name=name_format and name_format % i or None,
         **self.overflow_kwargs,
       )
       self._threads.append(_ThreadWorker(chunk_thread, 0))
       chunk_thread.start()
+      i += 1
 
 
 # Handle abrupt exit

src/thread/utils/algorithm.py

@@ -8,10 +8,10 @@
   |_ b.py
 """
 
-from typing import List, Sequence, Any
+from typing import List, Tuple
 
 
-def chunk_split(dataset: Sequence[Any], number_of_chunks: int) -> List[List[Any]]:
+def chunk_split(dataset_length: int, number_of_chunks: int) -> List[Tuple[int, int]]:
   """
   Splits a dataset into balanced chunks
 
@@ -21,33 +21,32 @@ def chunk_split(dataset: Sequence[Any], number_of_chunks: int) -> List[List[Any]
 
   Parameters
   ----------
-  :param dataset: This should be the dataset you want to split into chunks
+  :param dataset_length: This should be the length of the dataset you want to split into chunks
   :param number_of_chunks: The should be the number of chunks it will attempt to split into
 
 
   Returns
   -------
-  :returns list[list[Any]]: The split dataset
+  :returns list[tuple[int, int]]: The chunked dataset slices
 
   Raises
   ------
   AssertionError: The number of chunks specified is larger than the dataset size
   """
-  length = len(dataset)
   assert (
-    length >= number_of_chunks
+    dataset_length >= number_of_chunks
   ), 'The number of chunks specified is larger than the dataset size'
 
-  chunk_count = length // number_of_chunks
-  overflow = length % number_of_chunks
+  chunk_count = dataset_length // number_of_chunks
+  overflow = dataset_length % number_of_chunks
 
   i = 0
   split = []
-  while i < length:
+  while i < dataset_length:
     chunk_length = chunk_count + int(overflow > 0)
     b = i + chunk_length
 
-    split.append(dataset[i:b])
+    split.append((i, b))
     overflow -= 1
     i = b
 

tests/test_algorithm.py

@@ -0,0 +1,45 @@
+import random
+from src.thread.utils import algorithm
+
+
+def test_chunking_1():
+  assert algorithm.chunk_split(5, 1) == [(0, 5)]
+
+
+def test_chunking_2():
+  assert algorithm.chunk_split(5, 2) == [(0, 3), (3, 5)]
+
+
+def test_chunking_3():
+  assert algorithm.chunk_split(100, 8) == [
+    (0, 13),
+    (13, 26),
+    (26, 39),
+    (39, 52),
+    (52, 64),
+    (64, 76),
+    (76, 88),
+    (88, 100),
+  ]
+
+
+def test_chunking_dynamic():
+  dataset_length = random.randint(400, int(10e6))
+  thread_count = random.randint(2, 100)
+
+  expected_chunk_length_low = dataset_length // thread_count
+  expected_chunk_high = dataset_length % thread_count
+
+  i = 0
+  heap = []
+  while i < dataset_length:
+    chunk_length = expected_chunk_length_low + int(expected_chunk_high > 0)
+    b = i + chunk_length
+
+    heap.append((i, b))
+    expected_chunk_high -= 1
+    i = b
+
+  assert (
+    algorithm.chunk_split(dataset_length, thread_count) == heap
+  ), f'\nLength: {dataset_length}\nThreads: {thread_count}\nExpected: {heap}\nActual: {algorithm.chunk_split(dataset_length, thread_count)}'