From d35a3e31c08d60f68468d9b014c047a28f6d3ba1 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 17:19:55 -0400 Subject: [PATCH 1/3] fix(compile): apply dynamo settings per call -- serving threads saw a limit of 8 (#47) install_compiled_blocks set torch._dynamo.config.recompile_limit and suppress_errors once, on the installing thread. Dynamo's config overrides are ContextVars, so fits and predicts on engine/worker threads and asyncio.to_thread saw the defaults: a recompile limit of 8, hit at the first graph past warmup's 54 (as on prod's four A40 workers), after which dynamo stops compiling the blocks for the process -- every later fit ran eagerly -- and compile errors raised instead of falling back. with_compile_settings wraps the compiled forward in config.patch (one shared ConfigPatch; its saved state is itself a ContextVar stack, so it is safe across threads and nested calls; ~4 us per call). The settings no longer leak to other torch.compile users in the process. A100, every call on a separate thread (bench_stalls --call-thread): 3500x500 fit 2.31 -> 1.61 s, 3500x1000 5.49 -> 3.50 s; graphs past warmup now compile once instead of never. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 3 +- ROADMAP.md | 4 +- .../compile-call-thread-a100-after.jsonl | 22 +++++ .../compile-call-thread-a100-before.jsonl | 22 +++++ benchmarks/bench_stalls.py | 87 ++++++++++++++----- src/tabctx/backends/tabicl_compile.py | 42 ++++++--- tests/unit/test_tabicl_compile.py | 52 +++++++++++ 7 files changed, 195 insertions(+), 37 deletions(-) create mode 100644 benchmarks/baselines/compile-call-thread-a100-after.jsonl create mode 100644 benchmarks/baselines/compile-call-thread-a100-before.jsonl create mode 100644 tests/unit/test_tabicl_compile.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f1e336b..35320d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,8 +5,9 @@ - **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. +- **Compiled TabICL fell back to eager in serving** (#47). `install_compiled_blocks` set dynamo's `recompile_limit` and `suppress_errors` once, but dynamo's config overrides are per thread (a `ContextVar`), and fits and predicts run on other threads (engine and worker threads, `asyncio.to_thread`), which saw dynamo's defaults: a limit of 8. Prod's four A40 workers each hit `recompile_limit (8)` at the first graph past warmup's 54, after which dynamo stops compiling the blocks for the life of the process, so every later fit ran eagerly, even shapes warmup had compiled; and a compile error would have failed a tenant's fit instead of falling back to eager. The settings are now applied around each compiled call (`tabicl_compile.with_compile_settings`, ~4 us), so they hold on any thread and no longer change other `torch.compile` users in the process. A100, calls on a separate thread (`benchmarks/bench_stalls.py --call-thread`; `baselines/compile-call-thread-a100-{before,after}.jsonl`): 3,500 x 500 fits 2.31 -> 1.61 s, 3,500 x 1,000 5.49 -> 3.50 s. Shapes warmup doesn't cover now compile once per worker as documented (a first ~30k-row fit takes 5-25 s) instead of running eagerly forever. ### 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). +- `benchmarks/bench_stalls.py` (fixed-cost stalls across table sizes, with compiled-graph counts per call; `--tasks`, `--call-thread` to call from another thread as serving does), `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] - 2026-10-18 ### Added diff --git a/ROADMAP.md b/ROADMAP.md index fbb4f5b..f370936 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -123,7 +123,9 @@ no FP8). Ray Serve, 6 users on 4x A100: 5.78 -> 6.99 predicts/s. Dynamic shapes needed tabicl's ssmax log(n) computed symbolically (math.log on the row count recompiled every new table). ~50 graphs (fp16 path - only) compile in warmup (~130 s per worker). CUDA graphs aren't worth it: the GPU half + only) compile in warmup (~130 s per worker); dynamo's settings are + applied per call, since its config overrides are per thread and + serving calls from other threads (#47). CUDA graphs aren't worth it: the GPU half is GPU-bound, not launch-bound. TabPFN-3.5's own compile path (`enable_torch_compile`) measured 1.3-1.5x faster fits, ~30% faster predicts and ~22% lower peak, but it still recompiles for many new diff --git a/benchmarks/baselines/compile-call-thread-a100-after.jsonl b/benchmarks/baselines/compile-call-thread-a100-after.jsonl new file mode 100644 index 0000000..75004a1 --- /dev/null +++ b/benchmarks/baselines/compile-call-thread-a100-after.jsonl @@ -0,0 +1,22 @@ +{"event": "built", "models": ["tabicl-v2"], "device": "cuda", "build_s": 109.1, "graphs_after_warmup": 54} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x1", "fit_s": 0.164, "fit_graphs": 0, "predicts": [{"s": 0.055, "graphs": 0}, {"s": 0.064, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x1", "fit_s": 0.182, "fit_graphs": 0, "predicts": [{"s": 0.053, "graphs": 0}, {"s": 0.052, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x1", "fit_s": 0.289, "fit_graphs": 0, "predicts": [{"s": 0.057, "graphs": 0}, {"s": 0.062, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x1", "fit_s": 6.466, "fit_graphs": 2, "predicts": [{"s": 0.06, "graphs": 0}, {"s": 0.056, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x10", "fit_s": 0.27, "fit_graphs": 0, "predicts": [{"s": 0.056, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x10", "fit_s": 0.25, "fit_graphs": 0, "predicts": [{"s": 0.058, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x10", "fit_s": 0.473, "fit_graphs": 0, "predicts": [{"s": 0.056, "graphs": 0}, {"s": 0.054, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x10", "fit_s": 1.831, "fit_graphs": 0, "predicts": [{"s": 0.074, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x100", "fit_s": 0.285, "fit_graphs": 0, "predicts": [{"s": 0.063, "graphs": 0}, {"s": 0.064, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x100", "fit_s": 0.436, "fit_graphs": 0, "predicts": [{"s": 0.078, "graphs": 0}, {"s": 0.061, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x100", "fit_s": 1.021, "fit_graphs": 0, "predicts": [{"s": 0.066, "graphs": 0}, {"s": 0.061, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x100", "fit_s": 14.638, "fit_graphs": 3, "predicts": [{"s": 0.073, "graphs": 0}, {"s": 0.063, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x500", "fit_s": 0.706, "fit_graphs": 0, "predicts": [{"s": 0.105, "graphs": 0}, {"s": 0.104, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x500", "fit_s": 1.613, "fit_graphs": 0, "predicts": [{"s": 0.105, "graphs": 0}, {"s": 0.107, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x500", "fit_s": 7.416, "fit_graphs": 1, "predicts": [{"s": 0.113, "graphs": 0}, {"s": 0.109, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x500", "fit_s": 209.757, "fit_graphs": 0, "predicts": [{"s": 0.124, "graphs": 0}, {"s": 0.108, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x1000", "fit_s": 1.55, "fit_graphs": 0, "predicts": [{"s": 0.172, "graphs": 0}, {"s": 0.173, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x1000", "fit_s": 3.497, "fit_graphs": 0, "predicts": [{"s": 0.168, "graphs": 0}, {"s": 0.17, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x1000", "fit_s": 133.227, "fit_graphs": 0, "predicts": [{"s": 0.183, "graphs": 0}, {"s": 0.176, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x1000", "fit_s": 578.252, "fit_graphs": 0, "predicts": [{"s": 0.175, "graphs": 0}, {"s": 0.172, "graphs": 0}]} +{"event": "done", "graphs_total": 60} diff --git a/benchmarks/baselines/compile-call-thread-a100-before.jsonl b/benchmarks/baselines/compile-call-thread-a100-before.jsonl new file mode 100644 index 0000000..3311699 --- /dev/null +++ b/benchmarks/baselines/compile-call-thread-a100-before.jsonl @@ -0,0 +1,22 @@ +{"event": "built", "models": ["tabicl-v2"], "device": "cuda", "build_s": 107.3, "graphs_after_warmup": 54} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x1", "fit_s": 0.168, "fit_graphs": 0, "predicts": [{"s": 0.057, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x1", "fit_s": 0.185, "fit_graphs": 0, "predicts": [{"s": 0.054, "graphs": 0}, {"s": 0.053, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x1", "fit_s": 0.265, "fit_graphs": 0, "predicts": [{"s": 0.056, "graphs": 0}, {"s": 0.058, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x1", "fit_s": 0.95, "fit_graphs": 0, "predicts": [{"s": 0.061, "graphs": 0}, {"s": 0.052, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x10", "fit_s": 0.199, "fit_graphs": 0, "predicts": [{"s": 0.06, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x10", "fit_s": 0.242, "fit_graphs": 0, "predicts": [{"s": 0.065, "graphs": 0}, {"s": 0.055, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x10", "fit_s": 0.462, "fit_graphs": 0, "predicts": [{"s": 0.057, "graphs": 0}, {"s": 0.053, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x10", "fit_s": 1.965, "fit_graphs": 0, "predicts": [{"s": 0.065, "graphs": 0}, {"s": 0.054, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x100", "fit_s": 0.269, "fit_graphs": 0, "predicts": [{"s": 0.067, "graphs": 0}, {"s": 0.065, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x100", "fit_s": 0.435, "fit_graphs": 0, "predicts": [{"s": 0.066, "graphs": 0}, {"s": 0.067, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x100", "fit_s": 1.007, "fit_graphs": 0, "predicts": [{"s": 0.065, "graphs": 0}, {"s": 0.062, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x100", "fit_s": 4.557, "fit_graphs": 0, "predicts": [{"s": 0.072, "graphs": 0}, {"s": 0.061, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x500", "fit_s": 0.704, "fit_graphs": 0, "predicts": [{"s": 0.102, "graphs": 0}, {"s": 0.104, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x500", "fit_s": 2.307, "fit_graphs": 0, "predicts": [{"s": 0.103, "graphs": 0}, {"s": 0.104, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x500", "fit_s": 5.864, "fit_graphs": 0, "predicts": [{"s": 0.109, "graphs": 0}, {"s": 0.103, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x500", "fit_s": 207.812, "fit_graphs": 0, "predicts": [{"s": 0.125, "graphs": 0}, {"s": 0.11, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "1100x1000", "fit_s": 1.546, "fit_graphs": 0, "predicts": [{"s": 0.169, "graphs": 0}, {"s": 0.17, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "3500x1000", "fit_s": 5.49, "fit_graphs": 0, "predicts": [{"s": 0.168, "graphs": 0}, {"s": 0.167, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "10000x1000", "fit_s": 133.124, "fit_graphs": 0, "predicts": [{"s": 0.183, "graphs": 0}, {"s": 0.174, "graphs": 0}]} +{"model": "tabicl-v2", "task": "classification", "shape": "30000x1000", "fit_s": 569.824, "fit_graphs": 0, "predicts": [{"s": 0.174, "graphs": 0}, {"s": 0.175, "graphs": 0}]} +{"event": "done", "graphs_total": 54} diff --git a/benchmarks/bench_stalls.py b/benchmarks/bench_stalls.py index b8d3384..09b1fdd 100644 --- a/benchmarks/bench_stalls.py +++ b/benchmarks/bench_stalls.py @@ -41,6 +41,35 @@ def timed(fn, *args, **kwargs): return round(time.perf_counter() - t0, 3), compiled_graphs() - graphs +class _OnThread: + """The engine, with fit/predict/evict made on a pool thread.""" + + def __init__(self, engine, pool) -> None: + self._engine, self._pool = engine, pool + + def __getattr__(self, name): + method = getattr(self._engine, name) + return lambda *a, **kw: self._pool.submit(method, *a, **kw).result() + + +def stalls(engine, backend, task, rows, features, args, rng) -> dict: + """One fit of a rows x features table, then repeated predicts against it.""" + X = rng.normal(size=(rows, features)).astype(np.float32) + signal = X[:, 0] + 0.5 * X[:, 1 % features] + y = (signal > 0).astype(int) if task == "classification" else signal + X_test = rng.normal(size=(args.test_rows, features)).astype(np.float32) + dataset_id = f"stall-{backend}-{task}-{features}x{rows}" + fit_s, fit_graphs = timed( + engine.fit, X, y, task=task, 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) + return {"fit_s": fit_s, "fit_graphs": fit_graphs, "predicts": predicts} + + def main() -> int: p = argparse.ArgumentParser() p.add_argument("--models", nargs="+", default=["tabicl", "tabpfn"]) @@ -48,9 +77,22 @@ def main() -> int: 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( + "--tasks", + nargs="+", + default=["classification"], + choices=["classification", "regression"], + ) + p.add_argument( + "--call-thread", + action="store_true", + help="make every call on another thread than the one that built the " + "engine, as serving does (thread-local settings show up only then)", + ) p.add_argument("--out") args = p.parse_args() + from tabctx.errors import AdmissionRejected from tabctx.serve.factory import ServeSettings, build_engine t0 = time.perf_counter() @@ -68,33 +110,32 @@ def main() -> int: ), flush=True, ) + if args.call_thread: + from concurrent.futures import ThreadPoolExecutor + + pool = ThreadPoolExecutor(1, thread_name_prefix="call") + engine = _OnThread(engine, pool) 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) + for task in args.tasks: + for features in args.features: + for rows in args.rows: + shape = { + "model": backend, + "task": task, + "shape": f"{rows}x{features}", + } + try: + record = { + **shape, + **stalls(engine, backend, task, rows, features, args, rng), + } + except AdmissionRejected as e: + record = {**shape, "skipped": str(e)[:160]} + 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) diff --git a/src/tabctx/backends/tabicl_compile.py b/src/tabctx/backends/tabicl_compile.py index 5ea43b9..d8eb506 100644 --- a/src/tabctx/backends/tabicl_compile.py +++ b/src/tabctx/backends/tabicl_compile.py @@ -27,12 +27,24 @@ path (tables of 1,024+ rows) is compiled: smaller tables run in fp32 in ~0.1 s and would double the graphs. ~50 graphs compile in the backend's warmup at replica start (~130 s per worker on an A100, workers in -parallel; Inductor's cache barely shortens it, tracing dominates). Calls warmup doesn't cover still compile inside a request, -once per worker: very large tables reach 1-3 more graphs (tabicl batches -them differently), ~10 s each, and other unusual shapes may too. +parallel; Inductor's cache barely shortens it, tracing dominates). Calls +warmup doesn't cover still compile inside a request, once per worker: +tables of ~30k rows hand the blocks differently strided inputs (1-7 more +graphs, 5-25 s on an A100 for the first such fit, 10,000 x 500 one more), +and other unusual shapes may too. Failure. Compilation errors fall back to eager (dynamo's suppress_errors), so a missing compiler toolchain costs speed, not availability. + +Settings. Both dynamo settings these blocks need (suppress_errors, and a +recompile limit above the ~50 graphs) are applied around each compiled +call, not set at install: dynamo's config overrides are per thread (a +ContextVar), and calls run on other threads than the one that installs +the blocks -- engine and worker threads, asyncio.to_thread -- which saw +dynamo's defaults: a limit of 8, so serving workers stopped compiling and +ran shapes warmup missed eagerly, and compile errors raised in tenants' +fits instead of falling back. Scoped to these blocks, they also don't +change other torch.compile users in the process. """ from __future__ import annotations @@ -63,16 +75,10 @@ def install_compiled_blocks(mode: str = "default") -> bool: forward = MultiheadAttentionBlock.forward if getattr(forward, "_tabctx_compiled", False): return True - # Process-wide settings: they apply to any other torch.compile user in - # this process too (suppress_errors makes every compile failure fall - # back to eager; warn_if_not_compiled reports it for these blocks). - config: Any = torch._dynamo.config - devices = max(1, torch.cuda.device_count()) - config.recompile_limit = max( - config.recompile_limit, _RECOMPILE_LIMIT_PER_DEVICE * devices + compiled = with_compile_settings( + torch.compile(forward, dynamic=True, mode=mode), + devices=max(1, torch.cuda.device_count()), ) - config.suppress_errors = True - compiled = torch.compile(forward, dynamic=True, mode=mode) # Only under autocast: tabicl turns fp16 autocast on for tables of # 1,024+ rows, where the fused kernels pay off; smaller tables (fp32, @@ -89,6 +95,18 @@ def dispatch(self: Any, *args: Any, **kwargs: Any) -> Any: return True +def with_compile_settings(fn: Any, devices: int) -> Any: + """fn, run under the dynamo settings these blocks need, on whichever + thread or context calls it (see the module docstring): compile errors + fall back to eager (warn_if_not_compiled reports it), and the recompile + limit covers every graph on every device.""" + import torch._dynamo + + config: Any = torch._dynamo.config + limit = max(config.recompile_limit, _RECOMPILE_LIMIT_PER_DEVICE * devices) + return config.patch(recompile_limit=limit, suppress_errors=True)(fn) + + def warn_if_not_compiled() -> bool: """After warmup: warn if no graph was compiled (compile errors fall back to eager silently, losing the speedup) or a block hit dynamo's limit. diff --git a/tests/unit/test_tabicl_compile.py b/tests/unit/test_tabicl_compile.py new file mode 100644 index 0000000..ab74521 --- /dev/null +++ b/tests/unit/test_tabicl_compile.py @@ -0,0 +1,52 @@ +"""Dynamo settings for tabicl's compiled blocks (backends/tabicl_compile.py).""" + +import asyncio +import threading + +import pytest + +from tabctx.backends.tabicl_compile import with_compile_settings + +config = pytest.importorskip("torch._dynamo").config + + +def _settings() -> tuple[int, bool]: + return config.recompile_limit, config.suppress_errors + + +def test_settings_apply_on_any_thread_or_context(): + """Dynamo's config overrides are per thread (a ContextVar): set once at + install, serving threads saw a recompile limit of 8 and compile errors + that raised. The settings now travel with each call.""" + wrapped = with_compile_settings(_settings, devices=2) + seen = {} + thread = threading.Thread(target=lambda: seen.update(thread=wrapped())) + thread.start() + thread.join() + seen["to_thread"] = asyncio.run(asyncio.to_thread(wrapped)) + assert seen == {"thread": (512, True), "to_thread": (512, True)} + + +def test_settings_stay_scoped_to_the_call(): + default = _settings() + with_compile_settings(_settings, devices=1)() + assert _settings() == default + + +def test_nested_and_concurrent_calls_restore_their_own_settings(): + """One wrapper serves every block call: nested calls and calls on other + threads each restore what they found.""" + default = _settings() + inner = with_compile_settings(_settings, devices=1) + outer = with_compile_settings(lambda: (inner(), _settings()), devices=1) + results: list = [] + threads = [ + threading.Thread(target=lambda: results.append(outer())) for _ in range(8) + ] + for t in threads: + t.start() + for t in threads: + t.join() + assert results == [((256, True), (256, True))] * 8 + assert outer() == ((256, True), (256, True)) + assert _settings() == default From 4ba58be4631f765efd5a41fe0b7fa7f1c82680dc Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 17:27:02 -0400 Subject: [PATCH 2/3] fix(compile): review fixes -- raise accumulated_recompile_limit, set at install too - Dynamo also caps a code object's graphs at accumulated_recompile_limit (256), checked first, so the per-device recompile_limit above 256 had no effect; both are raised now (compile_settings). - with_compile_settings also sets them for the installing thread: on a torch whose config is process-wide, every call's shared patch then saves and restores the same values, so concurrent calls can't restore the defaults under one another. - Tests: an end-to-end CPU regression (eager backend: 10 graphs on one thread, then a new shape on another must compile -- fails with install-time-only settings), restore-on-raise, both limits. - CHANGELOG: past the limit dynamo compiles no new graph (warmed shapes keep theirs); the 3,500-row gains came from graphs built by the sweep's earlier 30k-row fits. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + src/tabctx/backends/tabicl_compile.py | 59 ++++++++++----- tests/unit/test_tabicl_compile.py | 105 +++++++++++++++++++++----- 3 files changed, 127 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 35320d3..ba4f1e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - **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. - **Compiled TabICL fell back to eager in serving** (#47). `install_compiled_blocks` set dynamo's `recompile_limit` and `suppress_errors` once, but dynamo's config overrides are per thread (a `ContextVar`), and fits and predicts run on other threads (engine and worker threads, `asyncio.to_thread`), which saw dynamo's defaults: a limit of 8. Prod's four A40 workers each hit `recompile_limit (8)` at the first graph past warmup's 54, after which dynamo stops compiling the blocks for the life of the process, so every later fit ran eagerly, even shapes warmup had compiled; and a compile error would have failed a tenant's fit instead of falling back to eager. The settings are now applied around each compiled call (`tabicl_compile.with_compile_settings`, ~4 us), so they hold on any thread and no longer change other `torch.compile` users in the process. A100, calls on a separate thread (`benchmarks/bench_stalls.py --call-thread`; `baselines/compile-call-thread-a100-{before,after}.jsonl`): 3,500 x 500 fits 2.31 -> 1.61 s, 3,500 x 1,000 5.49 -> 3.50 s. Shapes warmup doesn't cover now compile once per worker as documented (a first ~30k-row fit takes 5-25 s) instead of running eagerly forever. +- **Compiled TabICL stopped compiling in serving** (#47). `install_compiled_blocks` set dynamo's `recompile_limit` and `suppress_errors` once, but dynamo's config overrides are per thread (a `ContextVar`, torch 2.13+), and fits and predicts run on other threads (engine and worker threads, `asyncio.to_thread`), which saw dynamo's defaults: a limit of 8. Prod's four A40 workers each hit `recompile_limit (8)` at the first graph past warmup's 54, after which dynamo compiles no new graph of the blocks for the life of the process: every shape warmup missed ran eagerly from then on (warmed shapes kept their graphs), and a compile error would have failed a tenant's fit instead of falling back to eager. The settings are now applied around each compiled call too (`tabicl_compile.with_compile_settings`, ~4 us), so they hold on any thread, and `accumulated_recompile_limit` (256, which capped the per-device limit) is raised with them. A100, calls on a separate thread (`benchmarks/bench_stalls.py --call-thread`; `baselines/compile-call-thread-a100-{before,after}.jsonl`): shapes outside warmup's families run compiled once their graph exists -- in that sweep, 3,500 x 500 fits 2.31 -> 1.61 s and 3,500 x 1,000 5.49 -> 3.50 s, on graphs first built by its earlier 30k-row fits. Such a graph now compiles once per worker (a first ~30k-row fit takes 5-25 s) instead of never. ### Added - `benchmarks/bench_stalls.py` (fixed-cost stalls across table sizes, with compiled-graph counts per call; `--tasks`, `--call-thread` to call from another thread as serving does), `benchmarks/profile_predict.py` (cProfile of one predict), `benchmarks/calibrate_predict.py` and `benchmarks/gen_predict_grid.py` (predict calibration), `benchmarks/sweep_gateway.py` (the table-size sweep through the Vector gateway). diff --git a/src/tabctx/backends/tabicl_compile.py b/src/tabctx/backends/tabicl_compile.py index d8eb506..09221e6 100644 --- a/src/tabctx/backends/tabicl_compile.py +++ b/src/tabctx/backends/tabicl_compile.py @@ -36,15 +36,17 @@ Failure. Compilation errors fall back to eager (dynamo's suppress_errors), so a missing compiler toolchain costs speed, not availability. -Settings. Both dynamo settings these blocks need (suppress_errors, and a -recompile limit above the ~50 graphs) are applied around each compiled -call, not set at install: dynamo's config overrides are per thread (a -ContextVar), and calls run on other threads than the one that installs -the blocks -- engine and worker threads, asyncio.to_thread -- which saw -dynamo's defaults: a limit of 8, so serving workers stopped compiling and -ran shapes warmup missed eagerly, and compile errors raised in tenants' -fits instead of falling back. Scoped to these blocks, they also don't -change other torch.compile users in the process. +Settings. The dynamo settings these blocks need (suppress_errors, and +recompile limits above the ~50 graphs) are applied around each compiled +call as well as at install: dynamo's config overrides are per thread (a +ContextVar, as of torch 2.13), and calls run on other threads than the one +that installs the blocks -- engine and worker threads, asyncio.to_thread +-- which saw dynamo's defaults: a limit of 8. Past it dynamo compiles no +new graph of the block for the life of the process (existing ones still +run), so serving workers ran every shape warmup missed eagerly, and compile +errors raised in tenants' fits instead of falling back. Set at install too, +so a torch whose config is process-wide gets them everywhere, and every +call's patch saves and restores the same values there. """ from __future__ import annotations @@ -57,7 +59,9 @@ # Graphs per block, per device (stage x fp16/fp32 x cached/uncached x ..., # ~50 measured; guards include the device, so an in-process engine over # several GPUs needs that many per GPU): well above dynamo's default of 8, -# past which it silently runs the rest eagerly. +# past which it silently runs new shapes eagerly. Dynamo also caps every +# code object's graphs at accumulated_recompile_limit (256 by default), so +# that is raised to match. _RECOMPILE_LIMIT_PER_DEVICE = 256 @@ -96,21 +100,40 @@ def dispatch(self: Any, *args: Any, **kwargs: Any) -> Any: def with_compile_settings(fn: Any, devices: int) -> Any: - """fn, run under the dynamo settings these blocks need, on whichever - thread or context calls it (see the module docstring): compile errors - fall back to eager (warn_if_not_compiled reports it), and the recompile - limit covers every graph on every device.""" + """fn, run under compile_settings(devices) on whichever thread or + context calls it; also set for the installing thread (or the whole + process, where dynamo's config isn't per thread). + """ import torch._dynamo config: Any = torch._dynamo.config - limit = max(config.recompile_limit, _RECOMPILE_LIMIT_PER_DEVICE * devices) - return config.patch(recompile_limit=limit, suppress_errors=True)(fn) + settings = compile_settings(devices) + for name, value in settings.items(): + setattr(config, name, value) + return config.patch(settings)(fn) + + +def compile_settings(devices: int) -> dict[str, Any]: + """The dynamo settings the compiled blocks run under (see the module + docstring): compile errors fall back to eager (warn_if_not_compiled + reports it), and both recompile limits cover every graph on every + device.""" + import torch._dynamo + + config: Any = torch._dynamo.config + limit = _RECOMPILE_LIMIT_PER_DEVICE * devices + return { + "recompile_limit": max(config.recompile_limit, limit), + "accumulated_recompile_limit": max(config.accumulated_recompile_limit, limit), + "suppress_errors": True, + } def warn_if_not_compiled() -> bool: """After warmup: warn if no graph was compiled (compile errors fall back - to eager silently, losing the speedup) or a block hit dynamo's limit. - Returns whether compilation looks healthy.""" + to eager silently, losing the speedup). Returns whether compilation + looks healthy. A recompile limit hit later shows only in the log + (dynamo's 'hit config.recompile_limit' warning).""" from torch._dynamo.utils import counters graphs = counters["stats"]["unique_graphs"] diff --git a/tests/unit/test_tabicl_compile.py b/tests/unit/test_tabicl_compile.py index ab74521..da3f660 100644 --- a/tests/unit/test_tabicl_compile.py +++ b/tests/unit/test_tabicl_compile.py @@ -1,52 +1,117 @@ """Dynamo settings for tabicl's compiled blocks (backends/tabicl_compile.py).""" import asyncio +import contextvars import threading import pytest -from tabctx.backends.tabicl_compile import with_compile_settings +from tabctx.backends.tabicl_compile import compile_settings, with_compile_settings -config = pytest.importorskip("torch._dynamo").config +torch = pytest.importorskip("torch") +dynamo = pytest.importorskip("torch._dynamo") +config = dynamo.config -def _settings() -> tuple[int, bool]: - return config.recompile_limit, config.suppress_errors +def _settings() -> tuple[int, int, bool]: + return ( + config.recompile_limit, + config.accumulated_recompile_limit, + config.suppress_errors, + ) + + +def _fresh(fn): + """Run fn in an empty context, as a new thread would: the module-level + settings with_compile_settings applies for its caller don't leak into + what each test checks.""" + return contextvars.Context().run(fn) + + +def test_settings_raise_both_limits_for_every_device(): + settings = _fresh(lambda: compile_settings(devices=2)) + assert settings == { + "recompile_limit": 512, + "accumulated_recompile_limit": 512, + "suppress_errors": True, + } def test_settings_apply_on_any_thread_or_context(): """Dynamo's config overrides are per thread (a ContextVar): set once at install, serving threads saw a recompile limit of 8 and compile errors that raised. The settings now travel with each call.""" - wrapped = with_compile_settings(_settings, devices=2) + wrapped = _fresh(lambda: with_compile_settings(_settings, devices=2)) seen = {} thread = threading.Thread(target=lambda: seen.update(thread=wrapped())) thread.start() thread.join() seen["to_thread"] = asyncio.run(asyncio.to_thread(wrapped)) - assert seen == {"thread": (512, True), "to_thread": (512, True)} + assert seen == {"thread": (512, 512, True), "to_thread": (512, 512, True)} -def test_settings_stay_scoped_to_the_call(): - default = _settings() - with_compile_settings(_settings, devices=1)() - assert _settings() == default +def test_settings_are_restored_after_a_call_even_when_it_raises(): + def default_then_call(wrapped) -> tuple: + before = _settings() + with pytest.raises(RuntimeError): + wrapped() + return before, _settings() + + def boom(): + raise RuntimeError + + def run(): + wrapped = with_compile_settings(boom, devices=1) + return _settings(), contextvars.Context().run(default_then_call, wrapped) + + installer, (before, after) = _fresh(run) + assert installer == (256, 256, True) # set for the installing thread + assert before == after # another context: untouched by the failed call def test_nested_and_concurrent_calls_restore_their_own_settings(): """One wrapper serves every block call: nested calls and calls on other - threads each restore what they found.""" - default = _settings() - inner = with_compile_settings(_settings, devices=1) - outer = with_compile_settings(lambda: (inner(), _settings()), devices=1) + threads each see the settings and leave their context as they found it.""" + inner = _fresh(lambda: with_compile_settings(_settings, devices=1)) + outer = _fresh( + lambda: with_compile_settings(lambda: (inner(), _settings()), devices=1) + ) results: list = [] - threads = [ - threading.Thread(target=lambda: results.append(outer())) for _ in range(8) - ] + + def call() -> None: + before = _settings() + results.append((outer(), _settings() == before)) + + threads = [threading.Thread(target=call) for _ in range(8)] for t in threads: t.start() for t in threads: t.join() - assert results == [((256, True), (256, True))] * 8 - assert outer() == ((256, True), (256, True)) - assert _settings() == default + assert results == [(((256, 256, True), (256, 256, True)), True)] * 8 + + +def test_new_graphs_compile_on_a_serving_thread_past_dynamo_default_limit(): + """The bug end to end, on CPU: more graphs than dynamo's default limit + of 8 compiled on one thread, then a new shape on another thread. It + compiled nothing there before (limit 8, hit for the life of the + process); with the settings carried by the call it compiles.""" + from torch._dynamo.utils import counters + + def f(x): + return x * 2 + 1 + + dynamo.reset() + compiled = _fresh( + lambda: with_compile_settings( + torch.compile(f, backend="eager", dynamic=False), devices=1 + ) + ) + for n in range(1, 11): # 10 graphs: past the default limit of 8 + compiled(torch.ones(n)) + graphs = counters["stats"]["unique_graphs"] + + thread = threading.Thread(target=lambda: compiled(torch.ones(50))) + thread.start() + thread.join() + assert counters["stats"]["unique_graphs"] == graphs + 1 + dynamo.reset() From c3c4c11eee7e877ffdd179e3ddca116e9cb0b0a4 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 17:27:30 -0400 Subject: [PATCH 3/3] docs(changelog): drop the superseded compile entry from the rebase Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba4f1e1..91b00ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,6 @@ - **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. -- **Compiled TabICL fell back to eager in serving** (#47). `install_compiled_blocks` set dynamo's `recompile_limit` and `suppress_errors` once, but dynamo's config overrides are per thread (a `ContextVar`), and fits and predicts run on other threads (engine and worker threads, `asyncio.to_thread`), which saw dynamo's defaults: a limit of 8. Prod's four A40 workers each hit `recompile_limit (8)` at the first graph past warmup's 54, after which dynamo stops compiling the blocks for the life of the process, so every later fit ran eagerly, even shapes warmup had compiled; and a compile error would have failed a tenant's fit instead of falling back to eager. The settings are now applied around each compiled call (`tabicl_compile.with_compile_settings`, ~4 us), so they hold on any thread and no longer change other `torch.compile` users in the process. A100, calls on a separate thread (`benchmarks/bench_stalls.py --call-thread`; `baselines/compile-call-thread-a100-{before,after}.jsonl`): 3,500 x 500 fits 2.31 -> 1.61 s, 3,500 x 1,000 5.49 -> 3.50 s. Shapes warmup doesn't cover now compile once per worker as documented (a first ~30k-row fit takes 5-25 s) instead of running eagerly forever. - **Compiled TabICL stopped compiling in serving** (#47). `install_compiled_blocks` set dynamo's `recompile_limit` and `suppress_errors` once, but dynamo's config overrides are per thread (a `ContextVar`, torch 2.13+), and fits and predicts run on other threads (engine and worker threads, `asyncio.to_thread`), which saw dynamo's defaults: a limit of 8. Prod's four A40 workers each hit `recompile_limit (8)` at the first graph past warmup's 54, after which dynamo compiles no new graph of the blocks for the life of the process: every shape warmup missed ran eagerly from then on (warmed shapes kept their graphs), and a compile error would have failed a tenant's fit instead of falling back to eager. The settings are now applied around each compiled call too (`tabicl_compile.with_compile_settings`, ~4 us), so they hold on any thread, and `accumulated_recompile_limit` (256, which capped the per-device limit) is raised with them. A100, calls on a separate thread (`benchmarks/bench_stalls.py --call-thread`; `baselines/compile-call-thread-a100-{before,after}.jsonl`): shapes outside warmup's families run compiled once their graph exists -- in that sweep, 3,500 x 500 fits 2.31 -> 1.61 s and 3,500 x 1,000 5.49 -> 3.50 s, on graphs first built by its earlier 30k-row fits. Such a graph now compiles once per worker (a first ~30k-row fit takes 5-25 s) instead of never. ### Added - `benchmarks/bench_stalls.py` (fixed-cost stalls across table sizes, with compiled-graph counts per call; `--tasks`, `--call-thread` to call from another thread as serving does), `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).