From 0bc66515ad56f73df05b5133c44d8a431f9dbfe5 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 14:47:22 -0400 Subject: [PATCH 1/7] feat: multi-GPU replicas (MultiDeviceEngine) + workload benchmark One TabctxEngine per device behind the same API. Fits go where they fit without eviction on the least active device (else most headroom); predicts route to the owning device; spill per device. Backends measure and run on their own cuda:i. TABCTX_DEVICES=auto|list. Per-device /readyz. benchmarks/bench_workload.py: concurrent users, mixed table sizes. ROADMAP: multi-GPU and CUDA-level track (#30-#37). Co-Authored-By: Claude Opus 5.5 --- ROADMAP.md | 32 ++ benchmarks/bench_workload.py | 215 ++++++++++ src/tabctx/backends/_cuda.py | 27 ++ src/tabctx/backends/tabicl.py | 54 +-- src/tabctx/backends/tabpfn.py | 52 +-- src/tabctx/batching.py | 5 +- src/tabctx/engine.py | 12 + src/tabctx/multidevice.py | 234 +++++++++++ src/tabctx/serve/app.py | 56 ++- src/tabctx/serve/factory.py | 217 ++++++++--- .../integration/test_multi_device_service.py | 70 ++++ tests/unit/test_multidevice.py | 368 ++++++++++++++++++ tests/unit/test_serve_factory.py | 115 +++++- 13 files changed, 1338 insertions(+), 119 deletions(-) create mode 100644 benchmarks/bench_workload.py create mode 100644 src/tabctx/backends/_cuda.py create mode 100644 src/tabctx/multidevice.py create mode 100644 tests/integration/test_multi_device_service.py create mode 100644 tests/unit/test_multidevice.py diff --git a/ROADMAP.md b/ROADMAP.md index 60e3336..2ade128 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -45,6 +45,38 @@ gets deleted when it empties. 4. **Docs site** ([#9](https://github.com/VectorInstitute/tabctx/issues/9)): README carries the full story today; a docs site matters as external adoption starts. +## Throughput and table size: multi-GPU and CUDA-level work + +The LLM serving engines (vLLM, SGLang, TensorRT-LLM) got their gains first +from memory management and scheduling (paged KV, continuous batching, +prefix reuse) and then from CUDA kernels. tabctx already has the +scheduling layer (context cache, usage-aware admission, evict-ahead, +spill, same-context coalescing). This track adds multi-GPU replicas and +the kernel work. Everything is developed and tested on a GKE GPU node +first, then validated on prod's A40s (Ampere sm_86: no FlashAttention-3, +no FP8). + +5. **Multi-GPU replica** ([#30](https://github.com/VectorInstitute/tabctx/issues/30), in progress): one engine per device, + fits placed where they fit without eviction on the least active + device, predicts routed to the owning device. +6. **Profiling harness** ([#31](https://github.com/VectorInstitute/tabctx/issues/31)): per-stage time and peak-memory breakdown; + picks and sizes everything below. Hypothesis to confirm: the fit + peak is per-cell activations in the column/row stages, not attention. +7. **Per-device calibration grids** ([#37](https://github.com/VectorInstitute/tabctx/issues/37)): admission on non-A100 cards + falls back to the conservative formula and is the first ceiling a + large-table user hits. +8. **Tiled/fused column+row stage kernels** ([#32](https://github.com/VectorInstitute/tabctx/issues/32)): cut the fit peak, raising + the per-GPU table ceiling (FlashAttention-style tiling; chunked + prefill for the context KV). +9. **Paged context KV + cross-context varlen predict** ([#33](https://github.com/VectorInstitute/tabctx/issues/33)): throughput with + many users (PagedAttention + continuous batching analogue). +10. **INT8 context KV** ([#34](https://github.com/VectorInstitute/tabctx/issues/34)): about 2x contexts per GPU. +11. **Context parallelism across GPUs** ([#35](https://github.com/VectorInstitute/tabctx/issues/35)): tables larger than one GPU + (ring attention over the replica's devices). +12. **CUDA graphs for small predicts** ([#36](https://github.com/VectorInstitute/tabctx/issues/36)). + +Triton first; CUDA C++/CUTLASS where profiling shows Triton falls short. + ## Standing process note Every significant fix in this project came from testing against real diff --git a/benchmarks/bench_workload.py b/benchmarks/bench_workload.py new file mode 100644 index 0000000..94f2ea9 --- /dev/null +++ b/benchmarks/bench_workload.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +"""Mixed-workload benchmark: several concurrent users, small AND large +tables, the way a shared deployment is actually used. + +Each simulated user loops for --duration seconds: fit a table drawn from a +size mix (mostly small, some medium, occasionally large), then run a few +predicts against it, then move to a new table. Users run concurrently, +each with its own tenant and session, through TabctxClient (so the +gateway's header contract is exercised exactly). + +Reports, per size class: fit and predict latency percentiles, 413s +(admission rejections), and errors; overall predict throughput; and, on a +multi-GPU replica, how contexts and activity were spread across devices +(from /readyz). Compare runs before/after a change with --save. + +Usage: + bench_workload.py --base-url http://127.0.0.1:8000/tabctx \ + --users 6 --duration 120 --session-header x-tabctx-session-id \ + --save benchmarks/baselines/workload-a40x4.json +""" + +from __future__ import annotations + +import argparse +import json +import random +import statistics +import sys +import threading +import time +from dataclasses import dataclass, field +from datetime import UTC, datetime + +import numpy as np + +from tabctx.client import TabctxClient +from tabctx.errors import AdmissionRejected, TabctxError + +# Tables at least this many cells go through the CSV upload path, the way +# a real large-table user sends them; smaller ones are inline JSON. +UPLOAD_CELLS = 200_000 + +# (name, weight, rows, features). Defaults follow the single-L4 sweep: +# "large" is near what one GPU admits, "small" is a typical user table. +DEFAULT_MIX = [ + ("small", 0.6, 2_000, 20), + ("medium", 0.3, 10_000, 50), + ("large", 0.1, 40_000, 100), +] + + +@dataclass +class ClassStats: + fit_s: list[float] = field(default_factory=list) + predict_s: list[float] = field(default_factory=list) + rejected: int = 0 + errors: list[str] = field(default_factory=list) + + +def make_table(rows: int, features: int, rng: np.random.Generator): + X = rng.normal(size=(rows, features)).astype(np.float32) + y = (X[:, 0] + 0.5 * X[:, 1] > 0).astype(int).astype(str).tolist() + return X, y + + +def to_csv(X: np.ndarray, y: list[str]) -> bytes: + header = ",".join([f"f{i}" for i in range(X.shape[1])] + ["target"]) + lines = [header] + lines += [ + ",".join(f"{v:.5g}" for v in row) + f",{label}" + for row, label in zip(X, y, strict=True) + ] + return ("\n".join(lines) + "\n").encode() + + +def fit_table( + client: TabctxClient, X: np.ndarray, y: list[str], dataset_id: str, model +): + if X.size >= UPLOAD_CELLS: + upload_id = client.upload_csv(to_csv(X, y), dataset_id) + return client.fit_uploaded(upload_id, dataset_id, model=model) + return client.fit(X.tolist(), y, dataset_id=dataset_id, model=model) + + +def pct(values: list[float], p: float) -> float | None: + if not values: + return None + ordered = sorted(values) + return round(ordered[min(len(ordered) - 1, int(p / 100 * len(ordered)))], 3) + + +def user_loop(user: int, args, mix, stats: dict[str, ClassStats], lock, stop_at): + rng = np.random.default_rng(args.seed + user) + pick = random.Random(args.seed + user) + client = TabctxClient( + args.base_url, + tenant_id=f"{args.tenant}-{user}", + session_header=args.session_header, + timeout_s=args.timeout, + max_retries=10, + ) + n = 0 + while time.monotonic() < stop_at: + name, _, rows, features = pick.choices(mix, weights=[m[1] for m in mix])[0] + X, y = make_table(rows, features, rng) + dataset_id = f"wl-{user}-{n}" + n += 1 + t0 = time.monotonic() + try: + fit_table(client, X, y, dataset_id, args.model) + except AdmissionRejected: + with lock: + stats[name].rejected += 1 + continue + except (TabctxError, OSError) as e: + with lock: + stats[name].errors.append(f"fit: {type(e).__name__}: {e}"[:200]) + continue + fit_s = time.monotonic() - t0 + X_test, _ = make_table(args.test_rows, features, rng) + predicts = [] + for _ in range(args.predicts): + t0 = time.monotonic() + try: + client.predict(dataset_id, X_test.tolist()) + predicts.append(time.monotonic() - t0) + except (TabctxError, OSError) as e: + with lock: + stats[name].errors.append(f"predict: {type(e).__name__}: {e}"[:200]) + break + with lock: + stats[name].fit_s.append(fit_s) + stats[name].predict_s.extend(predicts) + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--base-url", required=True) + p.add_argument("--users", type=int, default=6) + p.add_argument("--duration", type=float, default=120) + p.add_argument("--predicts", type=int, default=5, help="per fitted table") + p.add_argument("--test-rows", type=int, default=200) + p.add_argument("--model", default=None) + p.add_argument("--tenant", default="bench") + p.add_argument("--session-header", default="x-session-id") + p.add_argument("--timeout", type=float, default=900) + p.add_argument("--seed", type=int, default=0) + p.add_argument( + "--mix", + default=None, + help="JSON list of [name, weight, rows, features], e.g. '[[\"s\",1,1000,10]]'", + ) + p.add_argument("--save", default=None, help="write JSON results here") + args = p.parse_args() + mix = [tuple(m) for m in json.loads(args.mix)] if args.mix else DEFAULT_MIX + + stats = {m[0]: ClassStats() for m in mix} + lock = threading.Lock() + started = time.monotonic() + stop_at = started + args.duration + threads = [ + threading.Thread(target=user_loop, args=(u, args, mix, stats, lock, stop_at)) + for u in range(args.users) + ] + for t in threads: + t.start() + for t in threads: + t.join() + elapsed = time.monotonic() - started + + ready = TabctxClient(args.base_url, timeout_s=args.timeout).ready() + total_predicts = sum(len(s.predict_s) for s in stats.values()) + result = { + "timestamp": datetime.now(UTC).isoformat(), + "base_url": args.base_url, + "users": args.users, + "duration_s": round(elapsed, 1), + "mix": mix, + "predict_throughput_per_s": round(total_predicts / elapsed, 2), + "classes": { + name: { + "fits": len(s.fit_s), + "fit_p50_s": pct(s.fit_s, 50), + "fit_p95_s": pct(s.fit_s, 95), + "predicts": len(s.predict_s), + "predict_p50_s": pct(s.predict_s, 50), + "predict_p95_s": pct(s.predict_s, 95), + "rejected_413": s.rejected, + "errors": len(s.errors), + "first_errors": s.errors[:3], + } + for name, s in stats.items() + }, + "devices": ready.get("devices"), + "cache_stats": ready.get("cache_stats"), + "gpu_capacity_bytes": ready.get("gpu_capacity_bytes"), + } + print(json.dumps(result, indent=2)) + if args.save: + with open(args.save, "w") as f: + json.dump(result, f, indent=2) + n_errors = sum(len(s.errors) for s in stats.values()) + fit_times = [v for s in stats.values() for v in s.fit_s] + print( + f"\n{args.users} users, {elapsed:.0f}s: {sum(len(s.fit_s) for s in stats.values())} " + f"fits (median {statistics.median(fit_times) if fit_times else 0:.2f}s), " + f"{total_predicts} predicts ({result['predict_throughput_per_s']}/s), " + f"{sum(s.rejected for s in stats.values())} rejected, {n_errors} errors", + file=sys.stderr, + ) + return 1 if n_errors else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/tabctx/backends/_cuda.py b/src/tabctx/backends/_cuda.py new file mode 100644 index 0000000..737ab83 --- /dev/null +++ b/src/tabctx/backends/_cuda.py @@ -0,0 +1,27 @@ +"""Device helpers shared by the torch backends. + +A backend bound to e.g. ``cuda:2`` must measure and allocate on THAT +device: torch's memory counters default to the *current* device, and +library internals (tabicl/tabpfn) allocate scratch tensors there too. +Multi-GPU replicas (see tabctx/multidevice.py) run one backend instance +per device, so every GPU call is wrapped in ``device_scope``. +""" + +from __future__ import annotations + +import contextlib +from typing import Any + + +def is_cuda(device: str) -> bool: + return device == "cuda" or device.startswith("cuda:") + + +def device_scope(device: str) -> contextlib.AbstractContextManager[Any]: + """Make ``device`` torch's current CUDA device for the block; a no-op + on CPU.""" + if not is_cuda(device): + return contextlib.nullcontext() + import torch + + return torch.cuda.device(torch.device(device)) diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py index 1849835..b1eb4be 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -53,6 +53,7 @@ import numpy as np +from tabctx.backends._cuda import device_scope, is_cuda from tabctx.errors import BackendComputeError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -132,12 +133,14 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: # shared across fits after the first (see __init__). model = self._make_estimator(task) - measuring = self._device == "cuda" - before_bytes = torch.cuda.memory_allocated() if measuring else 0 + device = self._device + measuring = is_cuda(device) + before_bytes = torch.cuda.memory_allocated(device) if measuring else 0 if measuring: - torch.cuda.reset_peak_memory_stats() + torch.cuda.reset_peak_memory_stats(device) try: - model.fit(X, y) + with device_scope(device): + model.fit(X, y) except torch.cuda.OutOfMemoryError as e: self._last_context_bytes = None self._last_fit_peak_bytes = None @@ -152,10 +155,10 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: # magnitude (activations, ensemble intermediates). # None on CPU; the engine falls back to the formula estimator. if measuring: - after_bytes = torch.cuda.memory_allocated() + after_bytes = torch.cuda.memory_allocated(device) self._last_context_bytes = max(0, after_bytes - before_bytes) self._last_fit_peak_bytes = max( - 0, torch.cuda.max_memory_allocated() - before_bytes + 0, torch.cuda.max_memory_allocated(device) - before_bytes ) else: self._last_context_bytes = None @@ -190,27 +193,31 @@ def predict( model = payload try: - if hasattr(model, "predict_proba"): - if return_proba: - proba = model.predict_proba(X_test) - class_idx = np.argmax(proba, axis=1) - classes = [str(c) for c in model.classes_] - predictions = [classes[i] for i in class_idx] - return PredictOutcome( - predictions=predictions, - probabilities=proba.tolist(), - classes=classes, - ) - predictions = model.predict(X_test).tolist() - return PredictOutcome(predictions=predictions) - # Regressor: no predict_proba to derive from. - predictions = model.predict(X_test).tolist() - return PredictOutcome(predictions=predictions) + with device_scope(self._device): + return self._predict(model, X_test, return_proba) except torch.cuda.OutOfMemoryError as e: raise BackendComputeError(f"CUDA OOM during predict(): {e}") from e finally: self._unpin_shared_module(model) + def _predict( + self, model: Any, X_test: ArrayLike, return_proba: bool + ) -> PredictOutcome: + if hasattr(model, "predict_proba"): + if return_proba: + proba = model.predict_proba(X_test) + class_idx = np.argmax(proba, axis=1) + classes = [str(c) for c in model.classes_] + predictions = [classes[i] for i in class_idx] + return PredictOutcome( + predictions=predictions, + probabilities=proba.tolist(), + classes=classes, + ) + return PredictOutcome(predictions=model.predict(X_test).tolist()) + # Regressor: no predict_proba to derive from. + return PredictOutcome(predictions=model.predict(X_test).tolist()) + def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: del n_train, n_features # Real measurement from the fit() call this is queried after (see @@ -255,7 +262,8 @@ def loads_payload(self, data: bytes) -> Any: # Cold restore on a fresh replica: borrow a throwaway # estimator's load path to populate the shared backbone. probe = self._make_estimator(task) - probe._load_model() + with device_scope(self._device): + probe._load_model() self._stash_backbone(task, probe) for attr, value in self._shared_backbones[task].items(): setattr(model, attr, value) diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py index ecc948a..39bce44 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -35,6 +35,7 @@ import numpy as np +from tabctx.backends._cuda import device_scope, is_cuda from tabctx.errors import BackendComputeError, InvalidInputError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -85,12 +86,14 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: import torch model = self._make_estimator(task) - measuring = self._device == "cuda" - before_bytes = torch.cuda.memory_allocated() if measuring else 0 + device = self._device + measuring = is_cuda(device) + before_bytes = torch.cuda.memory_allocated(device) if measuring else 0 if measuring: - torch.cuda.reset_peak_memory_stats() + torch.cuda.reset_peak_memory_stats(device) try: - model.fit(np.asarray(X, dtype=float), np.asarray(y)) + with device_scope(device): + model.fit(np.asarray(X, dtype=float), np.asarray(y)) except torch.cuda.OutOfMemoryError as e: self._last_context_bytes = None self._last_fit_peak_bytes = None @@ -105,10 +108,10 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: # backends/tabicl.py fit() for the full rationale. if measuring: self._last_context_bytes = max( - 0, torch.cuda.memory_allocated() - before_bytes + 0, torch.cuda.memory_allocated(device) - before_bytes ) self._last_fit_peak_bytes = max( - 0, torch.cuda.max_memory_allocated() - before_bytes + 0, torch.cuda.max_memory_allocated(device) - before_bytes ) else: self._last_context_bytes = None @@ -120,27 +123,32 @@ def predict( ) -> PredictOutcome: import torch - model = payload X_arr = np.asarray(X_test, dtype=float) try: - if hasattr(model, "predict_proba"): - # Single pass regardless of return_proba (protocol - # requirement): derive labels from probabilities. - proba = model.predict_proba(X_arr) - class_idx = np.argmax(proba, axis=1) - classes = [str(c) for c in model.classes_] - predictions = [classes[i] for i in class_idx] - if return_proba: - return PredictOutcome( - predictions=predictions, - probabilities=proba.tolist(), - classes=classes, - ) - return PredictOutcome(predictions=predictions) - return PredictOutcome(predictions=model.predict(X_arr).tolist()) + with device_scope(self._device): + return self._predict(payload, X_arr, return_proba) except torch.cuda.OutOfMemoryError as e: raise BackendComputeError(f"CUDA OOM during predict(): {e}") from e + def _predict( + self, model: Any, X_arr: np.ndarray, return_proba: bool + ) -> PredictOutcome: + if hasattr(model, "predict_proba"): + # Single pass regardless of return_proba (protocol + # requirement): derive labels from probabilities. + proba = model.predict_proba(X_arr) + class_idx = np.argmax(proba, axis=1) + classes = [str(c) for c in model.classes_] + predictions = [classes[i] for i in class_idx] + if return_proba: + return PredictOutcome( + predictions=predictions, + probabilities=proba.tolist(), + classes=classes, + ) + return PredictOutcome(predictions=predictions) + return PredictOutcome(predictions=model.predict(X_arr).tolist()) + def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: del n_train, n_features # Real measurement from the fit() this is queried after; None on diff --git a/src/tabctx/batching.py b/src/tabctx/batching.py index 6fcb590..2639df1 100644 --- a/src/tabctx/batching.py +++ b/src/tabctx/batching.py @@ -65,6 +65,7 @@ from tabctx.engine import TabctxEngine from tabctx.errors import InvalidInputError +from tabctx.multidevice import MultiDeviceEngine from tabctx.types import ArrayLike, PredictOutcome @@ -92,7 +93,9 @@ class CoalescingPredictor: when requested) for its own rows, in its own row order. """ - def __init__(self, engine: TabctxEngine, window_s: float = 0.005) -> None: + def __init__( + self, engine: TabctxEngine | MultiDeviceEngine, window_s: float = 0.005 + ) -> None: if window_s < 0: raise ValueError(f"window_s must be >= 0, got {window_s}") self._engine = engine diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py index 0e8865e..e21c42a 100644 --- a/src/tabctx/engine.py +++ b/src/tabctx/engine.py @@ -134,6 +134,18 @@ def default_backend(self) -> str: def estimator_for(self, backend: str | None = None) -> MemoryEstimator: return self._estimators[backend or self._default_backend] + def admission_headroom_bytes(self, backend: str | None = None) -> int: + """Transient headroom a fit could use right now, given what the + cache holds resident (the quantity fit admission bounds).""" + return self.estimator_for(backend).admission_headroom_bytes( + self._cache.stats().used_bytes + ) + + def holds(self, dataset_id: str) -> bool: + """Whether the context is resident in this engine's cache (a + spilled copy doesn't count; get() would reload it).""" + return dataset_id in self._cache.dataset_ids() + def _resolve_backend(self, backend: str | None) -> tuple[str, TabularICLBackend]: name = backend or self._default_backend b = self._backends.get(name) diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py new file mode 100644 index 0000000..72da8e9 --- /dev/null +++ b/src/tabctx/multidevice.py @@ -0,0 +1,234 @@ +"""MultiDeviceEngine: one replica serving from several GPUs. + +Neither TabICL nor TabPFN can split one training context across GPUs +(TabICL is single-device; TabPFN's device lists only parallelize ensemble +members, and not in its kv-cache fit mode). So a multi-GPU replica runs +one complete TabctxEngine per device -- its own backend instances bound +to that device, its own cache, lock and spill tier -- and this class +routes between them: + +- fit() places each context with a rule aimed at the usual workload -- + a few concurrent users, each fitting a table (small or large) and then + predicting against it repeatedly: + 1. Devices where the table's estimated peak fits WITHOUT evicting + anything are preferred; among them, the least active one wins + (fewest calls in flight, then fewest started recently), ties to + the most headroom. Each device serializes its own GPU work, so + spreading active users across devices is what lets them run + concurrently. + 2. If no device can take it without eviction (a large table, or a + full replica), it goes to the device with the most headroom, + evicting the fewest contexts. + Each device's own usage-aware admission gate then has the final say, + exactly as on a single GPU. A table's ceiling is therefore one + device's memory; what more devices buy is N x the resident cache and + concurrency. +- predict() runs on whichever device holds the context. Ownership is + tracked on fit; a context the map doesn't know (e.g. spilled before a + restart) is found by asking each device, which reloads it from that + device's own spill tier. +- Re-fitting a dataset_id drops it from every device first, so a stale + copy on another device can never answer a later predict. + +Placement is a snapshot taken outside the per-device locks: two +concurrent fits may pick the same device, in which case the second simply +evicts-ahead there like any single-GPU fit would. Serializing placement +would cost the concurrency this class exists for. + +Same public surface as TabctxEngine, so the serving layer and batching +use either interchangeably. +""" + +from __future__ import annotations + +import contextlib +import threading +import time +from collections import deque +from collections.abc import Callable, Iterator + +from tabctx.engine import TabctxEngine, _validate_fit_input +from tabctx.errors import DatasetNotFoundError +from tabctx.memory.estimator import MemoryEstimator +from tabctx.types import ArrayLike, EngineStats, PredictOutcome, Task + +# How long a started call counts toward a device's recent activity. +ACTIVITY_WINDOW_S = 30.0 + + +class MultiDeviceEngine: + def __init__( + self, + shards: list[TabctxEngine], + devices: list[str], + activity_window_s: float = ACTIVITY_WINDOW_S, + clock: Callable[[], float] = time.monotonic, + ) -> None: + if len(shards) < 2 or len(shards) != len(devices): + raise ValueError("need one engine per device, and at least two devices") + names = {tuple(s.backend_names) for s in shards} + defaults = {s.default_backend for s in shards} + if len(names) != 1 or len(defaults) != 1: + raise ValueError("every device must serve the same models") + self._shards = shards + self._devices = devices + self._owner: dict[str, int] = {} + self._owner_lock = threading.Lock() + self._window_s = activity_window_s + self._clock = clock + self._inflight = [0] * len(shards) + self._recent: list[deque[float]] = [deque() for _ in shards] + self._activity_lock = threading.Lock() + + @property + def devices(self) -> list[str]: + return list(self._devices) + + @property + def backend_names(self) -> list[str]: + return self._shards[0].backend_names + + @property + def default_backend(self) -> str: + return self._shards[0].default_backend + + def estimator_for(self, backend: str | None = None) -> MemoryEstimator: + return self._shards[0].estimator_for(backend) + + def admission_headroom_bytes(self, backend: str | None = None) -> int: + """The best any one device offers: a fit lands on a single device.""" + return max(s.admission_headroom_bytes(backend) for s in self._shards) + + @contextlib.contextmanager + def _active(self, index: int) -> Iterator[None]: + now = self._clock() + with self._activity_lock: + self._inflight[index] += 1 + self._recent[index].append(now) + try: + yield + finally: + with self._activity_lock: + self._inflight[index] -= 1 + + def _activity(self) -> list[tuple[int, int]]: + """(calls in flight, calls started within the window) per device.""" + cutoff = self._clock() - self._window_s + with self._activity_lock: + for recent in self._recent: + while recent and recent[0] < cutoff: + recent.popleft() + return [(self._inflight[i], len(r)) for i, r in enumerate(self._recent)] + + def _place(self, estimated_bytes: int, backend: str | None) -> int: + headrooms = [s.admission_headroom_bytes(backend) for s in self._shards] + activity = self._activity() + indices = range(len(self._shards)) + fits = [i for i in indices if headrooms[i] >= estimated_bytes] + if fits: + return min(fits, key=lambda i: (activity[i], -headrooms[i], i)) + return max(indices, key=lambda i: (headrooms[i], -i)) + + def _drop_everywhere(self, dataset_id: str) -> None: + with self._owner_lock: + self._owner.pop(dataset_id, None) + for shard in self._shards: + shard.evict(dataset_id) + + def fit( + self, + X: ArrayLike, + y: ArrayLike, + task: Task = "classification", + dataset_id: str | None = None, + backend: str | None = None, + feature_names: list[str] | None = None, + ) -> str: + n_features = _validate_fit_input(X, y) + if dataset_id is not None: + self._drop_everywhere(dataset_id) + estimated = self.estimator_for(backend).estimate_bytes(len(X), 0, n_features) + index = self._place(estimated, backend) + with self._active(index): + resolved = self._shards[index].fit( + X, + y, + task=task, + dataset_id=dataset_id, + backend=backend, + feature_names=feature_names, + ) + with self._owner_lock: + self._owner[resolved] = index + return resolved + + def _search_order(self, dataset_id: str) -> list[int]: + with self._owner_lock: + owner = self._owner.get(dataset_id) + if owner is not None: + return [owner] + resident = [i for i, s in enumerate(self._shards) if s.holds(dataset_id)] + return resident + [i for i in range(len(self._shards)) if i not in resident] + + def predict( + self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False + ) -> PredictOutcome: + last_error: DatasetNotFoundError | None = None + for index in self._search_order(dataset_id): + try: + with self._active(index): + outcome = self._shards[index].predict( + dataset_id, X_test, return_proba=return_proba + ) + except DatasetNotFoundError as e: + last_error = e + continue + with self._owner_lock: + self._owner[dataset_id] = index + return outcome + with self._owner_lock: + self._owner.pop(dataset_id, None) + assert last_error is not None + raise last_error + + def feature_names(self, dataset_id: str) -> list[str] | None: + for index in self._search_order(dataset_id): + names = self._shards[index].feature_names(dataset_id) + if names is not None: + return names + return None + + def fit_predict( + self, + X: ArrayLike, + y: ArrayLike, + X_test: ArrayLike, + task: Task = "classification", + return_proba: bool = False, + ) -> PredictOutcome: + dataset_id = self.fit(X, y, task) + try: + return self.predict(dataset_id, X_test, return_proba=return_proba) + finally: + self.evict(dataset_id) + + def evict(self, dataset_id: str) -> None: + self._drop_everywhere(dataset_id) + + def stats(self) -> EngineStats: + per_device = self.device_stats() + return EngineStats( + n_cached_contexts=sum(s.n_cached_contexts for s in per_device), + used_bytes=sum(s.used_bytes for s in per_device), + free_bytes=sum(s.free_bytes for s in per_device), + capacity_bytes=sum(s.capacity_bytes for s in per_device), + ) + + def device_stats(self) -> list[EngineStats]: + return [s.stats() for s in self._shards] + + def device_activity(self) -> list[dict[str, int]]: + return [ + {"in_flight": inflight, "recent_calls": recent} + for inflight, recent in self._activity() + ] diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py index 4a2fd86..bf1e02b 100644 --- a/src/tabctx/serve/app.py +++ b/src/tabctx/serve/app.py @@ -44,6 +44,7 @@ UploadTooLargeError, ) from tabctx.batching import CoalescingPredictor +from tabctx.multidevice import MultiDeviceEngine from tabctx.serve.affinity import resolve_dataset_id, session_id_from_headers from tabctx.serve.csv_io import parse_features_csv, parse_train_csv from tabctx.serve.factory import ServeSettings, build_engine @@ -214,6 +215,8 @@ def __init__(self) -> None: self._estimator = built.estimator self._backend_name = built.backend.name # default model id self._spill_store = built.spill_store + self._spill_stores = built.spill_stores + self._devices = built.devices self._device = built.device self._gpu_capacity_bytes = built.gpu_capacity_bytes # Same-context predict coalescing (see batching.py): concurrent @@ -508,7 +511,8 @@ def limits(self): the admission gate at representative feature counts, so it tightens/loosens as the adaptive estimator learns real costs.""" used_bytes = self._engine.stats().used_bytes - headroom = self._estimator.admission_headroom_bytes(used_bytes) + # The best any one device offers: a fit lands on a single device. + headroom = self._engine.admission_headroom_bytes() max_rows_by_features = {} for n_features in (10, 50, 100, 200, 500): lo, hi = 0, 2_000_000 @@ -528,13 +532,15 @@ def limits(self): "cache_used_bytes": used_bytes, "gpu_memory_fraction": self._settings.gpu_memory_fraction, "gpu_capacity_bytes": self._gpu_capacity_bytes, + "devices": list(self._devices), "max_admissible_train_rows_by_feature_count": max_rows_by_features, "note": ( "Admission limits are per-replica, usage-aware (they shrink " "as this replica's cache fills and recover as contexts are " "evicted), and adaptive: they loosen for shapes similar to " "ones this replica has measured. A rejected fit returns 413 " - "with the same numbers." + "with the same numbers. With several GPUs, a table must fit " + "on one of them: the limits are the best single device's." ), } @@ -557,6 +563,15 @@ def readyz(self): "allocated_mb": torch.cuda.memory_allocated() / 1e6, "reserved_mb": torch.cuda.memory_reserved() / 1e6, } + if len(self._devices) > 1: + real_gpu_memory = { + d: { + "allocated_mb": torch.cuda.memory_allocated(d) / 1e6, + "reserved_mb": torch.cuda.memory_reserved(d) / 1e6, + } + for d in self._devices + if d.startswith("cuda") + } return { "status": "ready", "device": self._device, @@ -577,9 +592,44 @@ def readyz(self): "engine_calls": self._predictor.engine_calls, }, "uploads": self._uploads.stats(), - "spill": self._spill_store.stats() if self._spill_store else None, + "spill": self._spill_stats(), "real_gpu_memory": real_gpu_memory, + "devices": self._device_breakdown(), } + def _spill_stats(self) -> dict | None: + if self._spill_store is not None: + return self._spill_store.stats() + if not self._spill_stores: + return None + per_device = [s.stats() for s in self._spill_stores] + totals: dict = {} + for stats in per_device: + for key, value in stats.items(): + if isinstance(value, (int, float)) and not isinstance(value, bool): + totals[key] = totals.get(key, 0) + value + return {**totals, "per_device": per_device} + + def _device_breakdown(self) -> list[dict] | None: + """Per-device cache use and activity on a multi-GPU replica (None + on a single device, where cache_stats already says it all).""" + if not isinstance(self._engine, MultiDeviceEngine): + return None + return [ + { + "device": device, + "n_cached_contexts": stats.n_cached_contexts, + "used_bytes": stats.used_bytes, + "capacity_bytes": stats.capacity_bytes, + **activity, + } + for device, stats, activity in zip( + self._engine.devices, + self._engine.device_stats(), + self._engine.device_activity(), + strict=True, + ) + ] + app = TabctxService.bind() diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py index 9de3ad9..c3d4b42 100644 --- a/src/tabctx/serve/factory.py +++ b/src/tabctx/serve/factory.py @@ -35,7 +35,14 @@ (see serve/uploads.py). - ``TABCTX_SPILL_DIR`` (default unset = spillover off) and ``TABCTX_SPILL_CAPACITY_BYTES`` (default 50GiB): disk spillover tier - for capacity-evicted contexts (see cache/spill.py). + for capacity-evicted contexts (see cache/spill.py). With several + devices, each gets its own subdirectory and the capacity is per device. +- ``TABCTX_DEVICES``: ``"auto"`` (default: every visible CUDA device, or + the CPU when there is none) or a comma-separated list such as + ``cuda:0,cuda:1``. Two or more devices build a MultiDeviceEngine (see + tabctx/multidevice.py): one engine per device, fits placed on the + device with the most headroom. ``TABCTX_GPU_MEMORY_FRACTION`` applies + to each device. """ from __future__ import annotations @@ -58,6 +65,7 @@ DEFAULT_HARD_CEILING_BYTES, MemoryEstimator, ) +from tabctx.multidevice import MultiDeviceEngine from tabctx.serve.uploads import DEFAULT_MAX_UPLOAD_BYTES, DEFAULT_TTL_S BACKEND_ENV_VAR = "TABCTX_BACKEND" @@ -68,6 +76,7 @@ SPILL_DIR_ENV_VAR = "TABCTX_SPILL_DIR" SPILL_CAPACITY_ENV_VAR = "TABCTX_SPILL_CAPACITY_BYTES" UPLOAD_TTL_S_ENV_VAR = "TABCTX_UPLOAD_TTL_S" +DEVICES_ENV_VAR = "TABCTX_DEVICES" BackendKind = Literal["tabicl", "tabpfn", "fake"] KvCacheMode = Literal["kv", "repr", "off"] @@ -85,6 +94,8 @@ class ServeSettings: upload_ttl_s: float = DEFAULT_TTL_S spill_dir: str | None = None spill_capacity_bytes: int = DEFAULT_SPILL_CAPACITY_BYTES + # Empty = "auto" (see resolve_devices). + devices: tuple[str, ...] = () @classmethod def from_env(cls) -> ServeSettings: @@ -151,6 +162,7 @@ def from_env(cls) -> ServeSettings: raise ValueError(f"{SPILL_CAPACITY_ENV_VAR} must be an int") from e if spill_capacity <= 0: raise ValueError(f"{SPILL_CAPACITY_ENV_VAR} must be positive") + devices = _parse_devices(os.environ.get(DEVICES_ENV_VAR, "auto")) return cls( backends=backends, gpu_memory_fraction=fraction, @@ -160,20 +172,41 @@ def from_env(cls) -> ServeSettings: upload_ttl_s=upload_ttl_s, spill_dir=spill_dir, spill_capacity_bytes=spill_capacity, + devices=devices, ) +def _parse_devices(raw: str) -> tuple[str, ...]: + raw = raw.strip().lower() + if raw in ("", "auto"): + return () + devices = tuple(d.strip() for d in raw.split(",") if d.strip()) + for d in devices: + index = d.removeprefix("cuda:") + if d not in ("cpu", "cuda") and not (d.startswith("cuda:") and index.isdigit()): + raise ValueError( + f"{DEVICES_ENV_VAR}={raw!r}: {d!r} is not 'cpu', 'cuda' or 'cuda:'" + ) + if len(set(devices)) != len(devices): + raise ValueError(f"{DEVICES_ENV_VAR} lists a device twice") + return devices + + @dataclass(frozen=True) class BuiltEngine: - engine: TabctxEngine + engine: TabctxEngine | MultiDeviceEngine # Keyed by model/backend name; `default` names the first-listed model. estimators: dict[str, MemoryEstimator] backends: dict[str, TabularICLBackend] default: str device: str spill_store: object | None = None - # Real device memory the budget was derived from; None = defaults. + # Real device memory the budget was derived from (per device; the + # smallest when several); None = defaults. gpu_capacity_bytes: int | None = None + devices: tuple[str, ...] = () + # One per device when there are several (spill_store is then None). + spill_stores: tuple = () # Single-model conveniences (the common deployment): @property @@ -212,16 +245,47 @@ def _preloaded_observations( return fit_grid, predict_grid -def detect_gpu_capacity_bytes() -> int | None: - """Total memory of the visible CUDA device, or None when torch is - absent or sees no GPU (the fake backend, CPU hosts).""" +def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | None: + """Total memory of the given CUDA devices (the smallest, since one + estimator budgets them all), or None when torch is absent or sees no + GPU (the fake backend, CPU hosts).""" try: import torch except ImportError: return None if not torch.cuda.is_available(): return None - return int(torch.cuda.get_device_properties(0).total_memory) + cuda = [torch.device(d) for d in devices if d.startswith("cuda")] or [ + torch.device("cuda") + ] + return min( + int(torch.cuda.get_device_properties(d.index or 0).total_memory) for d in cuda + ) + + +def _visible_cuda_count() -> int: + try: + import torch + except ImportError: + return 0 + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + + +def resolve_devices(settings: ServeSettings) -> tuple[str, ...]: + """The devices to build one engine for each. Explicit settings win; + "auto" means every visible GPU. A single auto-detected GPU stays + "cuda" (the pre-multi-device layout); the fake backend never touches + a GPU, so auto means one CPU engine for it.""" + if settings.devices: + return settings.devices + if set(settings.backends) == {"fake"}: + return ("cpu",) + n = _visible_cuda_count() + if n == 0: + return ("cpu",) + if n == 1: + return ("cuda",) + return tuple(f"cuda:{i}" for i in range(n)) def build_estimator( @@ -258,91 +322,122 @@ def build_estimator( def _build_real_backend( - kind: BackendKind, settings: ServeSettings -) -> tuple[TabularICLBackend, str]: # pragma: no cover + kind: BackendKind, settings: ServeSettings, device: str +) -> TabularICLBackend: # pragma: no cover # Requires torch + tabicl/tabpfn (GPU or licensed weights), which -- # like backends/tabicl.py and backends/tabpfn.py themselves -- are not # part of any unit- or integration-test CI extra; see those modules' # docstrings. Isolated in its own function so the pragma doesn't also # exclude the (fully unit-tested) fake-backend and dispatch logic in # _build_backend below. - import torch - - device = "cuda" if torch.cuda.is_available() else "cpu" if kind == "tabpfn": from tabctx.backends.tabpfn import TabPFNBackend # TABCTX_KV_CACHE maps onto TabPFN's fit_mode (see backends/tabpfn.py). - return TabPFNBackend(device=device, cache_mode=settings.kv_cache), device + return TabPFNBackend(device=device, cache_mode=settings.kv_cache) from tabctx.backends.tabicl import TabICLBackend kv_cache: bool | str = False if settings.kv_cache == "off" else settings.kv_cache - return TabICLBackend(device=device, kv_cache=kv_cache), device + return TabICLBackend(device=device, kv_cache=kv_cache) def _build_backend( - kind: BackendKind, settings: ServeSettings -) -> tuple[TabularICLBackend, str]: - """Returns (backend, device). Imports torch/tabicl/tabpfn only on the - paths that need them, so the fake backend runs with core deps alone.""" + kind: BackendKind, settings: ServeSettings, device: str = "cpu" +) -> TabularICLBackend: + """One backend bound to `device`. Imports torch/tabicl/tabpfn only on + the paths that need them, so the fake backend runs with core deps + alone.""" if kind == "fake": from tabctx.backends.fake import FakeBackend - return FakeBackend(), "cpu (fake backend)" - return _build_real_backend(kind, settings) + return FakeBackend() + return _build_real_backend(kind, settings, device) + + +def _build_spill_store( + settings: ServeSettings, + backends: dict[str, TabularICLBackend], + subdir: str | None, +) -> DiskSpillStore | None: + if not settings.spill_dir: + return None + # Backends may provide backbone-aware serialization; pickle is the + # default for those that don't (see cache/spill.py). Restores use the + # serializers of the device that spilled, so a reloaded context comes + # back onto that device. + serializers = { + name: (b.dumps_payload, b.loads_payload) + for name, b in backends.items() + if hasattr(b, "dumps_payload") + } + path = os.path.join(settings.spill_dir, subdir) if subdir else settings.spill_dir + return DiskSpillStore( + path, + capacity_bytes=settings.spill_capacity_bytes, + serializers=serializers, + ) def build_engine(settings: ServeSettings | None = None) -> BuiltEngine: settings = settings or ServeSettings.from_env() - backends: dict[str, TabularICLBackend] = {} - estimators: dict[str, MemoryEstimator] = {} - device = "cpu" + devices = resolve_devices(settings) # Budget against the real device, not the calibration card's size; # the fake backend never touches a GPU even when one is visible. gpu_capacity = ( - None if settings.backends == ("fake",) else detect_gpu_capacity_bytes() + None + if set(settings.backends) == {"fake"} + else detect_gpu_capacity_bytes(devices) ) - for kind in settings.backends: - backend, device = _build_backend(kind, settings) - backends[backend.name] = backend - estimators[backend.name] = build_estimator(settings, kind, gpu_capacity) - # Backends are keyed by model id (e.g. "tabicl-v2"), not by kind - # ("tabicl"); insertion order follows settings.backends, so the - # first-listed kind's model is the default. - default = next(iter(backends)) - - spill_store = None - if settings.spill_dir: - # Backends may provide backbone-aware serialization; pickle is - # the default for those that don't (see cache/spill.py). - serializers = { - name: (b.dumps_payload, b.loads_payload) - for name, b in backends.items() - if hasattr(b, "dumps_payload") - } - spill_store = DiskSpillStore( - settings.spill_dir, - capacity_bytes=settings.spill_capacity_bytes, - serializers=serializers, + multi = len(devices) > 1 + estimators: dict[str, MemoryEstimator] = {} + shards: list[TabctxEngine] = [] + spill_stores: list[DiskSpillStore | None] = [] + first_backends: dict[str, TabularICLBackend] = {} + for device in devices: + backends: dict[str, TabularICLBackend] = {} + for kind in settings.backends: + backend = _build_backend(kind, settings, device) + backends[backend.name] = backend + # One estimator per model, shared by every device: they are + # identical cards budgeted alike, and observations learned on + # one device are valid on the others. + if backend.name not in estimators: + estimators[backend.name] = build_estimator(settings, kind, gpu_capacity) + # Backends are keyed by model id (e.g. "tabicl-v2"), not by kind + # ("tabicl"); insertion order follows settings.backends, so the + # first-listed kind's model is the default. + default = next(iter(backends)) + spill_store = _build_spill_store( + settings, backends, device.replace(":", "_") if multi else None ) - # ONE cache and budget shared by every model on the device -- the - # whole point of serving them behind one endpoint (see engine.py). - cache = ContextCacheManager( - capacity_bytes=estimators[default].ceiling_bytes(), spill_store=spill_store - ) - engine = TabctxEngine( - backends=backends, - estimators=estimators, - cache=cache, - default_backend=default, + # ONE cache and budget shared by every model on a device -- the + # whole point of serving them behind one endpoint (see engine.py). + cache = ContextCacheManager( + capacity_bytes=estimators[default].ceiling_bytes(), spill_store=spill_store + ) + shards.append( + TabctxEngine( + backends=backends, + estimators={name: estimators[name] for name in backends}, + cache=cache, + default_backend=default, + ) + ) + spill_stores.append(spill_store) + first_backends = first_backends or backends + engine: TabctxEngine | MultiDeviceEngine = ( + MultiDeviceEngine(shards, list(devices)) if multi else shards[0] ) + fake_only = set(settings.backends) == {"fake"} return BuiltEngine( engine=engine, - estimators=estimators, - backends=backends, - default=default, - device=device, - spill_store=spill_store, + estimators={name: estimators[name] for name in first_backends}, + backends=first_backends, + default=next(iter(first_backends)), + device="cpu (fake backend)" if fake_only else ",".join(devices), + spill_store=None if multi else spill_stores[0], gpu_capacity_bytes=gpu_capacity, + devices=devices, + spill_stores=tuple(s for s in spill_stores if s is not None) if multi else (), ) diff --git a/tests/integration/test_multi_device_service.py b/tests/integration/test_multi_device_service.py new file mode 100644 index 0000000..dc18cf2 --- /dev/null +++ b/tests/integration/test_multi_device_service.py @@ -0,0 +1,70 @@ +"""The serving layer on a multi-device replica (fake backend: no GPU). + +TABCTX_DEVICES=cuda:0,cuda:1 with the fake backend builds a real +MultiDeviceEngine without touching CUDA, so this exercises the exact +HTTP paths a multi-GPU replica serves: placement across devices, routing +predicts back to them, and the per-device /readyz and /limits views. +""" + +import pytest + +ray = pytest.importorskip("ray") +httpx = pytest.importorskip("httpx") + +PREFIX = "/multi" + + +@pytest.fixture(scope="module") +def multi_device_service(two_replica_service): + from ray import serve + + from tabctx.serve.app import TabctxService + + serve.run( + TabctxService.options( + num_replicas=1, + ray_actor_options={ + "runtime_env": { + "env_vars": { + "TABCTX_BACKEND": "fake", + "TABCTX_DEVICES": "cuda:0,cuda:1", + } + } + }, + ).bind(), + name="multi-device", + route_prefix=PREFIX, + ) + yield two_replica_service + PREFIX + serve.delete("multi-device") + + +def test_contexts_spread_over_devices_and_predict_from_them(multi_device_service): + with httpx.Client(base_url=multi_device_service, timeout=60.0) as client: + ready = client.get("/readyz").json() + assert [d["device"] for d in ready["devices"]] == ["cuda:0", "cuda:1"] + + for i in range(6): + r = client.post( + "/v1/tabctx/fit", + json={ + "train_X": [[float(i), 1.0], [2.0, 3.0]], + "train_y": [f"y{i}", f"y{i}"], + "dataset_id": f"md-{i}", + }, + ) + assert r.status_code == 200, r.text + for i in range(6): + r = client.post( + "/v1/tabctx/predict", + json={"dataset_id": f"md-{i}", "test_X": [[0.0, 0.0]]}, + ) + assert r.status_code == 200, r.text + assert r.json()["predictions"] == [f"y{i}"] + + devices = client.get("/readyz").json()["devices"] + assert sum(d["n_cached_contexts"] for d in devices) == 6 + assert all(d["n_cached_contexts"] > 0 for d in devices) + + limits = client.get("/v1/tabctx/limits").json() + assert limits["devices"] == ["cuda:0", "cuda:1"] diff --git a/tests/unit/test_multidevice.py b/tests/unit/test_multidevice.py new file mode 100644 index 0000000..e80e259 --- /dev/null +++ b/tests/unit/test_multidevice.py @@ -0,0 +1,368 @@ +"""Unit tests for MultiDeviceEngine: one replica, several GPUs. + +Placement is tested against a deterministic estimator (bytes = cells x +BYTES_PER_CELL) and a fake backend whose resident size follows the same +rule, so every headroom number in these tests is exact. +""" + +import threading +import time + +import pytest + +from tabctx.backends.fake import FakeBackend +from tabctx.cache.manager import ContextCacheManager +from tabctx.cache.spill import DiskSpillStore +from tabctx.engine import TabctxEngine +from tabctx.errors import AdmissionRejected, DatasetNotFoundError, InvalidInputError +from tabctx.multidevice import MultiDeviceEngine + +BYTES_PER_CELL = 10 +CAPACITY = 10_000 # per device: 1,000 cells resident or in flight + + +class CellEstimator: + """Peak and resident cost are both cells x BYTES_PER_CELL; headroom is + whatever the device's cache doesn't hold.""" + + def __init__(self, capacity: int = CAPACITY) -> None: + self.capacity = capacity + self.observations = 0 + + def estimate_bytes(self, n_train: int, n_test: int, n_features: int) -> int: + return (n_train + n_test) * n_features * BYTES_PER_CELL + + def admit(self, n_train: int, n_test: int, n_features: int) -> bool: + return self.estimate_bytes(n_train, n_test, n_features) <= self.capacity + + def ceiling_bytes(self) -> int: + return self.capacity + + def admission_headroom_bytes(self, used_bytes: int) -> int: + return max(0, self.capacity - used_bytes) + + def confidence(self) -> str: + return "test" + + def record_observation( + self, n_train: int, n_features: int, real_bytes: int + ) -> None: + self.observations += 1 + + +class SizedFake(FakeBackend): + """Resident size = cells x BYTES_PER_CELL, like the estimator.""" + + def __init__(self, device: str, **kw) -> None: + super().__init__(**kw) + self.device = device + self._last_bytes = None + + def fit(self, X, y, task): + self._last_bytes = len(X) * len(X[0]) * BYTES_PER_CELL + return super().fit(X, y, task) + + def context_bytes_hint(self, n_train, n_features): + return self._last_bytes + + +class FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +def table(rows: int, features: int = 1, label: str = "a"): + return [[float(i)] * features for i in range(rows)], [label] * rows + + +def build(n_devices=2, spill_root=None, clock=None, **fake_kw): + devices = [f"cuda:{i}" for i in range(n_devices)] + estimator = CellEstimator() + shards, fakes = [], [] + for device in devices: + fake = SizedFake(device, name="m", **fake_kw) + spill = ( + DiskSpillStore(str(spill_root / device.replace(":", "_")), 10**9) + if spill_root + else None + ) + shards.append( + TabctxEngine( + backends={"m": fake}, + estimators={"m": estimator}, + cache=ContextCacheManager(CAPACITY, spill_store=spill), + default_backend="m", + ) + ) + fakes.append(fake) + engine = MultiDeviceEngine(shards, devices, clock=clock or FakeClock()) + return engine, shards, fakes + + +def where(shards, dataset_id): + return [i for i, s in enumerate(shards) if s.holds(dataset_id)] + + +class TestConstruction: + def test_needs_two_or_more_devices(self): + _, shards, _ = build(2) + with pytest.raises(ValueError): + MultiDeviceEngine(shards[:1], ["cuda:0"]) + with pytest.raises(ValueError): + MultiDeviceEngine(shards, ["cuda:0"]) + + def test_devices_must_serve_the_same_models(self): + _, shards, _ = build(2) + other = TabctxEngine( + backends={"x": FakeBackend(name="x")}, + estimators={"x": CellEstimator()}, + cache=ContextCacheManager(CAPACITY), + default_backend="x", + ) + with pytest.raises(ValueError, match="same models"): + MultiDeviceEngine([shards[0], other], ["cuda:0", "cuda:1"]) + + def test_exposes_the_single_engine_surface(self): + engine, _, _ = build(3) + assert engine.devices == ["cuda:0", "cuda:1", "cuda:2"] + assert engine.backend_names == ["m"] + assert engine.default_backend == "m" + assert engine.admission_headroom_bytes() == CAPACITY + + +class TestPlacementSmallTables: + def test_concurrent_users_spread_across_devices(self): + engine, shards, _ = build(4) + for user in range(4): + engine.fit(*table(10), dataset_id=f"user{user}") + placed = [where(shards, f"user{u}")[0] for u in range(4)] + assert sorted(placed) == [0, 1, 2, 3] + + def test_recently_active_device_is_avoided(self): + clock = FakeClock() + engine, shards, _ = build(2, clock=clock) + engine.fit(*table(10), dataset_id="busy") + busy = where(shards, "busy")[0] + for _ in range(5): + engine.predict("busy", [[0.0]]) + engine.fit(*table(10), dataset_id="newcomer") + assert where(shards, "newcomer") == [1 - busy] + + def test_activity_expires_after_the_window(self): + clock = FakeClock() + engine, shards, _ = build(2, clock=clock) + engine.fit(*table(10), dataset_id="old") + first = where(shards, "old")[0] + clock.now += 31 + # Activity has expired, so the tie breaks on headroom: the empty + # device wins over the one holding "old". + engine.fit(*table(10), dataset_id="new") + assert where(shards, "new") == [1 - first] + + def test_in_flight_work_outranks_recent_history(self): + clock = FakeClock() + engine, shards, fakes = build(2, clock=clock) + # Device 1 has more recent calls, device 0 has one in flight. + engine._recent[1].extend([clock.now] * 5) + engine._inflight[0] = 1 + engine.fit(*table(10), dataset_id="d") + assert where(shards, "d") == [1] + + +class TestPlacementLargeTables: + def test_large_table_goes_where_it_fits_without_eviction(self): + clock = FakeClock() + engine, shards, _ = build(2, clock=clock) + engine.fit(*table(600), dataset_id="big0") # 6,000 B on one device + clock.now += 60 + engine.fit(*table(500), dataset_id="big1") # only fits on the other + assert where(shards, "big0") != where(shards, "big1") + assert shards[0].stats().n_cached_contexts == 1 + assert shards[1].stats().n_cached_contexts == 1 + + def test_when_nothing_fits_it_evicts_on_the_emptiest_device_only(self): + clock = FakeClock() + engine, shards, _ = build(2, clock=clock) + engine.fit(*table(700), dataset_id="a") # device X: 7,000 used + clock.now += 60 + engine.fit(*table(400), dataset_id="b") # device Y: 4,000 used + clock.now += 60 + # 800 cells fits nowhere without eviction; Y has more headroom, + # so only "b" is evicted and "a" is untouched. + engine.fit(*table(800), dataset_id="c") + assert where(shards, "a") != [] and where(shards, "b") == [] + assert where(shards, "c") != where(shards, "a") + + def test_table_larger_than_any_device_is_rejected_without_evictions(self): + engine, shards, _ = build(2) + engine.fit(*table(100), dataset_id="keep0") + engine.fit(*table(100), dataset_id="keep1") + with pytest.raises(AdmissionRejected): + engine.fit(*table(2000), dataset_id="huge") + assert where(shards, "keep0") and where(shards, "keep1") + + def test_small_tables_do_not_block_a_later_large_table(self): + clock = FakeClock() + engine, shards, _ = build(4, clock=clock) + for u in range(8): + engine.fit(*table(50), dataset_id=f"small{u}") + clock.now += 1 + clock.now += 60 + engine.fit(*table(900), dataset_id="large") + # 8 small contexts spread 2 per device leave 9,000 B on every + # device, so the large table needs no eviction anywhere. + assert all(where(shards, f"small{u}") for u in range(8)) + assert len(where(shards, "large")) == 1 + + +class TestRouting: + def test_predict_runs_on_the_owning_device(self): + engine, shards, fakes = build(2) + engine.fit(*table(10, label="x"), dataset_id="d") + owner = where(shards, "d")[0] + assert engine.predict("d", [[0.0]]).predictions == ["x"] + assert fakes[owner].predict_calls == 1 + assert fakes[1 - owner].predict_calls == 0 + + def test_unknown_dataset_is_not_found(self): + engine, _, _ = build(2) + with pytest.raises(DatasetNotFoundError): + engine.predict("nope", [[0.0]]) + + def test_refit_moves_the_context_and_leaves_no_stale_copy(self): + clock = FakeClock() + engine, shards, _ = build(2, clock=clock) + engine.fit(*table(10, label="v1"), dataset_id="d") + first = where(shards, "d")[0] + # Make the first device look busy so the re-fit lands elsewhere. + engine._recent[first].extend([clock.now] * 10) + engine.fit(*table(10, label="v2"), dataset_id="d") + assert where(shards, "d") == [1 - first] + assert engine.predict("d", [[0.0]]).predictions == ["v2"] + + def test_evict_drops_every_copy(self): + engine, shards, _ = build(2) + engine.fit(*table(10), dataset_id="d") + engine.evict("d") + assert where(shards, "d") == [] + with pytest.raises(DatasetNotFoundError): + engine.predict("d", [[0.0]]) + + def test_generated_dataset_ids_are_routed(self): + engine, _, _ = build(2) + dataset_id = engine.fit(*table(10, label="g")) + assert engine.predict(dataset_id, [[0.0]]).predictions == ["g"] + + def test_feature_names_follow_the_context(self): + engine, _, _ = build(2) + engine.fit(*table(5, features=2), dataset_id="d", feature_names=["p", "q"]) + assert engine.feature_names("d") == ["p", "q"] + assert engine.feature_names("missing") is None + + def test_fit_predict_leaves_nothing_cached(self): + engine, shards, _ = build(2) + out = engine.fit_predict(*table(10, label="z"), [[0.0]]) + assert out.predictions == ["z"] + assert engine.stats().n_cached_contexts == 0 + + def test_invalid_input_is_rejected_before_placement(self): + engine, _, _ = build(2) + with pytest.raises(InvalidInputError): + engine.fit([[1.0], [2.0]], ["a"], dataset_id="bad") + assert engine.device_activity() == [ + {"in_flight": 0, "recent_calls": 0}, + {"in_flight": 0, "recent_calls": 0}, + ] + + +class TestSpill: + def test_capacity_eviction_spills_to_the_device_and_restores_there(self, tmp_path): + clock = FakeClock() + engine, shards, fakes = build(2, spill_root=tmp_path, clock=clock) + engine.fit(*table(600, label="cold"), dataset_id="cold") # idle tie: dev 0 + clock.now += 60 + engine.fit(*table(600), dataset_id="other") # fits only on dev 1 + assert (where(shards, "cold"), where(shards, "other")) == ([0], [1]) + clock.now += 60 + # 900 cells fits nowhere without eviction; headroom ties, so dev 0 + # takes it and spills "cold" to dev 0's own spill tier. + engine.fit(*table(900), dataset_id="big") + assert where(shards, "big") == [0] and where(shards, "cold") == [] + assert (tmp_path / "cuda_0").exists() + # Restoring "cold" brings it back onto dev 0 (spilling "big" there) + # and never disturbs dev 1. + assert engine.predict("cold", [[0.0]]).predictions == ["cold"] + assert where(shards, "cold") == [0] + assert where(shards, "other") == [1] + assert fakes[1].predict_calls == 0 + + def test_after_a_restart_spilled_contexts_are_found_on_their_device(self, tmp_path): + engine, shards, _ = build(2, spill_root=tmp_path) + engine.fit(*table(10, label="warm"), dataset_id="d") + owner = where(shards, "d")[0] + shards[owner]._cache.make_room(CAPACITY) # spill everything there + assert where(shards, "d") == [] + # Fresh engine over the same spill dirs = a restarted replica. + restarted, new_shards, new_fakes = build(2, spill_root=tmp_path) + assert restarted.predict("d", [[0.0]]).predictions == ["warm"] + assert where(new_shards, "d") == [owner] + assert new_fakes[owner].predict_calls == 1 + + +class TestStats: + def test_stats_sum_over_devices(self): + engine, _, _ = build(3) + engine.fit(*table(10), dataset_id="a") + engine.fit(*table(20), dataset_id="b") + total = engine.stats() + per = engine.device_stats() + assert total.n_cached_contexts == 2 + assert total.used_bytes == 300 == sum(s.used_bytes for s in per) + assert total.capacity_bytes == 3 * CAPACITY + + def test_headroom_is_the_best_single_device(self): + engine, _, _ = build(2) + engine.fit(*table(300), dataset_id="a") + assert engine.admission_headroom_bytes() == CAPACITY + + +class TestConcurrency: + def test_work_on_different_devices_runs_in_parallel(self): + engine, shards, fakes = build(2, predict_delay_s=0.3) + engine.fit(*table(10), dataset_id="a") + engine.fit(*table(10), dataset_id="b") + assert where(shards, "a") != where(shards, "b") + threads = [ + threading.Thread(target=engine.predict, args=(d, [[0.0]])) for d in "ab" + ] + start = time.monotonic() + for t in threads: + t.start() + for t in threads: + t.join() + assert time.monotonic() - start < 0.5 # ~0.3s, not ~0.6s serialized + + def test_concurrent_fits_spread_and_all_succeed(self): + engine, shards, _ = build(4, fit_delay_s=0.05) + errors = [] + + def fit(user): + try: + engine.fit(*table(20), dataset_id=f"u{user}") + except Exception as e: # pragma: no cover - surfaced below + errors.append(e) + + threads = [threading.Thread(target=fit, args=(u,)) for u in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + counts = [s.stats().n_cached_contexts for s in shards] + assert sum(counts) == 16 + assert min(counts) >= 2 # spread, not piled on one device + for u in range(16): + assert engine.predict(f"u{u}", [[0.0]]).predictions == ["a"] diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py index f6c71e1..f69a09b 100644 --- a/tests/unit/test_serve_factory.py +++ b/tests/unit/test_serve_factory.py @@ -9,6 +9,7 @@ from tabctx.serve.factory import ( BACKEND_ENV_VAR, BATCH_WINDOW_MS_ENV_VAR, + DEVICES_ENV_VAR, GPU_MEMORY_FRACTION_ENV_VAR, KV_CACHE_ENV_VAR, MAX_UPLOAD_BYTES_ENV_VAR, @@ -169,17 +170,18 @@ def test_none_capacity_keeps_the_a100_defaults(self): def test_build_engine_uses_detected_capacity_for_real_backends(self, monkeypatch): monkeypatch.setattr( "tabctx.serve.factory._build_real_backend", - lambda kind, settings: (FakeBackend(name=f"{kind}-model"), "cuda"), + lambda kind, settings, device: FakeBackend(name=f"{kind}-model"), ) monkeypatch.setattr( - "tabctx.serve.factory.detect_gpu_capacity_bytes", lambda: 24 * 1024**3 + "tabctx.serve.factory.detect_gpu_capacity_bytes", + lambda devices: 24 * 1024**3, ) - built = build_engine(ServeSettings(backends=("tabicl",))) + built = build_engine(ServeSettings(backends=("tabicl",), devices=("cuda",))) assert built.gpu_capacity_bytes == 24 * 1024**3 assert built.estimator._fallback.gpu_capacity_bytes == 24 * 1024**3 def test_fake_backend_never_probes_the_device(self, monkeypatch): - def boom(): + def boom(*_): raise AssertionError("fake backend must not touch torch/cuda") monkeypatch.setattr("tabctx.serve.factory.detect_gpu_capacity_bytes", boom) @@ -229,9 +231,11 @@ def test_default_model_is_the_first_listed_kind_by_model_id(self, monkeypatch): # model, resolved through that mapping. monkeypatch.setattr( "tabctx.serve.factory._build_real_backend", - lambda kind, settings: (FakeBackend(name=f"{kind}-model"), "cpu"), + lambda kind, settings, device: FakeBackend(name=f"{kind}-model"), + ) + built = build_engine( + ServeSettings(backends=("tabpfn", "fake"), devices=("cpu",)) ) - built = build_engine(ServeSettings(backends=("tabpfn", "fake"))) assert built.default == "tabpfn-model" assert built.engine.default_backend == "tabpfn-model" assert built.engine.backend_names == ["tabpfn-model", "fake"] @@ -246,6 +250,7 @@ def test_settings_defaults_match_from_env_defaults(self, monkeypatch): UPLOAD_TTL_S_ENV_VAR, SPILL_DIR_ENV_VAR, SPILL_CAPACITY_ENV_VAR, + DEVICES_ENV_VAR, ): monkeypatch.delenv(var, raising=False) assert ServeSettings.from_env() == ServeSettings() @@ -268,11 +273,13 @@ def test_dispatches_non_fake_kinds_to_the_real_backend_path(self, monkeypatch): calls = [] monkeypatch.setattr( "tabctx.serve.factory._build_real_backend", - lambda kind, settings: calls.append((kind, settings)) or (object(), "cpu"), + lambda kind, settings, device: ( + calls.append((kind, settings, device)) or object() + ), ) settings = ServeSettings(backends=("tabicl",)) - _build_backend("tabicl", settings) - assert calls == [("tabicl", settings)] + _build_backend("tabicl", settings, "cuda:1") + assert calls == [("tabicl", settings, "cuda:1")] class TestCalibrationPreload: @@ -314,3 +321,93 @@ def test_missing_calibration_module_degrades_to_no_preload(self, monkeypatch): monkeypatch.setitem(sys.modules, "tabctx.memory.calibration_tabicl_a100", None) est = build_estimator(ServeSettings(backends=("tabicl",))) assert "0 preloaded calibration measurement" in est.confidence() + + +class TestDevices: + def test_auto_is_the_default(self, monkeypatch): + monkeypatch.delenv(DEVICES_ENV_VAR, raising=False) + assert ServeSettings.from_env().devices == () + monkeypatch.setenv(DEVICES_ENV_VAR, " AUTO ") + assert ServeSettings.from_env().devices == () + + def test_explicit_list_is_parsed(self, monkeypatch): + monkeypatch.setenv(DEVICES_ENV_VAR, "cuda:0, cuda:2") + assert ServeSettings.from_env().devices == ("cuda:0", "cuda:2") + + @pytest.mark.parametrize("raw", ["gpu0", "cuda:x", "cuda:0,cuda:0", "tpu"]) + def test_bad_lists_are_rejected(self, monkeypatch, raw): + monkeypatch.setenv(DEVICES_ENV_VAR, raw) + with pytest.raises(ValueError, match=DEVICES_ENV_VAR): + ServeSettings.from_env() + + @pytest.mark.parametrize( + ("visible", "expected"), + [ + (0, ("cpu",)), + (1, ("cuda",)), + (4, ("cuda:0", "cuda:1", "cuda:2", "cuda:3")), + ], + ) + def test_auto_uses_every_visible_gpu(self, monkeypatch, visible, expected): + from tabctx.serve.factory import resolve_devices + + monkeypatch.setattr("tabctx.serve.factory._visible_cuda_count", lambda: visible) + assert resolve_devices(ServeSettings(backends=("tabicl",))) == expected + + def test_explicit_devices_win_and_fake_auto_is_cpu(self, monkeypatch): + from tabctx.serve.factory import resolve_devices + + monkeypatch.setattr("tabctx.serve.factory._visible_cuda_count", lambda: 4) + assert resolve_devices( + ServeSettings(backends=("tabicl",), devices=("cuda:1",)) + ) == ("cuda:1",) + assert resolve_devices(ServeSettings(backends=("fake",))) == ("cpu",) + + def test_single_device_builds_a_plain_engine(self): + from tabctx.engine import TabctxEngine + + built = build_engine(ServeSettings(backends=("fake",))) + assert type(built.engine) is TabctxEngine + assert built.devices == ("cpu",) + + def test_several_devices_build_one_engine_each(self, monkeypatch, tmp_path): + from tabctx.multidevice import MultiDeviceEngine + + created = [] + monkeypatch.setattr( + "tabctx.serve.factory._build_real_backend", + lambda kind, settings, device: ( + created.append(device) or FakeBackend(name=f"{kind}-model") + ), + ) + monkeypatch.setattr( + "tabctx.serve.factory.detect_gpu_capacity_bytes", + lambda devices: 48 * 1024**3, + ) + built = build_engine( + ServeSettings( + backends=("tabicl", "tabpfn"), + devices=("cuda:0", "cuda:1", "cuda:2"), + spill_dir=str(tmp_path), + ) + ) + assert isinstance(built.engine, MultiDeviceEngine) + # Every model gets a backend instance on every device. + assert created == ["cuda:0", "cuda:0", "cuda:1", "cuda:1", "cuda:2", "cuda:2"] + assert built.engine.backend_names == ["tabicl-model", "tabpfn-model"] + assert built.default == "tabicl-model" + # Estimators are shared across devices (identical cards). + shard_estimators = { + id(s.estimator_for("tabicl-model")) for s in built.engine._shards + } + assert len(shard_estimators) == 1 + # Each device has its own cache, budgeted per device, and spill dir. + caps = [s.capacity_bytes for s in built.engine.device_stats()] + assert len(set(caps)) == 1 and caps[0] == built.estimator.ceiling_bytes() + assert built.spill_store is None and len(built.spill_stores) == 3 + assert sorted(p.name for p in tmp_path.iterdir()) == [ + "cuda_0", + "cuda_1", + "cuda_2", + ] + assert built.device == "cuda:0,cuda:1,cuda:2" From 24dd3d2b4beb929085c5abfcabd840ce7cc57645 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 14:47:45 -0400 Subject: [PATCH 2/7] docs(roadmap): PyTorch-level first, custom kernels only where profiles demand Co-Authored-By: Claude Opus 5.5 --- ROADMAP.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ROADMAP.md b/ROADMAP.md index 2ade128..3a79a9f 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -75,7 +75,13 @@ no FP8). (ring attention over the replica's devices). 12. **CUDA graphs for small predicts** ([#36](https://github.com/VectorInstitute/tabctx/issues/36)). -Triton first; CUDA C++/CUTLASS where profiling shows Triton falls short. +PyTorch-level first: tiling and chunking, `torch.compile` (generated Triton +kernels, CUDA graphs via `reduce-overhead`), FlexAttention for paged and +cross-context attention, SDPA, and `torch.distributed` context +parallelism. Hand-written Triton/CUDA only where a profile (#31) shows a +gap those can't close -- the models change too often (TabPFN v3 -> v3.5, +TabICL v1 -> v2 within months) for version-locked kernels to be the +default. ## Standing process note From 7ceebb96b6e9ab264ab17d97ad801d15c90709c4 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 14:51:12 -0400 Subject: [PATCH 3/7] docs: multi-GPU README section, 0.12.0 changelog (unreleased), bump version Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 14 ++++ README.md | 20 ++++++ benchmarks/sweep_admission.py | 120 ++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- uv.lock | 2 +- 5 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 benchmarks/sweep_admission.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1808793..136fce9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,20 @@ ## [Unreleased] +## [0.12.0] - unreleased +### 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 owning device; spilled contexts are found on their device after a restart; re-fitting an id drops every copy first. Work on different devices runs concurrently. `/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. +- `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread. +- `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed). +- ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first. + +### Fixed +- The TabICL and TabPFN backends measured memory on, and let library internals allocate on, torch's *current* CUDA device rather than their own. Harmless on one GPU; on `cuda:i` it would have charged every context to `cuda:0`. They now measure on their device and run inside its scope. + +### Validated +- Unit (multi-device placement, routing, spill restore, concurrency) and integration tests (the serve app on two fake devices through Ray). +- On a local fake 3-device replica with 5 concurrent users: with spill on, 0 failed predicts at ~100 predicts/s. Without spill, capacity eviction cost 12 predicts on three devices versus 115 on one. + ## [0.11.0] - 2026-10-17 ### Changed - **The TabPFN model id follows the checkpoint actually loaded.** It was hardcoded to `tabpfn-3`, but `tabpfn` 9.x loads TabPFN-3.5 by default (a different model under its own license, `tabpfn-3-5-license-v1.0`), so a default install served TabPFN-3.5 named `tabpfn-3`. The id is now derived from `tabpfn`'s resolved `model_version` (`TABPFN_MODEL_VERSION`): `v3.5` -> `tabpfn-3.5` (the 9.x default), `v3` -> `tabpfn-3`, `v3.5-fast` -> `tabpfn-3.5-fast`. **Breaking for callers that send `"model": "tabpfn-3"` to a deployment on the default version**: pin `TABPFN_MODEL_VERSION=v3` to keep serving TabPFN-3 under its old id, or switch callers to `tabpfn-3.5`. diff --git a/README.md b/README.md index 83b2712..8586a35 100644 --- a/README.md +++ b/README.md @@ -200,6 +200,26 @@ Details worth knowing: torch -- `tests/integration/test_multi_replica_affinity.py` uses it to prove the multi-replica contract on a laptop. +## Several GPUs in one replica + +A replica that can see several GPUs serves from all of them +(`TABCTX_DEVICES=auto`, the default, or an explicit `cuda:0,cuda:1,...`): +one engine per GPU, with its own backends, cache and spill directory, +behind the same API. No client changes and no routing header needed. + +- **Placement.** Each fit goes to a GPU where the table fits *without + evicting* anyone, preferring the least active GPU (calls in flight, + then calls in the last 30s). If no GPU can take it without eviction, + it goes to the GPU with the most headroom. Concurrent users spread + across GPUs, and a large table gets a whole GPU's headroom. +- **Predicts** run on the GPU that holds the context, including after a + spill and restore, or a replica restart. +- **Limits.** Neither TabICL nor TabPFN can split one context across + GPUs, so a table's ceiling is one GPU's memory. What more GPUs buy is + N x the resident cache and N fits/predicts in flight at once. +- `/readyz` reports per-device contexts, bytes, activity and real GPU + memory; `/v1/tabctx/limits` reports the best single GPU's headroom. + ## Large tables: upload, then fit by reference Inline JSON tables are convenient to ~10^4 rows and painful beyond. diff --git a/benchmarks/sweep_admission.py b/benchmarks/sweep_admission.py new file mode 100644 index 0000000..cf8206a --- /dev/null +++ b/benchmarks/sweep_admission.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +"""Table-size sweep: how large a table does a deployment admit, and how +long do fit and predict take on the way up? + +For each feature count, fits tables of increasing row count until the +admission gate answers 413 (or --max-rows), timing fit and a fixed-size +predict at every step -- the method behind the single-L4 "fit latency +vs table size" chart. Each feature count re-fits ONE dataset_id, so +earlier steps never hold memory that would shrink later steps' headroom. +Tables at or above UPLOAD_CELLS go through the CSV upload path. + +Usage: + sweep_admission.py --base-url http://127.0.0.1:8000/tabctx \ + --session-header x-tabctx-session-id --model tabicl-v2 \ + --save benchmarks/baselines/sweep-l4x4.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from datetime import UTC, datetime + +import numpy as np +from bench_workload import UPLOAD_CELLS, fit_table, make_table + +from tabctx.client import TabctxClient +from tabctx.errors import AdmissionRejected, TabctxError + +ROWS = [ + 100, 250, 500, 1_000, 2_000, 3_500, 5_000, 7_500, 10_000, 15_000, 20_000, + 30_000, 45_000, 65_000, 100_000, 150_000, 200_000, 300_000, 500_000, +] # fmt: skip + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--base-url", required=True) + p.add_argument("--features", type=int, nargs="+", default=[10, 100, 500, 1000]) + p.add_argument("--max-rows", type=int, default=500_000) + p.add_argument("--test-rows", type=int, default=100) + p.add_argument("--model", default=None) + p.add_argument("--tenant", default="sweep") + p.add_argument("--session-header", default="x-session-id") + p.add_argument("--timeout", type=float, default=1800) + p.add_argument("--save", default=None) + args = p.parse_args() + + client = TabctxClient( + args.base_url, + tenant_id=args.tenant, + session_header=args.session_header, + timeout_s=args.timeout, + ) + rng = np.random.default_rng(0) + points = [] + for features in args.features: + dataset_id = f"sweep-{args.model or 'default'}-f{features}" + for rows in (r for r in ROWS if r <= args.max_rows): + X, y = make_table(rows, features, rng) + point = { + "n_features": features, + "n_rows": rows, + "upload": X.size >= UPLOAD_CELLS, + } + t0 = time.monotonic() + try: + fit_table(client, X, y, dataset_id, args.model) + except AdmissionRejected as e: + point.update( + status=413, + fit_s=round(time.monotonic() - t0, 3), + detail=str(e)[:300], + ) + points.append(point) + print(json.dumps(point), flush=True) + break + except (TabctxError, OSError) as e: + point.update(status="error", detail=f"{type(e).__name__}: {e}"[:300]) + points.append(point) + print(json.dumps(point), flush=True) + break + point["fit_s"] = round(time.monotonic() - t0, 3) + X_test, _ = make_table(args.test_rows, features, rng) + t0 = time.monotonic() + client.predict(dataset_id, X_test.tolist()) + point.update(status=200, predict_s=round(time.monotonic() - t0, 3)) + points.append(point) + print(json.dumps(point), flush=True) + + limits = client.limits() + result = { + "timestamp": datetime.now(UTC).isoformat(), + "base_url": args.base_url, + "model": args.model, + "points": points, + "max_admitted_rows": { + str(f): max( + ( + q["n_rows"] + for q in points + if q["n_features"] == f and q["status"] == 200 + ), + default=0, + ) + for f in args.features + }, + "limits": limits, + } + print(json.dumps(result["max_admitted_rows"])) + if args.save: + with open(args.save, "w") as f: + json.dump(result, f, indent=2) + return 1 if any(q["status"] == "error" for q in points) else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index f694b3d..5d60ba5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "tabctx" -version = "0.11.0" +version = "0.12.0" description = "Multi-tenant context caching and serving for tabular in-context-learning foundation models" readme = "README.md" license = "Apache-2.0" diff --git a/uv.lock b/uv.lock index e39fbb7..01125c2 100644 --- a/uv.lock +++ b/uv.lock @@ -3201,7 +3201,7 @@ wheels = [ [[package]] name = "tabctx" -version = "0.11.0" +version = "0.12.0" source = { editable = "." } dependencies = [ { name = "numpy" }, From 4c00471d757649061ff1f9badf80d6a3a1836d0f Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 14:58:09 -0400 Subject: [PATCH 4/7] bench: profiling harness (#31), copy bandwidth (#39), GKE GPU devbox Co-Authored-By: Claude Opus 5.5 --- benchmarks/bench_copy.py | 98 +++++++++++ benchmarks/gke/devbox.sh | 143 +++++++++++++++ benchmarks/profile_fit_predict.py | 279 ++++++++++++++++++++++++++++++ 3 files changed, 520 insertions(+) create mode 100644 benchmarks/bench_copy.py create mode 100755 benchmarks/gke/devbox.sh create mode 100644 benchmarks/profile_fit_predict.py diff --git a/benchmarks/bench_copy.py b/benchmarks/bench_copy.py new file mode 100644 index 0000000..4a3cf0c --- /dev/null +++ b/benchmarks/bench_copy.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Context copy bandwidth on a GPU node (#39): what moving a context costs. + +Fit/predict disaggregation and rebalancing contexts between GPUs both pay +for a copy. This measures, for buffer sizes spanning real contexts (MBs to +tens of GB): GPU->GPU (peer, when available), GPU->pinned host, +pinned host->GPU, and GPU->pageable host (what pickling a context to the +spill tier starts with). + +Usage: bench_copy.py --sizes-gib 0.25 1 4 16 --out copy-a100x4.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time + +import torch + +GiB = 1024**3 + + +def bandwidth(fn, nbytes: int, devices: list[torch.device], repeats: int) -> float: + fn() # warm + for d in devices: + torch.cuda.synchronize(d) + t0 = time.perf_counter() + for _ in range(repeats): + fn() + for d in devices: + torch.cuda.synchronize(d) + return nbytes * repeats / (time.perf_counter() - t0) / GiB + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--sizes-gib", type=float, nargs="+", default=[0.25, 1, 4, 16]) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--out", default=None) + args = p.parse_args() + + n = torch.cuda.device_count() + gpus = [torch.device(f"cuda:{i}") for i in range(n)] + result = { + "gpus": [torch.cuda.get_device_name(g) for g in gpus], + "peer_access": { + f"{i}->{j}": torch.cuda.can_device_access_peer(i, j) + for i in range(n) + for j in range(n) + if i != j + }, + "gib_per_s": [], + } + for size in args.sizes_gib: + nbytes = int(size * GiB) + src = torch.empty(nbytes, dtype=torch.uint8, device=gpus[0]) + pinned = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) + row = {"size_gib": size} + row["gpu_to_pinned_host"] = bandwidth( + lambda s=src, h=pinned: h.copy_(s, non_blocking=True), + nbytes, + gpus[:1], + args.repeats, + ) + row["pinned_host_to_gpu"] = bandwidth( + lambda s=src, h=pinned: s.copy_(h, non_blocking=True), + nbytes, + gpus[:1], + args.repeats, + ) + row["gpu_to_pageable_host"] = bandwidth( + lambda s=src: s.cpu(), nbytes, gpus[:1], max(1, args.repeats // 2) + ) + if n > 1: + dst = torch.empty(nbytes, dtype=torch.uint8, device=gpus[1]) + row["gpu0_to_gpu1"] = bandwidth( + lambda s=src, d=dst: d.copy_(s, non_blocking=True), + nbytes, + gpus[:2], + args.repeats, + ) + dst = None + row = {k: round(v, 2) if isinstance(v, float) else v for k, v in row.items()} + result["gib_per_s"].append(row) + print(json.dumps(row), flush=True) + src = pinned = None + torch.cuda.empty_cache() + if args.out: + with open(args.out, "w") as f: + json.dump(result, f, indent=2) + print(json.dumps({k: v for k, v in result.items() if k != "gib_per_s"})) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/gke/devbox.sh b/benchmarks/gke/devbox.sh new file mode 100755 index 0000000..4e71960 --- /dev/null +++ b/benchmarks/gke/devbox.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +# GKE GPU devbox for profiling/optimizing tabctx on real GPUs (#31-#36). +# +# One GPU node pool and one long-running pod with the PROD image (same torch/CUDA +# as inference-platform), the staged tabctx weights (offline, no token), tabicl/ +# tabpfn at prod's pins, and this checkout of tabctx installed editable. Iterate +# with `sync` (copy the working tree in) and `exec` (run anything there). +# +# benchmarks/gke/devbox.sh up # create cluster + pod (MACHINES order) +# benchmarks/gke/devbox.sh sync # copy src/ benchmarks/ tests/ into the pod +# benchmarks/gke/devbox.sh exec -- CMD # run CMD in the pod (cwd /work/tabctx) +# benchmarks/gke/devbox.sh fetch PATH # copy /work/tabctx/PATH back to ./gke-out/ +# benchmarks/gke/devbox.sh down # delete the cluster +# +# The cluster carries purpose=gke-mtest + expiry_epoch (MAX_HOURS) so the org's orphan +# sweep deletes it even if `down` never runs. +set -euo pipefail + +PROJECT="${PROJECT:-agentic-ai-evaluation-bootcamp}" +CLUSTER="${CLUSTER:-gke-mtest-tabctx-devbox}" +STATE="${STATE:-$HOME/.cache/tabctx-devbox}" # remembers the zone/machine chosen +MACHINES="${MACHINES:-a2-highgpu-1g g2-standard-12}" # 1x A100-40GB, else 1x L4 +ZONES="${ZONES:-us-central1-a us-central1-b us-central1-c us-central1-f us-east4-a us-east4-c us-east1-b us-east1-c us-east1-d us-west1-a us-west1-b}" +MAX_HOURS="${MAX_HOURS:-12}" +IMAGE="${IMAGE:-rayproject/ray-llm:nightly.261006.670387-extra-py312-cu130}" +WEIGHTS_BUCKET="${WEIGHTS_BUCKET:-gs://inference-platform-gke-model-weights-bootcamp/tabctx}" +NS=default +POD=tabctx-devbox +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$HERE/../.." && pwd)" +mkdir -p "$STATE" + +accelerator_for() { + case "$1" in + a2-highgpu-1g) echo "type=nvidia-tesla-a100,count=1" ;; + a2-highgpu-4g) echo "type=nvidia-tesla-a100,count=4" ;; + g2-standard-12) echo "type=nvidia-l4,count=1" ;; + g2-standard-48) echo "type=nvidia-l4,count=4" ;; + *) echo "unknown machine $1" >&2; exit 1 ;; + esac +} +gpus_for() { accelerator_for "$1" | sed -E 's/.*count=([0-9]+).*/\1/'; } + +has_capacity() { # zone machine + local probe="${CLUSTER}-probe" + if gcloud compute instances create "$probe" --project="$PROJECT" --zone="$1" \ + --machine-type="$2" --maintenance-policy=TERMINATE --no-address \ + --image-family=debian-12 --image-project=debian-cloud --boot-disk-size=10GB \ + --labels="purpose=gke-mtest,test=tabctx-devbox" --quiet >/dev/null 2>"$STATE/probe.err"; then + gcloud compute instances delete "$probe" --project="$PROJECT" --zone="$1" --quiet >/dev/null 2>&1 + return 0 + fi + echo " $1/$2: $(grep -oE 'ZONE_RESOURCE_POOL_EXHAUSTED[A-Z_]*|[Qq]uota [^.]*|does not exist[^.]*' "$STATE/probe.err" | head -1)" + return 1 +} + +creds() { + # shellcheck disable=SC1091 + source "$STATE/env" + gcloud container clusters get-credentials "$CLUSTER" --zone="$ZONE" --project="$PROJECT" >/dev/null 2>&1 + export KUBECONFIG_CONTEXT + KUBECONFIG_CONTEXT="gke_${PROJECT}_${ZONE}_${CLUSTER}" +} +k() { kubectl --context "$KUBECONFIG_CONTEXT" "$@"; } + +up() { + if [ ! -f "$STATE/env" ]; then + for m in $MACHINES; do + for z in $ZONES; do + if has_capacity "$z" "$m"; then printf 'ZONE=%s\nMACHINE=%s\n' "$z" "$m" > "$STATE/env"; break 2; fi + done + done + [ -f "$STATE/env" ] || { echo "no capacity for any of: $MACHINES" >&2; exit 1; } + fi + # shellcheck disable=SC1091 + source "$STATE/env" + echo "==> $CLUSTER in $ZONE on $MACHINE" + if ! gcloud container clusters describe "$CLUSTER" --project="$PROJECT" --zone="$ZONE" >/dev/null 2>&1; then + gcloud container clusters create "$CLUSTER" --project="$PROJECT" --zone="$ZONE" \ + --release-channel=regular --num-nodes=1 --machine-type=e2-standard-4 \ + --labels="purpose=gke-mtest,test=tabctx-devbox,expiry_epoch=$(( $(date +%s) + MAX_HOURS * 3600 ))" + fi + creds + if ! gcloud container node-pools describe gpu --cluster="$CLUSTER" --project="$PROJECT" --zone="$ZONE" >/dev/null 2>&1; then + gcloud container node-pools create gpu --project="$PROJECT" --zone="$ZONE" --cluster="$CLUSTER" \ + --machine-type="$MACHINE" --accelerator="$(accelerator_for "$MACHINE"),gpu-driver-version=latest" \ + --num-nodes=1 --disk-type=pd-ssd --disk-size=300 --no-enable-autoupgrade --no-enable-autorepair + fi + k wait --for=condition=Ready node -l cloud.google.com/gke-nodepool=gpu --timeout=20m + local gpus; gpus=$(gpus_for "$MACHINE") + cat < tuple[int, int]: + rows, features = text.lower().split("x") + return int(rows), int(features) + + +def make_table(rows: int, features: int, task: str, seed: int): + rng = np.random.default_rng(seed) + X = rng.normal(size=(rows, features)).astype(np.float32) + signal = X[:, 0] + 0.5 * X[:, min(1, features - 1)] + y = (signal > 0).astype(int) if task == "classification" else signal + return X, y + + +def build_backend(kind: str, device: str, cache_mode: str): + if kind == "tabicl": + from tabctx.backends.tabicl import TabICLBackend + + kv: bool | str = False if cache_mode == "off" else cache_mode + return TabICLBackend(device=device, kv_cache=kv) + from tabctx.backends.tabpfn import TabPFNBackend + + return TabPFNBackend(device=device, cache_mode=cache_mode) + + +def stage_modules(kind: str, backend, payload, task: str) -> dict[str, torch.nn.Module]: + """The model's top-level stages, from a fitted payload.""" + if kind == "tabicl": + model = backend._shared_backbones[task]["model_"] + return { + name: getattr(model, name) + for name in ("col_embedder", "row_interactor", "icl_predictor") + if hasattr(model, name) + } + model = None + for attr in ("models_", "model_"): + found = getattr(payload, attr, None) + if isinstance(found, list | tuple) and found: + model = found[0] + elif isinstance(found, torch.nn.Module): + model = found + if model is not None: + break + if model is None: + executor = getattr(payload, "executor_", None) + models = getattr(executor, "models", None) or getattr(executor, "model", None) + model = models[0] if isinstance(models, list | tuple) else models + if not isinstance(model, torch.nn.Module): + return {} + return dict(model.named_children()) + + +class StageRecorder: + """CUDA-event timing and peak memory per stage via forward hooks.""" + + def __init__(self, device: str, modules: dict[str, torch.nn.Module]) -> None: + self.device = torch.device(device) + self.calls: dict[str, int] = defaultdict(int) + self.time_ms: dict[str, float] = defaultdict(float) + self.peak_bytes: dict[str, int] = defaultdict(int) + self._open: dict[str, tuple] = {} + self._handles = [] + for name, module in modules.items(): + self._handles.append(module.register_forward_pre_hook(self._pre(name))) + self._handles.append(module.register_forward_hook(self._post(name))) + + def _pre(self, name): + def hook(module, args): + torch.cuda.synchronize(self.device) + base = torch.cuda.memory_allocated(self.device) + torch.cuda.reset_peak_memory_stats(self.device) + start = torch.cuda.Event(enable_timing=True) + start.record() + self._open[name] = (start, base) + + return hook + + def _post(self, name): + def hook(module, args, output): + start, base = self._open.pop(name) + end = torch.cuda.Event(enable_timing=True) + end.record() + torch.cuda.synchronize(self.device) + self.calls[name] += 1 + self.time_ms[name] += start.elapsed_time(end) + peak = torch.cuda.max_memory_allocated(self.device) - base + self.peak_bytes[name] = max(self.peak_bytes[name], peak) + + return hook + + def reset(self) -> None: + self.calls.clear() + self.time_ms.clear() + self.peak_bytes.clear() + + def summary(self) -> dict: + return { + name: { + "calls": self.calls[name], + "cuda_ms": round(self.time_ms[name], 1), + "peak_mib": round(self.peak_bytes[name] / MiB, 1), + } + for name in self.calls + } + + def close(self) -> None: + for h in self._handles: + h.remove() + + +def timed(device: torch.device, fn): + torch.cuda.synchronize(device) + t0 = time.perf_counter() + out = fn() + torch.cuda.synchronize(device) + return out, time.perf_counter() - t0 + + +def profile_shape(args, backend, rows, features, recorder_box) -> dict: + device = torch.device(args.device) + X, y = make_table(rows + args.n_test, features, args.task, seed=rows * 7 + features) + X_train, y_train, X_test = X[:rows], y[:rows], X[rows:] + record: dict = { + "backend": args.backend, + "task": args.task, + "cache_mode": args.cache_mode, + "rows": rows, + "features": features, + "n_test": args.n_test, + "gpu": torch.cuda.get_device_name(device), + } + try: + recorder = recorder_box.get("recorder") + if recorder is not None: + recorder.reset() + payload, fit_s = timed(device, lambda: backend.fit(X_train, y_train, args.task)) + record["fit_s"] = round(fit_s, 3) + record["fit_peak_mib"] = round((backend.fit_peak_bytes_hint() or 0) / MiB, 1) + record["resident_mib"] = round( + (backend.context_bytes_hint(rows, features) or 0) / MiB, 1 + ) + if recorder is None and args.stages: + recorder = StageRecorder( + args.device, stage_modules(args.backend, backend, payload, args.task) + ) + recorder_box["recorder"] = recorder + # First fit loaded the weights; re-fit so stages are measured. + payload = None + torch.cuda.empty_cache() + payload, fit_s = timed( + device, lambda: backend.fit(X_train, y_train, args.task) + ) + record["fit_s"] = round(fit_s, 3) + record["fit_peak_mib"] = round( + (backend.fit_peak_bytes_hint() or 0) / MiB, 1 + ) + if recorder is not None: + record["fit_stages"] = recorder.summary() + recorder.reset() + base = torch.cuda.memory_allocated(device) + torch.cuda.reset_peak_memory_stats(device) + _, predict_s = timed(device, lambda: backend.predict(payload, X_test, False)) + record["predict_s"] = round(predict_s, 3) + record["predict_peak_mib"] = round( + (torch.cuda.max_memory_allocated(device) - base) / MiB, 1 + ) + _, warm_s = timed(device, lambda: backend.predict(payload, X_test[:10], False)) + record["predict10_s"] = round(warm_s, 4) + if recorder is not None: + record["predict_stages"] = recorder.summary() + recorder.reset() + record["status"] = "ok" + payload = None + except Exception as e: # noqa: BLE001 -- OOM and model limits are results + record["status"] = type(e).__name__ + record["error"] = str(e)[:300] + torch.cuda.empty_cache() + return record + + +def trace_shape(args, backend, rows, features) -> None: + from torch.profiler import ProfilerActivity, profile + + out_dir = Path(args.trace_dir) + out_dir.mkdir(parents=True, exist_ok=True) + X, y = make_table(rows + args.n_test, features, args.task, seed=1) + backend.fit(X[:rows], y[:rows], args.task) # warm + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + profile_memory=True, + record_shapes=True, + ) as prof: + payload = backend.fit(X[:rows], y[:rows], args.task) + backend.predict(payload, X[rows:], False) + stem = f"{args.backend}-{args.cache_mode}-{rows}x{features}" + prof.export_chrome_trace(str(out_dir / f"{stem}.trace.json")) + averages = prof.key_averages() + (out_dir / f"{stem}.top-cuda.txt").write_text( + averages.table(sort_by="cuda_time_total", row_limit=30) + ) + (out_dir / f"{stem}.top-mem.txt").write_text( + averages.table(sort_by="self_cuda_memory_usage", row_limit=30) + ) + print(f"trace written: {out_dir / stem}.*", file=sys.stderr) + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--backend", choices=["tabicl", "tabpfn"], required=True) + p.add_argument("--device", default="cuda:0") + p.add_argument( + "--task", choices=["classification", "regression"], default="classification" + ) + p.add_argument("--cache-mode", choices=["kv", "repr", "off"], default="kv") + p.add_argument("--shapes", nargs="+", type=parse_shape, required=True) + p.add_argument("--n-test", type=int, default=1000) + p.add_argument("--no-stages", dest="stages", action="store_false") + p.add_argument("--out", default=None, help="append JSON lines here") + p.add_argument("--trace-shape", type=parse_shape, default=None) + p.add_argument("--trace-dir", default="traces") + args = p.parse_args() + + backend = build_backend(args.backend, args.device, args.cache_mode) + warm_X, warm_y = make_table(200, 8, args.task, seed=0) + backend.fit(warm_X, warm_y, args.task) # load weights outside measurements + recorder_box: dict = {} + out = open(args.out, "a") if args.out else None # noqa: SIM115 + for rows, features in args.shapes: + record = profile_shape(args, backend, rows, features, recorder_box) + line = json.dumps(record) + print(line, flush=True) + if out: + out.write(line + "\n") + out.flush() + if args.trace_shape: + trace_shape(args, backend, *args.trace_shape) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From b4bfa30d080021481f830e7d971c38eaf7a22082 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 15:09:40 -0400 Subject: [PATCH 5/7] fix(multidevice): review findings on #38 - Routing reads lock-free snapshots (residency, spill index, used bytes): fits/lookups/stats no longer wait on another GPU work (fits were effectively serialized across devices). - Re-fit lands on its device before other copies are removed: a failed re-fit (413/OOM/unknown model) no longer destroys the old context. - Per-dataset_id locks: concurrent same-id fits leave one copy. - Owner map removed (stale entries, 404 during re-fit, unbounded growth). - Unknown model -> 422 before anything changes; activity pruned on every call; TABCTX_DEVICES rejects mixed/ambiguous/out-of-range lists. - Spill: capacity is the replica total split per device; layout-change warning; fix warm-restart index keyed by hex filename (pre-existing). - TabICL spill restore no longer reloads the checkpoint (pre-existing). - /readyz + legacy endpoint read the replica own devices; flat real_gpu_memory kept, per_device added; /limits per-device headroom. - Benchmarks: CSV encoding outside the timed fit, per-user table slots, guarded predict/ready, import works from any cwd. - Deterministic gate/barrier concurrency tests (7 fail on the old code). Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 7 +- README.md | 18 ++- benchmarks/bench_workload.py | 34 +++-- benchmarks/sweep_admission.py | 33 +++-- src/tabctx/backends/tabicl.py | 75 +++++++++-- src/tabctx/cache/manager.py | 28 +++++ src/tabctx/cache/spill.py | 14 ++- src/tabctx/engine.py | 22 +++- src/tabctx/multidevice.py | 159 ++++++++++++++++-------- src/tabctx/serve/app.py | 50 +++++--- src/tabctx/serve/factory.py | 94 +++++++++++--- tests/unit/test_multidevice.py | 207 +++++++++++++++++++++++++++---- tests/unit/test_serve_factory.py | 42 ++++++- tests/unit/test_spill.py | 16 +++ 14 files changed, 649 insertions(+), 150 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 136fce9..e32c21e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,16 +4,19 @@ ## [0.12.0] - unreleased ### 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 owning device; spilled contexts are found on their device after a restart; re-fitting an id drops every copy first. Work on different devices runs concurrently. `/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. +- **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. - `benchmarks/bench_workload.py`: concurrent users looping fit-then-predict over a small/medium/large table mix (large ones via CSV upload), with per-class latency percentiles, 413s, errors and the per-device spread. - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed). - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first. ### Fixed +- **Spill tier warm restart:** the index rebuilt at startup was keyed by the hex-encoded file name instead of the dataset_id, so lookups, explicit deletes and disk-pressure cleanup all missed contexts a previous process had spilled (restores still worked because `load()` reads by path). +- **TabICL spill restore reloaded the checkpoint:** tabicl's `__setstate__` calls `_load_model()`, which loaded the full checkpoint (and a second backbone onto the GPU, unaccounted) on every restore before tabctx swapped in the shared one. Spilled TabICL payloads now carry a task tag and restore straight onto the shared backbone; untagged blobs from older versions still load. +- `/readyz` and the legacy endpoint read GPU memory from torch's current device; they now read the replica's own device(s). - The TabICL and TabPFN backends measured memory on, and let library internals allocate on, torch's *current* CUDA device rather than their own. Harmless on one GPU; on `cuda:i` it would have charged every context to `cuda:0`. They now measure on their device and run inside its scope. ### Validated -- Unit (multi-device placement, routing, spill restore, concurrency) and integration tests (the serve app on two fake devices through Ray). +- Unit (multi-device placement, routing, spill restore, failures) and integration tests (the serve app on two fake devices through Ray). Concurrency tests force interleavings with gates and barriers rather than timing: a fit on an idle device completes while another device is mid-call, lookups and stats never wait on a busy device, concurrent same-id fits leave one copy, a predict during a re-fit is answered. Seven of them fail against the first version of this engine (found by review). - On a local fake 3-device replica with 5 concurrent users: with spill on, 0 failed predicts at ~100 predicts/s. Without spill, capacity eviction cost 12 predicts on three devices versus 115 on one. ## [0.11.0] - 2026-10-17 diff --git a/README.md b/README.md index 8586a35..99ab65e 100644 --- a/README.md +++ b/README.md @@ -217,8 +217,24 @@ behind the same API. No client changes and no routing header needed. - **Limits.** Neither TabICL nor TabPFN can split one context across GPUs, so a table's ceiling is one GPU's memory. What more GPUs buy is N x the resident cache and N fits/predicts in flight at once. +- **Concurrency.** Routing reads lock-free snapshots, so it never waits + on a GPU's work: fits and predicts on different GPUs run at the same + time. Fits and evicts of the *same* dataset_id are serialized, so a + retried fit can't leave two copies. A re-fit lands first and only then + replaces the old copy: a failed re-fit (413, OOM) leaves the previous + context serving. +- **The replica must be given the GPUs.** Under Ray Serve that means + `ray_actor_options.num_gpus: N`, or, for a `num_gpus: 0` actor on a + GPU pod (as in inference-platform's prod config), + `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1`. +- **Spill.** Each GPU spills to its own subdirectory (`cuda_0/`, ...); + `TABCTX_SPILL_CAPACITY_BYTES` is the replica's total, split evenly. + Changing between one and several GPUs changes the layout, and the + replica logs a warning about spilled files it won't restore. - `/readyz` reports per-device contexts, bytes, activity and real GPU - memory; `/v1/tabctx/limits` reports the best single GPU's headroom. + memory (`real_gpu_memory` keeps its flat fields, summed, plus + `per_device`); `/v1/tabctx/limits` reports the best single GPU's + headroom and each device's. ## Large tables: upload, then fit by reference diff --git a/benchmarks/bench_workload.py b/benchmarks/bench_workload.py index 94f2ea9..512b078 100644 --- a/benchmarks/bench_workload.py +++ b/benchmarks/bench_workload.py @@ -73,13 +73,22 @@ def to_csv(X: np.ndarray, y: list[str]) -> bytes: return ("\n".join(lines) + "\n").encode() -def fit_table( - client: TabctxClient, X: np.ndarray, y: list[str], dataset_id: str, model -): +def encode_table(X: np.ndarray, y: list[str]): + """Client-side encoding (CSV or JSON lists), done BEFORE timing a fit: + building a large CSV is seconds of GIL-bound work that is not the + server's.""" if X.size >= UPLOAD_CELLS: - upload_id = client.upload_csv(to_csv(X, y), dataset_id) + return ("csv", to_csv(X, y)) + return ("json", (X.tolist(), y)) + + +def fit_table(client: TabctxClient, encoded, dataset_id: str, model): + kind, data = encoded + if kind == "csv": + upload_id = client.upload_csv(data, dataset_id) return client.fit_uploaded(upload_id, dataset_id, model=model) - return client.fit(X.tolist(), y, dataset_id=dataset_id, model=model) + X, y = data + return client.fit(X, y, dataset_id=dataset_id, model=model) def pct(values: list[float], p: float) -> float | None: @@ -103,11 +112,16 @@ def user_loop(user: int, args, mix, stats: dict[str, ClassStats], lock, stop_at) while time.monotonic() < stop_at: name, _, rows, features = pick.choices(mix, weights=[m[1] for m in mix])[0] X, y = make_table(rows, features, rng) - dataset_id = f"wl-{user}-{n}" + encoded = encode_table(X, y) + # A user works with a few tables at a time: re-fitting a slot + # replaces its previous table, so the cache holds about + # users x tables_per_user contexts -- "fit once, predict many", + # not an ever-growing pile of abandoned contexts. + dataset_id = f"wl-{user}-{n % args.tables_per_user}" n += 1 t0 = time.monotonic() try: - fit_table(client, X, y, dataset_id, args.model) + fit_table(client, encoded, dataset_id, args.model) except AdmissionRejected: with lock: stats[name].rejected += 1 @@ -139,6 +153,7 @@ def main() -> int: p.add_argument("--users", type=int, default=6) p.add_argument("--duration", type=float, default=120) p.add_argument("--predicts", type=int, default=5, help="per fitted table") + p.add_argument("--tables-per-user", type=int, default=2) p.add_argument("--test-rows", type=int, default=200) p.add_argument("--model", default=None) p.add_argument("--tenant", default="bench") @@ -168,7 +183,10 @@ def main() -> int: t.join() elapsed = time.monotonic() - started - ready = TabctxClient(args.base_url, timeout_s=args.timeout).ready() + try: + ready = TabctxClient(args.base_url, timeout_s=args.timeout).ready() + except (TabctxError, OSError) as e: + ready = {"error": f"{type(e).__name__}: {e}"} total_predicts = sum(len(s.predict_s) for s in stats.values()) result = { "timestamp": datetime.now(UTC).isoformat(), diff --git a/benchmarks/sweep_admission.py b/benchmarks/sweep_admission.py index cf8206a..8980b56 100644 --- a/benchmarks/sweep_admission.py +++ b/benchmarks/sweep_admission.py @@ -22,12 +22,20 @@ import sys import time from datetime import UTC, datetime +from pathlib import Path import numpy as np -from bench_workload import UPLOAD_CELLS, fit_table, make_table -from tabctx.client import TabctxClient -from tabctx.errors import AdmissionRejected, TabctxError +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from bench_workload import ( # noqa: E402 + UPLOAD_CELLS, + encode_table, + fit_table, + make_table, +) + +from tabctx.client import TabctxClient # noqa: E402 +from tabctx.errors import AdmissionRejected, TabctxError # noqa: E402 ROWS = [ 100, 250, 500, 1_000, 2_000, 3_500, 5_000, 7_500, 10_000, 15_000, 20_000, @@ -65,9 +73,10 @@ def main() -> int: "n_rows": rows, "upload": X.size >= UPLOAD_CELLS, } + encoded = encode_table(X, y) t0 = time.monotonic() try: - fit_table(client, X, y, dataset_id, args.model) + fit_table(client, encoded, dataset_id, args.model) except AdmissionRejected as e: point.update( status=413, @@ -85,12 +94,20 @@ def main() -> int: point["fit_s"] = round(time.monotonic() - t0, 3) X_test, _ = make_table(args.test_rows, features, rng) t0 = time.monotonic() - client.predict(dataset_id, X_test.tolist()) - point.update(status=200, predict_s=round(time.monotonic() - t0, 3)) + try: + client.predict(dataset_id, X_test.tolist()) + point.update(status=200, predict_s=round(time.monotonic() - t0, 3)) + except (TabctxError, OSError) as e: + point.update( + status="predict_error", detail=f"{type(e).__name__}: {e}"[:300] + ) points.append(point) print(json.dumps(point), flush=True) - limits = client.limits() + try: + limits = client.limits() + except (TabctxError, OSError) as e: + limits = {"error": f"{type(e).__name__}: {e}"} result = { "timestamp": datetime.now(UTC).isoformat(), "base_url": args.base_url, @@ -113,7 +130,7 @@ def main() -> int: if args.save: with open(args.save, "w") as f: json.dump(result, f, indent=2) - return 1 if any(q["status"] == "error" for q in points) else 0 + return 1 if any(str(q["status"]).endswith("error") for q in points) else 0 if __name__ == "__main__": diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py index b1eb4be..30d9187 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -49,6 +49,7 @@ from __future__ import annotations +import threading from typing import Any import numpy as np @@ -57,6 +58,32 @@ from tabctx.errors import BackendComputeError from tabctx.types import ArrayLike, PredictOutcome, Task +# Set (per thread) while loads_payload unpickles: tabicl's __setstate__ +# calls _load_model(), which would otherwise torch.load the full checkpoint +# and move a second backbone onto the GPU on every spill restore. +_restoring = threading.local() +_TASK_TAGS = {"classification": b"C", "regression": b"R"} + + +def _install_restore_hook(cls: type) -> None: + """Wrap cls._load_model: while a restore is in progress on this + thread, attach the backend's shared backbone instead of loading the + checkpoint. Idempotent; other threads and normal loads are untouched.""" + if getattr(cls, "_tabctx_restore_hook", False): + return + original = cls._load_model + + def _load_model(self: Any) -> None: + shared = getattr(_restoring, "backbone", None) + if shared is None: + return original(self) + for attr, value in shared.items(): + setattr(self, attr, value) + return None + + cls._load_model = _load_model + cls._tabctx_restore_hook = True + class TabICLBackend: # The exact model id callers select via the API's `model` field -- @@ -92,6 +119,10 @@ def __init__(self, device: str | None = None, kv_cache: bool | str = "kv") -> No # per-call mutable state (its _cache pointer, InferenceManager # config) is never touched concurrently. self._shared_backbones: dict[str, dict[str, Any]] = {} + from tabicl import TabICLClassifier, TabICLRegressor + + _install_restore_hook(TabICLClassifier) + _install_restore_hook(TabICLRegressor) def _make_estimator(self, task: Task) -> Any: if task == "classification": @@ -236,35 +267,59 @@ def dumps_payload(self, payload: Any) -> bytes: """Pickle a fitted estimator for the disk spill tier WITHOUT the shared pretrained backbone (hundreds of MB, identical across all contexts, unpicklable anyway once _load_model is shadowed with a - lambda): strip the shared attributes, pickle, restore.""" + lambda): strip the shared attributes, pickle, restore. A one-byte + task tag leads the blob so loads_payload can attach the right + backbone during unpickling.""" import pickle + task: Task = ( + "classification" if hasattr(payload, "predict_proba") else "regression" + ) stripped = {} for attr in ("model_", "_load_model"): if attr in payload.__dict__: stripped[attr] = payload.__dict__.pop(attr) try: - return pickle.dumps(payload) + return _TASK_TAGS[task] + pickle.dumps(payload) finally: payload.__dict__.update(stripped) + def _ensure_backbone(self, task: Task) -> None: + if task in self._shared_backbones: + return + # Cold restore on a fresh replica: borrow a throwaway estimator's + # load path to populate the shared backbone. + probe = self._make_estimator(task) + with device_scope(self._device): + probe._load_model() + self._stash_backbone(task, probe) + def loads_payload(self, data: bytes) -> Any: """Inverse of dumps_payload: unpickle and re-attach this process's shared backbone (loading it first if this replica has never fit this task type -- e.g. right after a restart).""" import pickle - model = pickle.loads(data) + tagged = {tag: task for task, tag in _TASK_TAGS.items()} + task_hint = tagged.get(data[:1]) + if task_hint is None: + # Untagged blob (spilled by tabctx < 0.12): tabicl's + # __setstate__ reloads the checkpoint; the shared backbone + # replaces it below. + with device_scope(self._device): + model = pickle.loads(data) + else: + self._ensure_backbone(task_hint) + _restoring.backbone = self._shared_backbones[task_hint] + try: + with device_scope(self._device): + model = pickle.loads(data[1:]) + finally: + _restoring.backbone = None task: Task = ( "classification" if hasattr(model, "predict_proba") else "regression" ) - if task not in self._shared_backbones: - # Cold restore on a fresh replica: borrow a throwaway - # estimator's load path to populate the shared backbone. - probe = self._make_estimator(task) - with device_scope(self._device): - probe._load_model() - self._stash_backbone(task, probe) + self._ensure_backbone(task) for attr, value in self._shared_backbones[task].items(): setattr(model, attr, value) model._load_model = lambda: None diff --git a/src/tabctx/cache/manager.py b/src/tabctx/cache/manager.py index 04df88c..de50c59 100644 --- a/src/tabctx/cache/manager.py +++ b/src/tabctx/cache/manager.py @@ -169,6 +169,34 @@ def evict(self, dataset_id: str) -> None: if self._spill is not None: self._spill.delete(dataset_id) + # ---- lock-free snapshots --------------------------------------- + # The engine holds this cache's lock across whole GPU calls, so a + # reader that only needs bookkeeping (a multi-device router choosing + # a device, /readyz) must not take it: it would wait for that GPU's + # work to finish. Single int / dict-membership reads are atomic under + # the GIL; these are snapshots, and every mutation stays locked. + + def holds(self, dataset_id: str) -> bool: + """Resident in the primary tier (snapshot, lock-free).""" + return dataset_id in self._entries + + def has(self, dataset_id: str) -> bool: + """Resident or spilled (snapshot, lock-free).""" + return dataset_id in self._entries or ( + self._spill is not None and self._spill.contains(dataset_id) + ) + + def stats_nowait(self) -> EngineStats: + """stats() without the lock: fields may be from slightly different + instants, which is fine for placement and monitoring.""" + used = self._used_bytes + return EngineStats( + n_cached_contexts=len(self._entries), + used_bytes=used, + free_bytes=max(0, self._capacity_bytes - used), + capacity_bytes=self._capacity_bytes, + ) + def dataset_ids(self) -> list[str]: """Ids currently resident in the primary tier (not spilled).""" with self._lock: diff --git a/src/tabctx/cache/spill.py b/src/tabctx/cache/spill.py index c8d32a2..48a6ba3 100644 --- a/src/tabctx/cache/spill.py +++ b/src/tabctx/cache/spill.py @@ -75,7 +75,13 @@ def __init__( # first to go under disk pressure. self._index: dict[str, float] = {} for meta_path in self._dir.glob("*.meta.json"): - self._index[meta_path.name[: -len(".meta.json")]] = 0.0 + # File names are hex-encoded dataset_ids (see _paths); index by + # the id itself, like entries this process spills. + stem = meta_path.name[: -len(".meta.json")] + try: + self._index[bytes.fromhex(stem).decode()] = 0.0 + except ValueError: + continue # not one of ours def _paths(self, dataset_id: str) -> tuple[Path, Path]: # dataset_ids may contain ':' (tenant scoping) -- hex-encode for @@ -148,6 +154,12 @@ def load(self, dataset_id: str) -> CachedContext | None: feature_names=meta.get("feature_names"), ) + def contains(self, dataset_id: str) -> bool: + """Whether a spilled copy exists. Lock-free snapshot (a dict + membership read is atomic), for routing decisions that must not + wait on another thread's spill I/O.""" + return dataset_id in self._index + def delete(self, dataset_id: str) -> None: meta_path, payload_path = self._paths(dataset_id) payload_path.unlink(missing_ok=True) diff --git a/src/tabctx/engine.py b/src/tabctx/engine.py index e21c42a..1186cf5 100644 --- a/src/tabctx/engine.py +++ b/src/tabctx/engine.py @@ -136,15 +136,27 @@ def estimator_for(self, backend: str | None = None) -> MemoryEstimator: def admission_headroom_bytes(self, backend: str | None = None) -> int: """Transient headroom a fit could use right now, given what the - cache holds resident (the quantity fit admission bounds).""" + cache holds resident (the quantity fit admission bounds). + Lock-free snapshot: never waits on this engine's GPU work.""" return self.estimator_for(backend).admission_headroom_bytes( - self._cache.stats().used_bytes + self._cache.stats_nowait().used_bytes ) def holds(self, dataset_id: str) -> bool: - """Whether the context is resident in this engine's cache (a - spilled copy doesn't count; get() would reload it).""" - return dataset_id in self._cache.dataset_ids() + """Resident in this engine's cache (lock-free snapshot; a + spilled copy doesn't count).""" + return self._cache.holds(dataset_id) + + def has(self, dataset_id: str) -> bool: + """Resident or spilled here (lock-free snapshot).""" + return self._cache.has(dataset_id) + + def resolve_backend_name(self, backend: str | None) -> str: + """The model a fit would use; InvalidInputError if unknown.""" + return self._resolve_backend(backend)[0] + + def stats_nowait(self) -> EngineStats: + return self._cache.stats_nowait() def _resolve_backend(self, backend: str | None) -> tuple[str, TabularICLBackend]: name = backend or self._default_backend diff --git a/src/tabctx/multidevice.py b/src/tabctx/multidevice.py index 72da8e9..33080cd 100644 --- a/src/tabctx/multidevice.py +++ b/src/tabctx/multidevice.py @@ -23,17 +23,25 @@ exactly as on a single GPU. A table's ceiling is therefore one device's memory; what more devices buy is N x the resident cache and concurrency. -- predict() runs on whichever device holds the context. Ownership is - tracked on fit; a context the map doesn't know (e.g. spilled before a - restart) is found by asking each device, which reloads it from that - device's own spill tier. -- Re-fitting a dataset_id drops it from every device first, so a stale - copy on another device can never answer a later predict. - -Placement is a snapshot taken outside the per-device locks: two -concurrent fits may pick the same device, in which case the second simply -evicts-ahead there like any single-GPU fit would. Serializing placement -would cost the concurrency this class exists for. +- predict() runs on the device that holds the context, found from each + device's residency and spill index (so a context spilled before a + restart is found on, and restored to, its own device). + +Concurrency. Each device's engine holds its cache lock across its whole +GPU call, so this router never takes another device's lock to make a +decision: placement and lookup read lock-free snapshots (used bytes, +residency, spill index). Work on different devices therefore runs +concurrently, fits included. Per-dataset_id locks serialize fits and +evicts of the SAME id, so two concurrent fits of one id end with a single +copy. A re-fit happens on the chosen device first; only after it succeeds +are copies on other devices removed -- a failed re-fit (413, OOM) leaves +the previous context serving, as on a single GPU. A predict racing a +re-fit is answered by whichever copy it reaches (old or new), never 404. + +Placement is a snapshot: two concurrent fits of different ids may pick +the same device, in which case the second evicts-ahead there like any +single-GPU fit would. Serializing placement would cost the concurrency +this class exists for. Same public surface as TabctxEngine, so the serving layer and batching use either interchangeably. @@ -56,6 +64,35 @@ ACTIVITY_WINDOW_S = 30.0 +class _KeyedLocks: + """One lock per key, created on demand and dropped when unused.""" + + def __init__(self) -> None: + self._guard = threading.Lock() + self._locks: dict[str, list] = {} # key -> [lock, users] + + @contextlib.contextmanager + def hold(self, key: str) -> Iterator[None]: + with self._guard: + entry = self._locks.setdefault(key, [threading.Lock(), 0]) + entry[1] += 1 + try: + with entry[0]: + yield + finally: + with self._guard: + entry[1] -= 1 + if entry[1] == 0: + del self._locks[key] + + +def _not_found(dataset_id: str) -> DatasetNotFoundError: + return DatasetNotFoundError( + f"no cached context for dataset_id={dataset_id!r} -- it was never " + "fit, was evicted, or was lost to a replica restart" + ) + + class MultiDeviceEngine: def __init__( self, @@ -72,8 +109,7 @@ def __init__( raise ValueError("every device must serve the same models") self._shards = shards self._devices = devices - self._owner: dict[str, int] = {} - self._owner_lock = threading.Lock() + self._id_locks = _KeyedLocks() self._window_s = activity_window_s self._clock = clock self._inflight = [0] * len(shards) @@ -97,7 +133,16 @@ def estimator_for(self, backend: str | None = None) -> MemoryEstimator: def admission_headroom_bytes(self, backend: str | None = None) -> int: """The best any one device offers: a fit lands on a single device.""" - return max(s.admission_headroom_bytes(backend) for s in self._shards) + return max(self.device_headroom_bytes(backend)) + + def device_headroom_bytes(self, backend: str | None = None) -> list[int]: + return [s.admission_headroom_bytes(backend) for s in self._shards] + + # ---- activity -------------------------------------------------- + + def _prune(self, recent: deque[float], cutoff: float) -> None: + while recent and recent[0] < cutoff: + recent.popleft() @contextlib.contextmanager def _active(self, index: int) -> Iterator[None]: @@ -105,6 +150,7 @@ def _active(self, index: int) -> Iterator[None]: with self._activity_lock: self._inflight[index] += 1 self._recent[index].append(now) + self._prune(self._recent[index], now - self._window_s) try: yield finally: @@ -116,12 +162,13 @@ def _activity(self) -> list[tuple[int, int]]: cutoff = self._clock() - self._window_s with self._activity_lock: for recent in self._recent: - while recent and recent[0] < cutoff: - recent.popleft() + self._prune(recent, cutoff) return [(self._inflight[i], len(r)) for i, r in enumerate(self._recent)] - def _place(self, estimated_bytes: int, backend: str | None) -> int: - headrooms = [s.admission_headroom_bytes(backend) for s in self._shards] + # ---- placement and lookup (lock-free snapshots) ---------------- + + def _place(self, estimated_bytes: int, backend: str) -> int: + headrooms = self.device_headroom_bytes(backend) activity = self._activity() indices = range(len(self._shards)) fits = [i for i in indices if headrooms[i] >= estimated_bytes] @@ -129,11 +176,17 @@ def _place(self, estimated_bytes: int, backend: str | None) -> int: return min(fits, key=lambda i: (activity[i], -headrooms[i], i)) return max(indices, key=lambda i: (headrooms[i], -i)) - def _drop_everywhere(self, dataset_id: str) -> None: - with self._owner_lock: - self._owner.pop(dataset_id, None) - for shard in self._shards: - shard.evict(dataset_id) + def _locate(self, dataset_id: str) -> list[int]: + """Devices holding the id: resident first, then spilled copies.""" + resident = [i for i, s in enumerate(self._shards) if s.holds(dataset_id)] + spilled = [ + i + for i, s in enumerate(self._shards) + if i not in resident and s.has(dataset_id) + ] + return resident + spilled + + # ---- API ------------------------------------------------------- def fit( self, @@ -145,12 +198,25 @@ def fit( feature_names: list[str] | None = None, ) -> str: n_features = _validate_fit_input(X, y) - if dataset_id is not None: - self._drop_everywhere(dataset_id) - estimated = self.estimator_for(backend).estimate_bytes(len(X), 0, n_features) - index = self._place(estimated, backend) + name = self._shards[0].resolve_backend_name(backend) + estimated = self.estimator_for(name).estimate_bytes(len(X), 0, n_features) + if dataset_id is None: + return self._fit_on( + self._place(estimated, name), X, y, task, None, name, feature_names + ) + with self._id_locks.hold(dataset_id): + index = self._place(estimated, name) + resolved = self._fit_on(index, X, y, task, dataset_id, name, feature_names) + # Only now that the new context is in place: drop copies on + # other devices (a failed fit above leaves the old one serving). + for i in self._locate(dataset_id): + if i != index: + self._shards[i].evict(dataset_id) + return resolved + + def _fit_on(self, index, X, y, task, dataset_id, backend, feature_names) -> str: with self._active(index): - resolved = self._shards[index].fit( + return self._shards[index].fit( X, y, task=task, @@ -158,41 +224,23 @@ def fit( backend=backend, feature_names=feature_names, ) - with self._owner_lock: - self._owner[resolved] = index - return resolved - - def _search_order(self, dataset_id: str) -> list[int]: - with self._owner_lock: - owner = self._owner.get(dataset_id) - if owner is not None: - return [owner] - resident = [i for i, s in enumerate(self._shards) if s.holds(dataset_id)] - return resident + [i for i in range(len(self._shards)) if i not in resident] def predict( self, dataset_id: str, X_test: ArrayLike, return_proba: bool = False ) -> PredictOutcome: - last_error: DatasetNotFoundError | None = None - for index in self._search_order(dataset_id): + for index in self._locate(dataset_id): try: with self._active(index): - outcome = self._shards[index].predict( + return self._shards[index].predict( dataset_id, X_test, return_proba=return_proba ) - except DatasetNotFoundError as e: - last_error = e + except DatasetNotFoundError: + # Moved or evicted between the snapshot and the call. continue - with self._owner_lock: - self._owner[dataset_id] = index - return outcome - with self._owner_lock: - self._owner.pop(dataset_id, None) - assert last_error is not None - raise last_error + raise _not_found(dataset_id) def feature_names(self, dataset_id: str) -> list[str] | None: - for index in self._search_order(dataset_id): + for index in self._locate(dataset_id): names = self._shards[index].feature_names(dataset_id) if names is not None: return names @@ -213,7 +261,9 @@ def fit_predict( self.evict(dataset_id) def evict(self, dataset_id: str) -> None: - self._drop_everywhere(dataset_id) + with self._id_locks.hold(dataset_id): + for index in self._locate(dataset_id): + self._shards[index].evict(dataset_id) def stats(self) -> EngineStats: per_device = self.device_stats() @@ -225,7 +275,8 @@ def stats(self) -> EngineStats: ) def device_stats(self) -> list[EngineStats]: - return [s.stats() for s in self._shards] + """Lock-free: monitoring never waits on a device's GPU work.""" + return [s.stats_nowait() for s in self._shards] def device_activity(self) -> list[dict[str, int]]: return [ diff --git a/src/tabctx/serve/app.py b/src/tabctx/serve/app.py index bf1e02b..982f66a 100644 --- a/src/tabctx/serve/app.py +++ b/src/tabctx/serve/app.py @@ -264,8 +264,9 @@ def _legacy_predict_sync(self, req: LegacyPredictRequest) -> LegacyPredictRespon n_train, n_test = len(req.train_X), len(req.test_X) n_features = len(req.train_X[0]) if n_train else 0 - if cuda: - torch.cuda.reset_peak_memory_stats() + gpus = self._cuda_devices() if cuda else [] + for d in gpus: + torch.cuda.reset_peak_memory_stats(d) start = time.monotonic() try: outcome = self._engine.fit_predict( @@ -286,10 +287,13 @@ def _legacy_predict_sync(self, req: LegacyPredictRequest) -> LegacyPredictRespon latency_ms = (time.monotonic() - start) * 1000 gpu_memory = None - if cuda: + if gpus: + # The fit ran on one of the devices; the peak over all of them + # is that device's (the others had no new work from this call). gpu_memory = GpuMemory( - allocated_mb=torch.cuda.max_memory_allocated() / 1e6, - reserved_mb=torch.cuda.max_memory_reserved() / 1e6, + allocated_mb=max(torch.cuda.max_memory_allocated(d) for d in gpus) + / 1e6, + reserved_mb=max(torch.cuda.max_memory_reserved(d) for d in gpus) / 1e6, ) logger.info( "[%s] legacy predict train=%d test=%d feats=%d latency_ms=%.1f gpu=%s", @@ -533,6 +537,13 @@ def limits(self): "gpu_memory_fraction": self._settings.gpu_memory_fraction, "gpu_capacity_bytes": self._gpu_capacity_bytes, "devices": list(self._devices), + # Scopes: cache_used_bytes is the replica's total; + # memory_ceiling_bytes and the headrooms are per device. + "device_admission_headroom_bytes": ( + self._engine.device_headroom_bytes() + if isinstance(self._engine, MultiDeviceEngine) + else [headroom] + ), "max_admissible_train_rows_by_feature_count": max_rows_by_features, "note": ( "Admission limits are per-replica, usage-aware (they shrink " @@ -559,19 +570,21 @@ def readyz(self): # failing to actually release GPU memory on eviction (Python # dropping a reference doesn't guarantee PyTorch's caching # allocator returns the memory immediately). + per_device = { + d: { + "allocated_mb": torch.cuda.memory_allocated(d) / 1e6, + "reserved_mb": torch.cuda.memory_reserved(d) / 1e6, + } + for d in self._cuda_devices() + } + # Same flat shape as ever (summed over devices when several), + # plus the per-device breakdown on a multi-GPU replica. real_gpu_memory = { - "allocated_mb": torch.cuda.memory_allocated() / 1e6, - "reserved_mb": torch.cuda.memory_reserved() / 1e6, + "allocated_mb": sum(m["allocated_mb"] for m in per_device.values()), + "reserved_mb": sum(m["reserved_mb"] for m in per_device.values()), } - if len(self._devices) > 1: - real_gpu_memory = { - d: { - "allocated_mb": torch.cuda.memory_allocated(d) / 1e6, - "reserved_mb": torch.cuda.memory_reserved(d) / 1e6, - } - for d in self._devices - if d.startswith("cuda") - } + if len(per_device) > 1: + real_gpu_memory["per_device"] = per_device return { "status": "ready", "device": self._device, @@ -597,6 +610,11 @@ def readyz(self): "devices": self._device_breakdown(), } + def _cuda_devices(self) -> list[str]: + """The CUDA devices this replica serves from (just "cuda", the + current device, on an auto single-GPU replica).""" + return [d for d in self._devices if d.startswith("cuda")] + def _spill_stats(self) -> dict | None: if self._spill_store is not None: return self._spill_store.stats() diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py index c3d4b42..ff474bc 100644 --- a/src/tabctx/serve/factory.py +++ b/src/tabctx/serve/factory.py @@ -36,19 +36,25 @@ - ``TABCTX_SPILL_DIR`` (default unset = spillover off) and ``TABCTX_SPILL_CAPACITY_BYTES`` (default 50GiB): disk spillover tier for capacity-evicted contexts (see cache/spill.py). With several - devices, each gets its own subdirectory and the capacity is per device. + devices, each gets its own subdirectory (``cuda_0/``, ...) and an equal + share of the capacity (the setting is the total for the replica). - ``TABCTX_DEVICES``: ``"auto"`` (default: every visible CUDA device, or the CPU when there is none) or a comma-separated list such as ``cuda:0,cuda:1``. Two or more devices build a MultiDeviceEngine (see - tabctx/multidevice.py): one engine per device, fits placed on the - device with the most headroom. ``TABCTX_GPU_MEMORY_FRACTION`` applies - to each device. + tabctx/multidevice.py): one engine per device; a fit goes where it fits + without eviction on the least active device, else to the device with + the most headroom. ``TABCTX_GPU_MEMORY_FRACTION`` applies to each + device. The replica only sees the GPUs its process is given: under Ray + Serve that is ``ray_actor_options.num_gpus`` (or, for a ``num_gpus: 0`` + actor on a GPU pod, ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1``). """ from __future__ import annotations +import logging import os from dataclasses import dataclass +from pathlib import Path from typing import Literal from tabctx.backends.base import TabularICLBackend @@ -78,6 +84,8 @@ UPLOAD_TTL_S_ENV_VAR = "TABCTX_UPLOAD_TTL_S" DEVICES_ENV_VAR = "TABCTX_DEVICES" +logger = logging.getLogger(__name__) + BackendKind = Literal["tabicl", "tabpfn", "fake"] KvCacheMode = Literal["kv", "repr", "off"] BACKEND_KINDS: tuple[BackendKind, ...] = ("tabicl", "tabpfn", "fake") @@ -189,6 +197,17 @@ def _parse_devices(raw: str) -> tuple[str, ...]: ) if len(set(devices)) != len(devices): raise ValueError(f"{DEVICES_ENV_VAR} lists a device twice") + if len(devices) > 1: + if "cpu" in devices: + raise ValueError( + f"{DEVICES_ENV_VAR}={raw!r}: several devices must all be GPUs " + "(a CPU engine would get a GPU budget and an equal share of users)" + ) + if "cuda" in devices: + raise ValueError( + f"{DEVICES_ENV_VAR}={raw!r}: name GPUs by index (cuda:0,cuda:1); " + "a bare 'cuda' may be the same GPU as one listed by index" + ) return devices @@ -255,9 +274,9 @@ def detect_gpu_capacity_bytes(devices: tuple[str, ...] = ("cuda",)) -> int | Non return None if not torch.cuda.is_available(): return None - cuda = [torch.device(d) for d in devices if d.startswith("cuda")] or [ - torch.device("cuda") - ] + cuda = [torch.device(d) for d in devices if d.startswith("cuda")] + if not cuda: + return None return min( int(torch.cuda.get_device_properties(d.index or 0).total_memory) for d in cuda ) @@ -272,11 +291,26 @@ def _visible_cuda_count() -> int: def resolve_devices(settings: ServeSettings) -> tuple[str, ...]: - """The devices to build one engine for each. Explicit settings win; - "auto" means every visible GPU. A single auto-detected GPU stays - "cuda" (the pre-multi-device layout); the fake backend never touches - a GPU, so auto means one CPU engine for it.""" + """The devices to build one engine for each. Explicit settings win + (GPU indices are checked against what this process can see); "auto" + means every visible GPU. A single auto-detected GPU stays "cuda" (the + pre-multi-device layout); the fake backend never touches a GPU, so + auto means one CPU engine for it.""" if settings.devices: + if set(settings.backends) != {"fake"}: + visible = _visible_cuda_count() + for d in settings.devices: + index = d.removeprefix("cuda:") + if ( + d == "cuda" + and visible == 0 + or index.isdigit() + and int(index) >= visible + ): + raise ValueError( + f"{DEVICES_ENV_VAR} names {d}, but this process sees " + f"{visible} GPU(s)" + ) return settings.devices if set(settings.backends) == {"fake"}: return ("cpu",) @@ -355,10 +389,35 @@ def _build_backend( return _build_real_backend(kind, settings, device) +def _warn_on_spill_layout(spill_dir: str, multi: bool) -> None: + """Spill files from the other layout (one device writes to the root, + several to per-device subdirectories) are invisible to this process: + never restored, counted or cleaned. Say so instead of silently + leaking disk.""" + root = Path(spill_dir) + if not root.is_dir(): + return + stray = ( + list(root.glob("*.meta.json")) + if multi + else [p for p in root.glob("cuda_*/*.meta.json")] + ) + if stray: + logger.warning( + "%d spilled context(s) under %s were written by a %s-device replica " + "and will not be restored by this one (layout changed); remove them " + "to reclaim disk", + len(stray), + spill_dir, + "single" if multi else "multi", + ) + + def _build_spill_store( settings: ServeSettings, backends: dict[str, TabularICLBackend], subdir: str | None, + capacity_bytes: int, ) -> DiskSpillStore | None: if not settings.spill_dir: return None @@ -372,11 +431,7 @@ def _build_spill_store( if hasattr(b, "dumps_payload") } path = os.path.join(settings.spill_dir, subdir) if subdir else settings.spill_dir - return DiskSpillStore( - path, - capacity_bytes=settings.spill_capacity_bytes, - serializers=serializers, - ) + return DiskSpillStore(path, capacity_bytes=capacity_bytes, serializers=serializers) def build_engine(settings: ServeSettings | None = None) -> BuiltEngine: @@ -390,6 +445,8 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine: else detect_gpu_capacity_bytes(devices) ) multi = len(devices) > 1 + if settings.spill_dir: + _warn_on_spill_layout(settings.spill_dir, multi) estimators: dict[str, MemoryEstimator] = {} shards: list[TabctxEngine] = [] spill_stores: list[DiskSpillStore | None] = [] @@ -409,7 +466,10 @@ def build_engine(settings: ServeSettings | None = None) -> BuiltEngine: # first-listed kind's model is the default. default = next(iter(backends)) spill_store = _build_spill_store( - settings, backends, device.replace(":", "_") if multi else None + settings, + backends, + device.replace(":", "_") if multi else None, + settings.spill_capacity_bytes // len(devices), ) # ONE cache and budget shared by every model on a device -- the # whole point of serving them behind one endpoint (see engine.py). diff --git a/tests/unit/test_multidevice.py b/tests/unit/test_multidevice.py index e80e259..775e609 100644 --- a/tests/unit/test_multidevice.py +++ b/tests/unit/test_multidevice.py @@ -6,7 +6,6 @@ """ import threading -import time import pytest @@ -329,40 +328,194 @@ def test_headroom_is_the_best_single_device(self): assert engine.admission_headroom_bytes() == CAPACITY +class GatedFake(SizedFake): + """Blocks inside fit/predict until released, to force interleavings + deterministically (no wall-clock assertions).""" + + def __init__(self, device: str, **kw) -> None: + super().__init__(device, **kw) + self.fit_gate: threading.Event | None = None + self.predict_gate: threading.Event | None = None + self.entered = threading.Event() + self.barrier: threading.Barrier | None = None + + def fit(self, X, y, task): + if self.fit_gate is not None: + self.entered.set() + assert self.fit_gate.wait(5), "fit gate never released" + return super().fit(X, y, task) + + def predict(self, payload, X_test, return_proba=False): + if self.barrier is not None: + self.barrier.wait(5) # BrokenBarrierError if not concurrent + if self.predict_gate is not None: + self.entered.set() + assert self.predict_gate.wait(5), "predict gate never released" + return super().predict(payload, X_test, return_proba) + + +def build_gated(n_devices=2, spill_root=None): + devices = [f"cuda:{i}" for i in range(n_devices)] + estimator = CellEstimator() + shards, fakes = [], [] + for device in devices: + fake = GatedFake(device, name="m") + spill = ( + DiskSpillStore(str(spill_root / device.replace(":", "_")), 10**9) + if spill_root + else None + ) + shards.append( + TabctxEngine( + backends={"m": fake}, + estimators={"m": estimator}, + cache=ContextCacheManager(CAPACITY, spill_store=spill), + default_backend="m", + ) + ) + fakes.append(fake) + return MultiDeviceEngine(shards, devices, clock=FakeClock()), shards, fakes + + +def start(fn, *args, **kw): + box: dict = {} + + def run(): + try: + box["result"] = fn(*args, **kw) + except Exception as e: # noqa: BLE001 -- surfaced by the caller + box["error"] = e + + t = threading.Thread(target=run, daemon=True) + t.start() + return t, box + + class TestConcurrency: - def test_work_on_different_devices_runs_in_parallel(self): - engine, shards, fakes = build(2, predict_delay_s=0.3) + def test_predicts_on_different_devices_overlap(self): + engine, shards, fakes = build_gated(2) engine.fit(*table(10), dataset_id="a") engine.fit(*table(10), dataset_id="b") assert where(shards, "a") != where(shards, "b") - threads = [ - threading.Thread(target=engine.predict, args=(d, [[0.0]])) for d in "ab" - ] - start = time.monotonic() - for t in threads: - t.start() - for t in threads: - t.join() - assert time.monotonic() - start < 0.5 # ~0.3s, not ~0.6s serialized + barrier = threading.Barrier(2) + for f in fakes: + f.barrier = barrier # both predicts must be inside at once + threads = [start(engine.predict, d, [[0.0]]) for d in "ab"] + for t, box in threads: + t.join(10) + assert "error" not in box, box.get("error") + + def test_fit_on_an_idle_device_is_not_blocked_by_a_busy_one(self): + engine, shards, fakes = build_gated(2) + engine.fit(*table(10), dataset_id="busy") + busy = where(shards, "busy")[0] + fakes[busy].predict_gate = threading.Event() + t_busy, _ = start(engine.predict, "busy", [[0.0]]) + assert fakes[busy].entered.wait(5) + # Device `busy` is mid-GPU-call holding its lock. A fit (with an + # id, as the serve layer always passes) must still complete. + t_fit, box = start(engine.fit, *table(10), dataset_id="new") + t_fit.join(3) + assert not t_fit.is_alive(), "fit waited on another device's GPU call" + assert box["result"] == "new" and where(shards, "new") == [1 - busy] + fakes[busy].predict_gate.set() + t_busy.join(5) + + def test_lookups_and_stats_never_wait_on_a_busy_device(self): + engine, shards, fakes = build_gated(2) + engine.fit(*table(10), dataset_id="busy") + busy = where(shards, "busy")[0] + fakes[busy].predict_gate = threading.Event() + t_busy, _ = start(engine.predict, "busy", [[0.0]]) + assert fakes[busy].entered.wait(5) + for call in ( + lambda: engine.predict("missing", [[0.0]]), + engine.stats, + engine.device_stats, + engine.admission_headroom_bytes, + lambda: engine.feature_names("missing"), + ): + t, box = start(call) + t.join(3) + assert not t.is_alive(), f"{call} waited on a busy device" + assert isinstance(box.get("error"), DatasetNotFoundError | type(None)) + fakes[busy].predict_gate.set() + t_busy.join(5) + + def test_concurrent_fits_of_one_id_leave_a_single_copy(self): + engine, shards, fakes = build_gated(2) + for f in fakes: + f.fit_gate = threading.Event() + t_a, _ = start(engine.fit, *table(10, label="A"), dataset_id="same") + assert any(f.entered.wait(0.5) for f in fakes) + t_b, _ = start(engine.fit, *table(10, label="B"), dataset_id="same") + for f in fakes: + f.fit_gate.set() + for t in (t_a, t_b): + t.join(5) + assert len(where(shards, "same")) == 1 + assert engine.stats().n_cached_contexts == 1 + assert engine.predict("same", [[0.0]]).predictions[0] in ("A", "B") + + def test_predict_during_a_refit_is_answered_not_404(self): + clock = FakeClock() + engine, shards, fakes = build_gated(2) + engine._clock = clock + engine.fit(*table(10, label="old"), dataset_id="d") + first = where(shards, "d")[0] + engine._recent[first].extend([clock.now] * 10) # re-fit goes elsewhere + other = 1 - first + fakes[other].fit_gate = threading.Event() + t_fit, _ = start(engine.fit, *table(10, label="new"), dataset_id="d") + assert fakes[other].entered.wait(5) + assert engine.predict("d", [[0.0]]).predictions == ["old"] + fakes[other].fit_gate.set() + t_fit.join(5) + assert where(shards, "d") == [other] + assert engine.predict("d", [[0.0]]).predictions == ["new"] def test_concurrent_fits_spread_and_all_succeed(self): - engine, shards, _ = build(4, fit_delay_s=0.05) - errors = [] - - def fit(user): - try: - engine.fit(*table(20), dataset_id=f"u{user}") - except Exception as e: # pragma: no cover - surfaced below - errors.append(e) - - threads = [threading.Thread(target=fit, args=(u,)) for u in range(16)] - for t in threads: - t.start() - for t in threads: - t.join() - assert not errors + engine, shards, _ = build(4, fit_delay_s=0.02) + threads = [start(engine.fit, *table(20), dataset_id=f"u{u}") for u in range(16)] + for t, box in threads: + t.join(10) + assert "error" not in box, box.get("error") counts = [s.stats().n_cached_contexts for s in shards] assert sum(counts) == 16 assert min(counts) >= 2 # spread, not piled on one device for u in range(16): assert engine.predict(f"u{u}", [[0.0]]).predictions == ["a"] + + +class TestFailures: + def test_failed_refit_keeps_the_previous_context(self): + engine, shards, _ = build(2) + engine.fit(*table(10, label="kept"), dataset_id="d") + with pytest.raises(AdmissionRejected): + engine.fit(*table(2000), dataset_id="d") + assert engine.predict("d", [[0.0]]).predictions == ["kept"] + + def test_failed_refit_keeps_a_spilled_previous_context(self, tmp_path): + engine, shards, _ = build(2, spill_root=tmp_path) + engine.fit(*table(10, label="kept"), dataset_id="d") + owner = where(shards, "d")[0] + shards[owner]._cache.make_room(CAPACITY) # spill it + with pytest.raises(AdmissionRejected): + engine.fit(*table(2000), dataset_id="d") + assert engine.predict("d", [[0.0]]).predictions == ["kept"] + + def test_unknown_model_is_invalid_input_and_loses_nothing(self): + engine, _, _ = build(2) + engine.fit(*table(10, label="kept"), dataset_id="d") + with pytest.raises(InvalidInputError, match="unknown backend"): + engine.fit(*table(10), dataset_id="d", backend="xgboost") + assert engine.predict("d", [[0.0]]).predictions == ["kept"] + + def test_activity_history_stays_bounded_without_fits(self): + clock = FakeClock() + engine, _, _ = build(2, clock=clock) + engine.fit(*table(10), dataset_id="d") + for _ in range(500): + clock.now += 1 + engine.predict("d", [[0.0]]) + assert max(len(r) for r in engine._recent) <= 31 # one 30s window diff --git a/tests/unit/test_serve_factory.py b/tests/unit/test_serve_factory.py index f69a09b..2417d23 100644 --- a/tests/unit/test_serve_factory.py +++ b/tests/unit/test_serve_factory.py @@ -176,6 +176,7 @@ def test_build_engine_uses_detected_capacity_for_real_backends(self, monkeypatch "tabctx.serve.factory.detect_gpu_capacity_bytes", lambda devices: 24 * 1024**3, ) + monkeypatch.setattr("tabctx.serve.factory._visible_cuda_count", lambda: 1) built = build_engine(ServeSettings(backends=("tabicl",), devices=("cuda",))) assert built.gpu_capacity_bytes == 24 * 1024**3 assert built.estimator._fallback.gpu_capacity_bytes == 24 * 1024**3 @@ -334,7 +335,10 @@ def test_explicit_list_is_parsed(self, monkeypatch): monkeypatch.setenv(DEVICES_ENV_VAR, "cuda:0, cuda:2") assert ServeSettings.from_env().devices == ("cuda:0", "cuda:2") - @pytest.mark.parametrize("raw", ["gpu0", "cuda:x", "cuda:0,cuda:0", "tpu"]) + @pytest.mark.parametrize( + "raw", + ["gpu0", "cuda:x", "cuda:0,cuda:0", "tpu", "cuda,cuda:0", "cpu,cuda:0"], + ) def test_bad_lists_are_rejected(self, monkeypatch, raw): monkeypatch.setenv(DEVICES_ENV_VAR, raw) with pytest.raises(ValueError, match=DEVICES_ENV_VAR): @@ -384,11 +388,13 @@ def test_several_devices_build_one_engine_each(self, monkeypatch, tmp_path): "tabctx.serve.factory.detect_gpu_capacity_bytes", lambda devices: 48 * 1024**3, ) + monkeypatch.setattr("tabctx.serve.factory._visible_cuda_count", lambda: 4) built = build_engine( ServeSettings( backends=("tabicl", "tabpfn"), devices=("cuda:0", "cuda:1", "cuda:2"), spill_dir=str(tmp_path), + spill_capacity_bytes=30 * 1024**3, ) ) assert isinstance(built.engine, MultiDeviceEngine) @@ -405,9 +411,43 @@ def test_several_devices_build_one_engine_each(self, monkeypatch, tmp_path): caps = [s.capacity_bytes for s in built.engine.device_stats()] assert len(set(caps)) == 1 and caps[0] == built.estimator.ceiling_bytes() assert built.spill_store is None and len(built.spill_stores) == 3 + # The spill capacity setting is the replica's total, split evenly. + assert {s.stats()["spill_capacity_bytes"] for s in built.spill_stores} == { + 10 * 1024**3 + } assert sorted(p.name for p in tmp_path.iterdir()) == [ "cuda_0", "cuda_1", "cuda_2", ] assert built.device == "cuda:0,cuda:1,cuda:2" + + +class TestDeviceValidation: + def test_gpu_index_beyond_what_the_process_sees_is_rejected(self, monkeypatch): + from tabctx.serve.factory import resolve_devices + + monkeypatch.setattr("tabctx.serve.factory._visible_cuda_count", lambda: 2) + with pytest.raises(ValueError, match="sees 2 GPU"): + resolve_devices( + ServeSettings(backends=("tabicl",), devices=("cuda:0", "cuda:3")) + ) + + def test_explicit_cpu_never_gets_a_gpu_budget(self, monkeypatch): + from tabctx.serve.factory import detect_gpu_capacity_bytes + + torch = pytest.importorskip("torch") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + assert detect_gpu_capacity_bytes(("cpu",)) is None + + def test_spill_layout_change_is_reported(self, tmp_path, caplog): + (tmp_path / "abcd.meta.json").write_text("{}") + with caplog.at_level("WARNING"): + build_engine( + ServeSettings( + backends=("fake",), + devices=("cuda:0", "cuda:1"), + spill_dir=str(tmp_path), + ) + ) + assert "will not be restored" in caplog.text diff --git a/tests/unit/test_spill.py b/tests/unit/test_spill.py index 602b905..2905083 100644 --- a/tests/unit/test_spill.py +++ b/tests/unit/test_spill.py @@ -185,3 +185,19 @@ def test_end_to_end_predictions_survive_eviction(self, tmp_path): outcome = engine.predict("ds-a", [[9.0, 9.0]]) assert outcome.predictions == ["hot"] assert backend.fit_calls == fits_before, "restore must not re-fit" + + +def test_warm_restart_indexes_by_dataset_id(tmp_path): + """Regression: the restart index was keyed by the hex file stem, so + contains()/delete()/disk-pressure cleanup missed every entry a + predecessor process spilled.""" + from tabctx.cache.manager import CachedContext + from tabctx.cache.spill import DiskSpillStore + + first = DiskSpillStore(tmp_path, capacity_bytes=10**6) + ctx = CachedContext("t:alice:ds-1", "fake", "classification", 3, 2, {"p": 1}, 10) + assert first.spill(ctx) + restarted = DiskSpillStore(tmp_path, capacity_bytes=10**6) + assert restarted.contains("t:alice:ds-1") + restarted.delete("t:alice:ds-1") + assert list(tmp_path.iterdir()) == [] From 81e7797045b64df5121bf6cf485a5590adaf82ad Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 15:15:23 -0400 Subject: [PATCH 6/7] fix(backends): initialize CUDA before device-explicit memory stats reset_peak_memory_stats('cuda:N') raises 'Invalid device argument' before CUDA is initialized (the no-argument form initialized it implicitly), so a replica's first fit failed. Found by the GKE profiling devbox. Co-Authored-By: Claude Opus 5.5 --- src/tabctx/backends/_cuda.py | 12 ++++++++++++ src/tabctx/backends/tabicl.py | 3 ++- src/tabctx/backends/tabpfn.py | 3 ++- 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/tabctx/backends/_cuda.py b/src/tabctx/backends/_cuda.py index 737ab83..6dd7444 100644 --- a/src/tabctx/backends/_cuda.py +++ b/src/tabctx/backends/_cuda.py @@ -17,6 +17,18 @@ def is_cuda(device: str) -> bool: return device == "cuda" or device.startswith("cuda:") +def ensure_initialized(device: str) -> None: + """Initialize CUDA for a backend bound to a GPU. The memory-stat calls + that take an explicit device (reset_peak_memory_stats("cuda:1"), ...) + raise "Invalid device argument" before CUDA is initialized -- unlike + their no-argument forms, which initialize it via current_device(). A + replica's first fit would otherwise fail.""" + if is_cuda(device): + import torch + + torch.cuda.init() + + def device_scope(device: str) -> contextlib.AbstractContextManager[Any]: """Make ``device`` torch's current CUDA device for the block; a no-op on CPU.""" diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py index 30d9187..09eb830 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -54,7 +54,7 @@ import numpy as np -from tabctx.backends._cuda import device_scope, is_cuda +from tabctx.backends._cuda import device_scope, ensure_initialized, is_cuda from tabctx.errors import BackendComputeError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -105,6 +105,7 @@ def __init__(self, device: str | None = None, kv_cache: bool | str = "kv") -> No import torch self._device = device or ("cuda" if torch.cuda.is_available() else "cpu") + ensure_initialized(self._device) self._kv_cache = kv_cache self._last_context_bytes: int | None = None self._last_fit_peak_bytes: int | None = None diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py index 39bce44..f7329af 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -35,7 +35,7 @@ import numpy as np -from tabctx.backends._cuda import device_scope, is_cuda +from tabctx.backends._cuda import device_scope, ensure_initialized, is_cuda from tabctx.errors import BackendComputeError, InvalidInputError from tabctx.types import ArrayLike, PredictOutcome, Task @@ -69,6 +69,7 @@ def __init__(self, device: str | None = None, cache_mode: str = "kv") -> None: self.model_version = settings.tabpfn.model_version.value self.name = served_model_id(self.model_version) self._device = device or ("cuda" if torch.cuda.is_available() else "cpu") + ensure_initialized(self._device) self._fit_mode = _FIT_MODES[cache_mode] self._last_context_bytes: int | None = None self._last_fit_peak_bytes: int | None = None From 87e518000e26d028375bb23821550ae9beac8df9 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Sat, 17 Oct 2026 15:19:02 -0400 Subject: [PATCH 7/7] perf(tabpfn): share the model across contexts (~42x less GPU per context); spill without weights; warm up weights at start Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 +++ src/tabctx/backends/tabicl.py | 14 ++++++ src/tabctx/backends/tabpfn.py | 85 +++++++++++++++++++++++++++++++++++ src/tabctx/serve/factory.py | 18 +++++--- 4 files changed, 118 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e32c21e..b1a5637 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,13 @@ - `benchmarks/sweep_admission.py`: the table-size sweep (rows up to the first 413 per feature count, fit and predict timed). - ROADMAP: multi-GPU and CUDA-level track (#30-#37), PyTorch-level first. +### Changed +- **TabPFN contexts share the model: ~42x less GPU memory per context.** TabPFN builds a fresh model for every estimator unless its built-model cache is on, so every cached context held its own ~850 MB copy of the weights on the GPU (measured on an A100: 4 contexts of a 300x6 table took 3.4 GB). `TabPFNBackend` now enables that cache (`TABPFN_MODEL_CACHE_SIZE`, default 16 entries, keyed by device and task): one shared model per (device, task), about 1.7 GB fixed per GPU, and a context costs only its own state (~20 MB for that table). Interleaved contexts predict bit-identically to each context alone. +- **TabPFN spill blobs no longer contain the model.** A spilled TabPFN payload was ~900 MB, mostly weights, and a restore built another private copy on the GPU. References to the shared model are now pickled as tokens and re-attached on load: a 2,000x20 context spills as 99 MB and restores with identical predictions and no extra GPU memory. Older blobs still load. +- **Weights load at replica start** (`warmup()` on both backends, called by the factory) instead of inside the first tenant's fit, which reported the checkpoint as part of its context (a 200x6 TabPFN context was booked at 855 MB, TabICL's at 226 MB) and paid the load latency. + ### Fixed +- **First fit on a GPU failed** with "Invalid device argument": device-explicit memory-stat calls need CUDA initialized, which the no-argument forms used to do implicitly. Backends now initialize CUDA at construction. - **Spill tier warm restart:** the index rebuilt at startup was keyed by the hex-encoded file name instead of the dataset_id, so lookups, explicit deletes and disk-pressure cleanup all missed contexts a previous process had spilled (restores still worked because `load()` reads by path). - **TabICL spill restore reloaded the checkpoint:** tabicl's `__setstate__` calls `_load_model()`, which loaded the full checkpoint (and a second backbone onto the GPU, unaccounted) on every restore before tabctx swapped in the shared one. Spilled TabICL payloads now carry a task tag and restore straight onto the shared backbone; untagged blobs from older versions still load. - `/readyz` and the legacy endpoint read GPU memory from torch's current device; they now read the replica's own device(s). diff --git a/src/tabctx/backends/tabicl.py b/src/tabctx/backends/tabicl.py index 09eb830..2329fb8 100644 --- a/src/tabctx/backends/tabicl.py +++ b/src/tabctx/backends/tabicl.py @@ -250,6 +250,20 @@ def _predict( # Regressor: no predict_proba to derive from. return PredictOutcome(predictions=model.predict(X_test).tolist()) + def warmup(self) -> None: + """Load the model weights for both tasks now, with throwaway fits, + instead of inside a tenant's first fit: that fit would otherwise + report the weights as part of its context (resident bytes for a + tiny table were the whole checkpoint), overcharging the cache and + leaving the books wrong once it is evicted -- and would pay the + load latency.""" + rng = np.random.default_rng(0) + X = rng.normal(size=(32, 4)) + self.fit(X, (X[:, 0] > 0).astype(int), "classification") + self.fit(X, X[:, 0], "regression") + self._last_context_bytes = None + self._last_fit_peak_bytes = None + def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: del n_train, n_features # Real measurement from the fit() call this is queried after (see diff --git a/src/tabctx/backends/tabpfn.py b/src/tabctx/backends/tabpfn.py index f7329af..dfa86f3 100644 --- a/src/tabctx/backends/tabpfn.py +++ b/src/tabctx/backends/tabpfn.py @@ -31,6 +31,9 @@ from __future__ import annotations +import io +import os +import pickle from typing import Any import numpy as np @@ -52,6 +55,35 @@ def served_model_id(model_version: str) -> str: return f"tabpfn-{model_version.removeprefix('v')}" +# TabPFN builds a fresh model for every estimator unless its built-model +# cache is on, so every cached context held its own ~850 MB copy of the +# weights on the GPU (measured: 4 contexts of a 300x6 table = 3.4 GB). With +# the cache, estimators on one device share one model per task and a +# context costs only its own state (~20 MB for that table), with +# bit-identical predictions. Keyed by device and task, so multi-GPU +# replicas get one shared model per (device, task). +_BUILT_MODEL_CACHE_ENTRIES = "16" +_SPILL_TAG = b"P1" + + +class _SharedModelPickler(pickle.Pickler): + def __init__(self, file, shared: dict[int, tuple]) -> None: + super().__init__(file, protocol=pickle.HIGHEST_PROTOCOL) + self._shared = shared + + def persistent_id(self, obj: Any) -> tuple | None: + return self._shared.get(id(obj)) + + +class _SharedModelUnpickler(pickle.Unpickler): + def __init__(self, file, resolve) -> None: + super().__init__(file) + self._resolve = resolve + + def persistent_load(self, pid: Any) -> Any: + return self._resolve(pid) + + class TabPFNBackend: def __init__(self, device: str | None = None, cache_mode: str = "kv") -> None: import torch @@ -73,6 +105,9 @@ def __init__(self, device: str | None = None, cache_mode: str = "kv") -> None: self._fit_mode = _FIT_MODES[cache_mode] self._last_context_bytes: int | None = None self._last_fit_peak_bytes: int | None = None + # (task, index) -> the shared model module(s) on this device. + self._shared_models: dict[tuple[str, int], Any] = {} + os.environ.setdefault("TABPFN_MODEL_CACHE_SIZE", _BUILT_MODEL_CACHE_ENTRIES) def _make_estimator(self, task: Task) -> Any: if task == "classification": @@ -117,6 +152,8 @@ def fit(self, X: ArrayLike, y: ArrayLike, task: Task) -> Any: else: self._last_context_bytes = None self._last_fit_peak_bytes = None + for index, module in enumerate(getattr(model, "models_", None) or []): + self._shared_models.setdefault((task, index), module) return model def predict( @@ -150,6 +187,54 @@ def _predict( return PredictOutcome(predictions=predictions) return PredictOutcome(predictions=model.predict(X_arr).tolist()) + # ---- spillover serialization (see cache/spill.py) ---------------- + + def dumps_payload(self, payload: Any) -> bytes: + """Pickle a fitted estimator for the spill tier WITHOUT the shared + model (pickled whole, a TabPFN payload was ~900 MB of weights, and a + restore built a private copy again): every reference to a shared + module is written as a (task, index) token instead.""" + shared = { + id(m): ("tabpfn-model", task, i) + for (task, i), m in self._shared_models.items() + } + buf = io.BytesIO() + _SharedModelPickler(buf, shared).dump(payload) + return _SPILL_TAG + buf.getvalue() + + def loads_payload(self, data: bytes) -> Any: + """Inverse of dumps_payload, re-attaching this device's shared model + (loading it first after a restart). Untagged blobs from older + versions are plain pickles and still load.""" + if not data.startswith(_SPILL_TAG): + with device_scope(self._device): + return pickle.loads(data) + + def resolve(pid: Any) -> Any: + _, task, index = pid + if (task, index) not in self._shared_models: + self.warmup() + return self._shared_models[(task, index)] + + with device_scope(self._device): + return _SharedModelUnpickler( + io.BytesIO(data[len(_SPILL_TAG) :]), resolve + ).load() + + def warmup(self) -> None: + """Load the model weights for both tasks now, with throwaway fits, + instead of inside a tenant's first fit: that fit would otherwise + report the weights as part of its context (resident bytes for a + tiny table were the whole checkpoint), overcharging the cache and + leaving the books wrong once it is evicted -- and would pay the + load latency.""" + rng = np.random.default_rng(0) + X = rng.normal(size=(32, 4)) + self.fit(X, (X[:, 0] > 0).astype(int), "classification") + self.fit(X, X[:, 0], "regression") + self._last_context_bytes = None + self._last_fit_peak_bytes = None + def context_bytes_hint(self, n_train: int, n_features: int) -> int | None: del n_train, n_features # Real measurement from the fit() this is queried after; None on diff --git a/src/tabctx/serve/factory.py b/src/tabctx/serve/factory.py index ff474bc..2ee2e36 100644 --- a/src/tabctx/serve/factory.py +++ b/src/tabctx/serve/factory.py @@ -368,12 +368,20 @@ def _build_real_backend( from tabctx.backends.tabpfn import TabPFNBackend # TABCTX_KV_CACHE maps onto TabPFN's fit_mode (see backends/tabpfn.py). - return TabPFNBackend(device=device, cache_mode=settings.kv_cache) - - from tabctx.backends.tabicl import TabICLBackend + backend: TabularICLBackend = TabPFNBackend( + device=device, cache_mode=settings.kv_cache + ) + else: + from tabctx.backends.tabicl import TabICLBackend - kv_cache: bool | str = False if settings.kv_cache == "off" else settings.kv_cache - return TabICLBackend(device=device, kv_cache=kv_cache) + kv_cache: bool | str = ( + False if settings.kv_cache == "off" else settings.kv_cache + ) + backend = TabICLBackend(device=device, kv_cache=kv_cache) + # Weights load at replica start, not inside (and charged to) the first + # tenant's fit. + backend.warmup() + return backend def _build_backend(