Bound S3 dataset listings

src/gasbench/dataset/config.py

@@ -57,6 +57,7 @@ class BenchmarkDatasetConfig:
     hf_revision: Optional[str] = None
     hf_subfolders: Optional[List[str]] = None
 
+    s3_prefixes: Optional[List[str]] = None  # Bucket-relative literal key prefixes
     include_paths: Optional[List[str]] = None
     exclude_paths: Optional[List[str]] = None
     # Filter media entries inside ZIP/TAR archives. Unlike include_paths and
@@ -349,6 +350,12 @@ def validate_dataset_config(
                 f"Must be one of {valid_sources}"
             )
 
+    if config_dict.get("s3_prefixes") is not None:
+        prefixes = config_dict["s3_prefixes"]
+        if (config_dict.get("source") != "s3" or not isinstance(prefixes, list)
+                or not prefixes or any(not isinstance(p, str) or not p for p in prefixes)):
+            errors.append(f"Dataset '{dataset_name}': s3_prefixes requires S3 and nonempty key prefixes")
+
     numeric_fields = [
         "media_per_archive",
         "archives_per_dataset",
@@ -387,6 +394,7 @@ def _dataset_dict_to_config(d: dict, **overrides) -> BenchmarkDatasetConfig:
         "source": d.get("source", "huggingface"),
         "hf_revision": d.get("hf_revision"),
         "hf_subfolders": d.get("hf_subfolders"),
+        "s3_prefixes": d.get("s3_prefixes"),
         "media_per_archive": d.get("media_per_archive", 100),
         "archives_per_dataset": d.get("archives_per_dataset", 5),
         "include_paths": d.get("include_paths"),
@@ -501,6 +509,8 @@ def _obfuscate_holdout_names(
                 hf_subfolders,
             ]
         )
+        if d.s3_prefixes:
+            fingerprint += "|s3_prefixes=" + ",".join(sorted(set(d.s3_prefixes)))
         short_hash = hashlib.sha1(fingerprint.encode("utf-8")).hexdigest()[:8]
         new_name = f"{d.media_type}-{d.modality}-holdout-{short_hash}"
         obfuscated.append(replace(d, name=new_name, original_name=orig_name))

src/gasbench/dataset/download/core.py

@@ -160,6 +160,7 @@ def download_and_extract(
                         max_files=None,  # need full list to shuffle and pick from
                         hf_revision=hf_revision,
                         hf_subfolders=hf_subfolders,
+                    s3_prefixes=getattr(dataset, "s3_prefixes", None),
                     )
                 except DatasetAccessError as e:
                     logger.warning(f"Skipping {dataset.name}: {e}")
@@ -206,6 +207,7 @@ def download_and_extract(
                     max_files=max_files_to_list,
                     hf_revision=hf_revision,
                     hf_subfolders=hf_subfolders,
+                    s3_prefixes=getattr(dataset, "s3_prefixes", None),
                 )
             except DatasetAccessError as e:
                 logger.warning(f"Skipping dataset {dataset.name}: {e}")

src/gasbench/dataset/download/listing.py

@@ -66,6 +66,7 @@ def _list_remote_dataset_files(
     max_files: Optional[int] = None,
     hf_revision: Optional[str] = None,
     hf_subfolders: Optional[List[str]] = None,
+    s3_prefixes: Optional[List[str]] = None,
 ) -> List[str]:
     """List available files in a dataset, filtered by source_format and path patterns.
 
@@ -123,17 +124,11 @@ def _list_remote_dataset_files(
         if not needs_week_filter:
             files = files[:max_files]
     elif source == "s3":
-        files = list_s3_files(path=dataset_path, extension=source_format)
-        if include_paths:
-            files = [
-                f for f in files if any(path_seg in f for path_seg in include_paths)
-            ]
-        if exclude_paths:
-            files = [
-                f for f in files if not any(path_seg in f for path_seg in exclude_paths)
-            ]
-        if not needs_week_filter:
-            files = files[:max_files]
+        files = list_s3_files(
+            path=dataset_path, extension=source_format, prefixes=s3_prefixes,
+            include_paths=include_paths, exclude_paths=exclude_paths,
+            max_files=listing_max,
+        )
     else:  # hf - supports early termination natively
         files = list_hf_files(
             repo_id=dataset_path,

src/gasbench/dataset/utils/s3_utils.py

@@ -125,68 +125,66 @@ def _parse_s3_path(path: str) -> Tuple[str, str]:
     return bucket, prefix
 
 
-def list_s3_files(path: str, extension=None) -> List[str]:
-    """List files from an S3 bucket.
-
-    Args:
-        path: Path in format 'bucket-name/prefix/path'
-        extension: Filter files by extension(s) (e.g., '.parquet' or ['.tar', '.tar.gz'])
-                  Special value 'frames' to detect frame directories
-
-    Returns:
-        List of file keys (paths within the bucket) or directory paths if extension='frames'
+def list_s3_files(path: str, extension=None, *, prefixes=None,
+                  include_paths=None, exclude_paths=None, max_files=None) -> List[str]:
+    """List matching keys, using optional bucket-relative literal prefixes.
+
+    Prefixes may end mid-directory name (for sharded corpora). Substring filters
+    retain their existing semantics. Apply the cap only after all filters.
     """
