From b7cf4c3ea8ab733ec71a055447d9578dab06d38e Mon Sep 17 00:00:00 2001 From: Sankalp Date: Sun, 19 Jul 2026 17:17:35 +0530 Subject: [PATCH 1/4] Include TAR path in WebDataset read errors --- .../packaged_modules/webdataset/webdataset.py | 9 +++ tests/packaged_modules/test_webdataset.py | 77 +++++++++++++++++++ 2 files changed, 86 insertions(+) diff --git a/src/datasets/packaged_modules/webdataset/webdataset.py b/src/datasets/packaged_modules/webdataset/webdataset.py index 30bbe139394..ea5eacf23d8 100644 --- a/src/datasets/packaged_modules/webdataset/webdataset.py +++ b/src/datasets/packaged_modules/webdataset/webdataset.py @@ -1,6 +1,7 @@ import io import json import re +import tarfile from itertools import islice from typing import Any, Callable @@ -28,6 +29,14 @@ class WebDataset(datasets.GeneratorBasedBuilder): @classmethod def _get_pipeline_from_tar(cls, tar_path, tar_iterator): + try: + yield from cls._get_pipeline_from_tar_without_error_context(tar_path, tar_iterator) + except tarfile.ReadError as error: + # Avoid tracked_str.__repr__, which can include a presigned origin URL in the traceback. + raise tarfile.ReadError(f"Failed to read TAR archive {str(tar_path)!r}: {error}") from error + + @classmethod + def _get_pipeline_from_tar_without_error_context(cls, tar_path, tar_iterator): current_example = {} for filename, f in tar_iterator: example_key, field_name = base_plus_ext(filename) diff --git a/tests/packaged_modules/test_webdataset.py b/tests/packaged_modules/test_webdataset.py index 7eef6db566a..1230cb22dfa 100644 --- a/tests/packaged_modules/test_webdataset.py +++ b/tests/packaged_modules/test_webdataset.py @@ -1,4 +1,5 @@ import json +import re import tarfile from pathlib import Path @@ -6,6 +7,7 @@ from datasets import Audio, DownloadManager, Features, Image, List, Value, Video from datasets.packaged_modules.webdataset.webdataset import WebDataset +from datasets.utils.track import tracked_str from ..utils import ( require_numpy1_on_windows, @@ -101,6 +103,34 @@ def bad_wds_file(tmp_path, image_file, text_file): return str(filename) +@pytest.fixture +def corrupted_wds_file(tmp_path): + filename = tmp_path / "corrupted.tar" + filename.touch() + return str(filename) + + +@pytest.fixture +def valid_text_wds_file(tmp_path): + data_file = tmp_path / "valid.txt" + data_file.write_text("valid", encoding="utf-8") + filename = tmp_path / "valid.tar" + with tarfile.open(filename, "w") as tar: + tar.add(data_file, arcname="00000.txt") + return str(filename) + + +@pytest.fixture +def truncated_wds_file(tmp_path): + data_file = tmp_path / "data.txt" + data_file.write_bytes(b"x" * 1024) + filename = tmp_path / "truncated.tar" + with tarfile.open(filename, "w") as tar: + tar.add(data_file, arcname="00000.txt") + filename.write_bytes(filename.read_bytes()[:612]) + return str(filename) + + @pytest.fixture def tensor_wds_file(tmp_path, tensor_file): json_file = tmp_path / "data.json" @@ -305,6 +335,53 @@ def test_webdataset_errors_on_bad_file(bad_wds_file): webdataset._split_generators(DownloadManager()) +def test_webdataset_read_error_includes_tar_path_during_feature_inference(corrupted_wds_file): + webdataset = WebDataset(data_files={"train": [corrupted_wds_file]}) + + with pytest.raises(tarfile.ReadError, match=re.escape(corrupted_wds_file)) as raised: + webdataset._split_generators(DownloadManager()) + + assert isinstance(raised.value.__cause__, tarfile.ReadError) + + +def test_webdataset_read_error_identifies_corrupt_shard_after_valid_shard(valid_text_wds_file, corrupted_wds_file): + features = Features({"__key__": Value("string"), "__url__": Value("string"), "txt": Value("string")}) + webdataset = WebDataset(data_files={"train": [valid_text_wds_file, corrupted_wds_file]}, features=features) + split_generator = webdataset._split_generators(DownloadManager())[0] + generator = webdataset._generate_examples(**split_generator.gen_kwargs) + + _, first_example = next(generator) + assert first_example["__url__"] == valid_text_wds_file + + with pytest.raises(tarfile.ReadError, match=re.escape(corrupted_wds_file)) as raised: + next(generator) + + assert isinstance(raised.value.__cause__, tarfile.ReadError) + + +def test_webdataset_read_error_includes_tar_path_while_reading_member(truncated_wds_file): + features = Features({"__key__": Value("string"), "__url__": Value("string"), "txt": Value("string")}) + webdataset = WebDataset(data_files={"train": [truncated_wds_file]}, features=features) + split_generator = webdataset._split_generators(DownloadManager())[0] + + with pytest.raises(tarfile.ReadError, match=re.escape(truncated_wds_file)) as raised: + next(webdataset._generate_examples(**split_generator.gen_kwargs)) + + assert isinstance(raised.value.__cause__, tarfile.ReadError) + + +def test_webdataset_read_error_does_not_include_tracked_origin(corrupted_wds_file): + tar_path = tracked_str(corrupted_wds_file) + tar_path.set_origin("https://user:password@example.com/corrupted.tar?signature=SECRET") + tar_iterator = DownloadManager().iter_archive(str(tar_path)) + + with pytest.raises(tarfile.ReadError) as raised: + next(WebDataset._get_pipeline_from_tar(tar_path, tar_iterator)) + + assert corrupted_wds_file in str(raised.value) + assert "SECRET" not in str(raised.value) + + @require_pil def test_webdataset_with_features(image_wds_file): import PIL.Image From 5ccc7c6f908e48bf792becbb285f495842b3bf9c Mon Sep 17 00:00:00 2001 From: Sankalp Thakur Date: Wed, 22 Jul 2026 23:54:40 +0530 Subject: [PATCH 2/4] Include origin alongside path in WebDataset TAR ReadError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restore tracked_str origin (e.g. hf://…) in the re-raised ReadError so remote archive location remains visible for debugging, while still reporting the resolved local path. --- .../packaged_modules/webdataset/webdataset.py | 13 ++++++++-- tests/packaged_modules/test_webdataset.py | 24 +++++++++++++++---- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/src/datasets/packaged_modules/webdataset/webdataset.py b/src/datasets/packaged_modules/webdataset/webdataset.py index ea5eacf23d8..da336c471df 100644 --- a/src/datasets/packaged_modules/webdataset/webdataset.py +++ b/src/datasets/packaged_modules/webdataset/webdataset.py @@ -13,6 +13,7 @@ from datasets.features.features import cast_to_python_objects from datasets.filesystems import EXTENSION_TO_COMPRESSION_FS_FILE_CLS from datasets.utils.file_utils import xbasename +from datasets.utils.track import tracked_str logger = datasets.utils.logging.get_logger(__name__) @@ -32,8 +33,16 @@ def _get_pipeline_from_tar(cls, tar_path, tar_iterator): try: yield from cls._get_pipeline_from_tar_without_error_context(tar_path, tar_iterator) except tarfile.ReadError as error: - # Avoid tracked_str.__repr__, which can include a presigned origin URL in the traceback. - raise tarfile.ReadError(f"Failed to read TAR archive {str(tar_path)!r}: {error}") from error + # Include both the resolved local path and origin (e.g. hf://…) for debugging. + # Prefer str()/get_origin() over tracked_str.__repr__ so formatting stays explicit. + archive = str(tar_path) + if isinstance(tar_path, tracked_str): + origin = tar_path.get_origin() + if origin != archive: + raise tarfile.ReadError( + f"Failed to read TAR archive {archive!r} (origin={origin}): {error}" + ) from error + raise tarfile.ReadError(f"Failed to read TAR archive {archive!r}: {error}") from error @classmethod def _get_pipeline_from_tar_without_error_context(cls, tar_path, tar_iterator): diff --git a/tests/packaged_modules/test_webdataset.py b/tests/packaged_modules/test_webdataset.py index 1230cb22dfa..893b28bb8b2 100644 --- a/tests/packaged_modules/test_webdataset.py +++ b/tests/packaged_modules/test_webdataset.py @@ -370,16 +370,32 @@ def test_webdataset_read_error_includes_tar_path_while_reading_member(truncated_ assert isinstance(raised.value.__cause__, tarfile.ReadError) -def test_webdataset_read_error_does_not_include_tracked_origin(corrupted_wds_file): +def test_webdataset_read_error_includes_path_and_tracked_origin(corrupted_wds_file): + origin = "hf://datasets/org/name@main/data/corrupted.tar" tar_path = tracked_str(corrupted_wds_file) - tar_path.set_origin("https://user:password@example.com/corrupted.tar?signature=SECRET") + tar_path.set_origin(origin) tar_iterator = DownloadManager().iter_archive(str(tar_path)) with pytest.raises(tarfile.ReadError) as raised: next(WebDataset._get_pipeline_from_tar(tar_path, tar_iterator)) - assert corrupted_wds_file in str(raised.value) - assert "SECRET" not in str(raised.value) + message = str(raised.value) + assert corrupted_wds_file in message + assert origin in message + assert message.startswith(f"Failed to read TAR archive {corrupted_wds_file!r} (origin={origin}):") + + +def test_webdataset_read_error_omits_origin_when_same_as_path(corrupted_wds_file): + tar_path = tracked_str(corrupted_wds_file) + tar_path.set_origin(corrupted_wds_file) + tar_iterator = DownloadManager().iter_archive(str(tar_path)) + + with pytest.raises(tarfile.ReadError) as raised: + next(WebDataset._get_pipeline_from_tar(tar_path, tar_iterator)) + + message = str(raised.value) + assert corrupted_wds_file in message + assert "origin=" not in message @require_pil From 09d51b1a92ad23e89196e928a8b6857681f4f59a Mon Sep 17 00:00:00 2001 From: Sankalp Thakur Date: Thu, 23 Jul 2026 17:22:07 +0530 Subject: [PATCH 3/4] Sanitize origin URLs in WebDataset TAR ReadError Strip userinfo and query from tracked origins before embedding them in exceptions so presigned URLs and basic-auth credentials don't leak into logs/tracebacks while keeping scheme/host/path for debugging. --- .../packaged_modules/webdataset/webdataset.py | 19 ++++++++++++++++++- tests/packaged_modules/test_webdataset.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/datasets/packaged_modules/webdataset/webdataset.py b/src/datasets/packaged_modules/webdataset/webdataset.py index da336c471df..aee0a273aea 100644 --- a/src/datasets/packaged_modules/webdataset/webdataset.py +++ b/src/datasets/packaged_modules/webdataset/webdataset.py @@ -4,6 +4,7 @@ import tarfile from itertools import islice from typing import Any, Callable +from urllib.parse import urlsplit, urlunsplit import numpy as np import pyarrow as pa @@ -19,6 +20,21 @@ logger = datasets.utils.logging.get_logger(__name__) +def _sanitize_origin_for_error(origin: str) -> str: + """Strip userinfo and query from origin URLs before embedding in exceptions. + + Keeps scheme/host/path (e.g. hf://…) for debugging while avoiding leaking + basic-auth credentials or presigned URL query parameters into logs/tracebacks. + """ + parts = urlsplit(origin) + if not parts.scheme: + return origin + netloc = parts.netloc.rsplit("@", 1)[-1] if "@" in parts.netloc else parts.netloc + if netloc == parts.netloc and not parts.query: + return origin + return urlunsplit((parts.scheme, netloc, parts.path, "", parts.fragment)) + + class WebDataset(datasets.GeneratorBasedBuilder): DEFAULT_WRITER_BATCH_SIZE = 100 IMAGE_EXTENSIONS: list[str] # definition at the bottom of the script @@ -39,8 +55,9 @@ def _get_pipeline_from_tar(cls, tar_path, tar_iterator): if isinstance(tar_path, tracked_str): origin = tar_path.get_origin() if origin != archive: + safe_origin = _sanitize_origin_for_error(origin) raise tarfile.ReadError( - f"Failed to read TAR archive {archive!r} (origin={origin}): {error}" + f"Failed to read TAR archive {archive!r} (origin={safe_origin}): {error}" ) from error raise tarfile.ReadError(f"Failed to read TAR archive {archive!r}: {error}") from error diff --git a/tests/packaged_modules/test_webdataset.py b/tests/packaged_modules/test_webdataset.py index 893b28bb8b2..edf4bed46a1 100644 --- a/tests/packaged_modules/test_webdataset.py +++ b/tests/packaged_modules/test_webdataset.py @@ -385,6 +385,24 @@ def test_webdataset_read_error_includes_path_and_tracked_origin(corrupted_wds_fi assert message.startswith(f"Failed to read TAR archive {corrupted_wds_file!r} (origin={origin}):") +def test_webdataset_read_error_strips_userinfo_and_query_from_origin(corrupted_wds_file): + origin = "https://user:pass@example.com/bucket/file.tar?X-Amz-Signature=abc&Expires=1" + safe_origin = "https://example.com/bucket/file.tar" + tar_path = tracked_str(corrupted_wds_file) + tar_path.set_origin(origin) + tar_iterator = DownloadManager().iter_archive(str(tar_path)) + + with pytest.raises(tarfile.ReadError) as raised: + next(WebDataset._get_pipeline_from_tar(tar_path, tar_iterator)) + + message = str(raised.value) + assert corrupted_wds_file in message + assert safe_origin in message + assert "user:pass" not in message + assert "X-Amz-Signature" not in message + assert message.startswith(f"Failed to read TAR archive {corrupted_wds_file!r} (origin={safe_origin}):") + + def test_webdataset_read_error_omits_origin_when_same_as_path(corrupted_wds_file): tar_path = tracked_str(corrupted_wds_file) tar_path.set_origin(corrupted_wds_file) From 0307f14b8ba6a921310df8329cc4efa437ad5ed7 Mon Sep 17 00:00:00 2001 From: Sankalp Thakur Date: Sat, 25 Jul 2026 14:38:42 +0530 Subject: [PATCH 4/4] Move origin sanitization into tracked_str so all loaders benefit --- .../packaged_modules/webdataset/webdataset.py | 31 ++----------------- src/datasets/utils/track.py | 22 +++++++++++-- 2 files changed, 22 insertions(+), 31 deletions(-) diff --git a/src/datasets/packaged_modules/webdataset/webdataset.py b/src/datasets/packaged_modules/webdataset/webdataset.py index aee0a273aea..51c8f7caa2b 100644 --- a/src/datasets/packaged_modules/webdataset/webdataset.py +++ b/src/datasets/packaged_modules/webdataset/webdataset.py @@ -4,7 +4,6 @@ import tarfile from itertools import islice from typing import Any, Callable -from urllib.parse import urlsplit, urlunsplit import numpy as np import pyarrow as pa @@ -14,27 +13,11 @@ from datasets.features.features import cast_to_python_objects from datasets.filesystems import EXTENSION_TO_COMPRESSION_FS_FILE_CLS from datasets.utils.file_utils import xbasename -from datasets.utils.track import tracked_str logger = datasets.utils.logging.get_logger(__name__) -def _sanitize_origin_for_error(origin: str) -> str: - """Strip userinfo and query from origin URLs before embedding in exceptions. - - Keeps scheme/host/path (e.g. hf://…) for debugging while avoiding leaking - basic-auth credentials or presigned URL query parameters into logs/tracebacks. - """ - parts = urlsplit(origin) - if not parts.scheme: - return origin - netloc = parts.netloc.rsplit("@", 1)[-1] if "@" in parts.netloc else parts.netloc - if netloc == parts.netloc and not parts.query: - return origin - return urlunsplit((parts.scheme, netloc, parts.path, "", parts.fragment)) - - class WebDataset(datasets.GeneratorBasedBuilder): DEFAULT_WRITER_BATCH_SIZE = 100 IMAGE_EXTENSIONS: list[str] # definition at the bottom of the script @@ -49,17 +32,9 @@ def _get_pipeline_from_tar(cls, tar_path, tar_iterator): try: yield from cls._get_pipeline_from_tar_without_error_context(tar_path, tar_iterator) except tarfile.ReadError as error: - # Include both the resolved local path and origin (e.g. hf://…) for debugging. - # Prefer str()/get_origin() over tracked_str.__repr__ so formatting stays explicit. - archive = str(tar_path) - if isinstance(tar_path, tracked_str): - origin = tar_path.get_origin() - if origin != archive: - safe_origin = _sanitize_origin_for_error(origin) - raise tarfile.ReadError( - f"Failed to read TAR archive {archive!r} (origin={safe_origin}): {error}" - ) from error - raise tarfile.ReadError(f"Failed to read TAR archive {archive!r}: {error}") from error + # tar_path is a tracked_str whose repr appends its sanitized origin (e.g. hf://…), + # so every loader that reprs a tracked path gets the same safe error context. + raise tarfile.ReadError(f"Failed to read TAR archive {tar_path!r}: {error}") from error @classmethod def _get_pipeline_from_tar_without_error_context(cls, tar_path, tar_iterator): diff --git a/src/datasets/utils/track.py b/src/datasets/utils/track.py index ffd3434af79..941b79269c0 100644 --- a/src/datasets/utils/track.py +++ b/src/datasets/utils/track.py @@ -1,4 +1,20 @@ from collections.abc import Iterable, Iterator +from urllib.parse import urlsplit, urlunsplit + + +def _sanitize_origin(origin: str) -> str: + """Strip userinfo and query from origin URLs before embedding them in reprs/errors. + + Keeps scheme/host/path (e.g. hf://…) for debugging while avoiding leaking basic-auth + credentials or presigned-URL query parameters into logs and tracebacks. + """ + parts = urlsplit(origin) + if not parts.scheme: + return origin + netloc = parts.netloc.rsplit("@", 1)[-1] if "@" in parts.netloc else parts.netloc + if netloc == parts.netloc and not parts.query: + return origin + return urlunsplit((parts.scheme, netloc, parts.path, "", parts.fragment)) class tracked_str(str): @@ -12,10 +28,10 @@ def get_origin(self): return self.origins.get(super().__repr__(), str(self)) def __repr__(self) -> str: - if super().__repr__() not in self.origins or self.origins[super().__repr__()] == self: + origin = self.origins.get(super().__repr__()) + if origin is None or origin == self: return super().__repr__() - else: - return f"{str(self)} (origin={self.origins[super().__repr__()]})" + return f"{super().__repr__()} (origin={_sanitize_origin(origin)})" class tracked_list(list):