diff --git a/src/tabpfn_client/api_models.py b/src/tabpfn_client/api_models.py index f7ecc4f..70ce502 100644 --- a/src/tabpfn_client/api_models.py +++ b/src/tabpfn_client/api_models.py @@ -52,6 +52,14 @@ class ClassifierPredictParams(BaseModel): ) = None +class FitMode(str, Enum): + # Only the two values that gapi wires end-to-end. The server-side enum + # also has internal `low_memory` / `batched` modes, but those are not + # accepted by the public API — see `FitRequest._validate_fit_mode`. + FIT_PREPROCESSORS = "fit_preprocessors" + FIT_WITH_CACHE = "fit_with_cache" + + class ClassifierTabPFNConfig(BaseModel): n_estimators: int | None = None categorical_features_indices: list[int] | None = None @@ -65,6 +73,7 @@ class ClassifierTabPFNConfig(BaseModel): Annotated[Literal["autocast", "auto"] | str, Field(union_mode="left_to_right")] | None ) = None ignore_pretraining_limits: bool | None = None + fit_mode: Annotated[FitMode | UnknownEnum, Field(union_mode="left_to_right")] | None = None model_path: str | None = None balance_probabilities: bool | None = None @@ -161,6 +170,7 @@ class RegressorTabPFNConfig(BaseModel): Annotated[Literal["autocast", "auto"] | str, Field(union_mode="left_to_right")] | None ) = None ignore_pretraining_limits: bool | None = None + fit_mode: Annotated[FitMode | UnknownEnum, Field(union_mode="left_to_right")] | None = None model_path: str | None = None @@ -226,7 +236,7 @@ class FitRequest(BaseModel): ) async_mode: bool | None = Field( default=None, - description="Submit the fit and return immediately with status=pending; poll GET /tabpfn/fit/{fitted_train_set_id} for the outcome. Defaults to the blocking behaviour.", + description="Submit the fit and return immediately with status=pending; poll GET /fit/{fitted_train_set_id} for the outcome. Defaults to the blocking behaviour.", ) @@ -239,6 +249,7 @@ class FitStatusResponse(BaseModel): fitted_train_set_id: UUID status: Annotated[FitStatus | UnknownEnum, Field(union_mode="left_to_right")] error: str | None = None + error_code: str | None = None class GetModelLimitsResponse(BaseModel): diff --git a/src/tabpfn_client/estimator.py b/src/tabpfn_client/estimator.py index 811057e..be9666a 100644 --- a/src/tabpfn_client/estimator.py +++ b/src/tabpfn_client/estimator.py @@ -3,17 +3,20 @@ from __future__ import annotations +import json import logging import sys import time +from pathlib import Path from uuid import uuid4 from concurrent.futures import ThreadPoolExecutor -from typing import Any, Callable, Literal, cast, overload +from typing import Any, Callable, Literal, Protocol, cast, overload from typing_extensions import Self from uuid import UUID import numpy as np import pandas as pd +from pydantic import BaseModel, ValidationError from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin from sklearn.utils import column_or_1d from sklearn.utils.multiclass import check_classification_targets @@ -31,6 +34,8 @@ from tabpfn_client.utils import model_limit_from_version, model_version_from_path from tabpfn_client.service_wrapper import InferenceClient from tabpfn_client.api_models import ( + FitMode, + PredictionTask, RegressorTabPFNConfig, ClassifierTabPFNConfig, RegressorPredictParams, @@ -68,6 +73,143 @@ _VALID_THINKING_EFFORT_LEVELS = frozenset({"medium", "high"}) +# `FitMode` exposes only the two values gapi wires end-to-end: +# `fit_preprocessors` (the stateless default) and `fit_with_cache` (persist a +# server-side KV cache keyed by the returned fitted-train-set id, so later +# predicts against that id are served from the cache instead of re-fitting). + +# Bumped when the shape of the `_ModelHandle` schema changes in a way pydantic +# validation cannot otherwise catch (e.g. renamed fields, changed semantics). +_MODEL_HANDLE_VERSION = 1 + + +class _SavableEstimator(Protocol): + """Read-only structural type for the helpers below. Declares only what + they actually access on the estimator so the helpers don't need `Any`.""" + + model_id: str | UUID | None + + def get_params(self, deep: bool = False) -> dict[str, Any]: ... + + +class _ModelHandle(BaseModel): + """Serialized reference to a fitted (server-side) TabPFN model. + + Structural validation (missing keys, wrong types, etc.) is done by + pydantic on `load_model`; `format_version` remains as an explicit gate + for schema evolution that pydantic cannot see (e.g. semantic changes + to fields whose shape hasn't changed). + """ + + format_version: int + task: PredictionTask + params: dict[str, Any] + classes: list[Any] | None = None # classification only + + +def _resolve_train_set_id(estimator: _SavableEstimator) -> UUID | None: + """Resolve the server-side fitted-train-set id an estimator should predict against. + + A real `fit()` sets `fitted_train_set_id_` and always wins; otherwise fall + back to the constructor's `model_id` (the load-by-id path). Keeping this as + a resolver instead of syncing state at both mutation sites means `fit()` + never has to mutate `self.model_id`, so `get_params()` / `clone()` stay + sklearn-correct. + + Returns None if `model_id` is set but not a valid UUID — sklearn's + convention is that `__init__` does no validation, so bad ids must not + crash `__sklearn_is_fitted__`. The estimator degrades to "not fitted" + and `predict` surfaces the standard `NotFittedError`. + """ + fitted = getattr(estimator, "fitted_train_set_id_", None) + if fitted is not None: + return fitted + raw = estimator.model_id + if raw is None or isinstance(raw, UUID): + return raw + try: + return UUID(str(raw)) + except (ValueError, TypeError): + return None + + +def _read_handle_data(source: dict[str, Any] | str | Path) -> dict[str, Any]: + """Accept either an in-memory handle dict or a path to a JSON file + produced by `save_model()`.""" + if isinstance(source, dict): + return source + return json.loads(Path(source).read_text()) + + +def _build_model_handle( + estimator: _SavableEstimator, + task: PredictionTask, + path: str | Path | None, +) -> dict[str, Any] | Path: + """Shared implementation of `save_model()` for both estimators.""" + check_is_fitted(estimator) + params = estimator.get_params(deep=False) + # `client_options` is runtime-only (timeouts, auth/trace headers) and not + # JSON-serializable; drop it so the handle stays portable. + params.pop("client_options", None) + # Persist the *effective* id: a real `fit()` supersedes the constructor's + # `model_id`. Coerce to str for JSON. + params["model_id"] = str(_resolve_train_set_id(estimator)) + classes: list[Any] | None = None + if task is PredictionTask.CLASSIFICATION: + classes_attr = getattr(estimator, "classes_", None) + if classes_attr is None: + # A bare `model_id` classifier (never fit / never load_model'd) has + # no class labels to persist; fail clearly instead of AttributeError. + raise ValueError( + "Cannot save a classifier without `classes_`. Fit it first, or " + "reconstruct it via `load_model()`." + ) + classes = classes_attr.tolist() + handle = _ModelHandle( + format_version=_MODEL_HANDLE_VERSION, + task=task, + params=params, + classes=classes, + ) + dumped = handle.model_dump(mode="json", exclude_none=True) + if path is None: + return dumped + out = Path(path) + out.write_text(json.dumps(dumped, indent=2)) + return out + + +def _load_model_handle_for( + cls: type, + source: dict[str, Any] | str | Path, + task: PredictionTask, +) -> _ModelHandle: + """Read + validate a `save_model()` handle for the target estimator class. + + Pydantic catches structural problems (missing keys, wrong types, bad + enum values); `format_version` is checked separately as an explicit + schema-evolution gate; task mismatch is caught last with a clear error. + """ + raw = _read_handle_data(source) + try: + handle = _ModelHandle.model_validate(raw) + except ValidationError as exc: + raise ValueError( + f"Malformed model handle for {cls.__name__} " + f"(may be from an unsupported format version): {exc}" + ) from exc + if handle.format_version != _MODEL_HANDLE_VERSION: + raise ValueError( + f"Unsupported model handle version {handle.format_version!r}; " + f"expected {_MODEL_HANDLE_VERSION}." + ) + if handle.task is not task: + raise ValueError( + f"Cannot load a {handle.task.value!r} model into {cls.__name__}." + ) + return handle + class TabPFNModelSelection: """Base class for TabPFN model selection and path handling.""" @@ -179,6 +321,7 @@ def __init__( random_state: int | None = 0, inference_config: dict[str, Any] | None = None, categorical_features_indices: list[int] | None = None, + fit_mode: FitMode | None = None, # end: tabpfn_config paper_version: bool = False, thinking_mode: bool = False, @@ -186,6 +329,7 @@ def __init__( thinking_timeout_s: float | None = None, thinking_metric: str | None = None, client_options: ClientOptions | None = None, + model_id: str | UUID | None = None, ): """Construct a TabPFN classifier. @@ -277,8 +421,26 @@ def __init__( "roc_auc_ovr_micro", "roc_auc_ovr_weighted". Aliases "acc", "nll", "pac_score" are also accepted. + fit_mode: {"fit_preprocessors", "fit_with_cache"} or None, default=None + Controls what the server persists at fit time. None defers to the + server default, which is "fit_preprocessors". + "fit_preprocessors" fits only the preprocessing state, so every + predict re-runs the forward pass from the uploaded train set. + "fit_with_cache" additionally builds and persists a server-side KV + cache keyed by the resulting fitted-train-set id; later predicts + against that id (see `model_id` / `save_model` / `load_model`) are + 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 + Reference to an already-fitted train set on the server (the id + returned by a previous `fit()`, exposed via `save_model()`). When + set, the estimator is considered fitted and `predict`/`predict_proba` + run against that id without calling `fit()` again — this is what + lets a `fit_with_cache` model be reused across processes. A fresh + `fit()` call always supersedes it. Note that a classifier + constructed this way has no `classes_` unless it is restored via + `load_model()`. """ self.model_path = model_path self.categorical_features_indices = categorical_features_indices @@ -290,21 +452,29 @@ def __init__( self.inference_precision = inference_precision self.random_state = random_state self.inference_config = inference_config + self.fit_mode = fit_mode self.paper_version = paper_version self.thinking_mode = thinking_mode self.thinking_effort: ThinkingEffort | None = thinking_effort self.thinking_timeout_s = thinking_timeout_s self.thinking_metric = thinking_metric self.client_options = client_options or ClientOptions() + self.model_id = model_id self._last_trace_id = None - self._last_fitted_train_set_id = None self._last_train_X = None self._last_train_y = None self._last_meta = {} self._last_train_set_description = None self._fit_count = 0 + def __sklearn_is_fitted__(self) -> bool: + # `fitted_` is a legacy fittedness flag some tests set directly; it is + # kept as a fallback so those tests keep passing. + return _resolve_train_set_id(self) is not None or getattr( + self, "fitted_", False + ) + def fit( self, X: pd.DataFrame | np.ndarray, @@ -354,7 +524,7 @@ def fit_task() -> UUID: description=description, ) - self._last_fitted_train_set_id = cast(UUID, run_task(fit_task, "Fitting")) + self.fitted_train_set_id_ = cast(UUID, run_task(fit_task, "Fitting")) self._last_train_X = X_clean self._last_train_y = y self.fitted_ = True @@ -396,6 +566,11 @@ def _predict( # we capture the original user-provided values. predict_params = self._get_predict_params(locals()) + # A load-by-id estimator (via `model_id` / `load_model`) can reach + # `predict` without ever calling `fit()`, so `init()` (which authorizes + # the HTTP client) must run here too. It short-circuits after the first + # successful call. + init() check_is_fitted(self) tabpfn_config = self._get_tabpfn_config() @@ -424,14 +599,25 @@ def predict_task() -> PredictionResult: try: return InferenceClient.predict( X_clean, - fitted_train_set_id=cast(UUID, self._last_fitted_train_set_id), + fitted_train_set_id=cast(UUID, _resolve_train_set_id(self)), task_config=task_config, client_options=self.client_options, ) except NeedsRefittingError as exc: last_exc = exc refit_attempts += 1 - self._last_fitted_train_set_id = InferenceClient.fit( + if self._last_train_X is None or self._last_train_y is None: + # A load-by-id estimator (constructed via `model_id` / + # `load_model`) has no training data in memory to rebuild + # from, so we can't transparently recover here. + raise RuntimeError( + "The referenced fitted model is no longer available " + "on the server and this estimator has no in-memory " + "training data to refit (it was created via " + "`model_id` / `load_model`). Re-create it by calling " + "`fit(X, y)`." + ) from exc + self.fitted_train_set_id_ = InferenceClient.fit( self._last_train_X, self._last_train_y, task_config=task_config, @@ -493,6 +679,44 @@ def _validate_targets_and_classes(self, y) -> None: f"{URL_TABPFN_EXTENSIONS_GITHUB_MANY_CLASS_CODE}" ) + @overload + def save_model(self, path: None = None) -> dict[str, Any]: ... + @overload + def save_model(self, path: str | Path) -> Path: ... + def save_model(self, path: str | Path | None = None) -> dict[str, Any] | Path: + """Serialize a portable reference to the fitted (server-side) model. + + The handle captures the server-side fitted-train-set id, the estimator + hyperparameters, and the class labels — everything `load_model()` needs + to reconstruct an equivalent classifier. No training data leaves the + client; predictions run against the model referenced by the id, so this + is most useful with `fit_mode="fit_with_cache"`. + + Parameters + ---------- + path : str, Path or None, default=None + If given, the handle is written to this path as JSON and the path + is returned. If None, the handle dict is returned directly. + """ + return _build_model_handle(self, task=PredictionTask.CLASSIFICATION, path=path) + + @classmethod + def load_model(cls, handle: dict[str, Any] | str | Path) -> "TabPFNClassifier": + """Reconstruct a classifier from a `save_model()` handle (dict or path). + + The resulting estimator is already fitted (it references the server-side + model by id) and restores `classes_`, so it can `predict` without + calling `fit()` again. + """ + data = _load_model_handle_for(cls, handle, task=PredictionTask.CLASSIFICATION) + if data.classes is None: + raise ValueError( + f"Classifier handle is missing 'classes'; cannot reconstruct {cls.__name__}." + ) + est = cls(**data.params) + est.classes_ = np.asarray(data.classes) + return est + class TabPFNRegressor(RegressorMixin, BaseEstimator, TabPFNModelSelection): _AVAILABLE_MODELS = [ @@ -528,6 +752,7 @@ def __init__( random_state: int | None = 0, inference_config: dict[str, Any] | None = None, categorical_features_indices: list[int] | None = None, + fit_mode: FitMode | None = None, # end: tabpfn_config paper_version: bool = False, thinking_mode: bool = False, @@ -535,6 +760,7 @@ def __init__( thinking_timeout_s: float | None = None, thinking_metric: str | None = None, client_options: ClientOptions | None = None, + model_id: str | UUID | None = None, ): """Construct a TabPFN regressor. @@ -610,8 +836,24 @@ def __init__( Aliases "mse", "rmse", "mae", "mape", "smape" are also accepted. + fit_mode: {"fit_preprocessors", "fit_with_cache"} or None, default=None + Controls what the server persists at fit time. None defers to the + server default, which is "fit_preprocessors". + "fit_preprocessors" fits only the preprocessing state, so every + predict re-runs the forward pass from the uploaded train set. + "fit_with_cache" additionally builds and persists a server-side KV + cache keyed by the resulting fitted-train-set id; later predicts + against that id (see `model_id` / `save_model` / `load_model`) are + 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 + Reference to an already-fitted train set on the server (the id + returned by a previous `fit()`, exposed via `save_model()`). When + set, the estimator is considered fitted and `predict` runs against + that id without calling `fit()` again — this is what lets a + `fit_with_cache` model be reused across processes. A fresh `fit()` + call always supersedes it. """ self.model_path = model_path self.categorical_features_indices = categorical_features_indices @@ -622,21 +864,29 @@ def __init__( self.inference_precision = inference_precision self.random_state = random_state self.inference_config = inference_config + self.fit_mode = fit_mode self.paper_version = paper_version self.thinking_mode = thinking_mode self.thinking_effort: ThinkingEffort | None = thinking_effort self.thinking_timeout_s = thinking_timeout_s self.thinking_metric = thinking_metric self.client_options = client_options or ClientOptions() + self.model_id = model_id self._last_trace_id = None - self._last_fitted_train_set_id = None self._last_train_X = None self._last_train_y = None self._last_meta = {} self._last_train_set_description = None self._fit_count = 0 + def __sklearn_is_fitted__(self) -> bool: + # `fitted_` is a legacy fittedness flag some tests set directly; it is + # kept as a fallback so those tests keep passing. + return _resolve_train_set_id(self) is not None or getattr( + self, "fitted_", False + ) + def fit( self, X: pd.DataFrame | np.ndarray, @@ -687,7 +937,7 @@ def fit_task() -> UUID: description=description, ) - self._last_fitted_train_set_id = cast(UUID, run_task(fit_task, "Fitting")) + self.fitted_train_set_id_ = cast(UUID, run_task(fit_task, "Fitting")) self._last_train_X = X_clean self._last_train_y = y self.fitted_ = True @@ -734,6 +984,11 @@ def predict( # we capture the original user-provided values. predict_params = self._get_predict_params(locals()) + # A load-by-id estimator (via `model_id` / `load_model`) can reach + # `predict` without ever calling `fit()`, so `init()` (which authorizes + # the HTTP client) must run here too. It short-circuits after the first + # successful call. + init() check_is_fitted(self) tabpfn_config = self._get_tabpfn_config() @@ -762,14 +1017,25 @@ def predict_task() -> PredictionResult: try: return InferenceClient.predict( X_clean, - fitted_train_set_id=cast(UUID, self._last_fitted_train_set_id), + fitted_train_set_id=cast(UUID, _resolve_train_set_id(self)), task_config=task_config, client_options=self.client_options, ) except NeedsRefittingError as exc: last_exc = exc refit_attempts += 1 - self._last_fitted_train_set_id = InferenceClient.fit( + if self._last_train_X is None or self._last_train_y is None: + # A load-by-id estimator (constructed via `model_id` / + # `load_model`) has no training data in memory to rebuild + # from, so we can't transparently recover here. + raise RuntimeError( + "The referenced fitted model is no longer available " + "on the server and this estimator has no in-memory " + "training data to refit (it was created via " + "`model_id` / `load_model`). Re-create it by calling " + "`fit(X, y)`." + ) from exc + self.fitted_train_set_id_ = InferenceClient.fit( self._last_train_X, self._last_train_y, task_config=task_config, @@ -830,6 +1096,37 @@ def _validate_targets(self, y) -> None: if sum(pd.isnull(y_)) > 0: raise ValueError("Input y contains NaN.") + @overload + def save_model(self, path: None = None) -> dict[str, Any]: ... + @overload + def save_model(self, path: str | Path) -> Path: ... + def save_model(self, path: str | Path | None = None) -> dict[str, Any] | Path: + """Serialize a portable reference to the fitted (server-side) model. + + The handle captures the server-side fitted-train-set id and the + estimator hyperparameters — everything `load_model()` needs to + reconstruct an equivalent regressor. No training data leaves the + client; predictions run against the model referenced by the id, so this + is most useful with `fit_mode="fit_with_cache"`. + + Parameters + ---------- + path : str, Path or None, default=None + If given, the handle is written to this path as JSON and the path + is returned. If None, the handle dict is returned directly. + """ + return _build_model_handle(self, task=PredictionTask.REGRESSION, path=path) + + @classmethod + def load_model(cls, handle: dict[str, Any] | str | Path) -> "TabPFNRegressor": + """Reconstruct a regressor from a `save_model()` handle (dict or path). + + The resulting estimator is already fitted (it references the server-side + model by id), so it can `predict` without calling `fit()` again. + """ + data = _load_model_handle_for(cls, handle, task=PredictionTask.REGRESSION) + return cls(**data.params) + def _is_thinking_supported_model_path(model_path: str | None) -> bool: """Thinking is server-side supported only for v3 models (or the auto sentinel, diff --git a/tests/unit/test_fit_mode_and_model_id.py b/tests/unit/test_fit_mode_and_model_id.py new file mode 100644 index 0000000..f8cde15 --- /dev/null +++ b/tests/unit/test_fit_mode_and_model_id.py @@ -0,0 +1,227 @@ +"""Unit tests for the KV-cache client surface: + +- `fit_mode` reaching the server-side `tabpfn_config` on the wire, +- the `model_id` constructor param (Option 1) making an estimator "fitted" + and predictable without a local `fit()`, +- `save_model()` / `load_model()` round-trips (Option 3), +- the refit-fallback guard for load-by-id estimators. + +These stay hermetic by patching `InferenceClient` / `ServiceClient` rather +than talking to a mock server — the concern here is client-side wiring, not +the HTTP layer. +""" + +import unittest +from unittest.mock import patch +from uuid import UUID + +import numpy as np + +from tabpfn_client.api_models import FitMode +from tabpfn_client.client import NeedsRefittingError, PredictionResult, ServiceClient +from tabpfn_client.estimator import ( + TabPFNClassifier, + TabPFNRegressor, + _resolve_train_set_id, +) +from tabpfn_client.service_wrapper import InferenceClient + +_MODEL_ID = "00000000-0000-0000-0000-000000000abc" + + +def _regressor_prediction() -> PredictionResult: + return PredictionResult(y_pred={"mean": np.zeros(4)}, metadata={}) + + +class TestFitMode(unittest.TestCase): + def test_default_omits_fit_mode_on_wire(self): + # Default is None, so nothing is sent and the server applies its own + # default (fit_preprocessors) — matching the SDK's "None means unset" + # convention for server-backed config fields. + for Est in (TabPFNClassifier, TabPFNRegressor): + cfg = Est()._get_tabpfn_config() + dumped = cfg.model_dump(mode="json", exclude_none=True) + self.assertNotIn("fit_mode", dumped, Est.__name__) + + def test_fit_with_cache_reaches_wire_config(self): + for Est in (TabPFNClassifier, TabPFNRegressor): + cfg = Est(fit_mode=FitMode.FIT_WITH_CACHE)._get_tabpfn_config() + dumped = cfg.model_dump(mode="json", exclude_none=True) + self.assertEqual(dumped.get("fit_mode"), "fit_with_cache", Est.__name__) + + def test_fit_mode_forwarded_on_predict_task_config(self): + reg = TabPFNRegressor(fit_mode=FitMode.FIT_WITH_CACHE, model_id=_MODEL_ID) + with patch("tabpfn_client.estimator.init"): + with patch.object(ServiceClient, "get_model_limits", return_value=None): + with patch.object( + InferenceClient, "predict", return_value=_regressor_prediction() + ) as mock_predict: + reg.predict(np.random.randn(4, 3)) + task_config = mock_predict.call_args[1]["task_config"] + dumped = task_config.tabpfn_config.model_dump(mode="json", exclude_none=True) + self.assertEqual(dumped.get("fit_mode"), "fit_with_cache") + + +class TestModelIdConstructor(unittest.TestCase): + def test_model_id_marks_estimator_fitted(self): + for Est in (TabPFNClassifier, TabPFNRegressor): + est = Est(model_id=_MODEL_ID) + self.assertTrue(est.__sklearn_is_fitted__(), Est.__name__) + self.assertEqual(_resolve_train_set_id(est), UUID(_MODEL_ID)) + self.assertFalse(Est().__sklearn_is_fitted__(), Est.__name__) + + def test_predict_uses_model_id_without_fitting(self): + reg = TabPFNRegressor(model_id=_MODEL_ID) + with patch("tabpfn_client.estimator.init"): + with patch.object(ServiceClient, "get_model_limits", return_value=None): + with patch.object(InferenceClient, "fit") as mock_fit: + with patch.object( + InferenceClient, "predict", return_value=_regressor_prediction() + ) as mock_predict: + reg.predict(np.random.randn(4, 3)) + mock_fit.assert_not_called() + self.assertEqual( + mock_predict.call_args[1]["fitted_train_set_id"], UUID(_MODEL_ID) + ) + + def test_model_id_accepts_uuid_and_str(self): + self.assertEqual( + _resolve_train_set_id(TabPFNRegressor(model_id=UUID(_MODEL_ID))), + UUID(_MODEL_ID), + ) + + +class TestSaveLoadModel(unittest.TestCase): + def test_regressor_round_trip_dict(self): + reg = TabPFNRegressor(model_id=_MODEL_ID, n_estimators=4) + handle = reg.save_model() + self.assertEqual(handle["task"], "regression") + self.assertNotIn("client_options", handle["params"]) + self.assertEqual(handle["params"]["model_id"], _MODEL_ID) + + loaded = TabPFNRegressor.load_model(handle) + self.assertTrue(loaded.__sklearn_is_fitted__()) + self.assertEqual(_resolve_train_set_id(loaded), UUID(_MODEL_ID)) + self.assertEqual(loaded.get_params()["n_estimators"], 4) + + def test_classifier_round_trip_restores_classes(self): + clf = TabPFNClassifier(model_id=_MODEL_ID) + clf.classes_ = np.array([0, 1, 2]) + loaded = TabPFNClassifier.load_model(clf.save_model()) + self.assertTrue(loaded.__sklearn_is_fitted__()) + np.testing.assert_array_equal(loaded.classes_, np.array([0, 1, 2])) + + def test_round_trip_via_file(self): + import tempfile + from pathlib import Path + + reg = TabPFNRegressor(model_id=_MODEL_ID) + with tempfile.TemporaryDirectory() as d: + path = Path(d) / "model.json" + returned = reg.save_model(path) + self.assertEqual(Path(returned), path) + self.assertTrue(path.exists()) + loaded = TabPFNRegressor.load_model(path) + self.assertEqual(_resolve_train_set_id(loaded), UUID(_MODEL_ID)) + + def test_effective_id_beats_constructor_model_id(self): + # A real fit supersedes the constructor's model_id; save_model must + # persist the effective (post-fit) id, not the original. + reg = TabPFNRegressor(model_id=_MODEL_ID) + new_id = "00000000-0000-0000-0000-000000000fff" + reg.fitted_train_set_id_ = UUID(new_id) + self.assertEqual(reg.save_model()["params"]["model_id"], new_id) + + def test_load_model_task_mismatch_raises(self): + clf = TabPFNClassifier(model_id=_MODEL_ID) + clf.classes_ = np.array([0, 1]) + with self.assertRaises(ValueError): + TabPFNRegressor.load_model(clf.save_model()) + + def test_save_model_unfitted_raises(self): + from sklearn.exceptions import NotFittedError + + with self.assertRaises(NotFittedError): + TabPFNRegressor().save_model() + + def test_save_classifier_without_classes_raises(self): + # A bare model_id classifier is "fitted" (has an id) but has no + # classes_ to persist; save_model must fail clearly, not AttributeError. + clf = TabPFNClassifier(model_id=_MODEL_ID) + with self.assertRaises(ValueError) as ctx: + clf.save_model() + self.assertIn("classes_", str(ctx.exception)) + + def test_unsupported_handle_version_raises(self): + handle = TabPFNRegressor(model_id=_MODEL_ID).save_model() + handle["format_version"] = 999 + with self.assertRaises(ValueError): + TabPFNRegressor.load_model(handle) + + +class TestRefitGuard(unittest.TestCase): + def test_load_by_id_cannot_auto_refit(self): + reg = TabPFNRegressor(model_id=_MODEL_ID) # no in-memory train data + with patch("tabpfn_client.estimator.init"): + with patch.object(ServiceClient, "get_model_limits", return_value=None): + with patch.object(InferenceClient, "fit") as mock_fit: + with patch.object( + InferenceClient, "predict", side_effect=NeedsRefittingError() + ): + with self.assertRaises(RuntimeError) as ctx: + reg.predict(np.random.randn(4, 3)) + self.assertIn("no in-memory", str(ctx.exception).lower()) + mock_fit.assert_not_called() + + +class TestSklearnCompatibility(unittest.TestCase): + def test_clone_preserves_model_id_and_stays_fitted(self): + # `clone` calls `get_params()` and re-instantiates. `model_id` is a + # constructor param, so the clone must keep it (and thus still be + # considered fitted for the load-by-id path) — this covers the + # get_params/clone desync surfaced by cursor. + from sklearn.base import clone + + from typing import Any, cast as _cast + + for Est in (TabPFNClassifier, TabPFNRegressor): + original = Est(model_id=_MODEL_ID) + # `clone`'s overloaded return over a `type[C] | type[R]` union + # confuses pyright; runtime type is Est, but we cast to Any to + # avoid re-typing at every call site below. + cloned = _cast(Any, clone(original)) + self.assertEqual(cloned.get_params()["model_id"], _MODEL_ID, Est.__name__) + self.assertEqual( + _resolve_train_set_id(cloned), UUID(_MODEL_ID), Est.__name__ + ) + self.assertTrue(cloned.__sklearn_is_fitted__(), Est.__name__) + + def test_invalid_model_id_reports_unfitted_not_crash(self): + # Constructor stores `model_id` verbatim (sklearn convention: no + # validation in __init__). A bad id must not crash sklearn's + # fittedness check — it should degrade to "not fitted". + from sklearn.exceptions import NotFittedError + from sklearn.utils.validation import check_is_fitted + + for Est in (TabPFNClassifier, TabPFNRegressor): + est = Est(model_id="not-a-uuid") + self.assertFalse(est.__sklearn_is_fitted__(), Est.__name__) + with self.assertRaises(NotFittedError): + check_is_fitted(est) + + def test_predict_triggers_init_from_cold_state(self): + # A fresh process that only calls `load_model(...).predict(...)` never + # touches `fit()`, so `_predict` itself must call `init()` to authorize + # the HTTP client — otherwise the main cross-process reuse path 401s. + reg = TabPFNRegressor(model_id=_MODEL_ID) + with patch("tabpfn_client.estimator.init") as mock_init: + with patch.object(ServiceClient, "get_model_limits", return_value=None): + with patch.object( + InferenceClient, "predict", return_value=_regressor_prediction() + ): + reg.predict(np.random.randn(4, 3)) + mock_init.assert_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_tabpfn_classifier.py b/tests/unit/test_tabpfn_classifier.py index 38fa64c..ac82553 100644 --- a/tests/unit/test_tabpfn_classifier.py +++ b/tests/unit/test_tabpfn_classifier.py @@ -498,9 +498,7 @@ def test_only_allowed_parameters_passed_to_config(self): # Skip fitting classifier.fitted_ = True - classifier._last_fitted_train_set_id = UUID( - "00000000-0000-0000-0000-000000000000" - ) + classifier.fitted_train_set_id_ = UUID("00000000-0000-0000-0000-000000000000") test_X = np.random.randn(10, 5) diff --git a/tests/unit/test_tabpfn_regressor.py b/tests/unit/test_tabpfn_regressor.py index e746cfb..fa067f5 100644 --- a/tests/unit/test_tabpfn_regressor.py +++ b/tests/unit/test_tabpfn_regressor.py @@ -492,9 +492,7 @@ def test_only_allowed_parameters_passed_to_config(self): # Skip fitting regressor.fitted_ = True - regressor._last_fitted_train_set_id = UUID( - "00000000-0000-0000-0000-000000000000" - ) + regressor.fitted_train_set_id_ = UUID("00000000-0000-0000-0000-000000000000") test_X = np.random.randn(10, 5) @@ -602,9 +600,7 @@ def test_predict_single_quantile_returns_list_of_one_array(self): def test_predict_full_adds_criterion_with_optional_dependencies(self): regressor = TabPFNRegressor() regressor.fitted_ = True - regressor._last_fitted_train_set_id = UUID( - "00000000-0000-0000-0000-000000000000" - ) + regressor.fitted_train_set_id_ = UUID("00000000-0000-0000-0000-000000000000") regressor._last_train_X = np.random.randn(5, 2) regressor._last_train_y = np.random.randn(5) @@ -651,9 +647,7 @@ def __init__(self, borders): def test_predict_full_missing_optional_dependencies_logs_warning(self): regressor = TabPFNRegressor() regressor.fitted_ = True - regressor._last_fitted_train_set_id = UUID( - "00000000-0000-0000-0000-000000000000" - ) + regressor.fitted_train_set_id_ = UUID("00000000-0000-0000-0000-000000000000") regressor._last_train_X = np.random.randn(5, 2) regressor._last_train_y = np.random.randn(5)