+    bucket, base = _parse_s3_path(path)
+    base = base.rstrip("/") + "/" if base else ""
+    if prefixes is not None and (
+        not isinstance(prefixes, (list, tuple)) or not prefixes
+        or any(not isinstance(p, str) or not p or not p.startswith(base) for p in prefixes)
+    ):
+        raise ValueError("S3 prefixes must be nonempty keys within the dataset path")
+    # Remove overlaps so each object is enumerated once, in global key order.
+    selected = []
+    for prefix in sorted(set(prefixes or [base])):
+        if not any(prefix.startswith(parent) for parent in selected):
+            selected.append(prefix)
+    if max_files is not None and max_files < 0:
+        raise ValueError("max_files must be nonnegative")
+    if max_files == 0:
+        return []
+    frames = extension == "frames"
+    extensions = IMAGE_FILE_EXTENSIONS if frames else extension
+    if isinstance(extensions, str):
+        extensions = [extensions]
+    extensions = tuple(e.lower() for e in extensions) if extensions else ()
+    files, seen = [], set()
     try:
-        bucket, prefix = _parse_s3_path(path)
-        client = _get_s3_client()
-
-        files = []
-        paginator = client.get_paginator("list_objects_v2")
-
-        if prefix and not prefix.endswith("/"):
-            prefix = prefix + "/"
-
-        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
-            if "Contents" not in page:
-                continue
-
-            for obj in page["Contents"]:
-                key = obj["Key"]
-                if key.endswith("/"):
-                    continue
-                files.append(key)
-
-        if extension == "frames":
-            frame_dirs = set()
-            for f in files:
-                if any(f.lower().endswith(ext) for ext in IMAGE_FILE_EXTENSIONS):
-                    parent = os.path.dirname(f)
-                    if parent:
-                        frame_dirs.add(parent)
-
-            logger.info(f"Found {len(frame_dirs)} frame directories in S3 bucket {bucket}/{prefix}")
-            return sorted(list(frame_dirs))
-
-        if extension and files:
-            if isinstance(extension, (list, tuple, set)):
-                exts = tuple(e.lower() for e in extension)
-                files = [f for f in files if f.lower().endswith(exts)]
-            else:
-                ext_lower = extension.lower()
-                files = [f for f in files if f.lower().endswith(ext_lower)]
-
-        logger.info(f"Found {len(files)} files in S3 bucket {bucket}/{prefix}")
-        return files
-
+        paginator = _get_s3_client().get_paginator("list_objects_v2")
+        for prefix in selected:
+            for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
+                for obj in page.get("Contents", []):
+                    key = obj["Key"]
+                    if key.endswith("/") or (extensions and not key.lower().endswith(extensions)):
+                        continue
+                    candidate = os.path.dirname(key) if frames else key
+                    if not candidate or candidate in seen:
+                        continue
+                    if include_paths and not any(p in candidate for p in include_paths):
+                        continue
+                    if exclude_paths and any(p in candidate for p in exclude_paths):
+                        continue
+                    seen.add(candidate)
+                    files.append(candidate)
+                    # Frame parents may interleave in key order; keep their historical
+                    # sorted selection instead of truncating before the sort.
+                    if not frames and max_files is not None and len(files) >= max_files:
+                        return files
+        if frames:
+            files.sort()
+        return files[:max_files] if max_files is not None else files
     except NoCredentialsError:
         logger.error("S3 credentials not found or invalid")
         return []
     except ClientError as e:
         logger.error(f"Failed to list S3 files from {path}: {e}")
         return []
