fix(compile): apply dynamo settings per call -- serving threads saw a limit of 8 (#47)

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

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

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}

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}

benchmarks/bench_stalls.py

@@ -41,16 +41,58 @@ 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"])
     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(
+        "--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)