fix(dataset): detect failed streaming packing workers

swift/dataset/packing.py

@@ -3,6 +3,8 @@
 import multiprocessing as mp
 import torch.distributed as dist
 from itertools import chain
+from multiprocessing.connection import wait
+from queue import Empty
 from torch.utils.data import Dataset, IterableDataset
 from tqdm import tqdm
 from typing import Optional
@@ -250,7 +252,16 @@ def _put_data_in_queue(self, iterator) -> int:
     def _fetch_data_out_queue(self, last_res, num_samples):
         res = [None] * num_samples
         for _ in range(num_samples):
-            i, data = self._out_queue.get()
+            while True:
+                try:
+                    i, data = self._out_queue.get(timeout=1)
+                    break
+                except Empty:
+                    # Sentinels also work when a DataLoader process inherits the workers.
+                    if wait([worker.sentinel for worker in self.workers], timeout=0):
+                        raise RuntimeError(
+                            'A packing worker exited unexpectedly. Check the worker logs for the original error.'
+                        ) from None
             if not data:
                 continue
             res[i] = data if isinstance(data, list) else [data]

tests/general/test_packing_multiprocessing_context.py

@@ -23,6 +23,7 @@
 import shutil
 import sys
 import tempfile
+import time
 import torch
 import types
 import unittest
@@ -31,7 +32,7 @@
 
 import swift.utils.utils as utils_module
 from swift.dataset.packing import IterablePackingDataset, PackingDataset, _resolve_mp_context, _spawn_workers
-from swift.template.base import Template
+from swift.template.base import MaxLengthError, Template
 from swift.utils import get_external_files, import_external_file, patch_dataloader_external_plugins
 
 FORK_AVAILABLE = 'fork' in mp.get_all_start_methods()
@@ -60,6 +61,45 @@ def __init__(self):
         self._not_picklable = lambda x: x
 
 
+class FailingTemplate(FakeTemplate):
+
+    def encode(self, data, return_length=True):
+        if data.get('error') == 'exit':
+            os._exit(1)
+        if data.get('error') == 'length':
+            raise MaxLengthError('sample exceeds max_length')
+        if data.get('error') == 'invalid':
+            raise ValueError('invalid training sample')
+        if data.get('delay'):
+            time.sleep(data['delay'])
+        return super().encode(data, return_length=return_length)
+
+
+def _collect_packing_result(queue, context, rows, strict, dataloader_num_workers=0):
+    dataset = IterablePackingDataset(
+        FailingTemplate(), rows, strict=strict, packing_interval=4, multiprocessing_context=context)
+    if dataloader_num_workers:
+        iterator = iter(
+            DataLoader(
+                dataset,
+                batch_size=None,
+                num_workers=dataloader_num_workers,
+                multiprocessing_context='fork',
+                timeout=20))
+    else:
+        iterator = iter(dataset)
+    try:
+        queue.put(('ok', [row['input_ids'] for pack in iterator for row in pack]))
+    except Exception as exc:
+        queue.put(('error', type(exc).__name__, str(exc)))
+    finally:
+        if dataloader_num_workers:
+            iterator._shutdown_workers()
+        for worker in dataset.workers:
+            worker.terminate()
+            worker.join(timeout=5)
+
+
 class ListDataset:
     """Map-style dataset exposing ``dataset['lengths']`` like swift datasets do."""
 
@@ -347,6 +387,57 @@ def test_unpicklable_template_raises_under_spawn(self):
             _run_iter_packing('spawn', self.rows, template=UnpicklableTemplate())
 
 
+class TestIterablePackingWorkerErrors(unittest.TestCase):
+
+    def _result(self, rows, strict, context='spawn', dataloader_num_workers=0):
+        ctx = mp.get_context(context)
+        queue = ctx.Queue()
+        process = ctx.Process(
+            target=_collect_packing_result, args=(queue, context, rows, strict, dataloader_num_workers))
+        process.start()
+        try:
+            # Both the consumer and packing worker need time to start under spawn.
+            return queue.get(timeout=60)
+        finally:
+            process.join(timeout=5)
+            if process.is_alive():
+                process.terminate()
+                process.join(timeout=5)
+            queue.close()
+
+    def test_strict_encoding_error_reaches_consumer(self):
+        result = self._result([{'error': 'invalid'}], strict=True)
+        self.assertEqual(result[:2], ('error', 'RuntimeError'))
+        self.assertIn('worker', result[2].lower())
+
+    def test_unexpected_worker_exit_reaches_consumer(self):
+        result = self._result([{'error': 'exit'}], strict=False)
+        self.assertEqual(result[:2], ('error', 'RuntimeError'))
+
+    @unittest.skipUnless(FORK_AVAILABLE, 'fork is required to inherit packing worker handles')
+    def test_failed_worker_reaches_dataloader_consumer(self):
+        result = self._result([{'error': 'invalid'}], strict=True, dataloader_num_workers=1)
+        self.assertEqual(result[:2], ('error', 'RuntimeError'))
+        self.assertIn('packing worker exited unexpectedly', result[2])
+
+    @unittest.skipUnless(FORK_AVAILABLE, 'fork is required to inherit packing worker handles')
+    def test_live_worker_with_dataloader_consumer(self):
+        result = self._result([{'input_ids': [1, 2], 'delay': 2}], strict=True, dataloader_num_workers=1)
+        self.assertEqual(result, ('ok', [[1, 2]]))
+
+    def test_non_strict_error_still_skips_sample(self):
+        result = self._result([{'error': 'invalid'}, {'input_ids': [1, 2]}], strict=False)
+        self.assertEqual(result, ('ok', [[1, 2]]))
+
+    def test_strict_length_error_still_skips_sample(self):
+        result = self._result([{'error': 'length'}, {'input_ids': [1, 2]}], strict=True)
+        self.assertEqual(result, ('ok', [[1, 2]]))
+
+    def test_live_worker_is_allowed_to_finish(self):
+        result = self._result([{'input_ids': [1, 2], 'delay': 2}], strict=True)
+        self.assertEqual(result, ('ok', [[1, 2]]))
+
+
 class TestDataloaderContextInjection(unittest.TestCase):
     """The DataLoader-side injection is pure kwargs plumbing; assert it only activates when set."""