Merge pull request #141 from J08nY/feature/mp-billiard

Use billiard instead of concurrent.futures

2 files changed, 11 insertions, 8 deletions

diff --git a/requirements.txt b/requirements.txt
index 81d253f3..35a1cae8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 wheel==0.34.2
 beautifulsoup4==4.9.3
+billiard==3.6.4.0
 certifi==2020.11.8
 chardet==3.0.4
 click==7.1.2
@@ -31,4 +32,4 @@ importlib-resources==5.1.2
 scikit-learn
 jsonschema==3.2.0
 packaging
-html5lib
\ No newline at end of file
+html5lib
diff --git a/sec_certs/parallel_processing.py b/sec_certs/parallel_processing.py
index 912ef73b..c42c9c16 100644
--- a/sec_certs/parallel_processing.py
+++ b/sec_certs/parallel_processing.py
@@ -1,25 +1,27 @@
 from sec_certs.helpers import tqdm
-from concurrent.futures import ProcessPoolExecutor as ProcessPool, ThreadPoolExecutor as ThreadPool
+from multiprocessing.pool import ThreadPool
+from billiard.pool import Pool
 from typing import Callable, Iterable, Optional, Union
 import time
 
 
-def process_parallel(func: Callable, items: Iterable, max_workers: int,
+def process_parallel(func: Callable, items: Iterable, max_workers: int, callback: Optional[Callable] = None,
                      use_threading: bool = True, progress_bar: bool = True, unpack: bool = False,
                      progress_bar_desc: Optional[str] = None):
 
-    pool: Union[ProcessPool, ThreadPool] = ThreadPool(max_workers) if use_threading else ProcessPool(max_workers)
-    results = [pool.submit(func, *i) for i in items] if unpack else [pool.submit(func, i) for i in items]
+    pool: Union[Pool, ThreadPool] = ThreadPool(max_workers) if use_threading else Pool(max_workers)
+    results = [pool.apply_async(func, (*i,), callback=callback) for i in items] if unpack else [pool.apply_async(func, (i, ), callback=callback) for i in items]
 
     if progress_bar is True and items:
         bar = tqdm(total=len(results), desc=progress_bar_desc)
-        while not all(all_done := [x.done() for x in results]):
+        while not all(all_done := [x.ready() for x in results]):
             done_count = len(list(filter(lambda x: x, all_done)))
             bar.update(done_count - bar.n)
             time.sleep(1)
         bar.update(len(results) - bar.n)
         bar.close()
 
-    pool.shutdown()
+    pool.close()
+    pool.join()
 
-    return [r.result() for r in results]
+    return [r.get() for r in results]