fix(predict): classification labels are strings whatever return_proba is (#45)

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).
 

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
+    )

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:

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:

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())
 

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) ----------------

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

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
     )

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)