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.

cursor

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

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

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

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>

simo-prior · 2026-07-24T08:49:32Z

        self.inference_precision = inference_precision
        self.random_state = random_state
        self.inference_config = inference_config
+        self.fit_mode: FitMode | None = fit_mode


Nit: probably not needed to specify the type as it will infer it anyway.

simo-prior · 2026-07-24T08:54:15Z

+
+
+def _build_model_handle(
+    estimator: Any,


We could have a type for the estimator here.

simo-prior · 2026-07-24T09:02:26Z

+) -> dict[str, Any] | Path:
+    """Shared implementation of `save_model()` for both estimators."""
+    check_is_fitted(estimator)
+    params = estimator.get_params(deep=False)


Here you want to use: estimator._get_tabpfn_config(), type safe and you get exactly the params that end up in tabpfn (no client_options for instance). Make it public if you prefer.

The reason we use get_params is because _get_tabpfn_config returns wide-config fields only - it drops thinking_*, model_id etc., and save_model needs the full reproducibility set.

simo-prior · 2026-07-24T09:23:43Z

+        return None
+
+
+def _read_model_handle(handle: dict[str, Any] | str | Path) -> dict[str, Any]:


Nit: too much indirection here, I would avoid the extra function.

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:31:26Z

+    """Read + validate a `save_model()` handle for the target estimator class."""
+    data = _read_model_handle(handle)
+    version = data.get("format_version")
+    if version != _MODEL_HANDLE_VERSION:


I would avoid the concept of _MODEL_HANDLE_VERSION altogether. The handle can be malformed for a multitude of reasons and in fact this could pass even when params is missing. I'd make a pydantic class for the model handle where params is a **TabpfnConfig, and if we fail to deserialize it we give a "Malformed error...possibly no longer support format" + we explain why the deserialization failed.

simo-prior · 2026-07-24T09:33:09Z

+def _load_model_handle_for(
+    cls: type,
+    handle: dict[str, Any] | str | Path,
+    task: Literal["classification", "regression"],


PredictionTask from api_models for type safety?

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.

simo-prior · 2026-07-24T09:39:18Z

+            served from the cache instead of re-fitting.
        client_options : ClientOptions, default=None
            Client specific options (e.g. timeout, headers).
+        model_id : str, UUID or None, default=None


Safe to set the model_id directly? Shouldn't it be private and save/load are basically setter and getter?

I'd keep it private because sklearn requires init params to be public attributes with the same name as the constructor args. Making it _model_id would break get_params() and clone()

simo-prior · 2026-07-24T09:42:13Z

        self._validate_targets_and_classes(y)

        if Config.use_server:
            # Create a new sentry trace at every fit, provided that:


I wonder whether we should change trace_id on predict_count > 0?

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>

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

simo-prior reviewed Jul 24, 2026

View reviewed changes

Comment thread src/tabpfn_client/estimator.py Outdated

simo-prior reviewed Jul 24, 2026

View reviewed changes

		return None


		def _read_model_handle(handle: dict[str, Any] \| str \| Path) -> dict[str, Any]:

Uh oh!

Conversation

safaricd commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Change Description

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

simo-prior Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simo-prior Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simo-prior Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simo-prior left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

safaricd commented Jul 23, 2026 •

edited

Loading

simo-prior Jul 24, 2026 •

edited

Loading

simo-prior Jul 24, 2026 •

edited

Loading

simo-prior Jul 24, 2026 •

edited

Loading