fix(predict): size predict chunks with a measured predict cost model (#47)

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 
 ## [Unreleased]
+### 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.
+### 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).
 
 ## [0.12.0] - unreleased
 ### Added

README.md

@@ -393,7 +393,11 @@ pip install "tabctx[dev]"        # + test dependencies
   the memory budget so one oversized request can't crash the whole
   replica, the failure mode this library exists partly to prevent (the
   naive one-shot wrapper it replaces did crash this way; see
-  [CHANGELOG](CHANGELOG.md)).
+  [CHANGELOG](CHANGELOG.md)). Chunks are sized by a measured predict cost
+  model (`memory/predict.py`): against a kv context a predict's memory
+  follows its test rows and their width, not the training table, so a
+  100-row predict against a 10,000 x 500 context is one call. A chunk that
+  still OOMs is halved and retried.
 - **TabICL's kv-cache, actually on**: tabicl ships with its fit-time
   context cache disabled, which makes every `predict()` re-encode the
   whole training set. tabctx enables it by default

benchmarks/baselines/predict-a100/predict_a100.jsonl

@@ -0,0 +1,61 @@
+{"meta": {"device": "NVIDIA A100-SXM4-40GB", "torch": "2.13.0+cu130", "settings": {"backends": "('tabicl', 'tabpfn')", "gpu_memory_fraction": "1.0", "kv_cache": "kv", "kv_precision": "int8", "batch_window_ms": "5.0", "max_upload_bytes": "4294967296", "upload_ttl_s": "3600.0", "spill_dir": "None", "spill_capacity_bytes": "53687091200", "devices": "()", "offload_dir": "None", "worker_processes": "True", "compile_models": "True", "fast_preprocessing": "True"}}}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 28756480}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 290993664}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 46926848}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 114848256}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 168823808}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 189512192}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 110932480}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 1112098304}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 47056384}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 443168256}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 168887808}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 443168256}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 85357056}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 855861760}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 85520896}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 855687168}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 168967680}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 853568000}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 167436800}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 1678824448}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 167436800}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 1679818752}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 169127424}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 1679818752}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 413677056}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 4136960000}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 413677056}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 4136768000}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 413677056}
+{"model": "tabicl-v2", "n_train": 20000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 4136768000}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 62734336}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 605032448}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 106308096}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 605608960}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 10, "n_test": 100, "peak_predict_bytes": 389089792}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 10, "n_test": 1000, "peak_predict_bytes": 606250496}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 205309440}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 1918574592}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 205099008}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 1918835712}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 50, "n_test": 100, "peak_predict_bytes": 389824000}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 50, "n_test": 1000, "peak_predict_bytes": 1920798208}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 405837824}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 3575545856}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 405239296}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 3575804928}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 100, "n_test": 100, "peak_predict_bytes": 407184384}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 100, "n_test": 1000, "peak_predict_bytes": 3576409600}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 807400960}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 6754677248}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 807410688}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 6754936320}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 200, "n_test": 100, "peak_predict_bytes": 807975936}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 200, "n_test": 1000, "peak_predict_bytes": 6755895296}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 2008236544}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 6675193856}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 2008526336}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 6675992576}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 500, "n_test": 100, "peak_predict_bytes": 2009817088}
+{"model": "tabpfn-3.5", "n_train": 20000, "n_features": 500, "n_test": 1000, "peak_predict_bytes": 6676935680}

benchmarks/baselines/predict-a100/predict_a100_tabicl_wide.jsonl

@@ -0,0 +1,5 @@
+{"meta": {"device": "NVIDIA A100-SXM4-40GB", "torch": "2.13.0+cu130", "settings": {"backends": "('tabicl', 'tabpfn')", "gpu_memory_fraction": "1.0", "kv_cache": "kv", "kv_precision": "int8", "batch_window_ms": "5.0", "max_upload_bytes": "4294967296", "upload_ttl_s": "3600.0", "spill_dir": "None", "spill_capacity_bytes": "53687091200", "devices": "()", "offload_dir": "None", "worker_processes": "True", "compile_models": "True", "fast_preprocessing": "True"}}}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 1000, "n_test": 100, "peak_predict_bytes": 824076800}
+{"model": "tabicl-v2", "n_train": 1000, "n_features": 1000, "n_test": 1000, "peak_predict_bytes": 8241390592}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 1000, "n_test": 100, "peak_predict_bytes": 825769984}
+{"model": "tabicl-v2", "n_train": 5000, "n_features": 1000, "n_test": 1000, "peak_predict_bytes": 8241390592}

