From 61571c0169dc2680488a9ba18696153d5badc391 Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Tue, 14 Jul 2026 16:19:08 +0800 Subject: [PATCH] Fix RecursionError on Repeated Streaming Dataset Loads via Idempotent HfFileSystem Patching (#1754) --- .../msdatasets/utils/hf_datasets_util.py | 28 ++++--- tests/msdatasets/test_stream_load.py | 77 +++++++++++++++++++ 2 files changed, 96 insertions(+), 9 deletions(-) diff --git a/modelscope/msdatasets/utils/hf_datasets_util.py b/modelscope/msdatasets/utils/hf_datasets_util.py index 8c89ecc9..f211fc81 100644 --- a/modelscope/msdatasets/utils/hf_datasets_util.py +++ b/modelscope/msdatasets/utils/hf_datasets_util.py @@ -1185,6 +1185,8 @@ def load_dataset_with_ctx(*args, **kwargs): generate_from_dict_origin = features.generate_from_dict hf_fs_open_origin = HfFileSystem._open hf_fs_init_origin = HfFileSystem.__init__ + hf_fs_open_was_patched = hf_fs_open_origin is _hf_fs_open + hf_fs_init_was_patched = hf_fs_init_origin is _hf_fs_init_with_cookie # Apply patches config.HF_ENDPOINT = get_endpoint() @@ -1201,25 +1203,33 @@ def load_dataset_with_ctx(*args, **kwargs): if _HAS_SCRIPT_LOADING: HubDatasetModuleFactoryWithScript.get_module = get_module_with_script features.generate_from_dict = generate_from_dict_ms - _hf_fs_open_original = hf_fs_open_origin - HfFileSystem._open = _hf_fs_open - _hf_fs_init_original = hf_fs_init_origin - HfFileSystem.__init__ = _hf_fs_init_with_cookie + if not hf_fs_open_was_patched: + _hf_fs_open_original = hf_fs_open_origin + HfFileSystem._open = _hf_fs_open + if not hf_fs_init_was_patched: + _hf_fs_init_original = hf_fs_init_origin + HfFileSystem.__init__ = _hf_fs_init_with_cookie streaming = kwargs.get('streaming', False) + _streaming_dataset_returned = False + try: dataset_res = DatasetsWrapperHF.load_dataset(*args, **kwargs) + _streaming_dataset_returned = streaming yield dataset_res finally: _repo_tree_cache.clear() HubApi._dataset_id_type_cache.clear() - if not streaming: - HfFileSystem._open = hf_fs_open_origin - _hf_fs_open_original = None - HfFileSystem.__init__ = hf_fs_init_origin - _hf_fs_init_original = None + should_restore = not _streaming_dataset_returned + if should_restore: + if not hf_fs_open_was_patched: + HfFileSystem._open = hf_fs_open_origin + _hf_fs_open_original = None + if not hf_fs_init_was_patched: + HfFileSystem.__init__ = hf_fs_init_origin + _hf_fs_init_original = None config.HF_ENDPOINT = hf_endpoint_origin file_utils.get_from_cache = get_from_cache_origin diff --git a/tests/msdatasets/test_stream_load.py b/tests/msdatasets/test_stream_load.py index 0ce46887..87bbf096 100644 --- a/tests/msdatasets/test_stream_load.py +++ b/tests/msdatasets/test_stream_load.py @@ -1,6 +1,9 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import unittest +from unittest import mock + +from huggingface_hub.hf_file_system import HfFileSystem from modelscope import MsDataset from modelscope.utils.logger import get_logger @@ -11,6 +14,80 @@ logger = get_logger() class TestStreamLoad(unittest.TestCase): + @staticmethod + def _reset_hf_filesystem_patch(hf_datasets_util): + if (HfFileSystem._open is hf_datasets_util._hf_fs_open + and hf_datasets_util._hf_fs_open_original is not None): + HfFileSystem._open = hf_datasets_util._hf_fs_open_original + hf_datasets_util._hf_fs_open_original = None + if (HfFileSystem.__init__ is hf_datasets_util._hf_fs_init_with_cookie + and hf_datasets_util._hf_fs_init_original is not None): + HfFileSystem.__init__ = hf_datasets_util._hf_fs_init_original + hf_datasets_util._hf_fs_init_original = None + + def test_hf_filesystem_patch_idempotent_for_repeated_streaming_loads(self): + from modelscope.msdatasets.utils import hf_datasets_util + + hf_fs_open_before = HfFileSystem._open + hf_fs_init_before = HfFileSystem.__init__ + open_original_before = hf_datasets_util._hf_fs_open_original + init_original_before = hf_datasets_util._hf_fs_init_original + try: + self._reset_hf_filesystem_patch(hf_datasets_util) + with mock.patch.object( + hf_datasets_util.DatasetsWrapperHF, + 'load_dataset', + return_value=object()): + with hf_datasets_util.load_dataset_with_ctx(streaming=True): + pass + with hf_datasets_util.load_dataset_with_ctx(streaming=True): + pass + + self.assertIs(HfFileSystem._open, hf_datasets_util._hf_fs_open) + self.assertIsNot( + hf_datasets_util._hf_fs_open_original, + hf_datasets_util._hf_fs_open) + self.assertIs( + HfFileSystem.__init__, + hf_datasets_util._hf_fs_init_with_cookie) + self.assertIsNot( + hf_datasets_util._hf_fs_init_original, + hf_datasets_util._hf_fs_init_with_cookie) + finally: + HfFileSystem._open = hf_fs_open_before + HfFileSystem.__init__ = hf_fs_init_before + hf_datasets_util._hf_fs_open_original = open_original_before + hf_datasets_util._hf_fs_init_original = init_original_before + + def test_hf_filesystem_patch_restored_when_streaming_load_fails(self): + from modelscope.msdatasets.utils import hf_datasets_util + + hf_fs_open_before = HfFileSystem._open + hf_fs_init_before = HfFileSystem.__init__ + open_original_before = hf_datasets_util._hf_fs_open_original + init_original_before = hf_datasets_util._hf_fs_init_original + try: + self._reset_hf_filesystem_patch(hf_datasets_util) + hf_fs_open_clean = HfFileSystem._open + hf_fs_init_clean = HfFileSystem.__init__ + with mock.patch.object( + hf_datasets_util.DatasetsWrapperHF, + 'load_dataset', + side_effect=RuntimeError('load failed')): + with self.assertRaises(RuntimeError): + with hf_datasets_util.load_dataset_with_ctx(streaming=True): + pass + + self.assertIs(HfFileSystem._open, hf_fs_open_clean) + self.assertIs(HfFileSystem.__init__, hf_fs_init_clean) + self.assertIsNone(hf_datasets_util._hf_fs_open_original) + self.assertIsNone(hf_datasets_util._hf_fs_init_original) + finally: + HfFileSystem._open = hf_fs_open_before + HfFileSystem.__init__ = hf_fs_init_before + hf_datasets_util._hf_fs_open_original = open_original_before + hf_datasets_util._hf_fs_init_original = init_original_before + def setUp(self): ...