Enable fit_with_cache for predictions on API with K/V by safaricd · Pull Request #347 · PriorLabs/tabpfn-client

safaricd · 2026-07-23T19:29:04Z

Change Description

Enable fit_with_cache predictions - previously setting the fit_mode was blocked by the API.
Naming conventions - refactor to be closer with scikit-learn standards.

Root-cause fixes for the two cursor review findings on the KV-cache branch, plus follow-through on the CI failures that surfaced during the fix. - Rename `_last_fitted_train_set_id` -> `fitted_train_set_id_` to adopt sklearn's convention for learned attributes, and route reads through a new `_resolve_train_set_id(estimator)` helper. This dissolves the model_id / clone / set_params desync: `fit()` no longer mutates `self.model_id`, so `get_params()` / `clone()` stay sklearn-correct. - Call `init()` at the top of `_predict` so the cross-process `load_model().predict()` flow authorizes the HTTP client instead of 401-ing. `init()` is idempotent so it's free on the hot path. - Make the resolver swallow ValueError/TypeError from `_cast_fitted_id` so a bad `model_id` degrades to "not fitted" (surfacing sklearn's standard NotFittedError) rather than crashing `__sklearn_is_fitted__`. - Add `@overload`s to `save_model()` so the return type narrows on the `path` argument (dict when omitted, Path when given). Improves IDE ergonomics for real users and eliminates the union-narrowing pyright errors that were failing Trunk. - Use `FitMode.FIT_WITH_CACHE` (not the raw string) in tests and drop an unused test fixture. Ruff format on the touched files. - New targeted tests: clone preserves constructor `model_id` and stays fitted; predict from a cold state triggers `init()`; an invalid `model_id` reports unfitted rather than crashing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Drop `low_memory` and `batched` from the client-side FitMode enum. The server's public API (`FitRequest._validate_fit_mode`) only accepts `fit_preprocessors` and `fit_with_cache`; exposing the internal values in the client only invited runtime 400s. Type checkers now catch bad picks at construction time. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

simo-prior · 2026-07-24T09:38:36Z

+    crash `__sklearn_is_fitted__`. The estimator degrades to "not fitted"
+    and `predict` surfaces the standard `NotFittedError`.
+    """
+    fitted = getattr(estimator, "fitted_train_set_id_", None)


I don't think we still need fitted_train_set_id_, why not always use model_id?

I'd keep the fitted_train_set_id_ because of sklearn-compatibility. Specifically, if fit() writes to self.model_id, then clone(fitted_estimator) carries the fitted ID into the clone, which violates sklearn's clone contract - which should strip learned state.

@safaricd but cloning after load_model() would equally carry learned state to the new estimator. I think what we want is just one variable but rather a model_id_ to differentiate from hyperparams. This would solve the cloning problem and also other unexpected behaviors like setting estimator.model_id silently not working if we ever called fit() on the model.

simo-prior

@safaricd left comments, let me know.

Focused cleanups from the ENG-880 review that don't need discussion: - Inline `_cast_fitted_id` into `_resolve_train_set_id`; only one caller remains after the resolver refactor. - Drop redundant `: FitMode | None` annotations on `self.fit_mode`; the assignment already inherits the parameter type. - Use `PredictionTask` enum in `_build_model_handle` / `_load_model_handle_for` signatures and at both call sites, replacing the free-form `Literal["classification", "regression"]`. - Replace the dict-based handle with a pydantic `_ModelHandle` model: free structural validation on load with a clear error, `format_version` kept as an explicit schema-evolution gate (orthogonal to structural validation), classes typed as an optional list. `save_model()` still returns a plain dict / Path for JSON portability. - Introduce a small `_SavableEstimator` Protocol so the helpers declare what they actually read instead of taking `Any`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

cursor

Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.

There are 3 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit a27babb. Configure here.}

… model_id_: drop resolver, legacy fitted_ flag and constructor param; store handle id top-level; restore optional predict_row_pairs_budget

…ing-classes check into handle validation, document None-as-unset dump convention

ggprior

LGTM mostly,
ignore the stuff about save/load model -- if it's already there, fine
I think we should address leakage of server side internals
Also the fit_mode which looks auto-generated
Then this is GTG

ggprior · 2026-08-03T12:59:52Z


 _VALID_THINKING_EFFORT_LEVELS = frozenset({"medium", "high"})

+# `FitMode` exposes only the two values gapi wires end-to-end:


internal info which is not needed here?

ggprior · 2026-08-03T13:01:30Z

    ) = None


+class FitMode(str, Enum):


why do we declare enums like low_memory? clients cannot set that

ggprior · 2026-08-03T13:02:50Z

        self._fit_count = 0

+    def __sklearn_is_fitted__(self) -> bool:
+        return getattr(self, "model_id_", None) is not None


is this actually true (even if you don't fit with cache, is this parameter set correctly)

It's set on every fit, see: https://github.com/PriorLabs/tabpfn-client/pull/347/changes/BASE..9b85e7285c23d57535274f6567a9f42508e996d0#diff-33bad89c7747dd680e3135ec5d350edb8949a3e1452e672861f1dbb6a5ed9856R491

ggprior · 2026-08-03T13:05:09Z

+    @overload
+    def save_model(self, path: None = None) -> dict[str, Any]: ...
+    @overload
+    def save_model(self, path: str | Path) -> Path: ...


pretty cool! but isn't it conflicting with the server-side caching/dedup logic? Shouldn't we remove that first and then add save/load semantics afterwards?

I feel this is overloading this PR to be honest. I.e. save/load semantics are a separate concern from adding KV cache support, and should be tracked separately.
I think this needs a lot more work in terms of testing before shipping something here

ggprior · 2026-08-03T13:07:42Z

        self._last_train_set_description = None
        self._fit_count = 0

+    def __sklearn_is_fitted__(self) -> bool:


Why do we need the is fitted function here, it would be perfectly valid with underscore attribute convention as well
https://scikit-learn.org/stable/modules/generated/sklearn.utils.validation.check_is_fitted.html#sklearn.utils.validation.check_is_fitted

On the classifier classes_ is set before the remote fit call succeeds, so model_id_ might never get populated. I don't think the regressor needs it today, but it doesn't hurt to keep it on both for consistency.

Enable fit_with_cache for predictions on API with K/V

3d67a48

safaricd requested a review from a team as a code owner July 23, 2026 19:29

safaricd requested review from simo-prior and removed request for a team and simo-prior July 23, 2026 19:29

cursor Bot reviewed Jul 23, 2026

View reviewed changes

Comment thread src/tabpfn_client/estimator.py

Comment thread src/tabpfn_client/estimator.py Outdated

Fit with cache support

bc2e672

cursor Bot reviewed Jul 23, 2026

View reviewed changes

Comment thread src/tabpfn_client/estimator.py Outdated

safaricd and others added 2 commits July 23, 2026 23:26

Fix failing unit tests

4a9ca2f

safaricd commented Jul 23, 2026

View reviewed changes

Comment thread src/tabpfn_client/api_models.py Outdated

safaricd requested a review from simo-prior July 23, 2026 22:18