diff --git a/CHANGELOG.md b/CHANGELOG.md index c0f7969..74dd0c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] ### Added +- **scikit-learn estimators** (#63): `TabctxClassifier` and `TabctxRegressor` (`tabctx/estimators.py`, `from tabctx import ...`; `pip install "tabctx[client]"` adds scikit-learn and pyarrow). They take DataFrames, numpy arrays or lists of rows with text columns, dates and missing values as they are; labels are checked as scikit-learn checks them and come back exactly as given (sent as class indices); `predict_proba` follows `classes_`; clone, pipelines and model selection work. Test columns match the training columns by name when both have names, else by position (with a warning); a column that was text at fit stays text at predict, whatever its new values look like. Each table is one Parquet file (`tabctx/table_files.py`, float32 numbers, timestamps and durations as such); predicts over `max_request_bytes` (default 29 MiB, under the Vector gateway's 30 MiB) are split until each part fits, and a training table over it is refused before sending. The context id comes from the table, labels, model and task; a predict that finds its context gone (the server's "no cached context") fits it again from the kept training file and retries once, with the model it was fit with. +- **Client:** `TabctxClient(api_key=...)` (sent as `Authorization: Bearer`, for the Vector gateway); `base_url` and `api_key` default to `TABCTX_BASE_URL` and `TABCTX_API_KEY`. `fit_file()` / `predict_file()` send a table file (path, streamed, or bytes) in the request as multipart/form-data; `predict(model=...)` for gateways, which route predicts by model. Redirects are refused rather than followed (urllib would carry the key to the new host, and turn a POST into a GET). The key never appears in the client's repr and is dropped from pickles: an unpickled client (a saved model, a worker with `n_jobs > 1`) reads `TABCTX_API_KEY`, and a 401 without a key says so; copies (sklearn's clone) share the client. 403 maps to `PermissionError` like 401. - **A table file in the fit or predict request itself.** `/v1/tabctx/fit` and `/v1/tabctx/predict` take multipart/form-data as well as JSON: the same fields as form fields, the table as a `file` part (`curl -F dataset_id=churn -F target=churned -F file=@train.parquet`), so a file is one call instead of upload-then-reference. The file streams to the replica's disk under the upload size cap, is read, and is deleted. OpenAPI documents both body types. - **Parquet and compressed CSV.** Files (sent with the request or through `/v1/tabctx/upload`) may be CSV (plain, gzip or zstd) or Parquet, detected from the content (`serve/files.py`, pyarrow, replacing the numpy CSV parser). A 5,000 x 500 float table is 51 MB as JSON and 13 MB as Parquet; reading it takes 16 ms. Dates and times become seconds, stored from the training minimum so float32 keeps sub-second resolution (`TableSchema.offsets`, for any float64 column past 2^24); header names are stripped; decimal targets are numbers. Labels keep their text as written (`"1.0"`, `"01"`, `"true"` are classes of their own; `"NA"` can be a class among text labels); a missing value among numeric labels, in any spelling, is a 422 naming the row, and regression with text labels is a 422. Reading a CSV peaks at about 3x its size in host memory (arrow's parser; Parquet much less), which matters only for multi-GB direct uploads. `pyarrow` and `python-multipart` join the `serve` extra. - **Multipart bodies stream to disk under the upload cap** (`serve/multipart.py`): the file part goes straight to the upload writer as it arrives, so an oversized body stops at the cap (or is refused from its `Content-Length` before any is read), a good one is written once, and disk writes run off the event loop. Tenant checks run before the body is read. diff --git a/README.md b/README.md index f43b7bd..425ba71 100644 --- a/README.md +++ b/README.md @@ -29,11 +29,45 @@ endpoint**, chat-completions style -- requests pick a model by exact id one memory budget. A protocol seam (`backends/base.py`) keeps adding models cheap. +## Quickstart: as a scikit-learn model + +```bash +pip install "tabctx[client]" +export TABCTX_BASE_URL=https://proxy.vectorinstitute.ai # or your own deployment +export TABCTX_API_KEY=... # a gateway key, if it needs one +``` + +```python +from sklearn.datasets import fetch_openml +from sklearn.model_selection import train_test_split +from tabctx import TabctxClassifier + +X, y = fetch_openml("credit-g", version=1, as_frame=True, return_X_y=True) +X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0) + +clf = TabctxClassifier(model="tabpfn-3.5").fit(X_train, y_train) +clf.score(X_test, y_test) # the training table stays cached on a GPU +clf.predict_proba(X_test) +``` + +`TabctxClassifier` and `TabctxRegressor` (`tabctx/estimators.py`) take +DataFrames, arrays or lists of rows with text columns, dates and missing +values as they are, return labels as given, and work with clone, +pipelines and `cross_val_score`. Each table travels as one Parquet file; +large predicts are split under the request limit (`max_request_bytes`), +and a context the server dropped is fit again from the training file +the estimator keeps. The API key never appears in a repr and is left out +of pickles, so with `n_jobs > 1` or a saved model set `TABCTX_API_KEY`. + +The lower-level client speaks the HTTP API directly: + ```python client = TabctxClient("http://localhost:8000") client.models() # -> tabicl-v2, tabpfn-3.5 client.fit(X, y, dataset_id="churn", model="tabpfn-3.5") client.predict("churn", X_test) # served by tabpfn-3.5 +client.fit_file("train.parquet", target="churned", dataset_id="churn", model="tabicl-v2") +client.predict_file("churn", "test.csv", model="tabicl-v2") ``` Deploy one model (`TABCTX_BACKEND=tabicl`) or several @@ -90,7 +124,9 @@ deployment is exactly fast enough. *Figure: a request enters through Ray Serve's consistent-hash proxy (affinity on the session header) into a `TabctxService` replica, which -resolves the tenant, parses inline JSON or an uploaded CSV, and coalesces +resolves the tenant, reads inline JSON rows or a table file (CSV or +Parquet, sent with the request or uploaded first), encodes text columns +with the context's schema, and coalesces same-context predicts. With several GPUs, a `MultiDeviceEngine` places each fit on one GPU and routes predicts to the GPU holding the context, from lock-free snapshots; each GPU's engine runs in its own worker @@ -106,8 +142,9 @@ The unit everything revolves around is a **context**: one training table, encoded once by one model, addressed by `dataset_id` (scoped per tenant) and reused by every later `predict()` against it. A context carries its payload (the fitted estimator with its kv cache), the model -that fit it, its shape and, when it came from a CSV, its column names, -plus a measured byte cost. All models deployed on a device share one +that fit it, its shape, its schema (column names, and the codes of its +text columns, so new rows are encoded the same way), plus a measured +byte cost. All models deployed on a device share one cache and one memory budget, so admission control sees the whole device's usage regardless of which model a request picks. diff --git a/pyproject.toml b/pyproject.toml index ca83218..e1ce6eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ serve = [ "pyarrow>=14", "python-multipart>=0.0.9", ] +# The scikit-learn-style estimators (tabctx.estimators); pyarrow sends +# tables as Parquet (without it, gzip CSV). +client = ["scikit-learn>=1.6", "pyarrow>=14"] dev = ["pytest>=7.0", "pytest-cov", "httpx"] [dependency-groups] diff --git a/src/tabctx/__init__.py b/src/tabctx/__init__.py index 4827664..dd281df 100644 --- a/src/tabctx/__init__.py +++ b/src/tabctx/__init__.py @@ -24,6 +24,27 @@ ) from tabctx.types import EngineStats, PredictOutcome, Task +# The scikit-learn estimators need scikit-learn, which tabctx itself +# doesn't: imported on first use (`from tabctx import TabctxClassifier`). +_ESTIMATORS = ("TabctxClassifier", "TabctxRegressor") + + +def __getattr__(name: str): + if name in _ESTIMATORS: + try: + from tabctx import estimators + except ImportError as e: + raise ImportError( + f"{name} needs scikit-learn and pyarrow: pip install 'tabctx[client]'" + ) from e + return getattr(estimators, name) + raise AttributeError(f"module 'tabctx' has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted([*globals(), *_ESTIMATORS]) + + __all__ = [ "AdmissionRejected", "BackendComputeError", diff --git a/src/tabctx/client.py b/src/tabctx/client.py index f5e83bd..feeecb4 100644 --- a/src/tabctx/client.py +++ b/src/tabctx/client.py @@ -1,4 +1,5 @@ -"""Minimal HTTP client for a running tabctx deployment. +"""Minimal HTTP client for a tabctx deployment, directly or through a +gateway (Vector's: https://proxy.vectorinstitute.ai, with an API key). Exists mainly so callers can't get the serving contracts wrong: the session-affinity header (`x-session-id` == dataset_id, required for @@ -6,17 +7,27 @@ tenant header (`x-tabctx-tenant-id`; see serve/tenancy.py) are set automatically on every call. Pure stdlib (urllib) -- no dependency beyond tabctx's own errors, so `pip install tabctx` is enough to talk -to a deployment. +to a deployment. For scikit-learn-style estimators over this client, see +tabctx.estimators. Usage: from tabctx.client import TabctxClient - client = TabctxClient("http://localhost:8000", tenant_id="acme") - dataset_id = client.fit(X_train, y_train, dataset_id="churn-v1") - result = client.predict("churn-v1", X_test, return_proba=True) + # Or set TABCTX_BASE_URL and TABCTX_API_KEY and call TabctxClient(). + client = TabctxClient("https://proxy.vectorinstitute.ai", api_key=KEY) + dataset_id = client.fit(X_train, y_train, dataset_id="churn-v1", model="tabicl-v2") + result = client.predict("churn-v1", X_test, return_proba=True, model="tabicl-v2") result.predictions, result.probabilities, result.classes + # A table file (CSV, gzip/zstd CSV or Parquet; path or bytes) in one call: + client.fit_file("train.parquet", target="churned", dataset_id="churn-v1", model="tabicl-v2") + client.predict_file("churn-v1", "test.csv", model="tabicl-v2") + +The API key is sent as `Authorization: Bearer`. It never appears in the +client's repr and is left out when the client is pickled: an unpickled +client (a saved model, a joblib worker) reads TABCTX_API_KEY instead. + Server-side errors come back as the same tabctx exceptions the engine raises locally (InvalidInputError for 422, DatasetNotFoundError for 404, AdmissionRejected for 413, BackendComputeError for 507), so code can be @@ -34,7 +45,10 @@ import time import urllib.error import urllib.request +import uuid +from collections.abc import Iterator from dataclasses import dataclass +from pathlib import Path from typing import IO from tabctx.errors import ( @@ -49,6 +63,27 @@ from tabctx.serve.tenancy import TENANT_HEADER from tabctx.types import ArrayLike, Task +BASE_URL_ENV = "TABCTX_BASE_URL" +API_KEY_ENV = "TABCTX_API_KEY" + +# A table file: a path, or its bytes. +TableFile = str | os.PathLike | bytes | bytearray | memoryview + + +class _NoRedirects(urllib.request.HTTPRedirectHandler): + """urllib would follow a redirect with the Authorization header, to any + host, and turn a POST into a body-less GET; tabctx follows none.""" + + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None # urllib raises the 3xx as an HTTPError + + +_OPENER = urllib.request.build_opener(_NoRedirects) + + +def _open(req: urllib.request.Request, timeout: float): + return _OPENER.open(req, timeout=timeout) + class TabctxBackpressureError(TabctxError): """The owning replica is at capacity (HTTP 503). Retryable: back off @@ -74,27 +109,58 @@ class PredictResult: class TabctxClient: def __init__( self, - base_url: str, + base_url: str | None = None, tenant_id: str | None = None, timeout_s: float = 120.0, max_retries: int = 3, retry_backoff_s: float = 0.25, session_header: str = DEFAULT_SESSION_ID_HEADER, + api_key: str | None = None, ) -> None: - """max_retries applies only to 503 backpressure (safe to retry by + """base_url and api_key default to TABCTX_BASE_URL and + TABCTX_API_KEY. A gateway needs the key (and sets the tenant + itself); a deployment reached directly needs neither. + + max_retries applies only to 503 backpressure (safe to retry by construction); every other error propagates immediately. session_header: the affinity header name. Must match the deployment's RAY_SERVE_SESSION_ID_HEADER_KEY (default "x-session-id"); a deployment that renames it would otherwise silently lose replica affinity.""" + base_url = base_url or os.environ.get(BASE_URL_ENV) + if not base_url: + raise ValueError( + f"no deployment URL: pass base_url or set {BASE_URL_ENV} " + "(e.g. https://proxy.vectorinstitute.ai)" + ) self._base_url = base_url.rstrip("/") + self._api_key = api_key or os.environ.get(API_KEY_ENV) self._tenant_id = tenant_id self._timeout_s = timeout_s self._max_retries = max_retries self._retry_backoff_s = retry_backoff_s self._session_header = session_header + def __repr__(self) -> str: + return f"TabctxClient({self._base_url!r})" + + def __getstate__(self) -> dict: + # Keys stay out of pickles (saved models, joblib workers). + return {**self.__dict__, "_api_key": None} + + def __setstate__(self, state: dict) -> None: + self.__dict__.update(state) + self._api_key = os.environ.get(API_KEY_ENV) + + def __copy__(self) -> TabctxClient: + # Stateless between calls, so copies (sklearn's clone) share it -- + # and keep the key, which a pickle round trip would drop. + return self + + def __deepcopy__(self, memo: dict) -> TabctxClient: + return self + # ---- public API ------------------------------------------------------ def fit( @@ -171,30 +237,66 @@ def fit_uploaded( ) return resp["dataset_id"] + def fit_file( + self, + file: TableFile, + target: str | None = None, + dataset_id: str | None = None, + model: str | None = None, + task: Task | None = None, + ) -> str: + """fit() from a table file sent with the request: a path (streamed + from disk) or the file's bytes. Columns are the features plus the + label column, `target` (default: the last). Returns the dataset_id.""" + fields = { + "target": target, + "dataset_id": dataset_id, + "model": model, + "task": task, + } + resp = self._post_file("/v1/tabctx/fit", fields, file, session_id=dataset_id) + return resp["dataset_id"] + def predict( self, dataset_id: str, X_test: ArrayLike | None = None, return_proba: bool = False, test_upload_id: str | None = None, + model: str | None = None, ) -> PredictResult: - body: dict = {"dataset_id": dataset_id, "return_proba": return_proba} - if X_test is not None: - body["test_X"] = X_test - if test_upload_id is not None: - body["test_upload_id"] = test_upload_id - resp = self._post( - "/v1/tabctx/predict", - body, - session_id=dataset_id, + """model: needed through a gateway, which routes predicts by it; + a deployment reached directly knows it from the fit.""" + body: dict = _without_none( + { + "dataset_id": dataset_id, + "return_proba": return_proba, + "test_X": X_test, + "test_upload_id": test_upload_id, + "model": model, + } ) - return PredictResult( - predictions=resp["predictions"], - probabilities=resp.get("probabilities"), - classes=resp.get("classes"), - latency_ms=resp["latency_ms"], - served_by=resp.get("served_by"), + resp = self._post("/v1/tabctx/predict", body, session_id=dataset_id) + return _predict_result(resp) + + def predict_file( + self, + dataset_id: str, + file: TableFile, + return_proba: bool = False, + model: str | None = None, + ) -> PredictResult: + """predict() on a table file sent with the request (path or bytes). + When the training table was a file, columns are matched by name.""" + fields = { + "dataset_id": dataset_id, + "return_proba": "true" if return_proba else "false", + "model": model, + } + resp = self._post_file( + "/v1/tabctx/predict", fields, file, session_id=dataset_id ) + return _predict_result(resp) def fit_predict( self, @@ -243,15 +345,20 @@ def limits(self) -> dict: # ---- internals ------------------------------------------------------- def _get(self, path: str) -> dict: - return self._request(path, None, {}) + return self._request(path, None, self._auth()) - def _headers(self, session_id: str | None) -> dict[str, str]: - headers = {"Content-Type": "application/json"} + def _headers( + self, session_id: str | None, content_type: str = "application/json" + ) -> dict[str, str]: + headers = {"Content-Type": content_type} if session_id is not None: headers[self._session_header] = session_id if self._tenant_id is not None: headers[TENANT_HEADER] = self._tenant_id - return headers + return {**headers, **self._auth()} + + def _auth(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} def _request( self, path: str, data: bytes | IO[bytes] | None, headers: dict[str, str] @@ -267,16 +374,46 @@ def _request( method="POST" if data is not None else "GET", ) try: - with urllib.request.urlopen(req, timeout=self._timeout_s) as resp: + with _open(req, timeout=self._timeout_s) as resp: return json.loads(resp.read().decode()) except urllib.error.HTTPError as e: - raise self._map_status(e.code, self._detail(e)) from e + if 300 <= e.code < 400: + raise TabctxError( + f"HTTP {e.code}: the deployment redirected to " + f"{e.headers.get('Location')!r}; redirects aren't followed (they " + "would carry the API key elsewhere): set base_url to the " + "deployment's final address" + ) from e + error = self._map_status(e.code, self._detail(e)) + if isinstance(error, PermissionError) and not self._api_key: + error = PermissionError( + f"{error} -- no API key was sent: pass api_key or set " + f"{API_KEY_ENV} (a client unpickled in a worker, e.g. with " + "n_jobs > 1, reads it from the environment)" + ) + raise error from e def _post(self, path: str, body: dict, session_id: str | None) -> dict: data = json.dumps(body).encode() + return self._with_retries( + lambda: self._request(path, data, self._headers(session_id)) + ) + + def _post_file( + self, path: str, fields: dict, file: TableFile, session_id: str | None + ) -> dict: + form = _Multipart({k: v for k, v in fields.items() if v is not None}, file) + headers = { + **self._headers(session_id, form.content_type), + "Content-Length": str(form.length), + } + # A fresh body per attempt: a path is streamed again from the start. + return self._with_retries(lambda: self._request(path, form.chunks(), headers)) + + def _with_retries(self, send): for attempt in range(self._max_retries + 1): try: - return self._request(path, data, self._headers(session_id)) + return send() except TabctxBackpressureError: if attempt >= self._max_retries: raise @@ -309,8 +446,59 @@ def _map_status(code: int, detail: str) -> Exception: return AdmissionRejected(detail) if code == 507: return BackendComputeError(detail) - if code == 401: + if code in (401, 403): return PermissionError(detail) if code == 503: return TabctxBackpressureError(detail) return TabctxError(f"HTTP {code}: {detail}") + + +def _predict_result(resp: dict) -> PredictResult: + return PredictResult( + predictions=resp["predictions"], + probabilities=resp.get("probabilities"), + classes=resp.get("classes"), + latency_ms=resp["latency_ms"], + served_by=resp.get("served_by"), + ) + + +class _Multipart: + """A multipart/form-data body of text fields plus one `file` part, + with its length known up front, so urllib streams it (a path is read + in blocks, never whole into memory).""" + + _BLOCK = 1 << 20 + + def __init__(self, fields: dict[str, str], file: TableFile) -> None: + boundary = uuid.uuid4().hex + self.content_type = f"multipart/form-data; boundary={boundary}" + head = b"".join( + f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"' + f"\r\n\r\n{value}\r\n".encode() + for name, value in fields.items() + ) + self._head = ( + head + + ( + f"--{boundary}\r\nContent-Disposition: form-data; " + f'name="file"; filename="table"\r\n' + "Content-Type: application/octet-stream\r\n\r\n" + ).encode() + ) + self._tail = f"\r\n--{boundary}--\r\n".encode() + if isinstance(file, bytearray | memoryview): + file = bytes(file) + self._file = file + size = len(file) if isinstance(file, bytes) else os.path.getsize(file) + self.length = len(self._head) + size + len(self._tail) + + def chunks(self) -> Iterator[bytes]: + yield self._head + if isinstance(self._file, bytes): + yield self._file + else: + with open(Path(self._file), "rb") as f: + while block := f.read(self._BLOCK): + yield block + yield self._tail diff --git a/src/tabctx/estimators.py b/src/tabctx/estimators.py new file mode 100644 index 0000000..9f41ea4 --- /dev/null +++ b/src/tabctx/estimators.py @@ -0,0 +1,246 @@ +"""scikit-learn estimators backed by a tabctx deployment. + + from tabctx import TabctxClassifier + + clf = TabctxClassifier(model="tabpfn-3.5") # TABCTX_BASE_URL, TABCTX_API_KEY + clf.fit(X_train, y_train) # DataFrame or array; text and NaN fine + clf.predict_proba(X_test) + +They behave like any scikit-learn estimator (clone, pipelines, +cross_val_score), except that the work happens on the deployment: fit +sends the training table, which the server encodes once and keeps on a +GPU; predict sends only the new rows. X may be a pandas DataFrame, a +numpy array or a list of rows, with text columns, dates and missing +values (NaN/None) as they are: the server encodes text and handles +missing values itself. Test columns are matched to the training columns +by name when both tables have names (DataFrames), else by position. + +Each table travels as one Parquet file. Predicts larger than +max_request_bytes are split into several requests; a training table +larger than it is refused here, before it is sent. The training +context's id comes from the table, the labels, the model and the task, +so fitting the same data again reuses the id; if the server has dropped +the context (eviction, a restart), predict fits it again from the +training file the estimator keeps and retries once. + +With n_jobs > 1 (joblib workers) or a saved model, the client reaches +the worker without its API key, by design: set TABCTX_API_KEY in the +environment. + +Requires scikit-learn and pyarrow (`pip install "tabctx[client]"`). +""" + +from __future__ import annotations + +import hashlib +import warnings + +import numpy as np +from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin +from sklearn.utils.multiclass import check_classification_targets +from sklearn.utils.validation import check_is_fitted, column_or_1d + +from tabctx.client import PredictResult, TabctxClient +from tabctx.errors import DatasetNotFoundError +from tabctx.table_files import ( + Columns, + _missing, + columns, + is_frame, + take_rows, + to_parquet, +) +from tabctx.types import Task + +DEFAULT_MODEL = "tabpfn-3.5" +# The Vector gateway takes request bodies up to 30 MiB; this leaves room +# for the multipart envelope around the file. +MAX_REQUEST_BYTES = 29 * 1024 * 1024 +_TARGET = "__tabctx_target__" +# How the server says a context isn't there (engine._not_found); other +# 404s (an unknown model or route) must not trigger a re-fit. +_CONTEXT_GONE = "no cached context" + + +class _TabctxEstimator(BaseEstimator): + _task: Task + + def __init__( + self, + model: str = DEFAULT_MODEL, + client: TabctxClient | None = None, + max_request_bytes: int | None = MAX_REQUEST_BYTES, + ) -> None: + """model: a model id from client.models() ("tabpfn-3.5", + "tabicl-v2"). client: where to send the tables; default + TabctxClient(), configured by TABCTX_BASE_URL and TABCTX_API_KEY. + max_request_bytes: the largest table file to send in one request + (the Vector gateway takes 30 MiB bodies); None for no limit.""" + self.model = model + self.client = client + self.max_request_bytes = max_request_bytes + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.input_tags.allow_nan = True + tags.input_tags.string = True + tags.input_tags.categorical = True + tags.non_deterministic = True # served, not computed here + return tags + + def _client(self) -> TabctxClient: + return self.client if self.client is not None else TabctxClient() + + # ---- fit ----------------------------------------------------------- + + def _fit(self, X, labels: np.ndarray): + cols = columns(X) + if cols.n_rows != len(labels): + raise ValueError(f"X has {cols.n_rows} rows but y has {len(labels)} labels") + data = to_parquet(cols, (_TARGET, labels)) + limit = self.max_request_bytes + if limit is not None and len(data) > limit: + raise ValueError( + f"the training table is {len(data):,} bytes as Parquet; a request " + f"can be at most {limit:,} (max_request_bytes). " + "Fit on fewer rows." + ) + digest = hashlib.sha256(data) + digest.update(f"\0{self.model}\0{self._task}".encode()) + self.dataset_id_ = f"sklearn-{digest.hexdigest()[:32]}" + self.model_ = self.model + self.n_features_in_ = len(cols.names) + self._names = cols.names + self._named = cols.named + self._text = cols.text_names + if cols.named and all(isinstance(c, str) for c in X.columns): + self.feature_names_in_ = np.asarray(cols.names, dtype=object) + elif hasattr(self, "feature_names_in_"): + del self.feature_names_in_ + self._train_file = data + self._send_fit() + return self + + def _send_fit(self) -> None: + self._client().fit_file( + self._train_file, + target=_TARGET, + dataset_id=self.dataset_id_, + model=self.model_, + task=self._task, + ) + + # ---- predict ------------------------------------------------------- + + def _predict(self, X, return_proba: bool) -> list[PredictResult]: + check_is_fitted(self, "dataset_id_") + cols = self._test_columns(X) + return [self._predict_file(data, return_proba) for data in self._files(cols)] + + def _test_columns(self, X) -> Columns: + if is_frame(X) and self._named: + # By name, as the server matches them; training text stays text. + return columns(X, text=self._text) + positions = [self._names.index(n) for n in self._text] + cols = columns(X, positions=positions) + if len(cols.names) != self.n_features_in_: + raise ValueError( + f"X has {len(cols.names)} features, but {type(self).__name__} is " + f"expecting {self.n_features_in_} features as input." + ) + if is_frame(X) or self._named: + warnings.warn( + "X and the training table don't both have column names; their " + "columns are matched by position", + UserWarning, + stacklevel=3, + ) + return cols.renamed(self._names) + + def _files(self, cols: Columns) -> list[bytes]: + """The test table as Parquet files, each under max_request_bytes.""" + data = to_parquet(cols) + limit = self.max_request_bytes + if limit is None or len(data) <= limit: + return [data] + if cols.n_rows == 1: + raise ValueError( + f"one test row is {len(data):,} bytes as Parquet, over " + f"max_request_bytes ({limit:,})" + ) + half = cols.n_rows // 2 + return [ + *self._files(take_rows(cols, 0, half)), + *self._files(take_rows(cols, half, cols.n_rows)), + ] + + def _predict_file(self, data: bytes, return_proba: bool) -> PredictResult: + def send(): + return self._client().predict_file( + self.dataset_id_, data, return_proba=return_proba, model=self.model_ + ) + + try: + return send() + except DatasetNotFoundError as e: + if _CONTEXT_GONE not in str(e): + raise + # Evicted or lost in a restart: fit again from the kept file. + self._send_fit() + return send() + + +class TabctxClassifier(ClassifierMixin, _TabctxEstimator): + """Classification on a tabctx deployment. Labels may be of any one + type (strings, ints, bools); predict returns them as given to fit, and + predict_proba's columns follow classes_.""" + + _task = "classification" + + def fit(self, X, y): + if any(_missing(v) for v in np.ravel(np.asarray(y, dtype=object))): + raise ValueError("y has missing labels; every training row needs one") + y = column_or_1d(np.asarray(y, dtype=object) if _mixed(y) else y, warn=True) + check_classification_targets(y) + self.classes_, codes = np.unique(y, return_inverse=True) + # Labels travel as class indices, so they come back exactly as given. + return self._fit(X, codes.astype(np.int64)) + + def predict(self, X) -> np.ndarray: + codes = [int(p) for r in self._predict(X, False) for p in r.predictions] + return self.classes_[codes] + + def predict_proba(self, X) -> np.ndarray: + blocks = [] + for result in self._predict(X, True): + proba = np.zeros((len(result.probabilities), len(self.classes_))) + proba[:, [int(c) for c in result.classes]] = result.probabilities + blocks.append(proba) + return np.vstack(blocks) + + +class TabctxRegressor(RegressorMixin, _TabctxEstimator): + """Regression on a tabctx deployment.""" + + _task = "regression" + + def fit(self, X, y): + y = column_or_1d(y, warn=True).astype(np.float64) + if not np.isfinite(y).all(): + raise ValueError("y has missing or infinite values") + return self._fit(X, y) + + def predict(self, X) -> np.ndarray: + return np.array( + [p for r in self._predict(X, False) for p in r.predictions], + dtype=np.float64, + ) + + +def _mixed(y) -> bool: + """Labels of more than one type (1 and "a"), which numpy would turn + into text; kept as objects so np.unique refuses them.""" + values = list(np.ravel(np.asarray(y, dtype=object))) + return len({type(v) for v in values if v is not None}) > 1 and any( + isinstance(v, str) for v in values + ) diff --git a/src/tabctx/table_files.py b/src/tabctx/table_files.py new file mode 100644 index 0000000..87e7471 --- /dev/null +++ b/src/tabctx/table_files.py @@ -0,0 +1,177 @@ +"""Tables in memory (pandas DataFrames, numpy arrays, lists of rows) as +Parquet files to send to a deployment. + +Parquet keeps what the server needs to see: which columns are text (a +string column stays text even when it reads "007"), missing values as +nulls, and dates and durations as such. Each column is numbers or text, +by the server's rule (tabctx.schema): numbers when every present value +is a number, else text, unless the caller names it as text (a column +that was text in the training table must stay text at predict, even +when its new values all look like numbers). Numbers travel as float32, +the precision both models compute in; labels as given. + +Requires pyarrow (the `client` extra). +""" + +from __future__ import annotations + +import io +import numbers +from collections.abc import Collection, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np + + +@dataclass(frozen=True) +class Columns: + """A table as named columns: float32 (NaN = missing), datetime64 or + timedelta64, or object arrays of text and None.""" + + names: list[str] + arrays: list[np.ndarray] + named: bool # the names came from the caller (a DataFrame's columns) + + @property + def n_rows(self) -> int: + return len(self.arrays[0]) + + @property + def text_names(self) -> list[str]: + return [ + n for n, a in zip(self.names, self.arrays, strict=True) if a.dtype == object + ] + + def renamed(self, names: Sequence[str]) -> Columns: + return Columns(list(names), self.arrays, self.named) + + +def columns( + X: Any, text: Collection[str] = (), positions: Collection[int] = () +) -> Columns: + """X as columns. text / positions: columns (by name, or by position) + to keep as text whatever their values look like.""" + if is_frame(X): + names = [str(c) for c in X.columns] + named = True + raw = [X.iloc[:, j] for j in range(len(names))] + else: + arr = _as_2d(X) + names = [f"x{j}" for j in range(arr.shape[1])] + named = False + raw = [arr[:, j] for j in range(arr.shape[1])] + if not raw or not len(raw[0]): + raise ValueError("X has no rows or no columns") + duplicated = sorted({n for n in names if names.count(n) > 1}) + if duplicated: + raise ValueError(f"X repeats column name(s) {duplicated}") + arrays = [ + _column(col, as_text=names[j] in text or j in positions) + for j, col in enumerate(raw) + ] + return Columns(names, arrays, named) + + +def to_parquet(cols: Columns, labels: tuple[str, np.ndarray] | None = None) -> bytes: + """The columns (and a label column, as (name, values)) as Parquet.""" + import pyarrow as pa + import pyarrow.parquet as pq + + names, arrays = list(cols.names), list(cols.arrays) + if labels is not None: + name, values = labels + if name in names: + raise ValueError(f"X has a column named {name!r}, which is reserved") + names.append(name) + arrays.append(np.asarray(values)) + fields = [ + pa.array(list(a), pa.string()) if a.dtype == object else pa.array(a) + for a in arrays + ] + buf = io.BytesIO() + pq.write_table(pa.table(fields, names=names), buf, compression="zstd") + return buf.getvalue() + + +def is_frame(X: Any) -> bool: + return hasattr(X, "columns") and hasattr(X, "iloc") + + +def take_rows(cols: Columns, start: int, stop: int) -> Columns: + return Columns(cols.names, [a[start:stop] for a in cols.arrays], cols.named) + + +# ---- columns ---------------------------------------------------------------- + + +def _as_2d(X: Any) -> np.ndarray: + if hasattr(X, "toarray") or hasattr(X, "tocsr"): + raise ValueError("sparse X is not supported; pass a dense array or DataFrame") + # Lists keep each cell's own type ("007" stays text); arrays keep theirs. + arr = X if isinstance(X, np.ndarray) else np.asarray(X, dtype=object) + if arr.ndim != 2: + raise ValueError(f"X must be 2-dimensional, got shape {arr.shape}") + if arr.dtype.kind == "c": + raise ValueError("complex X is not supported") + return arr + + +def _column(values: Any, as_text: bool) -> np.ndarray: + if hasattr(values, "dtype") and hasattr(values, "to_numpy"): # a Series + values = _series_values(values, as_text) + if values.dtype.kind in "Mm" and not as_text: + return values # Parquet timestamps and durations; the server takes seconds + if values.dtype.kind in "biuf" and not as_text: + return values.astype(np.float32) + values = values.astype(object) if values.dtype != object else values + present = [v for v in values if not _missing(v)] + if not as_text and all(isinstance(v, numbers.Real) for v in present): + return np.array( + [np.nan if _missing(v) else float(v) for v in values], dtype=np.float32 + ) + out = np.empty(len(values), dtype=object) + out[:] = [None if _missing(v) else _text(v) for v in values] + return out + + +def _series_values(series, as_text: bool) -> np.ndarray: + from pandas.api.types import ( + is_bool_dtype, + is_datetime64_any_dtype, + is_numeric_dtype, + is_timedelta64_dtype, + ) + + if as_text: + return series.to_numpy(dtype=object) + if is_datetime64_any_dtype(series.dtype): + if getattr(series.dt, "tz", None) is not None: + series = series.dt.tz_convert("UTC").dt.tz_localize(None) + return series.to_numpy(dtype="datetime64[us]") + if is_timedelta64_dtype(series.dtype): + return series.to_numpy(dtype="timedelta64[us]") + if is_numeric_dtype(series.dtype) or is_bool_dtype(series.dtype): + # Nullable dtypes (Int64, boolean) included: pd.NA -> NaN. + return series.to_numpy(dtype=np.float32, na_value=np.nan) + return series.to_numpy(dtype=object) + + +def _missing(v: object) -> bool: + if v is None: + return True + try: + return bool(v != v) # NaN, NaT + except (TypeError, ValueError): # pandas.NA + return True + + +def _text(v: object) -> str: + # As the server compares a number with text: 3, not 3.0. + if isinstance(v, bool | np.bool_): + return str(bool(v)) + if isinstance(v, numbers.Integral): + return str(int(v)) + if isinstance(v, numbers.Real) and float(v).is_integer(): + return str(int(v)) + return str(v) # numpy floats print their shortest form, as Python's do diff --git a/tests/integration/test_estimators.py b/tests/integration/test_estimators.py new file mode 100644 index 0000000..35a9474 --- /dev/null +++ b/tests/integration/test_estimators.py @@ -0,0 +1,71 @@ +"""The scikit-learn estimators and the client's file methods against a +real local 2-replica deployment (fake backend: predicts the majority +class, or the mean target).""" + +import gzip + +import numpy as np +import pytest + +pytest.importorskip("ray") +pytest.importorskip("sklearn") +pd = pytest.importorskip("pandas") + +from sklearn.model_selection import cross_val_score # noqa: E402 +from sklearn.pipeline import make_pipeline # noqa: E402 +from sklearn.preprocessing import FunctionTransformer # noqa: E402 + +from tabctx import TabctxClassifier, TabctxRegressor # noqa: E402 +from tabctx.client import TabctxClient # noqa: E402 + + +def churn(n=40): + rng = np.random.default_rng(0) + X = pd.DataFrame( + { + "tenure": rng.integers(1, 60, n).astype(float), + "plan": rng.choice(["basic", "gold", None], n), + "spend": np.where(rng.random(n) < 0.2, np.nan, rng.normal(50, 10, n)), + } + ) + y = np.array(["stay"] * (n - n // 4) + ["leave"] * (n // 4)) + return X, y + + +def test_classifier_on_a_dataframe_with_text_and_missing(two_replica_service): + X, y = churn() + clf = TabctxClassifier(model="fake", client=TabctxClient(two_replica_service)) + clf.fit(X, y) + assert clf.predict(X.iloc[:5]).tolist() == ["stay"] * 5 + proba = clf.predict_proba(X[["spend", "plan", "tenure"]].iloc[:3]) + assert proba.shape == (3, 2) and list(clf.classes_) == ["leave", "stay"] + assert np.allclose(proba.sum(axis=1), 1.0) + assert clf.score(X, y) == 0.75 + + +def test_regressor_and_numpy_input(two_replica_service): + X = np.random.default_rng(1).normal(size=(30, 4)) + y = np.arange(30, dtype=float) + reg = TabctxRegressor(model="fake", client=TabctxClient(two_replica_service)) + pred = reg.fit(X, y).predict(X[:4]) + assert pred.shape == (4,) and np.allclose(pred, y.mean()) + + +def test_cross_validation_and_pipelines(two_replica_service): + X, y = churn(60) + client = TabctxClient(two_replica_service) + pipe = make_pipeline( + FunctionTransformer(lambda df: df.assign(tenure_years=df.tenure / 12)), + TabctxClassifier(model="fake", client=client), + ) + scores = cross_val_score(pipe, X, y, cv=3) + assert len(scores) == 3 and all(0 <= s <= 1 for s in scores) + + +def test_client_file_methods(two_replica_service, tmp_path): + client = TabctxClient(two_replica_service) + train = tmp_path / "train.csv.gz" + train.write_bytes(gzip.compress(b"a,b,label\n1,x,yes\n2,y,yes\n3,x,no\n")) + ds = client.fit_file(train, target="label", dataset_id="client-files", model="fake") + result = client.predict_file(ds, b"b,a\nx,4\n", return_proba=True) + assert result.predictions == ["yes"] and result.classes == ["no", "yes"] diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 7fd644e..b654df5 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1,6 +1,6 @@ """Unit tests for the pure-stdlib HTTP client (client.py). -urllib.request.urlopen is monkeypatched throughout -- no real network I/O, +tabctx.client._open is monkeypatched throughout -- no real network I/O, no running server required. This exercises exactly what the client promises callers: automatic affinity/tenant headers, response parsing, server-error -> tabctx-exception mapping, and 503 backpressure retry. @@ -71,9 +71,7 @@ def no_real_sleep(monkeypatch): def _install(monkeypatch, handler): """handler(req) -> _FakeResponse, or raises urllib.error.HTTPError.""" - monkeypatch.setattr( - urllib.request, "urlopen", lambda req, timeout=None: handler(req) - ) + monkeypatch.setattr("tabctx.client._open", lambda req, timeout=None: handler(req)) class TestMapStatus: @@ -561,3 +559,58 @@ def handler(req): client.fit([[1.0]], ["a"]) assert attempts["n"] == 1 assert no_real_sleep == [] + + +class TestRedirectsAndKeys: + def test_redirects_are_refused_and_the_key_stays_put(self, monkeypatch): + import http.server + import threading + + seen: list[dict] = [] + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 + seen.append(dict(self.headers)) + self.rfile.read(int(self.headers.get("Content-Length", 0))) + self.send_response(302) + self.send_header("Location", "http://127.0.0.1:1/elsewhere") + self.end_headers() + + do_POST = do_GET # noqa: N815 + + def log_message(self, *args): + pass + + server = http.server.HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + client = TabctxClient( + f"http://127.0.0.1:{server.server_port}", api_key="secret-key" + ) + with pytest.raises(TabctxError, match="redirect") as err: + client.models() + assert "secret-key" not in str(err.value) + with pytest.raises(TabctxError, match="redirect"): + client.fit_file(b"a,b\n1,2\n", dataset_id="d") + assert len(seen) == 2 # each request went to this server only + finally: + server.shutdown() + + def test_a_401_without_a_key_says_how_to_send_one(self, monkeypatch): + monkeypatch.delenv("TABCTX_API_KEY", raising=False) + + def handler(req): + raise _http_error(401, {"detail": "Missing API key"}) + + _install(monkeypatch, handler) + with pytest.raises(PermissionError, match="TABCTX_API_KEY"): + TabctxClient("http://example.test").models() + + +@pytest.mark.parametrize("make", [bytes, bytearray, memoryview]) +def test_multipart_bodies_from_bytes_like_files(make): + from tabctx.client import _Multipart + + form = _Multipart({"dataset_id": "d"}, make(b"a,b\n1,2\n")) + body = b"".join(form.chunks()) + assert len(body) == form.length and b"a,b\n1,2\n" in body diff --git a/tests/unit/test_estimators.py b/tests/unit/test_estimators.py new file mode 100644 index 0000000..c5bfb73 --- /dev/null +++ b/tests/unit/test_estimators.py @@ -0,0 +1,303 @@ +"""The scikit-learn estimators (tabctx/estimators.py), against a fake +client that serves fits and predicts in memory the way a deployment +does: it reads the files the estimators send and encodes them with the +server's own schema code.""" + +import copy +import pickle + +import numpy as np +import pytest + +pytest.importorskip("sklearn") +pytest.importorskip("pyarrow") +pd = pytest.importorskip("pandas") + +from sklearn.base import clone +from sklearn.model_selection import GridSearchCV, cross_val_score + +from tabctx import TabctxClassifier, TabctxRegressor +from tabctx.client import PredictResult, TabctxClient +from tabctx.errors import DatasetNotFoundError +from tabctx.schema import TableSchema +from tabctx.serve.files import read_table, read_training_table + + +class FakeDeployment: + """Predicts a row's class as round(its first encoded feature) % k (so + results depend on how the rows were encoded); regression returns the + first encoded feature.""" + + def __init__(self, tmp_path, classes_order=None): + self.tmp_path = tmp_path + self.contexts: dict[str, tuple] = {} + self.fits: list[dict] = [] + self.predict_calls: list[int] = [] + self.classes_order = classes_order + + def _path(self, data): + p = self.tmp_path / "f" + p.write_bytes(data) + return p + + def fit_file(self, file, target=None, dataset_id=None, model=None, task=None): + features, y = read_training_table(self._path(file), target, task) + schema, _ = TableSchema.learn(features.columns, features.names) + self.fits.append({"dataset_id": dataset_id, "model": model, "task": task}) + classes = sorted(set(y), key=int) if task == "classification" else None + self.contexts[dataset_id] = (task, schema, classes) + return dataset_id + + def predict_file(self, dataset_id, file, return_proba=False, model=None): + self.predict_calls.append(len(file)) + if dataset_id not in self.contexts: + raise DatasetNotFoundError( + f"no cached context for dataset_id={dataset_id!r}" + ) + task, schema, classes = self.contexts[dataset_id] + table = read_table(self._path(file), text_columns=schema.text_names) + X = schema.encode(schema.select(table.names, table.columns)) + if task == "regression": + return PredictResult(X[:, 0].tolist(), None, None, 1.0, None) + k = len(classes) + preds = [classes[int(round(v)) % k] if v == v else classes[0] for v in X[:, 0]] + order = self.classes_order or classes + proba = [[1.0 if c == p else 0.0 for c in order] for p in preds] + return PredictResult(preds, proba if return_proba else None, order, 1.0, None) + + +@pytest.fixture +def deployment(tmp_path): + return FakeDeployment(tmp_path) + + +class TestLabels: + def test_labels_of_any_type_come_back_as_given(self, deployment): + X = np.array([[0.0], [1.0], [2.0], [0.0]]) + for y in ( + ["cat", "dog", "emu", "cat"], + [True, False, True, True], + [10, 20, 30, 10], + ): + clf = TabctxClassifier(client=deployment).fit(X, y) + assert list(clf.classes_) == sorted(set(y)) + k = len(clf.classes_) + assert clf.predict(X).tolist() == [ + clf.classes_[int(v) % k] for v in X[:, 0] + ] + assert deployment.fits[-1] == { + "dataset_id": clf.dataset_id_, + "model": "tabpfn-3.5", + "task": "classification", + } + + def test_predict_proba_columns_follow_classes_(self, tmp_path): + # The server may list classes in any order ("10" before "2"). + deployment = FakeDeployment(tmp_path, classes_order=["2", "0", "1"]) + clf = TabctxClassifier(client=deployment).fit( + [[0.0], [1.0], [2.0]], ["a", "b", "c"] + ) + assert clf.predict_proba([[0.0], [2.0]]).tolist() == [ + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + + @pytest.mark.parametrize( + "y", + [ + [0, 1, np.nan, 1], + np.linspace(0, 1, 4), + [1, "a", 1, "a"], + ["a", None, "b", "a"], + ], + ) + def test_bad_class_labels_are_rejected_here(self, deployment, y): + with pytest.raises((ValueError, TypeError)): + TabctxClassifier(client=deployment).fit([[0.0]] * 4, y) + assert deployment.fits == [] + + @pytest.mark.parametrize( + "y", + [ + ["a", None, "b", "a"], + pd.array(["a", pd.NA, "b", "a"], dtype="string"), + pd.Categorical(["a", None, "b", "a"]), + ], + ) + def test_missing_labels_are_a_value_error(self, deployment, y): + with pytest.raises(ValueError, match="missing labels"): + TabctxClassifier(client=deployment).fit([[0.0]] * 4, y) + + def test_column_vector_labels_are_raveled(self, deployment): + with pytest.warns(Warning): + clf = TabctxClassifier(client=deployment).fit( + [[0.0], [1.0]], [["a"], ["b"]] + ) + assert list(clf.classes_) == ["a", "b"] + + def test_regression_labels_must_be_finite(self, deployment): + with pytest.raises(ValueError, match="missing or infinite"): + TabctxRegressor(client=deployment).fit([[0.0], [1.0]], [1.0, np.nan]) + + +class TestColumns: + def test_dataframes_with_text_and_missing_values(self, deployment): + train = pd.DataFrame({"n": [0.0, 1.0, np.nan], "plan": ["gold", None, "basic"]}) + clf = TabctxClassifier(client=deployment).fit(train, ["x", "y", "x"]) + assert list(clf.feature_names_in_) == ["n", "plan"] + assert clf.n_features_in_ == 2 + test = pd.DataFrame({"plan": ["gold", "basic"], "n": [1.0, 0.0]}) # any order + assert clf.predict(test).tolist() == ["y", "x"] + + def test_training_text_stays_text_when_test_values_look_numeric(self, deployment): + # "c" was text at fit; 0.1 and 2.5 at predict must match "0.1" and "2.5". + train = pd.DataFrame({"c": ["a", 0.1, "b", 2.5]}) + reg = TabctxRegressor(client=deployment).fit(train, [0.0, 1.0, 2.0, 3.0]) + codes = {"0.1": 0.0, "2.5": 1.0, "a": 2.0, "b": 3.0} + got = reg.predict(pd.DataFrame({"c": [0.1, 2.5]})) + assert got.tolist() == [codes["0.1"], codes["2.5"]] + with pytest.warns(UserWarning, match="by position"): + assert reg.predict(np.array([[2.5]], dtype=object)).tolist() == [ + codes["2.5"] + ] + + def test_dataframes_and_arrays_mix_by_position(self, deployment): + df = pd.DataFrame({"a": [0.0, 1.0], "b": [5.0, 6.0]}) + reg = TabctxRegressor(client=deployment).fit(df, [0.0, 1.0]) + with pytest.warns(UserWarning, match="by position"): + assert reg.predict(df.to_numpy()[:1]).tolist() == [0.0] + reg.fit(df.to_numpy(), [0.0, 1.0]) + assert not hasattr(reg, "feature_names_in_") + with pytest.warns(UserWarning, match="by position"): + assert reg.predict(df.iloc[1:]).tolist() == [1.0] + + def test_feature_count_is_checked(self, deployment): + reg = TabctxRegressor(client=deployment).fit([[0.0, 1.0]], [0.0]) + with pytest.raises( + ValueError, match="X has 3 features, but TabctxRegressor is expecting 2" + ): + reg.predict([[0.0, 1.0, 2.0]]) + + def test_non_string_column_names_set_no_feature_names(self, deployment): + reg = TabctxRegressor(client=deployment).fit(pd.DataFrame([[1.0, 2.0]]), [0.0]) + assert not hasattr(reg, "feature_names_in_") + + +class TestRequests: + def test_same_data_same_id_and_model_or_labels_change_it(self, deployment): + X, y = [[0.0], [1.0]], ["a", "b"] + a = TabctxClassifier(client=deployment).fit(X, y).dataset_id_ + assert TabctxClassifier(client=deployment).fit(X, y).dataset_id_ == a + assert ( + TabctxClassifier(model="tabicl-v2", client=deployment).fit(X, y).dataset_id_ + != a + ) + assert TabctxClassifier(client=deployment).fit(X, ["b", "a"]).dataset_id_ != a + + def test_a_lost_context_is_fit_again_once(self, deployment): + clf = TabctxClassifier(client=deployment).fit([[0.0], [1.0]], ["a", "b"]) + deployment.contexts.clear() # evicted, or the replica restarted + assert clf.predict([[1.0]]).tolist() == ["b"] + assert len(deployment.fits) == 2 + + def test_other_404s_are_not_refit(self, deployment): + clf = TabctxClassifier(client=deployment).fit([[0.0], [1.0]], ["a", "b"]) + + def unknown_model(*args, **kwargs): + raise DatasetNotFoundError("Model 'tabpfn-3.5' not found") + + deployment.predict_file = unknown_model + with pytest.raises(DatasetNotFoundError, match="Model"): + clf.predict([[1.0]]) + assert len(deployment.fits) == 1 + + def test_the_model_used_is_the_one_fit(self, deployment): + clf = TabctxClassifier(client=deployment).fit([[0.0], [1.0]], ["a", "b"]) + clf.set_params(model="tabicl-v2") # without fitting again + deployment.contexts.clear() + clf.predict([[0.0]]) + assert deployment.fits[-1]["model"] == "tabpfn-3.5" + + def test_large_predicts_are_split_under_the_request_limit(self, deployment): + reg = TabctxRegressor(client=deployment, max_request_bytes=40_000).fit( + [[0.0, 0.0]], [0.0] + ) + X = np.random.default_rng(0).normal(size=(5000, 2)) + got = reg.predict(X) + assert len(deployment.predict_calls) > 1 + assert max(deployment.predict_calls) <= 40_000 + assert np.allclose(got, X[:, 0].astype(np.float32)) + + def test_uneven_rows_are_split_until_each_fits(self, deployment): + rng = np.random.default_rng(0) + text = [None] * 1800 + [ + "".join(rng.choice(list("abcdef"), 2000)) for _ in range(200) + ] + X = pd.DataFrame({"t": text}) + reg = TabctxRegressor(client=deployment, max_request_bytes=100_000).fit( + X.iloc[1798:1802], [0.0, 1.0, 2.0, 3.0] + ) + assert len(reg.predict(X)) == 2000 + assert max(deployment.predict_calls) <= 100_000 + + def test_oversized_tables_are_refused_before_sending(self, deployment): + X = np.random.default_rng(0).normal(size=(2000, 20)) + with pytest.raises(ValueError, match="training table is"): + TabctxRegressor(client=deployment, max_request_bytes=10_000).fit(X, X[:, 0]) + assert deployment.fits == [] + reg = TabctxRegressor(client=deployment, max_request_bytes=None).fit(X, X[:, 0]) + with pytest.raises(ValueError, match="one test row"): + reg.set_params(max_request_bytes=10).predict(X[:2]) + + def test_unfitted_predict_raises(self, deployment): + from sklearn.exceptions import NotFittedError + + with pytest.raises(NotFittedError): + TabctxClassifier(client=deployment).predict([[0.0]]) + + +def test_works_with_sklearn_model_selection(deployment): + X = np.arange(40, dtype=float).reshape(20, 2) % 3 + y = np.array(["a", "b", "c", "a"] * 5) + assert len(cross_val_score(TabctxClassifier(client=deployment), X, y, cv=2)) == 2 + search = GridSearchCV( + TabctxClassifier(client=deployment), + {"model": ["tabpfn-3.5", "tabicl-v2"]}, + cv=2, + ) + assert search.fit(X, y).best_params_["model"] in ("tabpfn-3.5", "tabicl-v2") + client = TabctxClient("https://gw.example") + copied = clone(TabctxClassifier(model="tabicl-v2", client=client)) + assert copied.get_params() == { + "model": "tabicl-v2", + "client": client, + "max_request_bytes": 29 * 1024 * 1024, + } + + +class TestApiKey: + def test_sent_never_shown_never_pickled(self, monkeypatch): + monkeypatch.delenv("TABCTX_API_KEY", raising=False) + client = TabctxClient("https://gw.example", api_key="secret-key") + assert client._headers(None)["Authorization"] == "Bearer secret-key" + clf = TabctxClassifier(client=client) + assert "secret-key" not in repr(clf) and "secret-key" not in repr(client) + assert "secret-key" not in clf._repr_html_() + assert b"secret-key" not in pickle.dumps(clf) + # Copies keep the key (they share the client). + assert clone(clf).client._api_key == "secret-key" + assert copy.copy(client)._api_key == "secret-key" + + def test_an_unpickled_client_reads_the_environment(self, monkeypatch): + client = TabctxClient("https://gw.example", api_key="secret-key") + monkeypatch.setenv("TABCTX_API_KEY", "from-env") + assert pickle.loads(pickle.dumps(client))._api_key == "from-env" + + def test_defaults_come_from_the_environment(self, monkeypatch): + monkeypatch.setenv("TABCTX_BASE_URL", "https://gw.example/") + monkeypatch.setenv("TABCTX_API_KEY", "k") + client = TabctxClient() + assert client._base_url == "https://gw.example" and client._api_key == "k" + monkeypatch.delenv("TABCTX_BASE_URL") + with pytest.raises(ValueError, match="TABCTX_BASE_URL"): + TabctxClient() diff --git a/tests/unit/test_table_files.py b/tests/unit/test_table_files.py new file mode 100644 index 0000000..753ec18 --- /dev/null +++ b/tests/unit/test_table_files.py @@ -0,0 +1,120 @@ +"""In-memory tables -> Parquet (tabctx/table_files.py), checked by +reading them back the way the server does (serve/files.py).""" + +import datetime +import math + +import numpy as np +import pytest + +pa = pytest.importorskip("pyarrow") +pd = pytest.importorskip("pandas") + +from tabctx.serve.files import read_table, read_training_table +from tabctx.table_files import columns, take_rows, to_parquet + + +def read_back(tmp_path, cols, labels=None): + p = tmp_path / "t.parquet" + p.write_bytes(to_parquet(cols, labels)) + if labels is None: + table = read_table(p) + return table.names, table.columns + table, y = read_training_table(p, labels[0], None) + return table.names, table.columns, y + + +def cells(col): + return [None if isinstance(v, float) and math.isnan(v) else v for v in col.tolist()] + + +def test_dataframe_with_text_missing_and_pandas_dtypes(tmp_path): + df = pd.DataFrame( + { + "age": [31.5, np.nan, 40.0], + "visits": pd.array([1, None, 3], dtype="Int64"), + "member": [True, False, True], + "plan": pd.Categorical(["gold", None, "basic"]), + "city": ["Toronto", "Montréal, QC", None], + "zip": ["007", "010", None], + } + ) + cols = columns(df) + assert cols.named and cols.text_names == ["plan", "city", "zip"] + names, got = read_back(tmp_path, cols) + assert names == list(df.columns) + assert cells(got[0]) == [31.5, None, 40.0] + assert cells(got[1]) == [1.0, None, 3.0] + assert got[2].tolist() == [1.0, 0.0, 1.0] + assert list(got[3]) == ["gold", None, "basic"] + assert list(got[4]) == ["Toronto", "Montréal, QC", None] + assert list(got[5]) == ["007", "010", None] # Parquet strings stay text + + +def test_dates_and_durations_travel_as_such(tmp_path): + df = pd.DataFrame( + { + "when": pd.to_datetime(["1970-01-01 00:00:01.5", None]), + "utc": pd.to_datetime(["1970-01-01 01:00:00+01:00"] * 2), + "took": pd.to_timedelta(["1s", "2.5s"]), + } + ) + _, got = read_back(tmp_path, columns(df)) + assert cells(got[0]) == [1.5, None] + assert got[1].tolist() == [0.0, 0.0] + assert got[2].tolist() == [1.0, 2.5] + + +def test_rows_keep_each_cells_type(tmp_path): + rows = [[1, "a", None, "007"], [2.5, 3, "b", "010"]] + cols = columns(rows) + assert cols.names == ["x0", "x1", "x2", "x3"] and not cols.named + _, got = read_back(tmp_path, cols) + assert got[0].tolist() == [1.0, 2.5] + assert list(got[1]) == ["a", "3"] # a number in a text column is text + assert list(got[2]) == [None, "b"] + assert list(got[3]) == ["007", "010"] + + +def test_float32_arrays_travel_exactly(tmp_path): + X = np.random.default_rng(0).normal(size=(50, 3)).astype(np.float32) + _, got = read_back(tmp_path, columns(X)) + assert np.array_equal(np.column_stack(got), X) + + +def test_forced_text_keeps_numbers_as_written(tmp_path): + df = pd.DataFrame({"c": [0.1, 2.5, 16777217]}) + _, got = read_back(tmp_path, columns(df, text=["c"])) + assert list(got[0]) == ["0.1", "2.5", "16777217"] + _, got = read_back(tmp_path, columns(df.to_numpy(), positions=[0])) + assert list(got[0]) == ["0.1", "2.5", "16777217"] + + +def test_labels_column_keeps_full_precision(tmp_path): + y = np.array([123456789.123, 2.0]) + _, _, got = read_back(tmp_path, columns([[1.0], [2.0]]), ("__y__", y)) + assert got == [123456789.123, 2.0] + + +def test_bad_tables_are_rejected(): + with pytest.raises(ValueError, match="2-dimensional"): + columns([1.0, 2.0]) + with pytest.raises(ValueError, match="no rows"): + columns(np.empty((0, 3))) + with pytest.raises(ValueError, match="reserved"): + to_parquet(columns(pd.DataFrame({"y": [1.0]})), ("y", np.array([0]))) + with pytest.raises(ValueError, match="repeats"): + columns(pd.DataFrame([[1, 2]], columns=["a", "a"])) + with pytest.raises(ValueError, match="complex"): + columns(np.ones((2, 2), dtype=complex)) + + +def test_take_rows(): + cols = columns(np.arange(10.0).reshape(5, 2)) + assert take_rows(cols, 3, 5).arrays[1].tolist() == [7.0, 9.0] + + +def test_a_dated_object_column_is_text(): + # datetime objects in an object column (not a datetime dtype) are text. + cols = columns([[datetime.date(2020, 1, 1)], [None]]) + assert list(cols.arrays[0]) == ["2020-01-01", None] diff --git a/uv.lock b/uv.lock index b3c72b0..05ce588 100644 --- a/uv.lock +++ b/uv.lock @@ -691,7 +691,7 @@ name = "cuda-bindings" version = "13.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "cuda-pathfinder" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, @@ -722,43 +722,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] curand = [ - { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] [[package]] @@ -1500,7 +1500,7 @@ name = "mlx" version = "0.32.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, + { name = "mlx-metal" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/38/f7/5664e13821491559c953494e6705c818203d6880d713a0c22e12bead9564/mlx-0.32.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:77217798a2b036bae9f213b851d4cde4581893787c9964458b7d471f86036bd6", size = 619370, upload-time = "2026-08-25T10:31:44.235Z" }, @@ -1936,7 +1936,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cuda-nvrtc" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -1975,7 +1975,7 @@ name = "nvidia-cudnn-cu13" version = "9.24.0.43" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ca/30/7c257e3d5cb4fecb147b93895c66e29c93f8e76d74b45bb418ff0587c4ec/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a6812a554a1ff0413e9c52b84c26c050380649ab9615f9c16bded368ce9f421f", size = 650976863, upload-time = "2026-07-02T16:23:39.248Z" }, @@ -1987,7 +1987,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -2017,9 +2017,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -2031,7 +2031,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -3267,6 +3267,10 @@ dependencies = [ ] [package.optional-dependencies] +client = [ + { name = "pyarrow" }, + { name = "scikit-learn" }, +] dev = [ { name = "httpx" }, { name = "pytest" }, @@ -3306,18 +3310,20 @@ requires-dist = [ { name = "httpx", marker = "extra == 'dev'" }, { name = "jinja2", marker = "extra == 'serve'" }, { name = "numpy", specifier = ">=1.24" }, + { name = "pyarrow", marker = "extra == 'client'", specifier = ">=14" }, { name = "pyarrow", marker = "extra == 'serve'", specifier = ">=14" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, { name = "pytest-cov", marker = "extra == 'dev'" }, { name = "python-multipart", marker = "extra == 'serve'", specifier = ">=0.0.9" }, { name = "ray", extras = ["serve"], marker = "extra == 'serve'", specifier = ">=2.58.0" }, + { name = "scikit-learn", marker = "extra == 'client'", specifier = ">=1.6" }, { name = "tabicl", marker = "extra == 'tabicl'", specifier = ">=2.2.0" }, { name = "tabpfn", marker = "extra == 'tabpfn'" }, { name = "torch", marker = "extra == 'tabicl'" }, { name = "torch", marker = "extra == 'tabpfn'" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'serve'" }, ] -provides-extras = ["tabicl", "tabpfn", "serve", "dev"] +provides-extras = ["tabicl", "tabpfn", "serve", "client", "dev"] [package.metadata.requires-dev] dev = [