Repository: Netflix/metaflow
Pull request: Check for the first finished parallel proccess
This fixes https://github.com/Netflix/metaflow/issues/1316 to optimize the parallel map runs by checking for the first finished process instead of waiting for the first process before checking others.

Problem statement:
Optimize parallel_map by properly waiting on child pids
Currently `parallel_map` launches K subprocesses and then waits for the Kth child to finish before launching another subprocess. A more efficient approach would be to use epoll etc to wait for the first subprocess to finish.

Hints from issue discussion:

Commit message: Fix code format with black

Patch to review:
diff --git a/metaflow/multicore_utils.py b/metaflow/multicore_utils.py
index 1ac81039c12..55704623d5b 100644
--- a/metaflow/multicore_utils.py
+++ b/metaflow/multicore_utils.py
@@ -4,6 +4,7 @@
 from itertools import islice
 from multiprocessing import cpu_count
 from tempfile import NamedTemporaryFile
+import time
 
 try:
     # Python 2
@@ -68,8 +69,17 @@ def parallel_imap_unordered(func, iterable, max_parallel=None, dir=None):
 
     while pids:
 
-        pid, output_file = pids.pop()
-        if os.waitpid(pid, 0)[1]:
+        for idx, pid_info in enumerate(pids):
+            pid, output_file = pid_info
+            pid, exit_code = os.waitpid(pid, os.WNOHANG)
+            if pid:
+                pids.pop(idx)
+                break
+        else:
+            time.sleep(1)  # Wait a bit before re-checking
+            continue
+
+        if exit_code:
             raise MulticoreException("Child failed")
 
         with open(output_file, "rb") as f:
diff --git a/test/unit/test_multicore_utils.py b/test/unit/test_multicore_utils.py
new file mode 100644
index 00000000000..9b14a20b99a
--- /dev/null
+++ b/test/unit/test_multicore_utils.py
@@ -0,0 +1,12 @@
+from metaflow.multicore_utils import parallel_map
+
+
+def test_parallel_map():
+    assert parallel_map(lambda s: s.upper(), ["a", "b", "c", "d", "e", "f"]) == [
+        "A",
+        "B",
+        "C",
+        "D",
+        "E",
+        "F",
+    ]


Code review comments:
diff_hunk: @@ -68,8 +69,17 @@ def parallel_imap_unordered(func, iterable, max_parallel=None, dir=None):
 
     while pids:
 
-        pid, output_file = pids.pop()
-        if os.waitpid(pid, 0)[1]:
+        for idx, pid_info in enumerate(pids):
+            pid, output_file = pid_info
+            pid, exit_code = os.waitpid(pid, os.WNOHANG)
+            if pid:
+                pids.pop(idx)
+                break
+        else:
+            time.sleep(1)  # Wait a bit before re-checking
path: metaflow/multicore_utils.py
text: @user1:
After a quick discussion, someone brought up the fact that for very short tasks, this now makes it at least 1 second. Could we either skip the sleep (and just do tight polling) or do something a bit fancier (no sleep initially and then some sleep after a while)?

@author:
I did consider that as well but decided not to optimize prematurely as I am not sure if those very short tasks (less than 1 sec) would need to use this function?  If you are certain there are real scenarios where all tasks are done less than 1 sec and they would use this multi-process function for that, then I can optimize it. 

@author:
Reduced sleep to 0.1 secs as that should be small enough for any short tasks while not affecting the CPU as it checks.