From 2b647e9c914094ba1942b4cc544d6e38ba4fce95 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 13:50:17 -0400 Subject: [PATCH 1/3] docs: redraw architecture diagram for multi-GPU router, worker processes, host guard The figure predated 0.12.0: it showed one TabctxService replica per GPU share with a single engine lock, and had no MultiDeviceEngine, per-GPU worker processes, residency events, host-memory guard or offload tier. Redraw it from the code, update the caption, and fix the README's coalescing claim (light predicts now run beside fits). Co-Authored-By: Claude Opus 5.5 --- README.md | 29 +++-- docs/figures/architecture.svg | 207 +++++++++++++++++++--------------- 2 files changed, 139 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index ef5f683..90793ad 100644 --- a/README.md +++ b/README.md @@ -90,12 +90,17 @@ deployment is exactly fast enough. *Figure: a request enters through Ray Serve's consistent-hash proxy (affinity on the session header) into a `TabctxService` replica, which -resolves the tenant, parses inline JSON or an uploaded CSV, and hands the -table to the engine. The engine validates input, gates admission on an -estimated memory peak, keeps encoded contexts in a byte-budgeted LRU -cache, and dispatches to whichever backend fit the context. Evicted -contexts optionally spill to local disk and reload on their next use -(editable source: `docs/figures/architecture.svg`).* +resolves the tenant, parses inline JSON or an uploaded CSV, and coalesces +same-context predicts. With several GPUs, a `MultiDeviceEngine` places +each fit on one GPU and routes predicts to the GPU holding the context, +from lock-free snapshots; each GPU's engine runs in its own worker +process. A GPU's engine validates input, admits a fit on its estimated +memory need, keeps encoded contexts in a byte-budgeted LRU cache, and +dispatches to whichever backend fit the context. Evicted contexts +optionally spill to local disk and reload on their next use, and one +host-memory guard admits large fits against the node's RAM across all +workers. A single-GPU replica calls its one engine directly (editable +source: `docs/figures/architecture.svg`).* The unit everything revolves around is a **context**: one training table, encoded once by one model, addressed by `dataset_id` (scoped per @@ -418,8 +423,16 @@ pip install "tabctx[dev]" # + test dependencies - **Same-context coalescing** (`CoalescingPredictor`): concurrent requests against one cached context are packed into a single GPU call within a ~5ms window and split back per caller, amortizing per-call - overhead without ever running concurrent GPU work (the memory budget's - single-in-flight-call assumption is preserved). + overhead without adding concurrent GPU calls (each GPU still runs at + most one predict at a time, beside at most one fit). +- **Several GPUs per replica** (`MultiDeviceEngine`): one engine per GPU, + each in its own worker process, behind the single-engine API; see + [Several GPUs in one replica](#several-gpus-in-one-replica). +- **Host-memory admission** (`HostMemoryGuard`): large TabICL fits stage + tens of GiB in host RAM, so calls reserve their estimated host bytes + against one budget shared by every worker process on the node (first + come, first served), so two huge fits queue instead of getting the pod + OOM-killed with every tenant's context. ## Validated at scale (real A100-40GB, not simulated) diff --git a/docs/figures/architecture.svg b/docs/figures/architecture.svg index 3e2424a..41c5c8d 100644 --- a/docs/figures/architecture.svg +++ b/docs/figures/architecture.svg @@ -1,101 +1,130 @@ - - tabctx architecture: clients, serving layer, engine, backends, and context tiers + + tabctx architecture: clients, Ray Serve proxy, TabctxService replica, multi-GPU router, one worker process per GPU, context tiers, host-memory guard - + + - - CLIENTS - - TabctxClient (pure stdlib) - fit / predict / upload_csv / models / limits - sets affinity + tenant headers, maps errors - - Any HTTP client - x-session-id: dataset_id (replica affinity) - x-tabctx-tenant-id: tenant (isolation) - - In-process use - TabctxEngine directly, no server - (notebooks, batch jobs, tests) + + CLIENTS + + TabctxClient (pure stdlib) + fit / upload_csv / fit_uploaded / predict + sets session + tenant headers, maps errors + + Any HTTP client + x-session-id: dataset_id (replica affinity) + x-tabctx-tenant-id: tenant (isolation) + + In-process use, no server + TabctxEngine (one device) or + MultiDeviceEngine (several) - - - + + + + - - SERVING - - Ray Serve proxy: consistent-hash router on the session header - same dataset_id always reaches the same replica; strict (no fallback replica), 503 under backpressure - - - TabctxService replica (serve/app.py, FastAPI) -- one per GPU share - • POST /v1/tabctx/upload -> fit -> predict; GET /v1/models, /v1/tabctx/limits, /readyz - • tenant scoping (t:tenant:dataset_id), replica-local upload store, CSV parsing - • CoalescingPredictor: concurrent predicts on one context share a single GPU call - - serve/factory.py - env -> engine: TABCTX_BACKEND, - GPU fraction, kv mode, spill dir; - detects device memory - + + SERVING + + Ray Serve proxy: consistent-hash router on the session header + same dataset_id, same replica; strict (no fallback replica), 503 under backpressure + + + TabctxService replica (serve/app.py, FastAPI): one or more GPUs + • POST /v1/tabctx/upload, /fit, /predict; GET /v1/models, /limits, /residency, /readyz + • tenancy: every dataset_id scoped per tenant before it reaches an engine + • UploadStore: CSVs streamed to replica disk, single-use, TTL; JSON or CSV parsing + • CoalescingPredictor: concurrent predicts on one context share a single GPU call + + serve/factory.py (at start) + env -> engine: TABCTX_DEVICES, + _BACKEND, _WORKER_PROCESSES, + _KV_PRECISION, _COMPILE, + _SPILL_DIR, GPU memory fraction + measures each GPU's memory + and the pod's host-RAM budget + 1 GPU: one TabctxEngine, + no router, no worker process + - - ENGINE - - TabctxEngine: fit(X, y) -> dataset_id; predict(dataset_id, X_test); one lock serializes all GPU work per replica - input validation (422 before any GPU work) · per-backend dispatch by the context's backend_name · test-set chunking against live headroom - - - - - AdaptiveMemoryEstimator - admission: estimated fit PEAK must fit - 85% of device minus resident contexts - calibrated grid + learns from real fits; - infeasible fits 413 before any eviction - - ContextCacheManager - dataset_id -> CachedContext (payload, - backend, shape, feature names, bytes) - LRU eviction against a byte budget; - evict-ahead so a safe fit is never refused - - TabularICLBackend protocol - fit -> opaque payload; predict = ONE pass - reports measured resident + peak bytes - optional dumps/loads for the spill tier - (backends/base.py -- add a model here) - - + + ROUTER + + MultiDeviceEngine (multidevice.py): one replica's fits and predicts across its GPUs (2 or more) + fit -> PlacementPolicy (default LeastActiveThatFits): a GPU it fits on without eviction, least active first, else most headroom + predict -> the GPU holding the context (resident or spilled) · per-dataset_id locks · a re-fit lands before the old copy goes + decides from lock-free snapshots, never waits on a GPU · emits ResidencyEvents (placed / removed) -> /v1/tabctx/residency - - - MODELS - - tabicl-v2 (TabICLBackend) - kv cache on, shared backbone, GPU deltas - - tabpfn-3.5 (TabPFNBackend) - fit_with_cache; 422 on pretraining limits - - fake (FakeBackend) - no GPU/torch: CI, laptops, integration tests - one shared cache and memory budget per device; each context remembers the model that fit it + + one pipe per worker: calls out, a Snapshot (residency, stats, headroom) back with every reply - - - holds / evicts - TIERS - - GPU-resident contexts (encoded training tables, kv caches) - what the cache budgets and evicts; warm predicts skip the re-encode - - Disk spill tier (TABCTX_SPILL_DIR, opt-in) - evicted contexts reload on next use, no training data needed - + + WORKERS + + + + + Worker process per GPU (serve/workers.py): cuda:0 ... cuda:N-1, one TabctxEngine each + a ProcessShard proxies it in the replica; if a worker dies, its calls fail, placement skips it and the health check fails + + + TabctxEngine (engine.py): validates input (422 before any GPU work); one fit at a time on its GPU + predicts of up to 2,048 rows on a resident context run beside a fit; larger ones wait, chunked against live headroom + + + AdaptiveMemoryEstimator (per model) + admits a fit if resident + working floor + fits 85% of GPU minus resident contexts + calibrated grids, learns from real fits + never-feasible fits: 413 before evicting + + + ContextCacheManager (per GPU) + dataset_id -> CachedContext; all models + share one byte budget; LRU eviction + in-use contexts pinned; evicts ahead + to make room for an admitted fit + + + Backends (backends/base.py protocol) + tabicl-v2: kv cache, int8 KV, compiled + attention, GPU Yeo-Johnson preprocessing + tabpfn-3.5: fit_with_cache, shared model + fake: no GPU or torch (CI, laptops) + + + + + TIERS + + + holds / spills + + large TabICL fits + + GPU-resident contexts + encoded training tables, kv caches + warm predicts skip the re-encode + + Disk spill (TABCTX_SPILL_DIR, opt-in) + a subdirectory per GPU (cuda_0/, ...) + evicted contexts reload on next use + + TabICL offload (TABCTX_OFFLOAD_DIR) + memory-mapped stage outputs too big + for GPU and host RAM during a fit + + + + HOST + + + + HostMemoryGuard (memory/host.py): one host-RAM budget shared by every worker process on the node + calls reserve their estimated host bytes, first come first served; a few offloading fits at once; 503, not a hang, if its lock is stuck From c9113f48fd421c6e2e6b229de0e13ea1235ae7bb Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 13:51:27 -0400 Subject: [PATCH 2/3] docs: repo description in pyproject, 0.12.0 release date and status Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- README.md | 8 +++++--- pyproject.toml | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0347d54..e82e159 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## [Unreleased] -## [0.12.0] - unreleased +## [0.12.0] - 2026-10-18 ### Added - **Multi-GPU replicas** (#30). `MultiDeviceEngine` runs one `TabctxEngine` per device -- backends bound to `cuda:i`, its own cache, lock and spill subdirectory -- behind the single-engine API, built automatically when a replica sees several GPUs (`TABCTX_DEVICES=auto`, or an explicit list). Placement is aimed at the expected workload, a few concurrent users with small and large tables: a fit goes where its estimated peak fits without eviction, on the least active device (in-flight calls, then calls in the last 30s), else to the device with the most headroom. Predicts route to the device holding the context (resident first, then spilled), found from lock-free residency/spill snapshots, so routing never waits on another GPU's work and fits and predicts on different devices run concurrently. Per-dataset_id locks serialize fits and evicts of one id (no duplicate copies from retried fits); a re-fit lands on its device before other copies are removed, so a failed re-fit leaves the previous context serving. Unknown models are rejected (422) before anything changes. `TABCTX_SPILL_CAPACITY_BYTES` is the replica's total, split across devices' spill subdirectories; a spill layout change (one vs several devices) is logged. `TABCTX_DEVICES` rejects mixed CPU/GPU lists, a bare `cuda` among indexed GPUs, and indices the process can't see. `/readyz` gains a per-device breakdown (contexts, bytes, activity, real GPU memory); `/v1/tabctx/limits` reports the best single device. A single GPU builds exactly the engine it did before. - **One worker process per GPU** (#39, `TABCTX_WORKER_PROCESSES`, default on for multi-GPU replicas). Fits are host-bound, so threads of one process driving several GPUs contend for the GIL: on 4x A100, four concurrent fits ran at 2.1-3.1x aggregate from threads and 3.96-4.3x from processes (`benchmarks/bench_scaling.py`). `serve/workers.py` runs each device's engine in its own process behind a `ProcessShard` proxy; the router reads versioned residency/headroom snapshots that every reply carries, so its decisions stay lock-free. A dead worker fails its calls, is skipped by placement, and fails `/readyz` and Ray Serve's `check_health` so the replica restarts. Nothing depends on Ray. Under Ray Serve on 4x A100 with 6 mixed users: 2.58 -> 5.0 predicts/s, small-fit p50 1.6 -> 0.6 s. diff --git a/README.md b/README.md index 90793ad..52af243 100644 --- a/README.md +++ b/README.md @@ -466,9 +466,11 @@ sharing one A100-40GB** (v0.7.0, 2026-08-29, ## Status & Roadmap -v0.10.0: multi-model endpoint, calibrated peak-aware admission, disk -spillover and multi-replica routing have all shipped and been validated -on real GPUs (see [CHANGELOG.md](CHANGELOG.md)). Open roadmap items are +v0.12.0: multi-model endpoint, calibrated peak-aware admission, disk +spillover, multi-replica routing and multi-GPU replicas with a worker +process per GPU have all shipped and been validated on real GPUs (see +[CHANGELOG.md](CHANGELOG.md)); 0.12.0 serves on Vector's 4x A40 prod +node. Open roadmap items are tracked as GitHub issues labelled `roadmap`. **Read [ROADMAP.md](ROADMAP.md) before starting new work** -- it ranks what's next and why. diff --git a/pyproject.toml b/pyproject.toml index 5d60ba5..d34a3cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "tabctx" version = "0.12.0" -description = "Multi-tenant context caching and serving for tabular in-context-learning foundation models" +description = "Serving engine for tabular foundation models (TabICL, TabPFN): fit a table once, keep its encoded context on the GPU, and serve fast repeat predictions to many tenants, with memory-aware admission and multi-GPU placement." readme = "README.md" license = "Apache-2.0" requires-python = ">=3.12" From 13184896df54daa7cfd2f8557b73b1bd4aa2f9dc Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Tue, 20 Oct 2026 13:54:26 -0400 Subject: [PATCH 3/3] docs: architecture review fixes (admission rule, header rename, host guard scope, contrast) Co-Authored-By: Claude Opus 5.5 --- README.md | 8 ++++---- docs/figures/architecture.svg | 26 +++++++++++++------------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 52af243..ba31984 100644 --- a/README.md +++ b/README.md @@ -98,8 +98,8 @@ process. A GPU's engine validates input, admits a fit on its estimated memory need, keeps encoded contexts in a byte-budgeted LRU cache, and dispatches to whichever backend fit the context. Evicted contexts optionally spill to local disk and reload on their next use, and one -host-memory guard admits large fits against the node's RAM across all -workers. A single-GPU replica calls its one engine directly (editable +host-memory guard admits large fits against the pod's RAM across all +its workers. A single-GPU replica calls its one engine directly (editable source: `docs/figures/architecture.svg`).* The unit everything revolves around is a **context**: one training @@ -429,8 +429,8 @@ pip install "tabctx[dev]" # + test dependencies each in its own worker process, behind the single-engine API; see [Several GPUs in one replica](#several-gpus-in-one-replica). - **Host-memory admission** (`HostMemoryGuard`): large TabICL fits stage - tens of GiB in host RAM, so calls reserve their estimated host bytes - against one budget shared by every worker process on the node (first + tens of GiB in host RAM, so calls of 1 GiB or more reserve their + estimated host bytes against one budget shared by the pod's workers (first come, first served), so two huge fits queue instead of getting the pod OOM-killed with every tenant's context. diff --git a/docs/figures/architecture.svg b/docs/figures/architecture.svg index 41c5c8d..fcdb1ca 100644 --- a/docs/figures/architecture.svg +++ b/docs/figures/architecture.svg @@ -9,14 +9,14 @@ - CLIENTS + CLIENTS TabctxClient (pure stdlib) fit / upload_csv / fit_uploaded / predict sets session + tenant headers, maps errors Any HTTP client - x-session-id: dataset_id (replica affinity) + x-session-id (renamable): dataset_id, affinity x-tabctx-tenant-id: tenant (isolation) In-process use, no server @@ -29,14 +29,14 @@ - SERVING + SERVING Ray Serve proxy: consistent-hash router on the session header same dataset_id, same replica; strict (no fallback replica), 503 under backpressure TabctxService replica (serve/app.py, FastAPI): one or more GPUs - • POST /v1/tabctx/upload, /fit, /predict; GET /v1/models, /limits, /residency, /readyz + • POST /v1/tabctx/{upload,fit,predict}; GET /v1/models, /v1/tabctx/{limits,residency}, /readyz • tenancy: every dataset_id scoped per tenant before it reaches an engine • UploadStore: CSVs streamed to replica disk, single-use, TTL; JSON or CSV parsing • CoalescingPredictor: concurrent predicts on one context share a single GPU call @@ -53,7 +53,7 @@ - ROUTER + ROUTER MultiDeviceEngine (multidevice.py): one replica's fits and predicts across its GPUs (2 or more) fit -> PlacementPolicy (default LeastActiveThatFits): a GPU it fits on without eviction, least active first, else most headroom @@ -61,7 +61,7 @@ decides from lock-free snapshots, never waits on a GPU · emits ResidencyEvents (placed / removed) -> /v1/tabctx/residency - one pipe per worker: calls out, a Snapshot (residency, stats, headroom) back with every reply + one pipe per worker; every reply carries a Snapshot (residency, stats, headroom) WORKERS @@ -69,7 +69,7 @@ - Worker process per GPU (serve/workers.py): cuda:0 ... cuda:N-1, one TabctxEngine each + Worker process per GPU (serve/workers.py; threads if TABCTX_WORKER_PROCESSES=0): one TabctxEngine each a ProcessShard proxies it in the replica; if a worker dies, its calls fail, placement skips it and the health check fails @@ -78,16 +78,16 @@ AdaptiveMemoryEstimator (per model) - admits a fit if resident + working floor - fits 85% of GPU minus resident contexts + 413 only if a fit can't fit an empty + GPU's 85%; else evicts to make room calibrated grids, learns from real fits - never-feasible fits: 413 before evicting + checks real free GPU memory first ContextCacheManager (per GPU) dataset_id -> CachedContext; all models share one byte budget; LRU eviction - in-use contexts pinned; evicts ahead + in-use contexts evicted last; evicts ahead to make room for an admitted fit @@ -125,6 +125,6 @@ - HostMemoryGuard (memory/host.py): one host-RAM budget shared by every worker process on the node - calls reserve their estimated host bytes, first come first served; a few offloading fits at once; 503, not a hang, if its lock is stuck + HostMemoryGuard (memory/host.py): one host-RAM budget per replica pod, shared by its workers + calls needing 1 GiB+ reserve their estimate, first come first served; 503, not a hang, if its lock is stuck