From f92a23f6adf76b9a4994715f6aa0df659894d022 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 16:34:31 -0400 Subject: [PATCH 1/2] fix(predict): classification labels are strings whatever return_proba is (#45) TabICL returned predict()'s labels without probabilities (numbers for a JSON train_y) and strings from classes with them; TabPFN always returned strings. Both now go through backends.base.classification_outcome: one predict_proba pass, argmax over the class strings, probabilities and classes only when asked. The predictions field documents the contract. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + src/tabctx/backends/base.py | 22 ++++++++- src/tabctx/backends/tabicl.py | 14 ++---- src/tabctx/backends/tabpfn.py | 12 +---- src/tabctx/serve/app.py | 12 +++-- tests/unit/test_classification_outcome.py | 54 +++++++++++++++++++++++ 6 files changed, 90 insertions(+), 25 deletions(-) create mode 100644 tests/unit/test_classification_outcome.py diff --git a/CHANGELOG.md b/CHANGELOG.md index df51361..a3bb679 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. Both backends now derive labels from one `predict_proba` pass through `backends.base.classification_outcome`, so a classification's `predictions` are always strings (drawn from `classes`), and TabICL no longer runs `predict()`. 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/base.py b/src/tabctx/backends/base.py index 20c28a6..3e62104 100644 --- a/src/tabctx/backends/base.py +++ b/src/tabctx/backends/base.py @@ -7,11 +7,30 @@ from __future__ import annotations +from collections.abc import Iterable from typing import Any, Protocol +import numpy as np + from tabctx.types import ArrayLike, PredictOutcome, Task +def classification_outcome( + proba: Any, classes: Iterable[Any], return_proba: bool +) -> PredictOutcome: + """A classification predict's outcome from its probabilities (rows of + test samples, one column per class, in `classes` order): each label is + its class as a string, argmax over the row, whatever return_proba is, + 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 + ) + + class TabularICLBackend(Protocol): """A tabular in-context-learning model, wrapped for tabctx. @@ -37,7 +56,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 (classification_outcome).""" ... 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..359c2ed 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -64,6 +64,7 @@ ensure_initialized, is_cuda, ) +from tabctx.backends.base 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..ad31e31 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -44,6 +44,7 @@ ensure_initialized, is_cuda, ) +from tabctx.backends.base 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..f9aa2b2 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,15 @@ class LegacyPredictRequest(BaseModel): return_proba: bool = False +_PREDICTIONS_DOC = ( + "One per test row. Classification: the predicted class as a string, whatever " + "return_proba is (with return_proba, one of `classes`). 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 +152,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/unit/test_classification_outcome.py b/tests/unit/test_classification_outcome.py new file mode 100644 index 0000000..721a163 --- /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.base 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) From 26e28a38820097fb3ab2ce4c644799d2f62d39a7 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 16:50:50 -0400 Subject: [PATCH 2/2] fix(predict): review fixes -- helper in backends/_outcome.py, fake aligned - classification_outcome moves to backends/_outcome.py, keeping base.py protocol-only; proba typed ArrayLike. - The fake backend reports through it too, so its label is the argmax of its own probabilities (ties went to the first label seen). - The real-model kv test checks return_proba=False gives the same labels. - API doc and CHANGELOG wording. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- src/tabctx/backends/_outcome.py | 33 +++++++++++++++++++++++ src/tabctx/backends/base.py | 21 +-------------- src/tabctx/backends/fake.py | 20 +++++--------- src/tabctx/backends/tabicl.py | 2 +- src/tabctx/backends/tabpfn.py | 2 +- src/tabctx/serve/app.py | 5 ++-- tests/integration/test_tabicl_kv_real.py | 4 +++ tests/unit/test_classification_outcome.py | 2 +- 9 files changed, 52 insertions(+), 39 deletions(-) create mode 100644 src/tabctx/backends/_outcome.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a3bb679..f1e336b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +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. Both backends now derive labels from one `predict_proba` pass through `backends.base.classification_outcome`, so a classification's `predictions` are always strings (drawn from `classes`), and TabICL no longer runs `predict()`. Documented on the predict responses' `predictions` field. +- **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 3e62104..771210d 100644 --- a/src/tabctx/backends/base.py +++ b/src/tabctx/backends/base.py @@ -7,30 +7,11 @@ from __future__ import annotations -from collections.abc import Iterable from typing import Any, Protocol -import numpy as np - from tabctx.types import ArrayLike, PredictOutcome, Task -def classification_outcome( - proba: Any, classes: Iterable[Any], return_proba: bool -) -> PredictOutcome: - """A classification predict's outcome from its probabilities (rows of - test samples, one column per class, in `classes` order): each label is - its class as a string, argmax over the row, whatever return_proba is, - 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 - ) - - class TabularICLBackend(Protocol): """A tabular in-context-learning model, wrapped for tabctx. @@ -57,7 +38,7 @@ def predict( probabilities (e.g. argmax) rather than calling predict() and predict_proba() separately, since each call is a full forward pass over the cached context. A classification's labels are strings, - whatever return_proba is (classification_outcome).""" + 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 359c2ed..46aa8f8 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -64,7 +64,7 @@ ensure_initialized, is_cuda, ) -from tabctx.backends.base import classification_outcome +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 diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py index ad31e31..c0dd448 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -44,7 +44,7 @@ ensure_initialized, is_cuda, ) -from tabctx.backends.base import classification_outcome +from tabctx.backends._outcome import classification_outcome from tabctx.errors import BackendComputeError, InvalidInputError from tabctx.types import ArrayLike, PredictOutcome, Task diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py index f9aa2b2..365ff9b 100644 --- a/src/tabctx/serve/app.py +++ b/src/tabctx/serve/app.py @@ -91,8 +91,9 @@ class LegacyPredictRequest(BaseModel): _PREDICTIONS_DOC = ( - "One per test row. Classification: the predicted class as a string, whatever " - "return_proba is (with return_proba, one of `classes`). Regression: numbers." + "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." ) 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 index 721a163..a4e157a 100644 --- a/tests/unit/test_classification_outcome.py +++ b/tests/unit/test_classification_outcome.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from tabctx.backends.base import classification_outcome +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