safe map

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

src/megatron/bridge/data/datasets/utils.py

@@ -20,8 +20,10 @@
 import os
 import pickle
 import re
+import signal
 import time
 from functools import lru_cache, partial
+from queue import Empty
 from typing import Any, Callable, Optional, Pattern, Type
 
 import numpy as np
@@ -87,6 +89,89 @@ def build_index_from_memdata(fn, newline_int):
     return midx
 
 
+def safe_map(fn, iterable, workers=1, ctx="fork"):
+    """
+    Crash-resilient alternative to multiprocessing.Pool.map() that can handle
+    worker process crashes gracefully without hanging the entire operation.
+
+    This function provides robustness when processing large datasets where individual
+    workers might crash. Unlike Pool.map(), it won't hang indefinitely if a worker dies.
+
+    Args:
+        fn: Function to apply to each item
+        iterable: Items to process
+        workers: Number of worker processes
+        ctx: Multiprocessing context ("fork", "spawn", etc.)
+
+    Returns:
+        List of results (same order as iterable). Failed items are None with warning logged.
+    """
+    ctx = mp.get_context(ctx)
+    input_queue = ctx.Queue()
+    output_queue = ctx.Queue()
+    indexed_inputs = list(enumerate(iterable))
+    for job in indexed_inputs:
+        input_queue.put(job)
+    for _ in range(workers):
+        input_queue.put(None)  # poison pill
+
+    def worker_loop():
+        while True:
+            job = input_queue.get()
+            if job is None:
+                break
+            i, item = job
+            try:
+                result = fn(item)
+                output_queue.put((i, True, result, None))
+            except Exception as e:
+                output_queue.put((i, False, None, str(e)))
+
+    processes = [ctx.Process(target=worker_loop) for _ in range(workers)]
+    for p in processes:
+        p.start()
+
+    results = [None] * len(indexed_inputs)
+    seen_indices = set()
+    expected = len(indexed_inputs)
+    received = 0
+
+    # Collect whatever gets returned from live workers
+    while received < expected:
+        try:
+            i, success, result, err = output_queue.get(timeout=0.5)
+            seen_indices.add(i)
+            results[i] = result if success else None
+            if not success:
+                logger.warning(f"Item {i}: {err}")
+            received += 1
+        except Empty:
+            # Check if all workers are dead
+            if all(not p.is_alive() for p in processes):
+                logger.error("All workers exited before completing all tasks.")
+                break
+            continue
+
+    # Join and check for crashes
+    for p in processes:
+        p.join()
+        if p.exitcode is not None and p.exitcode < 0:
+            sig = -p.exitcode
+            try:
+                sig_name = signal.Signals(sig).name
+            except Exception:
+                sig_name = f"signal {sig}"
+            logger.warning(f"PID {p.pid} died from {sig_name}")
+
+    # Patch any missing results from crashed workers
+    for i in range(len(results)):
+        if i not in seen_indices:
+            logger.warning(f"No result for item {i}, likely crash")
+            results[i] = None
+
+    return results
+
+
 class _TextMemMapDataset(Dataset):
     """
     Allow per-line lazy access to multiple text files using numpy memmap.
@@ -568,17 +653,16 @@ def build_index_files(
     logger.info(f"Processing {len(dataset_paths)} data files using {workers} workers")
     # load all files into memmap
     start_time = time.time()
-    ctx = mp.get_context("fork")
-    with ctx.Pool(workers) as p:
-        build_status = p.map(
-            partial(
-                _build_memmap_index_files,
-                newline_int,
-                build_index_fn,
-                index_mapping_dir=index_mapping_dir,
-            ),
-            dataset_paths,
-        )
+    build_status = safe_map(
+        partial(
+            _build_memmap_index_files,
+            newline_int,
+            build_index_fn,
+            index_mapping_dir=index_mapping_dir,
+        ),
+        dataset_paths,
+        workers=workers,
+    )
 
     logger.info(
         f"Time building {sum(build_status)} / {len(build_status)} "

tests/unit_tests/data/datasets/test_utils.py

@@ -35,6 +35,7 @@
     _response_value_formater,
     build_index_from_memdata,
     handle_index,
+    safe_map,
 )
 
 
@@ -308,3 +309,68 @@ def test_build_memmap_index_files_with_msc_url(self):
 
             assert msc.Path(f"msc://default{temp_dir}/training.jsonl.idx.npy")
             assert msc.Path(f"msc://default{temp_dir}/training.jsonl.idx.info")
+
+
+class TestSafeMap:
+    """Test cases for crash-resilient safe_map function."""
+
+    def test_safe_map_basic_functionality(self):
+        """Test that safe_map works like normal map for successful cases."""
+
+        def square(x):
+            return x * x
+
+        items = [1, 2, 3, 4, 5]
+        result = safe_map(square, items, workers=2)
+
+        assert result == [1, 4, 9, 16, 25]
+
+    def test_safe_map_handles_exceptions(self):
+        """Test that safe_map handles exceptions gracefully."""
+
+        def process_with_error(x):
+            if x == 3:
+                raise ValueError("Simulated error")
+            return x * 2
+
+        items = [1, 2, 3, 4, 5]
+        result = safe_map(process_with_error, items, workers=2)
+
+        # Item 3 should be None (failed), others should succeed
+        assert result[0] == 2
+        assert result[1] == 4
+        assert result[2] is None  # Failed item
+        assert result[3] == 8
+        assert result[4] == 10
+
+    def test_safe_map_preserves_order(self):
+        """Test that safe_map preserves input order even with parallel execution."""
+
+        def identity(x):
+            return x
+
+        items = list(range(100))
+        result = safe_map(identity, items, workers=4)
+
+        assert result == items
+
+    def test_safe_map_with_single_worker(self):
+        """Test safe_map with workers=1 (sequential execution)."""
+
+        def double(x):
+            return x * 2
+
+        items = [1, 2, 3]
+        result = safe_map(double, items, workers=1)
+
+        assert result == [2, 4, 6]
+
+    def test_safe_map_empty_iterable(self):
+        """Test safe_map with empty input."""
+
+        def identity(x):
+            return x
+
+        result = safe_map(identity, [], workers=2)
+
+        assert result == []