Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-09-01 19:49:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix RecursionError on Repeated Streaming Dataset Loads via Idempotent HfFileSystem Patching (#1754)

Browse Source
This commit is contained in:
Xingjun.Wang
2026-07-14 16:19:08 +08:00
committed by GitHub
parent 9379504fd3
commit 61571c0169
2 changed files with 96 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
modelscope/msdatasets/utils/hf_datasets_util.py
View File

@@ -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

77
tests/msdatasets/test_stream_load.py
View File

@@ -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):
...