From 38e4cfb3ede6ab35fbfe3c5cde6e719f03064b13 Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Wed, 9 Sep 2026 04:17:22 +0800 Subject: [PATCH 1/2] fix(dataset): detect failed streaming packing workers --- swift/dataset/packing.py | 13 ++- .../test_packing_multiprocessing_context.py | 92 ++++++++++++++++++- 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/swift/dataset/packing.py b/swift/dataset/packing.py index 7bc6042536..de1bcd2b4c 100644 --- a/swift/dataset/packing.py +++ b/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] diff --git a/tests/general/test_packing_multiprocessing_context.py b/tests/general/test_packing_multiprocessing_context.py index c657ede9c5..c6eb617db5 100644 --- a/tests/general/test_packing_multiprocessing_context.py +++ b/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,56 @@ 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: + return queue.get(timeout=60 if dataloader_num_workers else 30) + 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.""" From b0f23f6656703086f10ff5ceb7a3dc85dd3c8c0a Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Wed, 9 Sep 2026 23:34:22 +0800 Subject: [PATCH 2/2] test(dataset): allow for spawned packing worker startup --- tests/general/test_packing_multiprocessing_context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/general/test_packing_multiprocessing_context.py b/tests/general/test_packing_multiprocessing_context.py index c6eb617db5..c37406bc32 100644 --- a/tests/general/test_packing_multiprocessing_context.py +++ b/tests/general/test_packing_multiprocessing_context.py @@ -396,7 +396,8 @@ def _result(self, rows, strict, context='spawn', dataloader_num_workers=0): target=_collect_packing_result, args=(queue, context, rows, strict, dataloader_num_workers)) process.start() try: - return queue.get(timeout=60 if dataloader_num_workers else 30) + # 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():