From 9712099ce8a219293e06626c4f74d60c2558312f Mon Sep 17 00:00:00 2001
From: Vaishnav88sk <vaishnavsk8804@gmail.com>
Date: Wed, 29 Apr 2026 20:44:37 +0530
Subject: [PATCH 1/2] Fix bare except clauses in multicore_utils and add
 comprehensive tests

Replace bare except: with except ImportError: for the pickle import
fallback, and except Exception: in the child process error handler.
Bare except clauses catch SystemExit and KeyboardInterrupt which can
mask real problems and interfere with process management.

Also expands test coverage from 1 test to 12 tests, covering:
- Basic parallel_map functionality and order preservation
- Edge cases (empty input, single element)
- max_parallel parameter
- Complex object serialization
- parallel_imap_unordered behavior
- Child process failure propagation
---
 metaflow/multicore_utils.py       |  4 +-
 test/unit/test_multicore_utils.py | 73 ++++++++++++++++++++++++++++++-
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/metaflow/multicore_utils.py b/metaflow/multicore_utils.py
index 165f9837e94..cc69c3a318f 100644
--- a/metaflow/multicore_utils.py
+++ b/metaflow/multicore_utils.py
@@ -22,7 +22,7 @@
 try:
     # Python 2
     import cPickle as pickle
-except:
+except ImportError:
     # Python 3
     import pickle
 
@@ -66,7 +66,7 @@ def _spawn(
                 with open(output_file, "wb") as f:
                     pickle.dump(ret, f, protocol=pickle.HIGHEST_PROTOCOL)
                 exit_code = 0
-            except:
+            except Exception:
                 # we must not let any exceptions escape this function
                 # which might trigger unintended side-effects
                 traceback.print_exc()
diff --git a/test/unit/test_multicore_utils.py b/test/unit/test_multicore_utils.py
index 9b14a20b99a..3ec3ef69f6d 100644
--- a/test/unit/test_multicore_utils.py
+++ b/test/unit/test_multicore_utils.py
@@ -1,7 +1,9 @@
-from metaflow.multicore_utils import parallel_map
+import pytest
 
+from metaflow.multicore_utils import MulticoreException, parallel_imap_unordered, parallel_map
 
-def test_parallel_map():
+
+def test_parallel_map_basic():
     assert parallel_map(lambda s: s.upper(), ["a", "b", "c", "d", "e", "f"]) == [
         "A",
         "B",
@@ -10,3 +12,70 @@ def test_parallel_map():
         "E",
         "F",
     ]
+
+
+def test_parallel_map_preserves_order():
+    result = parallel_map(lambda x: x * 2, range(10))
+    assert result == [0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
+
+
+def test_parallel_map_empty_input():
+    assert parallel_map(lambda x: x, []) == []
+
+
+def test_parallel_map_single_element():
+    assert parallel_map(lambda x: x + 1, [41]) == [42]
+
+
+def test_parallel_map_with_max_parallel():
+    result = parallel_map(lambda x: x**2, range(5), max_parallel=2)
+    assert result == [0, 1, 4, 9, 16]
+
+
+def test_parallel_map_with_large_dataset():
+    result = parallel_map(lambda x: x * 3, range(100))
+    assert result == [x * 3 for x in range(100)]
+
+
+def test_parallel_imap_unordered_basic():
+    results = list(parallel_imap_unordered(lambda x: x * 2, range(4)))
+    assert sorted(results) == [0, 2, 4, 6]
+
+
+def test_parallel_imap_unordered_empty():
+    assert list(parallel_imap_unordered(lambda x: x, [])) == []
+
+
+def test_parallel_imap_unordered_with_max_parallel():
+    results = list(parallel_imap_unordered(lambda x: x + 1, range(3), max_parallel=1))
+    assert sorted(results) == [1, 2, 3]
+
+
+def test_parallel_map_raises_on_child_failure():
+    def failing_func(x):
+        if x == 2:
+            raise ValueError("Child process failure")
+        return x
+
+    with pytest.raises(MulticoreException, match="Child failed"):
+        parallel_map(failing_func, [1, 2, 3])
+
+
+def test_parallel_map_returns_complex_objects():
+    def make_dict(x):
+        return {"key": x, "nested": {"value": x * 2}}
+
+    result = parallel_map(make_dict, [1, 2, 3])
+    assert result == [
+        {"key": 1, "nested": {"value": 2}},
+        {"key": 2, "nested": {"value": 4}},
+        {"key": 3, "nested": {"value": 6}},
+    ]
+
+
+def test_parallel_map_with_named_function():
+    def square(x):
+        return x * x
+
+    result = parallel_map(square, [1, 2, 3, 4, 5])
+    assert result == [1, 4, 9, 16, 25]

From d775ddb2261500f058c713519b4376414c3f9ca8 Mon Sep 17 00:00:00 2001
From: Vaishnav88sk <vaishnavsk8804@gmail.com>
Date: Sat, 2 May 2026 01:09:07 +0530
Subject: [PATCH 2/2] Change except Exception to except BaseException in child
 process handler

Based on greptile bot review: except Exception does not catch
KeyboardInterrupt, SystemExit, or GeneratorExit. Using
BaseException matches the original intent of catching all exceptions
while being explicit (unlike bare except:).
---
 metaflow/multicore_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metaflow/multicore_utils.py b/metaflow/multicore_utils.py
index cc69c3a318f..f1309fb7329 100644
--- a/metaflow/multicore_utils.py
+++ b/metaflow/multicore_utils.py
@@ -66,7 +66,7 @@ def _spawn(
                 with open(output_file, "wb") as f:
                     pickle.dump(ret, f, protocol=pickle.HIGHEST_PROTOCOL)
                 exit_code = 0
-            except Exception:
+            except BaseException:
                 # we must not let any exceptions escape this function
                 # which might trigger unintended side-effects
                 traceback.print_exc()