MIPLearn v0.3

miplearn/parallel.py

@@ -0,0 +1,32 @@
+# Modified version of: https://github.com/swansonk14/p_tqdm
+# Copyright (c) 2022 Kyle Swanson
+# MIT License
+
+from collections.abc import Sized
+from typing import Any, Callable, Generator, Iterable, List
+
+from pathos.multiprocessing import _ProcessPool as Pool
+from tqdm.auto import tqdm
+
+
+def _parallel(function: Callable, *iterables: Iterable, **kwargs: Any) -> Generator:
+    # Determine length of tqdm (equal to length of the shortest iterable or total kwarg)
+    total = kwargs.pop("total", None)
+    lengths = [len(iterable) for iterable in iterables if isinstance(iterable, Sized)]
+    length = total or (min(lengths) if lengths else None)
+
+    # Create parallel generator
+    num_cpus = kwargs.pop("num_cpus", 1)
+    maxtasksperchild = kwargs.pop("maxtasksperchild", 1)
+    chunksize = kwargs.pop("chunksize", 1)
+    with Pool(num_cpus, maxtasksperchild=maxtasksperchild) as pool:
+        for item in tqdm(
+            pool.imap_unordered(function, *iterables, chunksize=chunksize),
+            total=length,
+            **kwargs
+        ):
+            yield item
+
+
+def p_umap(function: Callable, *iterables: Iterable, **kwargs: Any) -> List[Any]:
+    return list(_parallel(function, *iterables, **kwargs))