Adaptive per-call test set row limit at predict time

src/tabpfn_client/estimator.py

@@ -66,6 +66,11 @@
 
 THINKING_TIMEOUT_MAX_S = 40 * 60
 
+# Prediction compute scales with n_train_rows * n_test_rows, so the API caps
+# their product. The effective per-call test row limit therefore shrinks as
+# the fitted training set grows (e.g. 1M training rows -> 250k test rows).
+PREDICT_ROW_PAIRS_BUDGET = 250_000 * 1_000_000
+
 _VALID_THINKING_EFFORT_LEVELS = frozenset({"medium", "high"})
 
 
@@ -404,7 +409,14 @@ def _predict(
             predict_params=predict_params,
         )
 
-        validate_test_set(X, output_type, tabpfn_config.model_path)
+        validate_test_set(
+            X,
+            output_type,
+            tabpfn_config.model_path,
+            train_rows=self._last_train_X.shape[0]
+            if self._last_train_X is not None
+            else None,
+        )
         X_clean = _clean_text_features(X)
 
         if (
@@ -742,7 +754,14 @@ def predict(
             predict_params=predict_params,
         )
 
-        validate_test_set(X, output_type, tabpfn_config.model_path)
+        validate_test_set(
+            X,
+            output_type,
+            tabpfn_config.model_path,
+            train_rows=self._last_train_X.shape[0]
+            if self._last_train_X is not None
+            else None,
+        )
         X_clean = _clean_text_features(X)
 
         if (
@@ -921,6 +940,7 @@ def validate_test_set(
     X: pd.DataFrame | np.ndarray,
     output_type: str | None,
     model_path: str | None = None,
+    train_rows: int | None = None,
 ):
     """Check the integrity of the test data."""
 
@@ -934,9 +954,13 @@ def validate_test_set(
         model_version = model_version_from_path(model_path)
         limit = model_limit_from_version(model_version, limits.model_limits)
 
-    if X.shape[0] > limit.test_set_max_rows:
+    max_rows = limit.test_set_max_rows
+    if train_rows:
+        max_rows = min(max_rows, PREDICT_ROW_PAIRS_BUDGET // train_rows)
+
+    if X.shape[0] > max_rows:
         raise ValueError(
-            f"The number of test rows ({X.shape[0]}) exceeds the maximum of {limit.test_set_max_rows}. "
+            f"The number of test rows ({X.shape[0]}) exceeds the maximum of {max_rows}. "
             "Split the test set across multiple calls to reduce the number of rows."
         )
     if X.shape[1] > limit.max_cols:

tests/unit/test_validate_test_set.py

@@ -0,0 +1,64 @@
+from typing import Any
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from tabpfn_client.api_models import GetModelLimitsResponse
+from tabpfn_client.client import ServiceClient
+from tabpfn_client.estimator import PREDICT_ROW_PAIRS_BUDGET, validate_test_set
+
+
+def _limits(test_set_max_rows: int = 1_000_000) -> GetModelLimitsResponse:
+    model_limit: dict[str, Any] = {
+        "train_set_max_rows": 1_000_000,
+        "train_set_max_cells": 100_000_000,
+        "test_set_max_rows": test_set_max_rows,
+        "test_set_max_cells": 100_000_000,
+        "test_set_max_rows_w_full_regression_output": 400,
+        "max_cols": 2_000,
+        "max_classes": 10,
+    }
+    return GetModelLimitsResponse.model_validate(
+        {
+            "default_model_version": "v3",
+            "max_model_limit": model_limit,
+            "model_limits": {"v3": model_limit},
+            "dataset_max_size_bytes": 100_000_000,
+        }
+    )
+
+
+def _X(n_rows: int) -> np.ndarray:
+    return np.zeros((n_rows, 1))
+
+
+def _validate(n_test_rows: int, train_rows: int | None = None):
+    with patch.object(ServiceClient, "get_model_limits", return_value=_limits()):
+        validate_test_set(_X(n_test_rows), None, train_rows=train_rows)
+
+
+def test_static_limit_applies_without_train_rows():
+    _validate(1_000_000)
+    with pytest.raises(ValueError, match="exceeds the maximum of 1000000"):
+        _validate(1_000_001)
+
+
+def test_adaptive_limit_shrinks_with_train_rows():
+    # 1M train rows -> 250k test rows per call
+    _validate(250_000, train_rows=1_000_000)
+    with pytest.raises(ValueError, match="exceeds the maximum of 250000"):
+        _validate(250_001, train_rows=1_000_000)
+
+
+def test_adaptive_limit_never_exceeds_static_limit():
+    # Small train sets leave the static cap binding
+    small_train = PREDICT_ROW_PAIRS_BUDGET // 1_000_000 // 10
+    _validate(1_000_000, train_rows=small_train)
+    with pytest.raises(ValueError, match="exceeds the maximum of 1000000"):
+        _validate(1_000_001, train_rows=small_train)
+
+
+def test_no_limits_available_skips_validation():
+    with patch.object(ServiceClient, "get_model_limits", return_value=None):
+        validate_test_set(_X(2_000_000), None, train_rows=1_000_000)