diff --git a/CHANGELOG.md b/CHANGELOG.md index 0347d54..8b89a47 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index ef5f683..a316067 100644 --- a/README.md +++ b/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 diff --git a/benchmarks/baselines/predict-a100/predict_a100.jsonl b/benchmarks/baselines/predict-a100/predict_a100.jsonl new file mode 100644 index 0000000..9237189 --- /dev/null +++ b/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} diff --git a/benchmarks/baselines/predict-a100/predict_a100_tabicl_wide.jsonl b/benchmarks/baselines/predict-a100/predict_a100_tabicl_wide.jsonl new file mode 100644 index 0000000..eaff898 --- /dev/null +++ b/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} diff --git a/benchmarks/baselines/predict-a100/predict_a100_tabpfn_wide.jsonl b/benchmarks/baselines/predict-a100/predict_a100_tabpfn_wide.jsonl new file mode 100644 index 0000000..1f40cf6 --- /dev/null +++ b/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} diff --git a/benchmarks/bench_stalls.py b/benchmarks/bench_stalls.py new file mode 100644 index 0000000..b8d3384 --- /dev/null +++ b/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()) diff --git a/benchmarks/calibrate_predict.py b/benchmarks/calibrate_predict.py new file mode 100644 index 0000000..fce8d62 --- /dev/null +++ b/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()) diff --git a/benchmarks/gen_predict_grid.py b/benchmarks/gen_predict_grid.py new file mode 100644 index 0000000..c551e24 --- /dev/null +++ b/benchmarks/gen_predict_grid.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +"""Generate the predict calibration module from calibrate_predict.py output +(JSON lines): one tuple of PredictPoint per model, for PredictCostModel (see +memory/predict.py). + +Usage: + gen_predict_grid.py predict_a100.jsonl predict_a100_wide.jsonl \ + --prefix A100_40GB --out src/tabctx/memory/calibration_predict_a100.py +""" + +from __future__ import annotations + +import argparse +import json +import sys + +# Served model id -> the name the factory keys calibrations by. +_KINDS = {"tabicl": "TABICL", "tabpfn": "TABPFN"} + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("inputs", nargs="+") + p.add_argument("--prefix", required=True) + p.add_argument("--out", required=True) + args = p.parse_args() + metas, rows = [], [] + for path in args.inputs: + with open(path) as f: + for line in f: + if not line.strip(): + continue + record = json.loads(line) + if "meta" in record: + metas.append(record["meta"]) + else: + rows.append(record) + if not metas: + sys.exit("no meta record in the inputs: run calibrate_predict.py with --out") + if any(m != metas[0] for m in metas[1:]): + sys.exit("inputs were measured on different devices or settings") + source = [f"{metas[0]['device']}, torch {metas[0]['torch']}"] + settings = metas[0]["settings"] + config = ", ".join( + f"{k}={settings[k]}" + for k in ("kv_cache", "kv_precision", "compile_models") + if k in settings + ) + lines = [ + '"""GENERATED predict calibration -- do not hand-edit.', + "", + "Produced by benchmarks/gen_predict_grid.py from", + "benchmarks/calibrate_predict.py on real hardware: each PredictPoint is", + "one predict's peak allocation (above what was allocated before it)", + "against a kv-cached context, under the serving configuration. Read by", + "serve/factory.py into PredictCostModel (memory/predict.py).", + "", + f"Measured on {'; '.join(source)} ({config}).", + '"""', + "", + "from __future__ import annotations", + "", + "from tabctx.memory.predict import PredictPoint", + ] + for kind, name in _KINDS.items(): + points = sorted( + (r for r in rows if r["model"].startswith(kind)), + key=lambda r: (r["n_features"], r["n_train"], r["n_test"]), + ) + if not points: + continue + lines += [ + "", + f"{args.prefix}_{name}_PREDICT_POINTS: tuple[PredictPoint, ...] = (", + ] + lines += [ + f" PredictPoint(n_train={r['n_train']}, n_test={r['n_test']}, " + f"n_features={r['n_features']}, peak_bytes={r['peak_predict_bytes']})," + for r in points + ] + lines.append(")") + with open(args.out, "w") as f: + f.write("\n".join(lines) + "\n") + print(f"wrote {len(rows)} points to {args.out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/gke/devbox.sh b/benchmarks/gke/devbox.sh index 075b63d..f8eece4 100755 --- a/benchmarks/gke/devbox.sh +++ b/benchmarks/gke/devbox.sh @@ -150,6 +150,15 @@ case "${1:-}" in sync) sync ;; exec) shift; [ "${1:-}" = "--" ] && shift; creds; k exec $POD -- bash -lc "cd /work/tabctx && $*" ;; fetch) creds; mkdir -p gke-out; k cp "$NS/$POD:/work/tabctx/$2" "gke-out/$(basename "$2")" ;; - down) creds; gcloud container clusters delete "$CLUSTER" --project="$PROJECT" --zone="$ZONE" --quiet --async; rm -f "$STATE/env" ;; + down) + # Only the zone, never credentials: a cluster whose creation failed (state + # ERROR) has none to fetch, and must still be deletable. + # shellcheck disable=SC1091 + [ -f "$STATE/env" ] && source "$STATE/env" + ZONE="${ZONE:-$(gcloud container clusters list --project="$PROJECT" \ + --filter="name=$CLUSTER" --format='value(location)')}" + [ -n "$ZONE" ] || { echo "no cluster $CLUSTER" >&2; exit 1; } + gcloud container clusters delete "$CLUSTER" --project="$PROJECT" --zone="$ZONE" --quiet --async + rm -f "$STATE/env" ;; *) sed -n 2,16p "$0"; exit 1 ;; esac diff --git a/benchmarks/profile_predict.py b/benchmarks/profile_predict.py new file mode 100644 index 0000000..d0400d1 --- /dev/null +++ b/benchmarks/profile_predict.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""Where one predict's time goes (#47): cProfile over a cached-context predict +built exactly as serving builds it (serve/factory.build_engine). + +Usage (on a GPU host): + TABCTX_BACKEND=tabicl profile_predict.py --model tabicl --shape 3500x500 +""" + +from __future__ import annotations + +import argparse +import cProfile +import pstats +import sys +import time + +import numpy as np + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--model", default="tabicl") + p.add_argument("--shape", default="3500x500") + p.add_argument("--test-rows", type=int, default=100) + p.add_argument("--top", type=int, default=40) + args = p.parse_args() + rows, features = (int(v) for v in args.shape.split("x")) + + from tabctx.serve.factory import ServeSettings, build_engine + + built = build_engine(ServeSettings.from_env()) + engine = built.engine + backend = next(b for b in built.backend_names if b.startswith(args.model)) + rng = np.random.default_rng(0) + X = rng.normal(size=(rows, features)).astype(np.float32) + y = (X[:, 0] > 0).astype(int) + X_test = rng.normal(size=(args.test_rows, features)).astype(np.float32) + engine.fit(X, y, dataset_id="p", backend=backend) + engine.predict("p", X_test) # first call: anything one-off happens here + t0 = time.perf_counter() + profiler = cProfile.Profile() + profiler.runcall(engine.predict, "p", X_test) + print(f"predict {time.perf_counter() - t0:.2f}s", flush=True) + stats = pstats.Stats(profiler).sort_stats("cumulative") + stats.print_stats(args.top) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/sweep_gateway.py b/benchmarks/sweep_gateway.py new file mode 100644 index 0000000..ae6ee4a --- /dev/null +++ b/benchmarks/sweep_gateway.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +"""Table-size sweep through the Vector inference gateway. + +The same steps as sweep_admission.py, but as a gateway user sees them: +bearer auth, the gateway's own tenant and session headers, and its request +body cap. A 413 can come from that cap or from tabctx's admission gate, so +each rejection records its detail. Comparing the output with a direct +sweep shows what the gateway adds to fit and predict latency, and where +its body cap stops a table before admission would. + +Usage: + VP_API_KEY=... sweep_gateway.py --model tabicl-v2 \ + --save benchmarks/baselines/sweep-gateway-a40x4-tabicl-v2.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +import urllib.error +import urllib.request +from datetime import UTC, datetime +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from bench_workload import encode_table, make_table # noqa: E402 +from sweep_admission import ROWS # noqa: E402 + + +class Gateway: + """The three tabctx endpoints, called the way a gateway client calls them.""" + + def __init__(self, base_url: str, api_key: str, model: str, timeout_s: float): + self._base = base_url.rstrip("/") + self._auth = {"Authorization": f"Bearer {api_key}"} + self._model = model + self._timeout = timeout_s + + def _send(self, path: str, body: bytes, headers: dict) -> tuple[int, dict]: + req = urllib.request.Request( + self._base + path, data=body, headers=self._auth | headers, method="POST" + ) + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp: + return resp.status, json.loads(resp.read()) + except urllib.error.HTTPError as e: + raw = e.read() + try: + return e.code, json.loads(raw) + except ValueError: + return e.code, {"detail": raw[:300].decode(errors="replace")} + + def _post_json(self, path: str, payload: dict) -> tuple[int, dict]: + body = json.dumps(payload | {"model": self._model}).encode() + return self._send(path, body, {"Content-Type": "application/json"}) + + def fit(self, encoded, dataset_id: str) -> tuple[int, dict]: + kind, data = encoded + if kind == "csv": + status, out = self._send( + "/tabctx/upload", + data, + { + "Content-Type": "text/csv", + "x-tabctx-model": self._model, + "x-tabctx-dataset-id": dataset_id, + }, + ) + if status != 200: + return status, out + payload = {"dataset_id": dataset_id, "train_upload_id": out["upload_id"]} + else: + X, y = data + payload = {"dataset_id": dataset_id, "train_X": X, "train_y": y} + return self._post_json("/tabctx/fit", payload) + + def predict(self, dataset_id: str, X_test: list) -> tuple[int, dict]: + return self._post_json( + "/tabctx/predict", {"dataset_id": dataset_id, "test_X": X_test} + ) + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--base-url", default="https://proxy.vectorinstitute.ai/v1") + p.add_argument("--model", required=True) + p.add_argument("--features", type=int, nargs="+", default=[10, 100, 500, 1000]) + p.add_argument("--max-rows", type=int, default=500_000) + p.add_argument("--test-rows", type=int, default=100) + p.add_argument("--timeout", type=float, default=1800) + p.add_argument("--save", default=None) + args = p.parse_args() + + gateway = Gateway(args.base_url, os.environ["VP_API_KEY"], args.model, args.timeout) + rng = np.random.default_rng(0) + points = [] + for features in args.features: + dataset_id = f"sweep-gw-{args.model}-f{features}" + for rows in (r for r in ROWS if r <= args.max_rows): + X, y = make_table(rows, features, rng) + encoded = encode_table(X, y) + body_bytes = len(encoded[1]) if encoded[0] == "csv" else None + point = { + "n_features": features, + "n_rows": rows, + "upload": encoded[0] == "csv", + "upload_bytes": body_bytes, + } + t0 = time.monotonic() + status, out = gateway.fit(encoded, dataset_id) + point["fit_s"] = round(time.monotonic() - t0, 3) + if status != 200: + point.update(status=status, detail=str(out.get("detail"))[:300]) + points.append(point) + print(json.dumps(point), flush=True) + break + X_test, _ = make_table(args.test_rows, features, rng) + t0 = time.monotonic() + status, out = gateway.predict(dataset_id, X_test.tolist()) + point["predict_s"] = round(time.monotonic() - t0, 3) + point["status"] = 200 if status == 200 else f"predict_{status}" + if status != 200: + point["detail"] = str(out.get("detail"))[:300] + points.append(point) + print(json.dumps(point), flush=True) + + result = { + "timestamp": datetime.now(UTC).isoformat(), + "base_url": args.base_url, + "model": args.model, + "points": points, + "max_admitted_rows": { + str(f): max( + ( + q["n_rows"] + for q in points + if q["n_features"] == f and q["status"] == 200 + ), + default=0, + ) + for f in args.features + }, + } + print(json.dumps(result["max_admitted_rows"])) + if args.save: + with open(args.save, "w") as f: + json.dump(result, f, indent=2) + bad = [q for q in points if q["status"] not in (200, 413)] + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/tabctx/chunking.py b/src/tabctx/chunking.py index 0e21cc9..2dc46af 100644 --- a/src/tabctx/chunking.py +++ b/src/tabctx/chunking.py @@ -18,7 +18,9 @@ from tabctx.memory.estimator import MemoryEstimator -_FALLBACK_CHUNK_ROWS = 2_000 +# One call's worth of test rows at most: the engine's LIGHT_PREDICT_ROWS, so +# a light predict (one that may run beside a fit) is a single call. +MAX_CHUNK_ROWS = 2_048 def choose_chunk_rows( @@ -28,19 +30,26 @@ def choose_chunk_rows( remaining_budget_bytes: int, min_chunk_rows: int = 1, ) -> int: - """Largest test-row chunk size such that estimate_bytes(n_train, - chunk_rows, n_features) fits within remaining_budget_bytes, found by - halving from a fallback starting point. Falls back to - _FALLBACK_CHUNK_ROWS (clamped to at least min_chunk_rows) if even the - smallest chunk doesn't fit -- the caller (engine) is responsible for - surfacing that as an admission failure rather than looping forever. + """Largest test-row chunk, halving from MAX_CHUNK_ROWS down to + min_chunk_rows, whose predict fits remaining_budget_bytes. + + What a call of min_chunk_rows costs is, beyond those rows, the call's + fixed cost, which no chunk size reduces. When that alone exceeds the + budget, smaller chunks would only multiply the calls (a 100-row predict + became 100 one-row calls, #47), so the fixed cost is taken as given and + only the rows' cost above it is sized against the budget: a chunk never + adds more than the budget on top of what the smallest call needs. A + chunk that still OOMs is halved by the engine. """ - candidate = _FALLBACK_CHUNK_ROWS + if min_chunk_rows < 1: + raise ValueError("min_chunk_rows must be at least 1") + floor_bytes = estimator.estimate_bytes(n_train, min_chunk_rows, n_features) + allowance = remaining_budget_bytes + if floor_bytes > remaining_budget_bytes: + allowance = floor_bytes + remaining_budget_bytes + candidate = MAX_CHUNK_ROWS while candidate > min_chunk_rows: - if ( - estimator.estimate_bytes(n_train, candidate, n_features) - <= remaining_budget_bytes - ): + if estimator.estimate_bytes(n_train, candidate, n_features) <= allowance: return candidate candidate //= 2 return min_chunk_rows diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py index 566013a..5cf8973 100644 --- a/src/tabctx/engine.py +++ b/src/tabctx/engine.py @@ -15,7 +15,7 @@ from tabctx.backends._cuda import release_cached_device_memory from tabctx.backends.base import TabularICLBackend from tabctx.cache.manager import CachedContext, ContextCacheManager -from tabctx.chunking import choose_chunk_rows, split_rows +from tabctx.chunking import MAX_CHUNK_ROWS, choose_chunk_rows, split_rows from tabctx.errors import ( AdmissionRejected, BackendComputeError, @@ -31,8 +31,9 @@ # #39): measured on an A100, a 100-row predict took ~0.1 s p50 during a # 40k x 100 fit (vs the fit's full 11 s), slowing the fit ~15%. Larger # predicts, and predicts that must first restore a spilled context, need -# device memory a running fit may be using, so they wait their turn. -LIGHT_PREDICT_ROWS = 2048 +# device memory a running fit may be using, so they wait their turn. One +# chunk's worth (chunking.MAX_CHUNK_ROWS), so a light predict is one call. +LIGHT_PREDICT_ROWS = MAX_CHUNK_ROWS # A predict whose context is evicted again (under memory pressure) between # its restore and getting the device tries this many times. _PREDICT_ATTEMPTS = 3 diff --git a/src/tabctx/memory/adaptive.py b/src/tabctx/memory/adaptive.py index 7dd4452..9e6d226 100644 --- a/src/tabctx/memory/adaptive.py +++ b/src/tabctx/memory/adaptive.py @@ -35,6 +35,7 @@ from dataclasses import dataclass from tabctx.memory.estimator import MemoryEstimator +from tabctx.memory.predict import PredictCostModel from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel # The n_test basis at which calibration predict-peaks were measured @@ -73,6 +74,7 @@ def __init__( working_floor: WorkingFloorModel | None = None, resident_rate: ResidentRateModel | None = None, resident_margin: float = 1.2, + predict_cost: PredictCostModel | None = None, ) -> None: """preloaded: factory-installed observations (e.g. a measured calibration grid -- see memory/calibration_data.py), consulted @@ -103,7 +105,13 @@ def __init__( resident_rate: with working_floor, extrapolates shapes no measurement dominates as (extrapolated resident x resident_margin) + floor -- a kv context grows linearly in rows -- instead of the - fallback formula. Learns from every fit.""" + fallback formula. Learns from every fit. + + predict_cost: for kv-cached contexts, whose predicts don't + re-encode the training table, bounds a predict's peak from test + rows x features alone (see memory/predict.py). Takes precedence + over preloaded_predict, whose dominance rule can't extrapolate to + wider or longer tables than were measured.""" self._fallback = fallback self._safety_margin = safety_margin self._preloaded: tuple[Observation, ...] = tuple(preloaded) @@ -114,6 +122,7 @@ def __init__( # degrades chunking to 1-row chunks (a 2000-row predict became # 2000 sequential GPU calls on a real A100 -- the bug this fixes). self._preloaded_predict: tuple[Observation, ...] = tuple(preloaded_predict) + self._predict_cost = predict_cost self._preloaded_margin = preloaded_margin self._transient_capacity_fraction = transient_capacity_fraction self._working_floor = working_floor @@ -198,6 +207,10 @@ def estimate_bytes(self, n_train: int, n_test: int, n_features: int) -> int: extrapolated = self._extrapolated_requirement(n_train, n_features) if extrapolated is not None: return extrapolated + elif self._predict_cost is not None: + # Predict-time (chunking) query against a kv context: its cost + # is mostly the test rows', not a re-encode of the training table. + return self._predict_cost.bytes(n_train, n_test, n_features) else: # Predict-time (chunking) query: use a measured predict peak # when a dominating one exists, scaled linearly in n_test diff --git a/src/tabctx/memory/calibration_predict_a100.py b/src/tabctx/memory/calibration_predict_a100.py new file mode 100644 index 0000000..c8296fa --- /dev/null +++ b/src/tabctx/memory/calibration_predict_a100.py @@ -0,0 +1,88 @@ +"""GENERATED predict calibration -- do not hand-edit. + +Produced by benchmarks/gen_predict_grid.py from +benchmarks/calibrate_predict.py on real hardware: each PredictPoint is +one predict's peak allocation (above what was allocated before it) +against a kv-cached context, under the serving configuration. Read by +serve/factory.py into PredictCostModel (memory/predict.py). + +Measured on NVIDIA A100-SXM4-40GB, torch 2.13.0+cu130 (kv_cache=kv, kv_precision=int8, compile_models=True). +""" + +from __future__ import annotations + +from tabctx.memory.predict import PredictPoint + +A100_40GB_TABICL_PREDICT_POINTS: tuple[PredictPoint, ...] = ( + PredictPoint(n_train=1000, n_test=100, n_features=10, peak_bytes=28756480), + PredictPoint(n_train=1000, n_test=1000, n_features=10, peak_bytes=290993664), + PredictPoint(n_train=5000, n_test=100, n_features=10, peak_bytes=46926848), + PredictPoint(n_train=5000, n_test=1000, n_features=10, peak_bytes=114848256), + PredictPoint(n_train=20000, n_test=100, n_features=10, peak_bytes=168823808), + PredictPoint(n_train=20000, n_test=1000, n_features=10, peak_bytes=189512192), + PredictPoint(n_train=1000, n_test=100, n_features=50, peak_bytes=110932480), + PredictPoint(n_train=1000, n_test=1000, n_features=50, peak_bytes=1112098304), + PredictPoint(n_train=5000, n_test=100, n_features=50, peak_bytes=47056384), + PredictPoint(n_train=5000, n_test=1000, n_features=50, peak_bytes=443168256), + PredictPoint(n_train=20000, n_test=100, n_features=50, peak_bytes=168887808), + PredictPoint(n_train=20000, n_test=1000, n_features=50, peak_bytes=443168256), + PredictPoint(n_train=1000, n_test=100, n_features=100, peak_bytes=85357056), + PredictPoint(n_train=1000, n_test=1000, n_features=100, peak_bytes=855861760), + PredictPoint(n_train=5000, n_test=100, n_features=100, peak_bytes=85520896), + PredictPoint(n_train=5000, n_test=1000, n_features=100, peak_bytes=855687168), + PredictPoint(n_train=20000, n_test=100, n_features=100, peak_bytes=168967680), + PredictPoint(n_train=20000, n_test=1000, n_features=100, peak_bytes=853568000), + PredictPoint(n_train=1000, n_test=100, n_features=200, peak_bytes=167436800), + PredictPoint(n_train=1000, n_test=1000, n_features=200, peak_bytes=1678824448), + PredictPoint(n_train=5000, n_test=100, n_features=200, peak_bytes=167436800), + PredictPoint(n_train=5000, n_test=1000, n_features=200, peak_bytes=1679818752), + PredictPoint(n_train=20000, n_test=100, n_features=200, peak_bytes=169127424), + PredictPoint(n_train=20000, n_test=1000, n_features=200, peak_bytes=1679818752), + PredictPoint(n_train=1000, n_test=100, n_features=500, peak_bytes=413677056), + PredictPoint(n_train=1000, n_test=1000, n_features=500, peak_bytes=4136960000), + PredictPoint(n_train=5000, n_test=100, n_features=500, peak_bytes=413677056), + PredictPoint(n_train=5000, n_test=1000, n_features=500, peak_bytes=4136768000), + PredictPoint(n_train=20000, n_test=100, n_features=500, peak_bytes=413677056), + PredictPoint(n_train=20000, n_test=1000, n_features=500, peak_bytes=4136768000), + PredictPoint(n_train=1000, n_test=100, n_features=1000, peak_bytes=824076800), + PredictPoint(n_train=1000, n_test=1000, n_features=1000, peak_bytes=8241390592), + PredictPoint(n_train=5000, n_test=100, n_features=1000, peak_bytes=825769984), + PredictPoint(n_train=5000, n_test=1000, n_features=1000, peak_bytes=8241390592), +) + +A100_40GB_TABPFN_PREDICT_POINTS: tuple[PredictPoint, ...] = ( + PredictPoint(n_train=1000, n_test=100, n_features=10, peak_bytes=62734336), + PredictPoint(n_train=1000, n_test=1000, n_features=10, peak_bytes=605032448), + PredictPoint(n_train=5000, n_test=100, n_features=10, peak_bytes=106308096), + PredictPoint(n_train=5000, n_test=1000, n_features=10, peak_bytes=605608960), + PredictPoint(n_train=20000, n_test=100, n_features=10, peak_bytes=389089792), + PredictPoint(n_train=20000, n_test=1000, n_features=10, peak_bytes=606250496), + PredictPoint(n_train=1000, n_test=100, n_features=50, peak_bytes=205309440), + PredictPoint(n_train=1000, n_test=1000, n_features=50, peak_bytes=1918574592), + PredictPoint(n_train=5000, n_test=100, n_features=50, peak_bytes=205099008), + PredictPoint(n_train=5000, n_test=1000, n_features=50, peak_bytes=1918835712), + PredictPoint(n_train=20000, n_test=100, n_features=50, peak_bytes=389824000), + PredictPoint(n_train=20000, n_test=1000, n_features=50, peak_bytes=1920798208), + PredictPoint(n_train=1000, n_test=100, n_features=100, peak_bytes=405837824), + PredictPoint(n_train=1000, n_test=1000, n_features=100, peak_bytes=3575545856), + PredictPoint(n_train=5000, n_test=100, n_features=100, peak_bytes=405239296), + PredictPoint(n_train=5000, n_test=1000, n_features=100, peak_bytes=3575804928), + PredictPoint(n_train=20000, n_test=100, n_features=100, peak_bytes=407184384), + PredictPoint(n_train=20000, n_test=1000, n_features=100, peak_bytes=3576409600), + PredictPoint(n_train=1000, n_test=100, n_features=200, peak_bytes=807400960), + PredictPoint(n_train=1000, n_test=1000, n_features=200, peak_bytes=6754677248), + PredictPoint(n_train=5000, n_test=100, n_features=200, peak_bytes=807410688), + PredictPoint(n_train=5000, n_test=1000, n_features=200, peak_bytes=6754936320), + PredictPoint(n_train=20000, n_test=100, n_features=200, peak_bytes=807975936), + PredictPoint(n_train=20000, n_test=1000, n_features=200, peak_bytes=6755895296), + PredictPoint(n_train=1000, n_test=100, n_features=500, peak_bytes=2008236544), + PredictPoint(n_train=1000, n_test=1000, n_features=500, peak_bytes=6675193856), + PredictPoint(n_train=5000, n_test=100, n_features=500, peak_bytes=2008526336), + PredictPoint(n_train=5000, n_test=1000, n_features=500, peak_bytes=6675992576), + PredictPoint(n_train=20000, n_test=100, n_features=500, peak_bytes=2009817088), + PredictPoint(n_train=20000, n_test=1000, n_features=500, peak_bytes=6676935680), + PredictPoint(n_train=1000, n_test=100, n_features=1000, peak_bytes=3082464768), + PredictPoint(n_train=1000, n_test=1000, n_features=1000, peak_bytes=6683149824), + PredictPoint(n_train=5000, n_test=100, n_features=1000, peak_bytes=3082723840), + PredictPoint(n_train=5000, n_test=1000, n_features=1000, peak_bytes=6683408896), +) diff --git a/src/tabctx/memory/predict.py b/src/tabctx/memory/predict.py new file mode 100644 index 0000000..7d169d3 --- /dev/null +++ b/src/tabctx/memory/predict.py @@ -0,0 +1,173 @@ +"""What a predict against a resident kv context needs on the device. + +A predict against a kv-cached context doesn't re-encode the training +table: the test rows run through the model and attend to cached keys and +values. Measured under the serving configuration (A100, compiled TabICL, +int8 kv; benchmarks/calibrate_predict.py), its peak has three parts: + +- a cost per test row that grows with the table's width (TabICLv2: ~1.1 MB + at 100 features, ~4.1 MB at 500, ~8.2 MB at 1,000), whatever the size of + the training table. TabPFN-3.5 batches test rows itself, so its peak at + 1,000 test rows stops growing past ~200 features (~6.7 GB); +- a fixed cost per call that also grows with width (TabPFN-3.5: ~3 GB at + 1,000 features for even 100 test rows); +- the cached context read back for attention: ~7 KB (TabICLv2) and ~17 KB + (TabPFN-3.5) per training row, visible when tables are narrow (an int8 + cache is dequantized a layer at a time). + +Sizing a predict's chunk with the fit model instead -- which scales with the +whole training table -- is wrong in the direction that hurts most: past some +table size no chunk "fits", chunking falls to one row, and a 100-row request +becomes 100 sequential model calls (14 s instead of 0.1 s on an A100, #47). + +``PredictCostModel`` is built from measured marginal costs, not a fitted +surface: fitted surfaces and fitted slopes underestimated held-out widths by +up to a third (tabicl's fp32 path for small contexts and TabPFN's own test +batching bend the curve), which a bound must not do. + +- ``row_bytes(f)``: per measured width, the extra peak per extra test row + (largest across context sizes). +- ``train_row_bytes``: the extra peak per extra training row, largest across + measured widths and test sizes. +- ``call_bytes(f)``: per measured width, the largest remainder of a measured + peak after the two above -- the fixed cost of a call at that width. + +Both width-dependent costs are made non-decreasing in width; a width between +measurements takes the next wider one's cost (as fit admission's dominance +rule does), and past the widest the widest's cost scales in proportion to +width. + +bytes = margin x (call_bytes(f) + n_test x row_bytes(f) + n_train x train_row_bytes). + +What is guaranteed and what is assumed: the bound is at or above every +measurement it was built from. Elsewhere it assumes each cost grows at most +linearly: in test rows and training rows beyond the largest measured (slopes +are the steepest between consecutive measured sizes), and affine in width +past the widest (a + b x f with a >= 0, for which cost(w) x f / w >= a + b x f +when f >= w). A cost that grows faster than that past the measured range +would be underestimated, so calibrate across the shapes a deployment +serves; the engine halves a chunk that OOMs regardless. +""" + +from __future__ import annotations + +import bisect +import math +from collections import defaultdict +from collections.abc import Iterable, Mapping +from dataclasses import dataclass + + +@dataclass(frozen=True) +class PredictPoint: + """One measured predict: its shape and its peak allocation above what + was allocated before the call.""" + + n_train: int + n_test: int + n_features: int + peak_bytes: int + + +class PredictCostModel: + """Upper bound on a kv-context predict's peak memory, from measurements. + + Needs predicts of two or more test sizes against the same context; + otherwise raises ValueError, so a deployment without a usable calibration + gets no model rather than a guess. + """ + + def __init__(self, points: Iterable[PredictPoint], margin: float = 1.1) -> None: + pts = [p for p in points if p.n_test > 0 and p.n_features > 0] + per_row = _largest(_marginals(pts, "n_test", ("n_features", "n_train"))) + if not per_row: + raise ValueError( + "PredictCostModel needs predicts of two test sizes against one context" + ) + self._row = _WidthCurve(per_row) + by_train = _marginals(pts, "n_train", ("n_features", "n_test")) + self._train_row_bytes = max( + (max(0.0, r) for rates in by_train.values() for r in rates), default=0.0 + ) + call: dict[int, float] = defaultdict(float) + for p in pts: + rest = p.peak_bytes - self._variable(p.n_train, p.n_test, p.n_features) + call[p.n_features] = max(call[p.n_features], rest) + self._call = _WidthCurve(call) + self._margin = margin + + @property + def train_row_bytes(self) -> float: + return self._train_row_bytes + + def row_bytes(self, n_features: int) -> float: + """Peak per test row at this width.""" + return self._row(n_features) + + def call_bytes(self, n_features: int) -> float: + """Fixed peak of one call at this width, whatever its rows.""" + return self._call(n_features) + + def bytes(self, n_train: int, n_test: int, n_features: int) -> int: + """Bound on the peak of a predict of n_test rows x n_features against + a context of n_train training rows.""" + raw = self._call(n_features) + self._variable(n_train, n_test, n_features) + return math.ceil(raw * self._margin) + + def _variable(self, n_train: int, n_test: int, n_features: int) -> float: + return ( + max(0, n_test) * self._row(n_features) + + max(0, n_train) * self._train_row_bytes + ) + + +class _WidthCurve: + """A measured cost per width as a bound for any width: non-decreasing, + the next wider measurement's cost between measurements, and scaled in + proportion to width past the widest.""" + + def __init__(self, by_width: Mapping[int, float]) -> None: + self._widths = sorted(by_width) + self._costs: list[float] = [] + running = 0.0 + for w in self._widths: + running = max(running, by_width[w]) + self._costs.append(running) + + def __call__(self, n_features: int) -> float: + i = bisect.bisect_left(self._widths, n_features) + if i < len(self._widths): + return self._costs[i] + return self._costs[-1] * n_features / self._widths[-1] + + +def _marginals( + pts: list[PredictPoint], vary: str, fixed: tuple[str, ...] +) -> dict[tuple, list[float]]: + """Extra peak per unit of `vary` between consecutive measured sizes, in + each group of measurements that holds `fixed` constant (a group with one + size has none).""" + groups: dict[tuple, list[PredictPoint]] = defaultdict(list) + for p in pts: + groups[tuple(getattr(p, f) for f in fixed)].append(p) + out: dict[tuple, list[float]] = {} + for key, group in groups.items(): + ordered = sorted(group, key=lambda p: getattr(p, vary)) + slopes = [ + (hi.peak_bytes - lo.peak_bytes) / (getattr(hi, vary) - getattr(lo, vary)) + for lo, hi in zip(ordered, ordered[1:], strict=False) + if getattr(hi, vary) > getattr(lo, vary) + ] + if slopes: + out[key] = slopes + return out + + +def _largest(by_group: dict[tuple, list[float]]) -> dict[int, float]: + """Per width (each group key's first element), the largest rate, and + never below zero: a negative slope between two measurements is noise (a + predict doesn't cost less for more rows).""" + per_width: dict[int, float] = {} + for (n_features, *_), rates in by_group.items(): + per_width[n_features] = max(per_width.get(n_features, 0.0), *rates) + return per_width diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py index 907ed4f..c4422d8 100644 --- a/src/tabctx/serve/factory.py +++ b/src/tabctx/serve/factory.py @@ -99,6 +99,7 @@ HostMemoryGuard, detect_host_budget_bytes, ) +from tabctx.memory.predict import PredictCostModel from tabctx.memory.requirement import ResidentRateModel, WorkingFloorModel from tabctx.multidevice import MultiDeviceEngine from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S @@ -388,6 +389,41 @@ def _requirement_calibration( ) +# Generated predict calibration (benchmarks/gen_predict_grid.py), per +# (model, kv precision, compiled) -- the settings that change a predict's +# memory, which only TabICL has: each model's PredictPoint tuple in this +# module. Measured on an A100 for the serving defaults; other configurations +# get no model (the dominance grid, else the fallback formula) rather than a +# borrowed one. Only kv-cached contexts: their predicts don't re-encode the +# training table ("repr" and "off" do). +_PREDICT_CALIBRATION_MODULE = "calibration_predict_a100" +_PREDICT_CALIBRATIONS: dict[tuple[str, str | None, bool], str] = { + ("tabicl", "int8", True): "A100_40GB_TABICL_PREDICT_POINTS", + ("tabpfn", None, False): "A100_40GB_TABPFN_PREDICT_POINTS", +} +# Backends whose kv_precision setting changes their memory profile. +_KV_PRECISION_KINDS = frozenset({"tabicl"}) + + +def _predict_cost_model( + kind: BackendKind, settings: ServeSettings +) -> PredictCostModel | None: + """The model's measured predict cost (memory/predict.py), which sizes + predict chunks, for the deployment's configuration; None without one.""" + if settings.kv_cache != "kv": + return None + key = ( + kind, + settings.kv_precision if kind in _KV_PRECISION_KINDS else None, + settings.compile_models and kind in _COMPILED_KINDS, + ) + name = _PREDICT_CALIBRATIONS.get(key) + if name is None: + return None + grid = importlib.import_module(f"tabctx.memory.{_PREDICT_CALIBRATION_MODULE}") + return PredictCostModel(getattr(grid, name)) + + def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | None: """Total memory of the given CUDA devices (the smallest, since one estimator budgets them all), or None when torch is absent or sees no @@ -496,6 +532,7 @@ def build_estimator( # Both are valid upper bounds; the tightest dominating one wins. preloaded=tuple(fit_grid) + tuple(requirements), preloaded_predict=predict_grid, + predict_cost=_predict_cost_model(kind, settings), working_floor=working_floor, resident_rate=ResidentRateModel( floor_points, width_dependent=kind in _WIDTH_DEPENDENT_RESIDENT diff --git a/tests/unit/test_chunking.py b/tests/unit/test_chunking.py index c27e470..58dc8dd 100644 --- a/tests/unit/test_chunking.py +++ b/tests/unit/test_chunking.py @@ -19,19 +19,12 @@ def test_split_rows_rejects_non_positive_chunk_size(): def test_choose_chunk_rows_smaller_for_larger_training_context(): est = PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION) - budget = est.ceiling_bytes() + # A budget the large context fits at 64 test rows (a halving of 2,048). + budget = est.estimate_bytes(90_000, 64, 200) small_context_chunk = choose_chunk_rows( est, n_train=500, n_features=10, remaining_budget_bytes=budget ) large_context_chunk = choose_chunk_rows( est, n_train=90_000, n_features=200, remaining_budget_bytes=budget ) - assert large_context_chunk < small_context_chunk - - -def test_choose_chunk_rows_respects_min_chunk_rows_floor(): - est = PowerLawMemoryEstimator(A100_40GB_TABICL_CALIBRATION) - chunk = choose_chunk_rows( - est, n_train=90_000, n_features=200, remaining_budget_bytes=1, min_chunk_rows=1 - ) - assert chunk == 1 + assert large_context_chunk == 64 < small_context_chunk diff --git a/tests/unit/test_predict_cost.py b/tests/unit/test_predict_cost.py new file mode 100644 index 0000000..c33cd65 --- /dev/null +++ b/tests/unit/test_predict_cost.py @@ -0,0 +1,212 @@ +"""PredictCostModel (memory/predict.py) and how predict chunking uses it (#47).""" + +import itertools + +import pytest + +from tabctx.chunking import choose_chunk_rows +from tabctx.memory.adaptive import AdaptiveMemoryEstimator +from tabctx.memory.calibration_data import A100_40GB_TABICL_CALIBRATION +from tabctx.memory.estimator import PowerLawMemoryEstimator +from tabctx.memory.predict import PredictCostModel, PredictPoint +from tabctx.serve.factory import ServeSettings, build_estimator + +GB = 10**9 +A100 = 42_404_806_656 + + +def _true_peak(n_train: int, n_test: int, n_features: int) -> int: + """A ground truth of the measured form: a fixed part plus, per test + row, a fixed part and a per-feature part, plus a cost per training row.""" + per_row = 200_000 + 8_000 * n_features + return 50_000_000 + n_test * per_row + 7_000 * n_train + + +def _measure(widths, n_trains=(1_000, 20_000), n_tests=(100, 1_000)): + return [ + PredictPoint(n, t, f, _true_peak(n, t, f)) + for f, n, t in itertools.product(widths, n_trains, n_tests) + ] + + +class TestPredictCostModel: + def test_bounds_every_measurement(self): + points = _measure((10, 100, 500)) + model = PredictCostModel(points, margin=1.0) + for p in points: + assert model.bytes(p.n_train, p.n_test, p.n_features) >= p.peak_bytes + + @pytest.mark.parametrize("n_features", [20, 150, 499, 1_000, 5_000]) + @pytest.mark.parametrize("n_train", [500, 50_000, 300_000]) + def test_never_underestimates_unmeasured_shapes(self, n_features, n_train): + """Between measured widths (next wider one's cost) and past the + widest (proportional scaling), for any context size and test rows.""" + model = PredictCostModel(_measure((10, 100, 500)), margin=1.0) + for n_test in (1, 100, 2_000): + truth = _true_peak(n_train, n_test, n_features) + assert model.bytes(n_train, n_test, n_features) >= truth + + def test_cost_follows_test_rows_not_the_training_table(self): + """The bug this model fixes: a predict's cost was sized with the + fit model, which grows with the training table.""" + model = PredictCostModel(_measure((10, 100, 500))) + small = model.bytes(1_000, 100, 500) + large_context = model.bytes(20_000, 100, 500) + assert large_context - small == pytest.approx( + 19_000 * model.train_row_bytes * 1.1, rel=1e-6 + ) + assert model.bytes(20_000, 1_000, 500) > 5 * model.bytes(20_000, 100, 500) + + def test_monotone_in_every_dimension(self): + model = PredictCostModel(_measure((10, 100, 500))) + for base in [(5_000, 100, 50), (5_000, 100, 700)]: + n_train, n_test, n_features = base + cost = model.bytes(*base) + assert model.bytes(n_train * 2, n_test, n_features) >= cost + assert model.bytes(n_train, n_test * 2, n_features) >= cost + assert model.bytes(n_train, n_test, n_features * 2) >= cost + + def test_a_narrower_width_never_costs_more_per_row(self): + """Non-decreasing in width even when a narrow measurement is costlier + (tabicl's fp32 path for small contexts measured so).""" + points = _measure((10, 100, 500)) + points += [ + PredictPoint(1_000, 100, 50, 10 * GB), + PredictPoint(1_000, 1_000, 50, 90 * GB), + ] + model = PredictCostModel(points) + assert model.row_bytes(100) >= model.row_bytes(50) + + def test_margin_scales_the_bound(self): + points = _measure((10, 100)) + assert PredictCostModel(points, margin=1.5).bytes( + 1_000, 10, 10 + ) == pytest.approx( + 1.5 * PredictCostModel(points, margin=1.0).bytes(1_000, 10, 10), rel=1e-6 + ) + + def test_rejects_calibration_without_two_test_sizes(self): + one_size = [PredictPoint(1_000, 100, f, GB) for f in (10, 100)] + with pytest.raises(ValueError): + PredictCostModel(one_size) + + +class TestShippedCalibration: + @pytest.mark.parametrize("kind", ["tabicl", "tabpfn"]) + def test_bounds_every_measured_predict(self, kind): + from tabctx.memory import calibration_predict_a100 as grid + + points = getattr(grid, f"A100_40GB_{kind.upper()}_PREDICT_POINTS") + model = PredictCostModel(points) + for p in points: + assert model.bytes(p.n_train, p.n_test, p.n_features) >= p.peak_bytes + + +class TestRaggedCalibration: + def test_a_width_measured_at_one_test_size_still_bounds_its_points(self): + """No per-row slope at that width (one size), but its point still + lies under the bound: the width takes the next wider row cost, and + what's left becomes its fixed call cost.""" + points = _measure((10, 500)) + [PredictPoint(1_000, 100, 100, 5 * GB)] + model = PredictCostModel(points, margin=1.0) + for p in points: + assert model.bytes(p.n_train, p.n_test, p.n_features) >= p.peak_bytes + + def test_negative_slopes_are_noise_not_savings(self): + points = _measure((10, 100)) + [ + PredictPoint(1_000, 100, 200, 3 * GB), + PredictPoint(1_000, 1_000, 200, 2 * GB), + ] + model = PredictCostModel(points, margin=1.0) + assert model.row_bytes(200) >= model.row_bytes(100) >= 0 + for p in points: + assert model.bytes(p.n_train, p.n_test, p.n_features) >= p.peak_bytes + + def test_narrower_than_any_measurement_takes_the_narrowest_cost(self): + model = PredictCostModel(_measure((10, 100))) + assert model.row_bytes(1) == model.row_bytes(10) + assert model.call_bytes(1) == model.call_bytes(10) + + def test_steepest_consecutive_slope_bounds_a_convex_curve(self): + """Three test sizes on a convex curve: the end-to-end slope would + undershoot the middle-to-end one.""" + peaks = {100: 1 * GB, 1_000: 2 * GB, 2_000: 6 * GB} + points = [PredictPoint(1_000, t, 50, b) for t, b in peaks.items()] + points += [PredictPoint(1_000, t, 10, b // 5) for t, b in peaks.items()] + model = PredictCostModel(points, margin=1.0) + assert model.row_bytes(50) == pytest.approx(4 * GB / 1_000) + + +class TestChunking: + class _Linear: + """estimate = fixed + per_row x rows, for exercising the sizing rule.""" + + def __init__(self, fixed: int, per_row: int) -> None: + self.fixed, self.per_row = fixed, per_row + + def estimate_bytes(self, n_train, n_test, n_features): + return self.fixed + self.per_row * n_test + + def test_largest_halving_that_fits(self): + est = self._Linear(fixed=0, per_row=1_000) + assert choose_chunk_rows(est, 1, 1, remaining_budget_bytes=600_000) == 512 + assert choose_chunk_rows(est, 1, 1, remaining_budget_bytes=10**12) == 2_048 + + def test_a_fixed_cost_over_budget_is_not_split_away(self): + """#47: splitting can't shrink a call's fixed cost, so only the rows + above it are sized against the budget, not split to one row each.""" + est = self._Linear(fixed=10 * GB, per_row=1_000_000) + # 1 GB of budget above the smallest call fits ~1,000 more rows. + assert choose_chunk_rows(est, 1, 1, remaining_budget_bytes=GB) == 512 + + def test_no_headroom_means_the_smallest_chunk(self): + """A full device (say, beside a fit at its peak): nothing beyond the + smallest call is added, slow but safe.""" + est = self._Linear(fixed=10 * GB, per_row=1_000) + assert choose_chunk_rows(est, 1, 1, remaining_budget_bytes=0) == 1 + + def test_rejects_a_zero_floor(self): + with pytest.raises(ValueError): + choose_chunk_rows(self._Linear(0, 1), 1, 1, 10, min_chunk_rows=0) + + def test_predict_queries_use_the_cost_model_and_fits_do_not(self): + model = PredictCostModel(_measure((10, 100, 500))) + fallback = PowerLawMemoryEstimator( + A100_40GB_TABICL_CALIBRATION, gpu_capacity_bytes=A100 + ) + with_model = AdaptiveMemoryEstimator(fallback=fallback, predict_cost=model) + without = AdaptiveMemoryEstimator(fallback=fallback) + assert with_model.estimate_bytes(3_500, 100, 500) == model.bytes( + 3_500, 100, 500 + ) + assert with_model.estimate_bytes(3_500, 0, 500) == without.estimate_bytes( + 3_500, 0, 500 + ) + + @pytest.mark.parametrize("kind", ["tabicl", "tabpfn"]) + @pytest.mark.parametrize("n_train", [3_500, 10_000, 100_000]) + def test_wide_context_predict_fits_in_one_call(self, kind, n_train): + """Regression (#47): a 100-row predict against a 3,500 x 500 context + ran as 100 one-row calls (14 s instead of 0.1 s on an A100), because + its memory was estimated with the fit model, which charges for the + whole training table and exceeded the device's headroom.""" + settings = ServeSettings(backends=(kind,)) + estimator = build_estimator(settings, kind, gpu_capacity_bytes=A100) + headroom = estimator.admission_headroom_bytes(0) + assert estimator.estimate_bytes(n_train, 100, 500) <= headroom + assert choose_chunk_rows(estimator, n_train, 500, headroom) >= 100 + # What it was: the same query on the fit model alone. + fit_model = build_estimator( + ServeSettings(backends=(kind,), kv_cache="repr"), kind, A100 + ) + assert fit_model.estimate_bytes(n_train, 100, 500) > headroom + + def test_factory_installs_the_model_for_its_calibrated_configuration(self): + def cost(**kw): + settings = ServeSettings(backends=("tabicl",), **kw) + return build_estimator(settings, "tabicl", A100)._predict_cost # noqa: SLF001 + + assert cost() is not None # serving defaults: kv, int8, compiled + assert cost(kv_cache="repr") is None + assert cost(kv_precision="native") is None + assert cost(compile_models=False) is None