[FIX] Cap listing request batches at 1000

openml/utils/_openml.py

@@ -27,6 +27,8 @@
 import openml._api_calls
 import openml.exceptions
 
+MAX_LISTING_BATCH_SIZE = 1_000
+
 # Avoid import cycles: https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
 if TYPE_CHECKING:
     from openml.base import OpenMLBase
@@ -249,7 +251,7 @@ def _list_all(  # noqa: C901
     *,
     limit: int | None = None,
     offset: int | None = None,
-    batch_size: int | None = 10_000,
+    batch_size: int | None = MAX_LISTING_BATCH_SIZE,
 ) -> list[_SizedT]:
     """Helper to handle paged listing requests.
 
@@ -263,7 +265,7 @@ def _list_all(  # noqa: C901
         Call listing, e.g. list_evaluations. Takes two positional
         arguments: batch_size and offset.
     batch_size : int, optional
-        The batch size to use for the listing call.
+        The batch size to use for the listing call. Requests are capped at 1,000 items.
     offset : int, optional
         The initial offset to use for the listing call.
     limit : int, optional
@@ -279,7 +281,7 @@ def _list_all(  # noqa: C901
     results: list[_SizedT] = []
 
     offset = offset if offset is not None else 0
-    batch_size = batch_size if batch_size is not None else 10_000
+    batch_size = batch_size if batch_size is not None else MAX_LISTING_BATCH_SIZE
 
     LIMIT = limit
     BATCH_SIZE_ORIG = batch_size
@@ -290,6 +292,8 @@ def _list_all(  # noqa: C901
     if not isinstance(BATCH_SIZE_ORIG, int):
         raise ValueError(f"'batch_size' should be an integer but got {BATCH_SIZE_ORIG}")
 
+    BATCH_SIZE_ORIG = min(BATCH_SIZE_ORIG, MAX_LISTING_BATCH_SIZE)
+
     if (LIMIT is not None) and (not isinstance(LIMIT, int)) and (not np.isinf(LIMIT)):
         raise ValueError(f"'limit' should be an integer or inf but got {LIMIT}")
 

tests/test_utils/test_utils.py

@@ -72,6 +72,21 @@ def test_list_all_with_multiple_batches(min_number_tasks_on_test_server):
     assert min_number_tasks_on_test_server <= sum(len(batch) for batch in batches)
 
 
+@pytest.mark.parametrize("batch_size", [None, 10_000])
+def test_list_all_caps_batch_size(batch_size):
+    calls = []
+
+    def listing_call(size, offset):
+        calls.append((size, offset))
+        return list(range(offset, min(offset + size, 2_500)))
+
+    kwargs = {} if batch_size is None else {"batch_size": batch_size}
+    batches = openml.utils._list_all(listing_call, limit=2_500, **kwargs)
+
+    assert calls == [(1_000, 0), (1_000, 1_000), (500, 2_000)]
+    assert sum(len(batch) for batch in batches) == 2_500
+
+
 @pytest.mark.test_server()
 def test_list_all_for_datasets(min_number_datasets_on_test_server):
     datasets = openml.datasets.list_datasets(