-    except Exception as e:
-        logger.error(f"Error listing S3 files from {path}: {e}")
-        return []
+
 
 
 def _get_s3_urls(path: str, filenames: List[str]) -> List[str]:

tests/unit/test_s3_listing_limits.py

@@ -0,0 +1,70 @@
+from types import SimpleNamespace
+
+import pytest
+
+from src.gasbench.dataset.download.listing import _list_remote_dataset_files
+from src.gasbench.dataset.utils import s3_utils
+
+
+def mock_pages(monkeypatch, pages):
+    calls = []
+    def paginate(**kwargs):
+        prefix = kwargs["Prefix"]
+        calls.append(prefix)
+        yield from pages(prefix)
+    monkeypatch.setattr(s3_utils, "_get_s3_client", lambda: SimpleNamespace(
+        get_paginator=lambda _: SimpleNamespace(paginate=paginate)))
+    return calls
+
+
+def test_prefixes_keep_partial_shard_names_deduplicate_and_stop(monkeypatch):
+    def pages(prefix):
+        yield {"Contents": [{"Key": prefix + "001/no.txt"},
+                            {"Key": prefix + "001/excluded.jpg"},
+                            {"Key": prefix + "002/keep.PNG"}]}
+        raise AssertionError("Fetched a page beyond the matching limit")
+    calls = mock_pages(monkeypatch, pages)
+    assert _list_remote_dataset_files(
+        "bucket/root", ["jpg", "png"], source="s3",
+        s3_prefixes=["root/corpus-", "root/corpus-001/", "root/later-"],
+        exclude_paths=["excluded"], max_files=1,
+    ) == ["root/corpus-002/keep.PNG"]
+    assert calls == ["root/corpus-"]
+
+
+def test_substring_filters_still_match_nested_keys(monkeypatch):
+    mock_pages(monkeypatch, lambda _: iter([
+        {"Contents": [{"Key": "root/other/no.jpg"}]},
+        {"Contents": [{"Key": "root/deep/corpus/a.jpg"}]},
+    ]))
+    assert _list_remote_dataset_files(
+        "bucket/root", "jpg", source="s3", include_paths=["corpus"], max_files=1
+    ) == ["root/deep/corpus/a.jpg"]
+
+
+def test_week_filter_runs_before_limit(monkeypatch):
+    mock_pages(monkeypatch, lambda _: iter([
+        {"Contents": [{"Key": "gasstation/old/a.jpg"}]},
+        {"Contents": [{"Key": "gasstation/desired/b.jpg"}]},
+    ]))
+    assert _list_remote_dataset_files(
+        "bucket/gasstation", "jpg", source="s3", target_week="desired", max_files=1,
+    ) == ["gasstation/desired/b.jpg"]
+
+
+@pytest.mark.parametrize("prefixes", [[], ["outside/"], ["root-other/"], [""]])
+def test_prefix_cannot_escape_dataset_path(prefixes):
+    with pytest.raises(ValueError, match="prefixes"):
+        s3_utils.list_s3_files("bucket/root", prefixes=prefixes)
+
+
+def test_prefix_config_is_loaded_and_changes_cache_identity():
+    from src.gasbench.dataset.config import _dataset_dict_to_config, _obfuscate_holdout_names
+    row = {"name": "corpus", "path": "bucket/root", "source": "s3",
+           "modality": "image", "media_type": "real", "source_format": "jpg"}
+    original = _dataset_dict_to_config(row)
+    narrowed = _dataset_dict_to_config({**row, "s3_prefixes": ["root/corpus-"]})
+    assert narrowed.s3_prefixes == ["root/corpus-"]
+    before, _ = _obfuscate_holdout_names([original])
+    after, _ = _obfuscate_holdout_names([narrowed])
+    assert before[0].name != after[0].name