benchmarks/baselines/predict-a100/predict_a100_tabpfn_wide.jsonl

@@ -0,0 +1,5 @@
+{"meta": {"device": "NVIDIA A100-SXM4-40GB", "torch": "2.13.0+cu130", "settings": {"backends": "('tabicl', 'tabpfn')", "gpu_memory_fraction": "1.0", "kv_cache": "kv", "kv_precision": "int8", "batch_window_ms": "5.0", "max_upload_bytes": "4294967296", "upload_ttl_s": "3600.0", "spill_dir": "None", "spill_capacity_bytes": "53687091200", "devices": "()", "offload_dir": "None", "worker_processes": "True", "compile_models": "True", "fast_preprocessing": "True"}}}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 1000, "n_test": 100, "peak_predict_bytes": 3082464768}
+{"model": "tabpfn-3.5", "n_train": 1000, "n_features": 1000, "n_test": 1000, "peak_predict_bytes": 6683149824}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 1000, "n_test": 100, "peak_predict_bytes": 3082723840}
+{"model": "tabpfn-3.5", "n_train": 5000, "n_features": 1000, "n_test": 1000, "peak_predict_bytes": 6683408896}

benchmarks/bench_stalls.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Fixed-cost stalls across table sizes (#47).
+
+Builds the engine exactly as serving does (serve/factory.build_engine, so
+warmup included), then walks a sequence of table shapes per model: one fit,
+then three predicts of the same rows against it. For every call it records
+wall time and how many graphs torch.compile built during it. A stall that
+recompiles once shows as a slow first call that the repeats don't share; one
+that recompiles on every call shows on all three; a stall without compiles
+points elsewhere.
+
+Usage (on a GPU host, with the serving environment, e.g. TABCTX_BACKEND):
+  TORCH_LOGS=recompiles bench_stalls.py --models tabicl tabpfn \
+      --features 10 100 500 --rows 100 500 2000 3500 10000 --out stalls.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+
+import numpy as np
+
+
+def compiled_graphs() -> int:
+    """Graphs torch.compile has built in this process so far (0 without it)."""
+    try:
+        from torch._dynamo.utils import counters
+    except ImportError:
+        return 0
+    return int(counters["stats"]["unique_graphs"])
+
+
+def timed(fn, *args, **kwargs):
+    """(seconds, graphs compiled) of one call."""
+    graphs = compiled_graphs()
+    t0 = time.perf_counter()
+    fn(*args, **kwargs)
+    return round(time.perf_counter() - t0, 3), compiled_graphs() - graphs
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--models", nargs="+", default=["tabicl", "tabpfn"])
+    p.add_argument("--features", type=int, nargs="+", default=[10, 100, 500])
+    p.add_argument("--rows", type=int, nargs="+", default=[100, 500, 2000, 3500, 10000])
+    p.add_argument("--test-rows", type=int, default=100)
+    p.add_argument("--repeats", type=int, default=3)
+    p.add_argument("--out")
+    args = p.parse_args()
+
+    from tabctx.serve.factory import ServeSettings, build_engine
+
+    t0 = time.perf_counter()
+    built = build_engine(ServeSettings.from_env())
+    engine = built.engine
+    print(
+        json.dumps(
+            {
+                "event": "built",
+                "models": built.backend_names,
+                "device": built.device,
+                "build_s": round(time.perf_counter() - t0, 1),
+                "graphs_after_warmup": compiled_graphs(),
+            }
+        ),
+        flush=True,
+    )
+    rng = np.random.default_rng(0)
+    records = []
+    for model in args.models:
+        backend = next(b for b in built.backend_names if b.startswith(model))
+        for features in args.features:
+            for rows in args.rows:
+                X = rng.normal(size=(rows, features)).astype(np.float32)
+                y = (X[:, 0] + 0.5 * X[:, 1 % features] > 0).astype(int)
+                X_test = rng.normal(size=(args.test_rows, features)).astype(np.float32)
+                dataset_id = f"stall-{model}-{features}x{rows}"
+                fit_s, fit_graphs = timed(
+                    engine.fit, X, y, dataset_id=dataset_id, backend=backend
+                )
+                predicts = []
+                for _ in range(args.repeats):
+                    s, graphs = timed(engine.predict, dataset_id, X_test)
+                    predicts.append({"s": s, "graphs": graphs})
+                engine.evict(dataset_id)
+                record = {
+                    "model": backend,
+                    "shape": f"{rows}x{features}",
+                    "fit_s": fit_s,
+                    "fit_graphs": fit_graphs,
+                    "predicts": predicts,
+                }
+                records.append(record)
+                print(json.dumps(record), flush=True)
+    if args.out:
+        with open(args.out, "w") as f:
+            f.writelines(json.dumps(r) + "\n" for r in records)
+    print(json.dumps({"event": "done", "graphs_total": compiled_graphs()}), flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

benchmarks/calibrate_predict.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""Predict-time GPU memory under the serving configuration (#47).
+
+For each model, fits a grid of (n_train, n_features) tables through the
+engine exactly as serving builds it (serve/factory.build_engine: compiled
+blocks, kv precision and all), then measures one predict's peak: the
+high-water mark of torch's allocator during the call, above what was
+allocated before it. The predict runs on the backend directly, so the
+engine's own chunking (what this calibrates) can't shrink it. Output is one
+JSON record per shape, the input to memory/predict.py's cost model.
+
+Usage (on a CUDA host):
+  TABCTX_BACKEND=tabicl,tabpfn calibrate_predict.py --out predict.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+import numpy as np
+
+
+def main() -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--models", nargs="+", default=["tabicl", "tabpfn"])
+    p.add_argument("--n-train", type=int, nargs="+", default=[1000, 5000, 20000])
+    p.add_argument("--features", type=int, nargs="+", default=[10, 50, 100, 200, 500])
+    p.add_argument("--n-test", type=int, nargs="+", default=[100, 1000])
+    p.add_argument("--out")
+    args = p.parse_args()
+
+    import torch
+
+    from tabctx.errors import AdmissionRejected
+    from tabctx.serve.factory import ServeSettings, build_engine
+
+    built = build_engine(ServeSettings.from_env())
+    engine = built.engine
+    device = torch.device(built.device if ":" in built.device else "cuda")
+    rng = np.random.default_rng(0)
+    records = []
+    for model in args.models:
+        name = next(b for b in built.backend_names if b.startswith(model))
+        backend = engine.backend_for(name)
+        for n_features in args.features:
+            for n_train in args.n_train:
+                X = rng.normal(size=(n_train, n_features)).astype(np.float32)
+                y = (X[:, 0] > 0).astype(int)
+                dataset_id = f"cal-{model}-{n_train}x{n_features}"
+                try:
+                    engine.fit(X, y, dataset_id=dataset_id, backend=name)
+                except AdmissionRejected as e:
+                    skip = {"model": name, "n_train": n_train, "n_features": n_features}
+                    print(json.dumps({**skip, "skipped": str(e)[:200]}), flush=True)
+                    continue
+                # The cached payload, as the engine would hand it to the
+                # backend (calibration needs the unchunked call).
+                payload = engine._cache.peek(dataset_id).payload  # noqa: SLF001
+                for n_test in args.n_test:
+                    X_test = rng.normal(size=(n_test, n_features)).astype(np.float32)
+                    torch.cuda.synchronize(device)
+                    base = torch.cuda.memory_allocated(device)
+                    torch.cuda.reset_peak_memory_stats(device)
+                    backend.predict(payload, X_test, return_proba=True)
+                    torch.cuda.synchronize(device)
+                    peak = torch.cuda.max_memory_allocated(device) - base
+                    record = {
+                        "model": name,
+                        "n_train": n_train,
+                        "n_features": n_features,
+                        "n_test": n_test,
+                        "peak_predict_bytes": int(peak),
+                    }
+                    records.append(record)
+                    print(json.dumps(record), flush=True)
+                payload = None
+                engine.evict(dataset_id)
+    meta = {
+        "device": torch.cuda.get_device_name(device),
+        "torch": torch.__version__,
+        "settings": {
+            k: str(v) for k, v in vars(ServeSettings.from_env()).items() if k != "token"
+        },
+    }
+    print(json.dumps({"event": "done", **meta}), flush=True)
+    if args.out:
+        with open(args.out, "w") as f:
+            f.write(json.dumps({"meta": meta}) + "\n")
+            f.writelines(json.dumps(r) + "\n" for r in records)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())