From 2224bcc0fb4a91c6a293629e112b8f8414ed34bc Mon Sep 17 00:00:00 2001 From: WeiweiZhang1 Date: Fri, 11 Sep 2026 15:44:38 +0800 Subject: [PATCH 1/3] add fineweb-edu dataset as calibration bakeup Signed-off-by: WeiweiZhang1 --- auto_round/calib_dataset.py | 130 +++++++++++++----- docs/step_by_step.md | 4 +- docs/step_by_step_CN.md | 5 +- .../utils/test_calib_dataset_helpers.py | 65 +++++++++ 4 files changed, 165 insertions(+), 39 deletions(-) diff --git a/auto_round/calib_dataset.py b/auto_round/calib_dataset.py index 9910b8d249..f2e7a8192e 100644 --- a/auto_round/calib_dataset.py +++ b/auto_round/calib_dataset.py @@ -18,6 +18,7 @@ import multiprocessing import os import random +import ssl import sys logging.getLogger("datasets").setLevel(logging.WARNING) @@ -32,6 +33,51 @@ CALIB_DATASETS = {} _GITHUB_CODE_CLEAN_MAX_DATASETS_VERSION = Version("3.6.0") +_FINEWEB_EDU_HF_DATASET = "HuggingFaceFW/fineweb-edu" +_FINEWEB_EDU_MODELSCOPE_DATASET = "AI-ModelScope/fineweb-edu" +_FINEWEB_EDU_CONFIG = "sample-10BT" + + +def _get_dataset_network_error(error): + """Return the network failure in an exception chain, if present.""" + visited = set() + current = error + while current is not None and id(current) not in visited: + visited.add(id(current)) + error_name = type(current).__name__.lower() + error_module = type(current).__module__.lower() + error_message = str(current).lower() + if isinstance(current, (ConnectionError, TimeoutError, ssl.SSLError)): + return current + if any(term in error_name for term in ("connection", "http", "proxy", "ssl", "timeout")) and any( + package in error_module + for package in ("datasets", "huggingface_hub", "httpcore", "httpx", "requests", "urllib3") + ): + return current + if any( + marker in error_message + for marker in ("connection timed out", "max retries exceeded", "proxy error", "ssl error") + ): + return current + current = current.__cause__ or current.__context__ + return None + + +def _warn_on_dataset_network_error(error, dataset_name): + """Give users an actionable cross-hub fallback for dataset network failures.""" + network_error = _get_dataset_network_error(error) + if network_error is None: + return + logger.warning( + "Failed to load calibration dataset %r because of an HTTP/proxy network issue: %s. " + "Check your network and proxy settings, or switch to FineWeb-Edu with `--dataset fineweb-edu`. " + "By default it loads from Hugging Face (%s); set `AR_USE_MODELSCOPE=1` and install `modelscope` " + "to load its ModelScope mirror (%s).", + dataset_name, + network_error, + _FINEWEB_EDU_HF_DATASET, + _FINEWEB_EDU_MODELSCOPE_DATASET, + ) def get_code_calibration_dataset(nsamples, datasets_version=None): @@ -193,22 +239,7 @@ def get_pile_dataset( tokenizer_function = get_tokenizer_function( tokenizer, seqlen, apply_chat_template=apply_chat_template, system_prompt=system_prompt ) - try: - calib_dataset = load_dataset("NeelNanda/pile-10k", split=split) - except Exception as e: - import ssl - - error_message = str(e) - # Check for proxy or SSL error - if "proxy" in error_message.lower() or isinstance(e, ssl.SSLError) or "SSL" in error_message.upper(): - logger.error( - f"Network error detected, please check proxy settings. " - f"Error: {error_message}. Or consider using a backup dataset by `pip install modelscope` " - f"and set '--dataset swift/pile-val-backup' in AutoRound API." - ) - else: - logger.error(f"Failed to load the dataset: {error_message}") - sys.exit(1) + calib_dataset = load_dataset("NeelNanda/pile-10k", split=split) calib_dataset = calib_dataset.shuffle(seed=seed) calib_dataset = calib_dataset.map( tokenizer_function, @@ -221,46 +252,65 @@ def get_pile_dataset( return calib_dataset -@register_dataset(["swift/pile-val-backup", "pile-val-backup"]) -def get_pile_val_dataset( +@register_dataset([_FINEWEB_EDU_HF_DATASET, _FINEWEB_EDU_MODELSCOPE_DATASET, "fineweb-edu"]) +def get_fineweb_edu_dataset( tokenizer, seqlen, - dataset_name="swift/pile-val-backup", + dataset_name="fineweb-edu", split=None, seed=42, apply_chat_template=False, system_prompt=None, ): - """Returns a dataloader for the specified dataset and split. + """Return a streaming FineWeb-Edu calibration dataset from either supported hub. Args: tokenizer: The tokenizer to be used for tokenization. seqlen: The maximum sequence length. - data_name: The name of the dataset. - split: The data split to be used (e.g., "train", "test", "validation"). + dataset_name: FineWeb-Edu alias or explicit Hugging Face/ModelScope repository. + split: The data split to use. FineWeb-Edu currently provides ``train``. seed: The random seed for shuffling the dataset. apply_chat_template: Whether to apply chat template in tokenization. Returns: - A dataloader for the specified dataset and split, using the provided tokenizer and sequence length. + A tokenized streaming dataset using the provided tokenizer and sequence length. """ - - split = "validation" + split = "train" if split is None else split + if isinstance(split, list): + if len(split) != 1: + raise ValueError("FineWeb-Edu supports only one split at a time.") + split = split[0] + if split != "train": + raise ValueError("FineWeb-Edu supports only the train split.") tokenizer_function = get_tokenizer_function( tokenizer, seqlen, apply_chat_template=apply_chat_template, system_prompt=system_prompt ) - from transformers.utils.versions import require_version - - require_version( - "modelscope", - "Loading 'swift/pile-val-backup' dataset requires modelscope to be installed, " "`pip install modelscope`", + use_modelscope = dataset_name == _FINEWEB_EDU_MODELSCOPE_DATASET or ( + dataset_name == "fineweb-edu" and envs.AR_USE_MODELSCOPE ) - from modelscope import MsDataset # pylint: disable=E0401 + if use_modelscope: + from transformers.utils.versions import require_version + + require_version( + "modelscope", + "Loading FineWeb-Edu from ModelScope requires `modelscope`; install it with `pip install modelscope`.", + ) + from modelscope import MsDataset # pylint: disable=E0401 - calib_dataset = MsDataset.load( - "swift/pile-val-backup", "default", split=split - ).to_iterable_dataset() # , use_streaming=True + calib_dataset = MsDataset.load( + _FINEWEB_EDU_MODELSCOPE_DATASET, + subset_name=_FINEWEB_EDU_CONFIG, + split=split, + use_streaming=True, + ) + else: + calib_dataset = load_dataset( + _FINEWEB_EDU_HF_DATASET, + name=_FINEWEB_EDU_CONFIG, + split=split, + streaming=True, + ) calib_dataset = calib_dataset.shuffle(seed=seed).take(10000) calib_dataset = calib_dataset.map(tokenizer_function, batched=True) @@ -1109,7 +1159,11 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n """ # Allow disabling subprocess mode via environment variable if envs.AR_DISABLE_DATASET_SUBPROCESS: - return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) + try: + return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) + except Exception as error: + _warn_on_dataset_network_error(error, dataset_name) + raise # Run preprocessing in a subprocess so all temporary memory is freed on exit. # The HuggingFace datasets cache is warmed up as a side effect. @@ -1139,7 +1193,11 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n # (Re-)load the dataset in the main process. When the subprocess # succeeded the HF datasets cache makes this almost instant. - return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) + try: + return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) + except Exception as error: + _warn_on_dataset_network_error(error, dataset_name) + raise def get_dataloader(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, bs=8, nsamples=512): diff --git a/docs/step_by_step.md b/docs/step_by_step.md index d8221a92cb..6f222913b8 100755 --- a/docs/step_by_step.md +++ b/docs/step_by_step.md @@ -63,7 +63,9 @@ pip install auto-round The [NeelNanda/pile-10k](https://huggingface.co/datasets/NeelNanda/pile-10k) in huggingface is adopted as the default calibration data and will be downloaded automatically from the datasets Hub. Other available datasets include: -- `swift/pile-val-backup` from modelscope for addressing HF network issue +- [`fineweb-edu`](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) from Hugging Face. If Hugging Face is + inaccessible, install `modelscope`, set `AR_USE_MODELSCOPE=1`, and use the same alias to load the + [ModelScope mirror](https://modelscope.cn/datasets/AI-ModelScope/fineweb-edu) - `BAAI/CCI3-HQ` for Chinese - `codeparrot/github-code-clean` for code - `HuggingFaceH4/ultrachat_200k` for chat data diff --git a/docs/step_by_step_CN.md b/docs/step_by_step_CN.md index 4dced5f010..071e418f92 100755 --- a/docs/step_by_step_CN.md +++ b/docs/step_by_step_CN.md @@ -61,10 +61,11 @@ pip install auto-round ## 2 准备标定数据集 ### 默认数据集 -**对于中国大陆用户推荐使用 ModelScope 中的 swift/pile-val-backup 以解决 Huggingface 不能访问的问题** +**如果无法访问 Hugging Face，建议安装 `modelscope`、设置 `AR_USE_MODELSCOPE=1`，并使用 `fineweb-edu` 数据集。** 默认标定数据集为 Hugging Face 上的 [NeelNanda/pile-10k](https://huggingface.co/datasets/NeelNanda/pile-10k) ，该数据集会自动从 Huggingface Hub 下载。同时也支持使用以下数据集： -- ModelScope 中的 `swift/pile-val-backup`：用于解决 HF 访问问题 +- Hugging Face 上的 [`fineweb-edu`](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)。如果无法访问 Hugging Face， + 可安装 `modelscope`、设置 `AR_USE_MODELSCOPE=1`，并使用同一别名从 [ModelScope 镜像](https://modelscope.cn/datasets/AI-ModelScope/fineweb-edu) 加载 - `BAAI/CCI3-HQ`：用于中文场景 - `codeparrot/github-code-clean`：用于代码场景 - `HuggingFaceH4/ultrachat_200k`：用于对话数据 diff --git a/test/unit/common/utils/test_calib_dataset_helpers.py b/test/unit/common/utils/test_calib_dataset_helpers.py index c210881431..4268525570 100644 --- a/test/unit/common/utils/test_calib_dataset_helpers.py +++ b/test/unit/common/utils/test_calib_dataset_helpers.py @@ -13,6 +13,9 @@ # limitations under the License. """Tests for the small pure helpers in ``auto_round/calib_dataset.py``.""" +import ssl +import sys +import types from unittest.mock import MagicMock import pytest @@ -63,6 +66,68 @@ class _Returned: CALIB_DATASETS.pop("_return_check_", None) +# --------------------------------------------------------------------------- +# network failures and FineWeb-Edu routing +# --------------------------------------------------------------------------- +class TestDatasetNetworkErrors: + def test_network_warning_recommends_fineweb_edu(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset + + warning = MagicMock() + monkeypatch.setattr(calib_dataset.logger, "warning", warning) + + calib_dataset._warn_on_dataset_network_error(ssl.SSLError("proxy unavailable"), "dataset") + + warning.assert_called_once() + message = warning.call_args.args[0] + assert "--dataset fineweb-edu" in message + assert "AR_USE_MODELSCOPE=1" in message + + +class TestFineWebEduDataset: + @staticmethod + def _streaming_dataset_mock(): + dataset = MagicMock() + dataset.shuffle.return_value = dataset + dataset.take.return_value = dataset + dataset.map.return_value = dataset + return dataset + + def test_alias_loads_huggingface_sample_by_default(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset + + dataset = self._streaming_dataset_mock() + load_dataset = MagicMock(return_value=dataset) + monkeypatch.setattr(calib_dataset, "load_dataset", load_dataset) + monkeypatch.setattr(calib_dataset.envs, "AR_USE_MODELSCOPE", False) + + result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128) + + assert result is dataset + load_dataset.assert_called_once_with( + "HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True + ) + dataset.shuffle.assert_called_once_with(seed=42) + dataset.take.assert_called_once_with(10000) + + def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset + + dataset = self._streaming_dataset_mock() + load = MagicMock(return_value=dataset) + modelscope = types.ModuleType("modelscope") + modelscope.MsDataset = types.SimpleNamespace(load=load) + monkeypatch.setitem(sys.modules, "modelscope", modelscope) + monkeypatch.setattr(calib_dataset.envs, "AR_USE_MODELSCOPE", True) + + result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128) + + assert result is dataset + load.assert_called_once_with( + "AI-ModelScope/fineweb-edu", subset_name="sample-10BT", split="train", use_streaming=True + ) + + # --------------------------------------------------------------------------- # _make_map_fingerprint # --------------------------------------------------------------------------- From 7467dd4aa88bad83297a5f8fc4a53021b78c6494 Mon Sep 17 00:00:00 2001 From: WeiweiZhang1 Date: Mon, 14 Sep 2026 10:32:50 +0800 Subject: [PATCH 2/3] fix CI Signed-off-by: WeiweiZhang1 --- test/unit/common/utils/test_calib_dataset_helpers.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test/unit/common/utils/test_calib_dataset_helpers.py b/test/unit/common/utils/test_calib_dataset_helpers.py index 4268525570..ff6ab882bc 100644 --- a/test/unit/common/utils/test_calib_dataset_helpers.py +++ b/test/unit/common/utils/test_calib_dataset_helpers.py @@ -93,13 +93,14 @@ def _streaming_dataset_mock(): dataset.map.return_value = dataset return dataset - def test_alias_loads_huggingface_sample_by_default(self, monkeypatch): + def test_alias_loads_huggingface_sample_without_modelscope(self, monkeypatch): import auto_round.calib_dataset as calib_dataset dataset = self._streaming_dataset_mock() load_dataset = MagicMock(return_value=dataset) monkeypatch.setattr(calib_dataset, "load_dataset", load_dataset) - monkeypatch.setattr(calib_dataset.envs, "AR_USE_MODELSCOPE", False) + monkeypatch.delenv("AR_USE_MODELSCOPE", raising=False) + monkeypatch.setitem(sys.modules, "modelscope", None) result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128) @@ -111,6 +112,8 @@ def test_alias_loads_huggingface_sample_by_default(self, monkeypatch): dataset.take.assert_called_once_with(10000) def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch): + import transformers.utils.versions as transformers_versions + import auto_round.calib_dataset as calib_dataset dataset = self._streaming_dataset_mock() @@ -118,7 +121,8 @@ def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch): modelscope = types.ModuleType("modelscope") modelscope.MsDataset = types.SimpleNamespace(load=load) monkeypatch.setitem(sys.modules, "modelscope", modelscope) - monkeypatch.setattr(calib_dataset.envs, "AR_USE_MODELSCOPE", True) + monkeypatch.setenv("AR_USE_MODELSCOPE", "1") + monkeypatch.setattr(transformers_versions, "require_version", MagicMock()) result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128) From 72fe29d2d57d1b45020d81804cd5558b6c618cb6 Mon Sep 17 00:00:00 2001 From: WeiweiZhang1 Date: Tue, 15 Sep 2026 11:19:02 +0800 Subject: [PATCH 3/3] refine fineweb calib, add more CI Signed-off-by: WeiweiZhang1 --- auto_round/calib_dataset.py | 94 +++++++++++++++---- .../core/test_calib_dataset_subprocess.py | 57 +++++++++++ .../utils/test_calib_dataset_helpers.py | 86 +++++++++++++++-- 3 files changed, 210 insertions(+), 27 deletions(-) diff --git a/auto_round/calib_dataset.py b/auto_round/calib_dataset.py index f2e7a8192e..49f59a8c8f 100644 --- a/auto_round/calib_dataset.py +++ b/auto_round/calib_dataset.py @@ -17,9 +17,11 @@ import logging import multiprocessing import os +import queue import random import ssl import sys +from importlib.metadata import PackageNotFoundError, version logging.getLogger("datasets").setLevel(logging.WARNING) @@ -36,6 +38,9 @@ _FINEWEB_EDU_HF_DATASET = "HuggingFaceFW/fineweb-edu" _FINEWEB_EDU_MODELSCOPE_DATASET = "AI-ModelScope/fineweb-edu" _FINEWEB_EDU_CONFIG = "sample-10BT" +_FINEWEB_EDU_MIN_CANDIDATES = 10000 +_DATASET_RESULT_SUCCESS = "success" +_DATASET_RESULT_ERROR = "error" def _get_dataset_network_error(error): @@ -56,28 +61,55 @@ def _get_dataset_network_error(error): return current if any( marker in error_message - for marker in ("connection timed out", "max retries exceeded", "proxy error", "ssl error") + for marker in ( + "cannot send a request, as the client has been closed", + "connection timed out", + "max retries exceeded", + "proxy error", + "ssl error", + ) ): return current current = current.__cause__ or current.__context__ return None -def _warn_on_dataset_network_error(error, dataset_name): - """Give users an actionable cross-hub fallback for dataset network failures.""" +def _fallback_to_fineweb_edu(error, tokenizer, seqlen, dataset_name, seed, nsamples): + """Switch to ModelScope FineWeb-Edu after a calibration dataset network failure.""" network_error = _get_dataset_network_error(error) if network_error is None: - return + raise error + try: + version("modelscope") + except PackageNotFoundError: + raise RuntimeError( + f"Failed to load calibration dataset {dataset_name!r} because of an HTTP/proxy network issue: " + f"{network_error}. Install ModelScope with `pip install modelscope` to enable automatic fallback " + "to FineWeb-Edu." + ) from error + + if dataset_name == _FINEWEB_EDU_MODELSCOPE_DATASET: + raise error + logger.warning( "Failed to load calibration dataset %r because of an HTTP/proxy network issue: %s. " - "Check your network and proxy settings, or switch to FineWeb-Edu with `--dataset fineweb-edu`. " - "By default it loads from Hugging Face (%s); set `AR_USE_MODELSCOPE=1` and install `modelscope` " - "to load its ModelScope mirror (%s).", + "Automatically switching to the ModelScope FineWeb-Edu mirror (%s).", dataset_name, network_error, - _FINEWEB_EDU_HF_DATASET, _FINEWEB_EDU_MODELSCOPE_DATASET, ) + return _get_dataset_impl(tokenizer, seqlen, _FINEWEB_EDU_MODELSCOPE_DATASET, seed, nsamples) + + +def _preprocess_dataset_in_subprocess(result_queue, tokenizer, seqlen, dataset_name, seed, nsamples): + """Run dataset preprocessing and report network failures to the parent process.""" + try: + _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) + except Exception as error: + network_error = _get_dataset_network_error(error) + result_queue.put((_DATASET_RESULT_ERROR, str(network_error) if network_error is not None else None)) + raise + result_queue.put((_DATASET_RESULT_SUCCESS, None)) def get_code_calibration_dataset(nsamples, datasets_version=None): @@ -261,6 +293,7 @@ def get_fineweb_edu_dataset( seed=42, apply_chat_template=False, system_prompt=None, + max_samples=_FINEWEB_EDU_MIN_CANDIDATES, ): """Return a streaming FineWeb-Edu calibration dataset from either supported hub. @@ -304,6 +337,11 @@ def get_fineweb_edu_dataset( split=split, use_streaming=True, ) + # ModelScope's streaming reader can keep background HTTP requests alive + # after iteration. Materialize the bounded sample before tokenization so + # downloads finish cleanly and shuffle remains deterministic. + calib_dataset = Dataset.from_list(list(calib_dataset.take(max_samples))) + calib_dataset = calib_dataset.shuffle(seed=seed) else: calib_dataset = load_dataset( _FINEWEB_EDU_HF_DATASET, @@ -311,7 +349,7 @@ def get_fineweb_edu_dataset( split=split, streaming=True, ) - calib_dataset = calib_dataset.shuffle(seed=seed).take(10000) + calib_dataset = calib_dataset.shuffle(seed=seed).take(max_samples) calib_dataset = calib_dataset.map(tokenizer_function, batched=True) return calib_dataset @@ -1050,15 +1088,21 @@ def concat_dataset_element(dataset): raise ValueError( f"Dataset '{name}' is not found. Please choose from the supported datasets: {filtered_keys}." ) - dataset = get_dataset( - tokenizer, - seqlen, + dataset_kwargs = dict( + tokenizer=tokenizer, + seqlen=seqlen, seed=seed, split=split, dataset_name=name, apply_chat_template=apply_chat_template, system_prompt=system_prompt, ) + if get_dataset is get_fineweb_edu_dataset: + dataset_kwargs["max_samples"] = max( + _FINEWEB_EDU_MIN_CANDIDATES, + data_lens.get(name, nsamples), + ) + dataset = get_dataset(**dataset_kwargs) if do_concat: dataset = concat_dataset_element(dataset) @@ -1162,13 +1206,13 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n try: return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) except Exception as error: - _warn_on_dataset_network_error(error, dataset_name) - raise + return _fallback_to_fineweb_edu(error, tokenizer, seqlen, dataset_name, seed, nsamples) # Run preprocessing in a subprocess so all temporary memory is freed on exit. # The HuggingFace datasets cache is warmed up as a side effect. logger.info("Preprocessing calibration dataset in a subprocess to avoid memory leaks...") + subprocess_network_error = None try: if os.name == "nt": raise OSError("fork is not available on Windows") @@ -1178,26 +1222,38 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n # threads). Use "spawn" on macOS, which is safe but requires pickling args. mp_context = "spawn" if sys.platform == "darwin" else "fork" ctx = multiprocessing.get_context(mp_context) + result_queue = ctx.Queue() p = ctx.Process( - target=_get_dataset_impl, - args=(tokenizer, seqlen, dataset_name, seed, nsamples), + target=_preprocess_dataset_in_subprocess, + args=(result_queue, tokenizer, seqlen, dataset_name, seed, nsamples), ) p.start() p.join() + try: + result_status, network_error_message = result_queue.get(timeout=1) + except queue.Empty: + result_status, network_error_message = None, None + result_queue.close() + result_queue.join_thread() if p.exitcode != 0: - raise RuntimeError(f"Dataset preprocessing subprocess exited with code {p.exitcode}") + if result_status == _DATASET_RESULT_ERROR and network_error_message is not None: + subprocess_network_error = ConnectionError(network_error_message) + else: + raise RuntimeError(f"Dataset preprocessing subprocess exited with code {p.exitcode}") except Exception as e: logger.warning(f"Subprocess dataset preprocessing failed ({e}), falling back to in-process mode.") + if subprocess_network_error is not None: + return _fallback_to_fineweb_edu(subprocess_network_error, tokenizer, seqlen, dataset_name, seed, nsamples) + # (Re-)load the dataset in the main process. When the subprocess # succeeded the HF datasets cache makes this almost instant. try: return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples) except Exception as error: - _warn_on_dataset_network_error(error, dataset_name) - raise + return _fallback_to_fineweb_edu(error, tokenizer, seqlen, dataset_name, seed, nsamples) def get_dataloader(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, bs=8, nsamples=512): diff --git a/test/unit/common/core/test_calib_dataset_subprocess.py b/test/unit/common/core/test_calib_dataset_subprocess.py index b28a3c5e5e..7c077cd3bd 100644 --- a/test/unit/common/core/test_calib_dataset_subprocess.py +++ b/test/unit/common/core/test_calib_dataset_subprocess.py @@ -7,6 +7,9 @@ import os +import pytest +from datasets import Dataset + class _FakeProcess: """Minimal subprocess stub: starts, joins, and exits cleanly.""" @@ -23,6 +26,17 @@ def join(self): exitcode = 0 +class _FakeQueue: + def get(self, timeout=None): + raise __import__("queue").Empty + + def close(self): + pass + + def join_thread(self): + pass + + def _fake_get_context(captured): """Return a factory that records the requested multiprocessing context name.""" @@ -31,6 +45,7 @@ def get_context(method): class _FakeCtx: Process = _FakeProcess + Queue = _FakeQueue return _FakeCtx() @@ -81,3 +96,45 @@ def test_windows_falls_back_to_inprocess(monkeypatch): cd.get_dataset(tokenizer=None, seqlen=512) assert inprocess_called, "in-process fallback should have been called on Windows" + + +def test_subprocess_network_error_falls_back_without_retrying_source(monkeypatch): + """A child-process network error must switch directly to FineWeb-Edu.""" + import auto_round.calib_dataset as cd + + class _NetworkErrorQueue: + def get(self, timeout=None): + return cd._DATASET_RESULT_ERROR, "simulated proxy failure" + + def close(self): + pass + + def join_thread(self): + pass + + class _NetworkErrorProcess(_FakeProcess): + exitcode = 1 + + class _NetworkErrorContext: + Process = _NetworkErrorProcess + Queue = _NetworkErrorQueue + + fallback_dataset = Dataset.from_dict({"input_ids": [[1]], "attention_mask": [[1]]}) + fallback = [] + + def fallback_to_fineweb(error, tokenizer, seqlen, dataset_name, seed, nsamples): + fallback.append((error, dataset_name)) + return fallback_dataset + + monkeypatch.setattr(cd.multiprocessing, "get_context", lambda method: _NetworkErrorContext()) + monkeypatch.setattr(cd.os, "name", "posix") + monkeypatch.setattr(cd.sys, "platform", "linux") + monkeypatch.setattr(cd.envs, "AR_DISABLE_DATASET_SUBPROCESS", False) + monkeypatch.setattr(cd, "_fallback_to_fineweb_edu", fallback_to_fineweb) + monkeypatch.setattr(cd, "_get_dataset_impl", lambda *args: pytest.fail("source retried")) + + result = cd.get_dataset(tokenizer=None, seqlen=128, dataset_name="source") + + assert result is fallback_dataset + assert isinstance(fallback[0][0], ConnectionError) + assert fallback[0][1] == "source" diff --git a/test/unit/common/utils/test_calib_dataset_helpers.py b/test/unit/common/utils/test_calib_dataset_helpers.py index ff6ab882bc..8b1abfebcc 100644 --- a/test/unit/common/utils/test_calib_dataset_helpers.py +++ b/test/unit/common/utils/test_calib_dataset_helpers.py @@ -70,24 +70,61 @@ class _Returned: # network failures and FineWeb-Edu routing # --------------------------------------------------------------------------- class TestDatasetNetworkErrors: - def test_network_warning_recommends_fineweb_edu(self, monkeypatch): + def test_network_error_without_modelscope_prompts_install(self, monkeypatch): import auto_round.calib_dataset as calib_dataset + monkeypatch.setattr(calib_dataset, "version", MagicMock(side_effect=calib_dataset.PackageNotFoundError)) + + with pytest.raises(RuntimeError, match="pip install modelscope"): + calib_dataset._fallback_to_fineweb_edu( + ssl.SSLError("proxy unavailable"), MagicMock(), 128, "dataset", 42, 1 + ) + + def test_network_error_with_modelscope_switches_to_fineweb_edu(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset + + fallback_dataset = MagicMock() + load_dataset = MagicMock(return_value=fallback_dataset) warning = MagicMock() + monkeypatch.setattr(calib_dataset, "version", MagicMock(return_value="1.0")) + monkeypatch.setattr(calib_dataset, "_get_dataset_impl", load_dataset) monkeypatch.setattr(calib_dataset.logger, "warning", warning) + tokenizer = MagicMock() + + result = calib_dataset._fallback_to_fineweb_edu( + ssl.SSLError("proxy unavailable"), tokenizer, 128, "dataset", 42, 2 + ) + + assert result is fallback_dataset + load_dataset.assert_called_once_with(tokenizer, 128, "AI-ModelScope/fineweb-edu", 42, 2) + assert "Automatically switching" in warning.call_args.args[0] + + def test_closed_http_client_is_treated_as_network_error(self): + import auto_round.calib_dataset as calib_dataset + + error = RuntimeError("Cannot send a request, as the client has been closed.") + + assert calib_dataset._get_dataset_network_error(error) is error + + def test_modelscope_fineweb_network_error_is_not_retried(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset - calib_dataset._warn_on_dataset_network_error(ssl.SSLError("proxy unavailable"), "dataset") + monkeypatch.setattr(calib_dataset, "version", MagicMock(return_value="1.0")) + load_dataset = MagicMock() + monkeypatch.setattr(calib_dataset, "_get_dataset_impl", load_dataset) + error = ConnectionError("ModelScope is unavailable") - warning.assert_called_once() - message = warning.call_args.args[0] - assert "--dataset fineweb-edu" in message - assert "AR_USE_MODELSCOPE=1" in message + with pytest.raises(ConnectionError, match="ModelScope is unavailable"): + calib_dataset._fallback_to_fineweb_edu(error, MagicMock(), 128, "AI-ModelScope/fineweb-edu", 42, 1) + + load_dataset.assert_not_called() class TestFineWebEduDataset: @staticmethod def _streaming_dataset_mock(): dataset = MagicMock() + dataset.n_shards = 1 dataset.shuffle.return_value = dataset dataset.take.return_value = dataset dataset.map.return_value = dataset @@ -116,7 +153,15 @@ def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch): import auto_round.calib_dataset as calib_dataset + streamed_samples = [ + { + "text": "sample text", + "input_ids": [1, 2, 3], + "attention_mask": [1, 1, 1], + } + ] dataset = self._streaming_dataset_mock() + dataset.take.return_value = streamed_samples load = MagicMock(return_value=dataset) modelscope = types.ModuleType("modelscope") modelscope.MsDataset = types.SimpleNamespace(load=load) @@ -124,12 +169,20 @@ def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch): monkeypatch.setenv("AR_USE_MODELSCOPE", "1") monkeypatch.setattr(transformers_versions, "require_version", MagicMock()) - result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128) + tokenizer = MagicMock( + side_effect=lambda texts, **kwargs: { + "input_ids": [[1, 2, 3] for _ in texts], + "attention_mask": [[1, 1, 1] for _ in texts], + } + ) + result = calib_dataset.get_fineweb_edu_dataset(tokenizer, 128) - assert result is dataset + assert isinstance(result, Dataset) + assert len(result) == 1 load.assert_called_once_with( "AI-ModelScope/fineweb-edu", subset_name="sample-10BT", split="train", use_streaming=True ) + dataset.take.assert_called_once_with(10000) # --------------------------------------------------------------------------- @@ -368,6 +421,23 @@ def test_fallback_when_template_fails(self): # _get_dataset_impl # --------------------------------------------------------------------------- class TestGetDatasetImpl: + def test_fineweb_edu_loader_receives_requested_sample_limit(self, monkeypatch): + import auto_round.calib_dataset as calib_dataset + + dataset = Dataset.from_dict( + { + "input_ids": [list(range(8))], + "attention_mask": [[1] * 8], + } + ) + loader = MagicMock(return_value=dataset) + monkeypatch.setattr(calib_dataset, "get_fineweb_edu_dataset", loader) + monkeypatch.setattr(calib_dataset, "CALIB_DATASETS", {"fineweb-edu": loader}) + + calib_dataset._get_dataset_impl(MagicMock(), 8, "fineweb-edu", nsamples=1) + + assert loader.call_args.kwargs["max_samples"] == calib_dataset._FINEWEB_EDU_MIN_CANDIDATES + def test_combines_sources_with_different_metadata_schemas(self, monkeypatch): import auto_round.calib_dataset as calib_dataset