diff --git a/CHANGELOG.md b/CHANGELOG.md index f1e336b..91b00ac 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 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 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..09221e6 100644 --- a/src/tabctx/backends/tabicl_compile.py +++ b/src/tabctx/backends/tabicl_compile.py @@ -27,12 +27,26 @@ 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. 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 @@ -45,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 @@ -63,16 +79,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,10 +99,41 @@ def dispatch(self: Any, *args: Any, **kwargs: Any) -> Any: return True +def with_compile_settings(fn: Any, devices: int) -> Any: + """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 + 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 new file mode 100644 index 0000000..da3f660 --- /dev/null +++ b/tests/unit/test_tabicl_compile.py @@ -0,0 +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 compile_settings, with_compile_settings + +torch = pytest.importorskip("torch") +dynamo = pytest.importorskip("torch._dynamo") +config = dynamo.config + + +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 = _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, 512, True), "to_thread": (512, 512, True)} + + +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 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 = [] + + 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, 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()