Files
modelscope/tests/hub/test_download_callback.py
Xingjun.Wang 9ab427b29a [Fix] Fix progress_callbacks for snapshot_download func (#1760)
* fix(msdatasets): make HfFileSystem monkey-patch idempotent for repeated streaming loads

When streaming=True the patches on HfFileSystem._open and HfFileSystem.__init__
are intentionally kept alive after load_dataset_with_ctx exits. A second call to
load_dataset_with_ctx then snaps the already-patched wrappers as 'originals',
causing _hf_fs_init_with_cookie / _hf_fs_open to call themselves recursively
until RecursionError: maximum recursion depth exceeded.

Fix: detect whether HfFileSystem is already patched before overwriting
_hf_fs_{open,init}_original, and skip re-applying the patch when it is already
in place.  The finally-block restores only patches that were applied in the
current invocation, leaving pre-existing patches intact.

Adds: test_hf_filesystem_patch_idempotent_for_repeated_streaming_loads

* fix(msdatasets): restore patches when streaming load fails; hermetic tests

Addresses both Gemini Code Assist review comments on PR #1754:

1. Bug fix – patch leak on failed streaming loads:
   Replace  guard with  flag.
   The flag is set True only after load_dataset() succeeds with streaming=True.
   Any exception in load_dataset (network, auth, invalid dataset, etc.) leaves
   _streaming_dataset_returned=False, so the finally block always restores all
   monkey-patches, preventing permanent global-state corruption.

2. Test improvement – hermetic test environment:
   Add _reset_hf_filesystem_patch() helper that strips pre-existing patches
   before each unit test, ensuring tests are independent of execution order.
   Add test_hf_filesystem_patch_restored_when_streaming_load_fails to cover
   the bug scenario introduced in fix 1.

* update requirements for hub

* feat(hub): restore progress_callbacks on snapshot_download shim (issue #1757)

The progress_callbacks parameter was dropped from snapshot_download after v1.38, breaking legacy code and GUI progress reporting. Re-expose it on modelscope.hub.snapshot_download (positioned before token per the v1.34 signature) and forward it to modelscope_hub.compat. Requires modelscope-hub >= 0.1.8.

Add network-free forwarding tests that patch the compat delegate so the real shim executes.

* fix lint

* feat(hub): add legacy-cache capability guard to snapshot_download shim

Warn once (thread-safe) when the installed modelscope-hub lacks pre-1.38 legacy cache auto-detection (DownloadManager._find_legacy_repo_dir, added in modelscope-hub>=0.1.7), so programmatic snapshot_download/dataset_snapshot_download callers are not silently downloading into the new layout while an old cache exists. Capability is probed, not reimplemented (kept single-sourced in modelscope-hub). Adds network-free tests for present/absent capability and fire-once behavior.
2026-07-21 17:57:49 +08:00

85 lines
2.6 KiB
Python

import tempfile
import unittest
from unittest import mock
from tqdm import tqdm
from modelscope import snapshot_download
from modelscope.hub import ProgressCallback
class NewProgressCallback(ProgressCallback):
all_files = set() # just for test
def __init__(self, filename: str, file_size: int):
super().__init__(filename, file_size)
self.progress = tqdm(total=file_size)
self.all_files.add(filename)
def update(self, size: int):
self.progress.update(size)
def end(self):
self.all_files.remove(self.filename)
assert self.progress.n == self.progress.total == self.file_size
self.progress.close()
class ProgressCallbackTest(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
def tearDown(self):
self.temp_dir.cleanup()
def test_progress_callback(self):
model_dir = snapshot_download(
'swift/test_lora',
progress_callbacks=[NewProgressCallback],
cache_dir=self.temp_dir.name)
print(f'model_dir: {model_dir}')
self.assertTrue(len(NewProgressCallback.all_files) == 0)
def test_empty_progress_callback(self):
model_dir = snapshot_download(
'swift/test_lora',
progress_callbacks=[],
cache_dir=self.temp_dir.name)
print(f'model_dir: {model_dir}')
class SnapshotDownloadForwardTest(unittest.TestCase):
"""Network-free tests: the shim forwards progress_callbacks to compat."""
# ``modelscope.hub.snapshot_download`` the attribute is shadowed by the
# re-exported function, so patch via the fully qualified module string.
_COMPAT_TARGET = \
'modelscope.hub.snapshot_download._compat_snapshot_download'
def test_progress_callbacks_forwarded_to_compat(self):
from modelscope.hub.snapshot_download import snapshot_download
with mock.patch(
self._COMPAT_TARGET, return_value='/tmp/snapshot') as m:
result = snapshot_download(
'owner/repo', progress_callbacks=[NewProgressCallback])
self.assertEqual(result, '/tmp/snapshot')
_, kwargs = m.call_args
self.assertEqual(kwargs['progress_callbacks'], [NewProgressCallback])
def test_progress_callbacks_default_none(self):
from modelscope.hub.snapshot_download import snapshot_download
with mock.patch(
self._COMPAT_TARGET, return_value='/tmp/snapshot') as m:
snapshot_download('owner/repo')
_, kwargs = m.call_args
self.assertIsNone(kwargs['progress_callbacks'])
if __name__ == '__main__':
unittest.main()