diff --git a/CHANGELOG.md b/CHANGELOG.md index df51361..f1e336b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed - **Predicts against wide or long contexts ran as one-row calls** (#47). Predict chunks were sized with the fit memory model, which grows with the training table; past ~3,500 x 500 no chunk "fit", chunking fell to one row, and a 100-row predict became 100 sequential model calls: 14.4 s (TabICLv2) and 6.3 s (TabPFN-3.5) instead of ~0.1 s on an A100, ~24 s and ~11 s on prod's A40s. Chunks are now sized by `PredictCostModel` (`memory/predict.py`), built from predict peaks measured under the serving configuration (`benchmarks/calibrate_predict.py`, A100, compiled TabICL, int8 kv; generated `memory/calibration_predict_a100.py`): a cost per test row and a fixed cost per call, both by width, plus a cost per training row. It bounds every measurement it was built from and assumes at most linear growth beyond them (affine in width past the widest, where it scales in proportion to width), and it applies only to the calibrated configuration (kv cache, int8, compiled TabICL; TabPFN). When a call's fixed cost alone exceeds the headroom, only the rows above it are sized against the budget instead of splitting to single rows, which only multiplied that fixed cost; with no headroom a predict still runs as the smallest chunks, and a real OOM still halves a chunk. `LIGHT_PREDICT_ROWS` is now `chunking.MAX_CHUNK_ROWS` (2,048), so a light predict is one call. - **A killed worker could look alive to `close()`**, which then kept its host-memory reservations (and failed `test_workers_close_together_within_one_grace_period` intermittently in CI). The reader thread and `close()` both reaped a dying worker; when two threads call waitpid at once, the loser finds no child and CPython reports the process still running until the winner records its exit code. Every wait, liveness check and kill of a worker is now serialized under one lock (waits on its process sentinel happen outside it, so health checks never block), so every caller sees the exit, and its code, once one does; a sentinel a forked grandchild still holds falls back to a non-blocking check. +- **TabICL classification labels changed type with `return_proba`** (#45): without probabilities it returned `predict()`'s labels (numbers for a JSON `train_y`), with them strings from `classes`; TabPFN always returned strings. TabICL now derives labels the way TabPFN already did, from one `predict_proba` pass, and both report them through `backends._outcome.classification_outcome`: a classification's `predictions` are always strings drawn from `classes`. (Integer labels sent as JSON numbers still come back in float form, `"1.0"`; see #45.) Documented on the predict responses' `predictions` field. ### Added - `benchmarks/bench_stalls.py` (fixed-cost stalls across table sizes, with compiled-graph counts per call), `benchmarks/profile_predict.py` (cProfile of one predict), `benchmarks/calibrate_predict.py` and `benchmarks/gen_predict_grid.py` (predict calibration), `benchmarks/sweep_gateway.py` (the table-size sweep through the Vector gateway). diff --git a/src/tabctx/backends/_outcome.py b/src/tabctx/backends/_outcome.py new file mode 100644 index 0000000..558392a --- /dev/null +++ b/src/tabctx/backends/_outcome.py @@ -0,0 +1,33 @@ +"""How a backend reports a classification predict (#45). + +Every backend derives a classification's labels from its probabilities, +in one pass (the protocol in backends/base.py), and reports them the same +way: as its classes' strings, whatever return_proba is. A backend that +returned predict()'s labels when probabilities weren't asked for gave +numbers one way and strings the other. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import numpy as np + +from tabctx.types import ArrayLike, PredictOutcome + + +def classification_outcome( + proba: ArrayLike, classes: Iterable[Any], return_proba: bool +) -> PredictOutcome: + """A classification predict's outcome from its probabilities (one row + per test sample, one column per class, in `classes` order): each label + is the row's most probable class as a string, so a client gets the same + labels, drawn from `classes`, either way.""" + names = [str(c) for c in classes] + predictions = [names[i] for i in np.argmax(proba, axis=1)] + if not return_proba: + return PredictOutcome(predictions=predictions) + return PredictOutcome( + predictions=predictions, probabilities=np.asarray(proba).tolist(), classes=names + ) diff --git a/src/tabctx/backends/base.py b/src/tabctx/backends/base.py index 20c28a6..771210d 100644 --- a/src/tabctx/backends/base.py +++ b/src/tabctx/backends/base.py @@ -37,7 +37,8 @@ def predict( inference pass regardless of return_proba -- derive predictions from probabilities (e.g. argmax) rather than calling predict() and predict_proba() separately, since each call is a full forward pass - over the cached context.""" + over the cached context. A classification's labels are strings, + whatever return_proba is (backends._outcome).""" ... def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: diff --git a/src/tabctx/backends/fake.py b/src/tabctx/backends/fake.py index d311944..02161d9 100644 --- a/src/tabctx/backends/fake.py +++ b/src/tabctx/backends/fake.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from typing import Any +from tabctx.backends._outcome import classification_outcome from tabctx.types import ArrayLike, PredictOutcome, Task @@ -22,7 +23,6 @@ class _FakePayload: task: Task n_train: int n_features: int - majority_class: str | None = None class_fractions: dict[str, float] | None = None mean_y: float | None = None @@ -62,7 +62,6 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: task=task, n_train=n_train, n_features=n_features, - majority_class=counts.most_common(1)[0][0] if counts else None, class_fractions={k: v / total for k, v in counts.items()}, ) mean_y = sum(float(v) for v in y) / len(y) if len(y) else 0.0 @@ -78,17 +77,12 @@ def predict( time.sleep(self._predict_delay_s) n_test = len(X_test) if payload.task == "classification": - predictions = [payload.majority_class] * n_test - probabilities = None - classes = None - if return_proba: - fractions = payload.class_fractions or {} - classes = sorted(fractions) - row = [fractions.get(c, 0.0) for c in classes] - probabilities = [row for _ in range(n_test)] - return PredictOutcome( - predictions=predictions, probabilities=probabilities, classes=classes - ) + fractions = payload.class_fractions or {} + classes = sorted(fractions) + if not classes: + return PredictOutcome(predictions=[None] * n_test) + row = [fractions[c] for c in classes] + return classification_outcome([row] * n_test, classes, return_proba) return PredictOutcome(predictions=[payload.mean_y] * n_test) def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py index 2fb7ece..46aa8f8 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -64,6 +64,7 @@ ensure_initialized, is_cuda, ) +from tabctx.backends._outcome import classification_outcome from tabctx.backends.tabicl_kv import KV_PRECISIONS, KVPrecision, compress_icl_cache from tabctx.errors import BackendComputeError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -517,17 +518,8 @@ def _predict( self, model: Any, X_test: ArrayLike, return_proba: bool ) -> PredictOutcome: if hasattr(model, "predict_proba"): - if return_proba: - proba = model.predict_proba(X_test) - class_idx = np.argmax(proba, axis=1) - classes = [str(c) for c in model.classes_] - predictions = [classes[i] for i in class_idx] - return PredictOutcome( - predictions=predictions, - probabilities=proba.tolist(), - classes=classes, - ) - return PredictOutcome(predictions=model.predict(X_test).tolist()) + proba = model.predict_proba(X_test) + return classification_outcome(proba, model.classes_, return_proba) # Regressor: no predict_proba to derive from. return PredictOutcome(predictions=model.predict(X_test).tolist()) diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py index 624e4dd..c0dd448 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -44,6 +44,7 @@ ensure_initialized, is_cuda, ) +from tabctx.backends._outcome import classification_outcome from tabctx.errors import BackendComputeError, InvalidInputError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -189,16 +190,7 @@ def _predict( # Single pass regardless of return_proba (protocol # requirement): derive labels from probabilities. proba = model.predict_proba(X_arr) - class_idx = np.argmax(proba, axis=1) - classes = [str(c) for c in model.classes_] - predictions = [classes[i] for i in class_idx] - if return_proba: - return PredictOutcome( - predictions=predictions, - probabilities=proba.tolist(), - classes=classes, - ) - return PredictOutcome(predictions=predictions) + return classification_outcome(proba, model.classes_, return_proba) return PredictOutcome(predictions=model.predict(X_arr).tolist()) # ---- spillover serialization (see cache/spill.py) ---------------- diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py index 958a6b7..365ff9b 100644 --- a/src/tabctx/serve/app.py +++ b/src/tabctx/serve/app.py @@ -29,7 +29,7 @@ from typing import Literal from fastapi import FastAPI, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, Field from ray import serve from ray.serve.config import RequestRouterConfig from starlette.concurrency import run_in_threadpool @@ -90,9 +90,16 @@ class LegacyPredictRequest(BaseModel): return_proba: bool = False +_PREDICTIONS_DOC = ( + "One per test row. Classification: the predicted class as a string, one of " + "the training labels (the `classes` returned with return_proba), whatever " + "return_proba is. Regression: numbers." +) + + class LegacyPredictResponse(BaseModel): request_id: str - predictions: list + predictions: list = Field(description=_PREDICTIONS_DOC) probabilities: list[list[float]] | None = None classes: list[str] | None = None n_train: int @@ -146,7 +153,7 @@ class TabctxPredictRequest(BaseModel): class TabctxPredictResponse(BaseModel): - predictions: list + predictions: list = Field(description=_PREDICTIONS_DOC) probabilities: list[list[float]] | None = None classes: list[str] | None = None n_test: int diff --git a/tests/integration/test_tabicl_kv_real.py b/tests/integration/test_tabicl_kv_real.py index 0770cbb..2b37b16 100644 --- a/tests/integration/test_tabicl_kv_real.py +++ b/tests/integration/test_tabicl_kv_real.py @@ -46,6 +46,10 @@ def test_compressed_context_predicts_like_native(data, precision): assert entries and all(isinstance(e, CompressedKVEntry) for e in entries) got = comp.predict(payload, X[1000:], True) assert got.predictions == want.predictions + # Without probabilities, the same labels, as strings from classes (#45). + plain = comp.predict(payload, X[1000:], False) + assert plain.predictions == got.predictions + assert set(plain.predictions) <= set(got.classes) assert ( np.abs(np.array(got.probabilities) - np.array(want.probabilities)).max() < 0.02 ) diff --git a/tests/unit/test_classification_outcome.py b/tests/unit/test_classification_outcome.py new file mode 100644 index 0000000..a4e157a --- /dev/null +++ b/tests/unit/test_classification_outcome.py @@ -0,0 +1,54 @@ +"""Classification predicts report the same labels, whatever return_proba is (#45).""" + +import numpy as np +import pytest + +from tabctx.backends._outcome import classification_outcome + +PROBA = np.array([[0.1, 0.7, 0.2], [0.6, 0.3, 0.1]]) +CLASSES = np.array([0.0, 1.0, 2.0]) # a JSON train_y of [0, 1, 2] parses as floats + + +@pytest.mark.parametrize("return_proba", [False, True]) +def test_labels_are_class_strings_either_way(return_proba): + outcome = classification_outcome(PROBA, CLASSES, return_proba) + assert outcome.predictions == ["1.0", "0.0"] + if return_proba: + assert outcome.classes == ["0.0", "1.0", "2.0"] + assert outcome.probabilities == PROBA.tolist() + else: + assert outcome.classes is None and outcome.probabilities is None + + +class _Classifier: + """A fitted estimator whose predict() must not be called: a second + forward pass, and the source of numeric labels (#45).""" + + classes_ = CLASSES + + def predict_proba(self, X): + return PROBA[: len(X)] + + def predict(self, X): + raise AssertionError("labels must come from predict_proba") + + +def _backend(module: str, cls: str): + backend_cls = getattr(pytest.importorskip(module), cls) + return object.__new__(backend_cls) # _predict needs no weights + + +@pytest.mark.parametrize( + ("module", "cls"), + [ + ("tabctx.backends.tabicl", "TabICLBackend"), + ("tabctx.backends.tabpfn", "TabPFNBackend"), + ], +) +def test_backends_return_the_same_labels_with_and_without_probabilities(module, cls): + backend = _backend(module, cls) + X = np.zeros((2, 3)) + plain = backend._predict(_Classifier(), X, False) + with_proba = backend._predict(_Classifier(), X, True) + assert plain.predictions == with_proba.predictions == ["1.0", "0.0"] + assert set(plain.predictions) <= set(with_proba.classes)