diff --git a/docs/tabicl_strong_feature_comparison.md b/docs/tabicl_strong_feature_comparison.md new file mode 100644 index 0000000..afd5dd4 --- /dev/null +++ b/docs/tabicl_strong_feature_comparison.md @@ -0,0 +1,180 @@ +# GBM (strong) vs TabICL (strong, reduced) on MIMIC + +**Date:** 2026-08-25 + +## Why this run exists + +Every prior TabICL row in this project's registry (`docs/experiments.md`) scored TabICL on the +**basic** 17-feature panel, while the GBM it was compared against used the **strong** 609-feature +panel. TabICL never saw the wide panel at all: a naive fit at 50,000 context rows, 609 features, +and the library's default of 8 estimators costs an estimated ~70 GB per `predict_proba` call, and +reliably OOM-killed the host (three confirmed kills on 2026-08-23). + +This is the first TabICL run on the strong panel. To make it fit in memory, three settings were +changed from this project's usual TabICL defaults: + +- `offload_mode="cpu"` (moves the library's large column-wise embedding tensor off GPU VRAM) +- `n_estimators` reduced from the library default of 8 to **1** +- the in-context row cap (`TABICL_MAX_ROWS`) cut from this project's usual 50,000 to **20,000** + +Two earlier attempts at the full-default configuration failed before this one worked: + +1. `offload_mode="disk"` with `batch_size=1` stalled in what looked like a disk I/O deadlock at + the standard 50,000-row, 8-estimator scale. Never resolved, just abandoned. +2. `offload_mode="cpu"` with all 12 models fit up front and held in memory for the whole predict + loop was kernel OOM-killed partway through (9 of 12 cells done, `shmem-rss` around 75 GB). This + was a repeat of a memory-accumulation problem this project had already hit and fixed once + before, not a new failure mode. + +The run reported here fits, scores, and drops one model at a time, which is what let all 12 cells +finish cleanly. + +## Setup + +- Checkpoint: `subset_run_v8_taskset_v3` (MIMIC, task_set v3, `model_kind=cbm`) +- GPU host: `odyssey-cbm-a100` +- GBM (strong) scores: the existing registered alerts run (`alerts_rows_v3.parquet`), fit on 30 + train shards, scored on 4 held-out shards +- TabICL (strong, reduced) scores: fit on 8 of those same 30 train shards (not all 30; TabICL's + in-context window caps out well below 30 shards' worth of rows regardless), scored on the same + 4 held-out shards +- Both models are scored on exactly the same rows and labels + +## The caveat that matters most + +Matching the feature set controls how much information each model gets per row. It does not +control the modeling approach itself. The GBM is gradient-boosted and hyperparameter-tuned per +task on the full training set (400 rounds, grid search over 4 configs). TabICL is zero-shot +in-context learning: no gradient descent on this data at all, just one forward pass conditioned +on a capped, subsampled 20,000-row context, with only 1 ensemble member instead of the library's +usual 8 (so the random feature/class permutation averaging TabICL normally relies on for +robustness is largely absent here). + +If TabICL still loses after matching features, that is not evidence TabICL "can't use" these +features. It is evidence that zero-shot, single-member, context-capped in-context learning is not +competitive with a fully-supervised, tuned GBM on this task family. That is a narrower and +different claim. + +## Can TabICL run at full capability instead of reduced? + +We tried. Here is what happened. + +The `offload_mode`/`disk_offload_dir` params in this project's code only cover one stage of +TabICL, the column-wise embedding step. That stage's disk offload works: a test with synthetic +data at full scale (50,000 rows, 609 features, 8 estimators, `disk_dtype` set to float16) wrote a +real file to disk, about 7.4 GB per estimator, and fit in 27 seconds with no memory error. + +Fitting was never the problem. Scoring is. TabICL rereads its entire context from disk on every +`predict_proba` call, with no caching by default. Scoring 200 rows against the full context took +728.9 seconds (about 12 minutes). Scoring 8,192 rows (this project's normal batch size) did not +finish inside a 15-minute timeout. + +At that rate, one MIMIC alert cell (about 48,000 held-out rows, 6+ batches) would take hours, and +the full 12-cell comparison would take a day or more of sequential compute. Not memory. Disk I/O +repeated on every call. + +One thing we did not try: `kv_cache="repr"`, which caches embeddings across calls and might avoid +the repeated disk reads. Flagged as a real next step, not assumed to work. + +Bottom line: full capability (8 estimators, 50,000-row context) fits and scores without crashing, +just not in a practical amount of time on this hardware. The reduced-config result below is what +was actually achievable, not a shortcut taken for convenience. + +## Update, 2026-08-26: full capability tested on new hardware + +The host was migrated to `a2-ultragpu-1g` (170 GB RAM, one A100-80GB, `odyssey-cbm-a100-ultra`, +`us-central1-a`), specifically to test whether more RAM, not disk, fixes this. It does, partly. + +RAM-resident full capability (`offload_mode="cpu"`, `n_estimators=8`, 50,000-row context, no +disk) works. It does not hang and does not crash. Peak RSS during a single (event, horizon) +fit and predict reached about 98 GB, comfortably inside the new 165 GB free RAM, with zero swap +use. This confirms the diagnosis: disk I/O was the wrong lever, not memory capacity. + +One full cell was run end to end: vasopressor_start@8h. Fit took 109 seconds for all 3 horizons +of this event together. Scoring the held-out rows for just the 8h horizon took 1,854 seconds +(31 minutes). Result: + +| Event | Horizon | n | GBM (strong) | TabICL (strong, reduced) | TabICL (strong, full) | +|---|---|---|---|---|---| +| Vasopressor start | 8h | 111,450 | 0.934 [0.916, 0.950] | 0.859 [0.832, 0.887] | 0.915 [0.896, 0.936] | + +Full capability closes most of the gap. The reduced config's 0.075 AUROC deficit against the +GBM drops to 0.019, and the confidence intervals now overlap (0.916-0.950 vs 0.896-0.936): at +full capability, this cell is statistically indistinguishable from the GBM, not a real loss. +The reduced config was a real handicap, not just a formality. + +But 31 minutes for one horizon of one event is real, and it does not fit in one session for all +12 cells. At the measured rate (about 0.0166 seconds per held-out row), scoring all 12 core +(event, horizon) cells in the table below would take on the order of 5 hours of sequential +compute, not counting fit time (negligible by comparison). This was not run to completion. The +process was stopped deliberately after this one cell, and the VM was stopped, rather than run +the full sweep unattended on an hourly-billed A100-80GB host without asking first. + +Where this leaves the table below: it is a real result, honestly obtained, and still the only +complete 12-cell comparison that exists. But the vasopressor@8h data point above shows it likely +understates TabICL's true performance across the board, not just for that one cell. Whether to +spend the roughly 5 hours of GPU time to get the real, full 12-cell table at full capability is +a decision for Amrit, not something to default into. + +## Results + +95% subject-clustered bootstrap confidence intervals, 1000 resamples +(`odyssey.inference.uncertainty.bootstrap_auroc`). + +| Event | Horizon | n | GBM (strong) AUROC | TabICL (strong, reduced) AUROC | Gap | Verdict | +|---|---|---|---|---|---|---| +| Acute kidney injury | 8h | 95,471 | 0.894 [0.881, 0.906] | 0.723 [0.702, 0.745] | 0.171 | real gap | +| Acute kidney injury | 24h | 84,147 | 0.845 [0.827, 0.861] | 0.714 [0.690, 0.736] | 0.131 | real gap | +| Acute kidney injury | 72h | 57,442 | 0.782 [0.758, 0.805] | 0.674 [0.646, 0.702] | 0.108 | real gap | +| Death | 8h | 136,850 | 0.953 [0.934, 0.969] | 0.941 [0.921, 0.958] | 0.012 | within noise | +| Death | 24h | 135,061 | 0.959 [0.947, 0.969] | 0.924 [0.900, 0.945] | 0.035 | real gap | +| Death | 72h | 130,818 | 0.940 [0.923, 0.954] | 0.903 [0.879, 0.926] | 0.036 | within noise (barely) | +| ICU admission | 8h | 85,838 | 0.968 [0.962, 0.974] | 0.957 [0.950, 0.964] | 0.011 | within noise | +| ICU admission | 24h | 74,839 | 0.954 [0.945, 0.963] | 0.940 [0.929, 0.950] | 0.014 | within noise | +| ICU admission | 72h | 48,971 | 0.931 [0.912, 0.946] | 0.907 [0.886, 0.926] | 0.024 | within noise | +| Vasopressor start | 8h | 111,450 | 0.934 [0.916, 0.950] | 0.859 [0.832, 0.887] | 0.075 | real gap | +| Vasopressor start | 24h | 98,722 | 0.914 [0.895, 0.933] | 0.848 [0.820, 0.877] | 0.066 | real gap | +| Vasopressor start | 72h | 67,385 | 0.883 [0.850, 0.913] | 0.812 [0.778, 0.845] | 0.070 | real gap | + +"Real gap" means the two confidence intervals do not overlap. "Within noise" means they do. + +## What this means + +Matching features does not close the gap. TabICL (strong, reduced) loses to the tuned GBM on all +12 cells, and the loss is statistically real on 7 of them. + +AKI has the largest and most consistent gap (0.108 to 0.171, real on all three horizons). This +matches an earlier finding in this project (`docs/experiments.md`, journal entry 52): the GBM's +edge on AKI comes from window aggregates and trend statistics it computes explicitly from the raw +values. TabICL sees the same raw features but has no equivalent way to aggregate them across time +in a single zero-shot forward pass. + +ICU admission is where TabICL comes closest to the GBM. That also matches the same earlier +finding: the GBM's edge on ICU admission concentrates in a small set of count features, which may +be easier for in-context learning to pick up directly from raw values than a window trend is. + +The remaining gap is probably a mix of two things this run cannot separate: the real zero-shot vs. +tuned-and-supervised difference, and the memory-driven reductions (1 estimator instead of 8, +20,000-row context instead of 50,000). A full-ensemble, full-context run was not achievable on +this host in a reasonable time. Treat this result as a lower bound on TabICL (strong)'s real +ceiling, not as its true performance. + +## Where things live + +- Comparison script (not committed, diagnostic only): `scripts/tabicl_strong_compare.py`, run on + the GPU host from a disposable git worktree (since removed) +- Code changes needed to make this run possible: `odyssey/inference/tabicl_baseline.py` (adds + `offload_mode`, `batch_size`, and `disk_offload_dir` passthrough to `TabICLClassifier`) +- Raw per-cell results, including full bootstrap output (mean, std, resample counts): + `compare_result.json`, pulled to `/tmp/compare_result.json` this session. Ask if you want it + moved somewhere durable. +- GBM (strong) scores: `~/runs/subset_run_v8_taskset_v3/alerts_rows_v3.parquet` on + `odyssey-cbm-a100` (already-registered run, unchanged by this work) +- Full-capability timing test (synthetic data, not committed, since removed from the host): + `~/tabicl_disk_test/probe.py`, a standalone script isolating just the fit/predict cost at + 50,000 rows x 609 features x 8 estimators with disk offload, independent of this project's + data-loading pipeline +- Full-capability real-data script (not committed, diagnostic only): `scripts/tabicl_strong_compare.py` + on `odyssey-cbm-a100-ultra` (`us-central1-a`), fits/scores one (event, horizon) cell at a time using + the existing `alerts_rows_v3.parquet` for GBM reference scores. The one completed cell's full + bootstrap output is in `~/tabicl_full_validate.json` on that host. diff --git a/odyssey/inference/tabicl_baseline.py b/odyssey/inference/tabicl_baseline.py index dfb4328..529253a 100644 --- a/odyssey/inference/tabicl_baseline.py +++ b/odyssey/inference/tabicl_baseline.py @@ -106,9 +106,37 @@ def estimate_peak_gb(n_context_rows: int, n_features: int, n_estimators: int) -> def check_inference_cost( - n_context_rows: int, n_features: int, n_estimators: int, *, context: str + n_context_rows: int, + n_features: int, + n_estimators: int, + *, + context: str, + offload_mode: str = "auto", + disk_offload_dir: Optional[str] = None, ) -> None: - """Raise before fitting a TabICL configuration that cannot then be scored.""" + """Raise before fitting a TabICL configuration that cannot then be scored. + + This GB estimate assumes the column-wise embedding tensor -- the + documented ``(n_estimators, n_rows, n_columns, embed_dim)`` memory + bottleneck -- stays resident (GPU or CPU RAM). That assumption holds + for ``offload_mode`` "gpu"/"cpu"/"auto": "auto"'s own fallback chain + ends in "CPU (swap as last resort)" (per tabicl's + ``_resolve_offload_mode``), and this host has zero configured swap + (measured 2026-08-25, ``free -h``), so a fit that would exceed the + budget under those modes is a real, unrecoverable kernel-OOM-kill + risk, not a conservative guess -- the gate stays a hard block there. + + ``offload_mode="disk"`` with a working ``disk_offload_dir`` changes + the actual constraint from resident memory to disk space and I/O + wall-clock, which this GB estimate does not model at all; blocking + on it here would be wrong in the other direction (refusing a fit + that would actually succeed, just slower). That combination skips + this check -- callers are responsible for checking disk headroom + themselves (e.g. a ``df`` check before fitting), since there is no + equivalent measured constant for disk throughput yet. + """ + if offload_mode == "disk" and disk_offload_dir: + return peak = estimate_peak_gb(n_context_rows, n_features, n_estimators) if peak > _MEMORY_BUDGET_GB: raise ValueError( @@ -120,8 +148,10 @@ def check_inference_cost( "OOM-kills, 2026-08-23, the last from a single 2000-row call). Use " "the basic feature set (what every completed TabICL run in this " "project has used: ~8 GB at the same context size), lower " - "TABICL_MAX_ROWS, or raise ODYSSEY_TABICL_MEMORY_BUDGET_GB knowing " - "the memory is real." + "TABICL_MAX_ROWS, raise ODYSSEY_TABICL_MEMORY_BUDGET_GB knowing " + "the memory is real, or pass offload_mode='disk' with a " + "disk_offload_dir (this check does not apply to that combination, " + "since the constraint becomes disk space/I-O, not RAM)." ) @@ -199,7 +229,7 @@ class TabICLBaselineModel: feature_set: str = "strong" n_features: int = 0 - params: dict[str, float] = field(default_factory=dict) + params: dict[str, Any] = field(default_factory=dict) all_nan_cols: np.ndarray | None = None """Boolean ``(n_features,)`` mask of columns that were entirely NaN in the fit-time context, or ``None`` if none were. tabicl's own @@ -253,6 +283,9 @@ def _fit_one_tabicl( n_estimators: int, device: str | None, cache: Optional[FitCache] = None, + offload_mode: str = "auto", + batch_size: Optional[int] = 8, + disk_offload_dir: Optional[str] = None, ) -> dict[float, TabICLBaselineModel]: """Fit one TabICL context per horizon for a single event. @@ -264,6 +297,15 @@ def _fit_one_tabicl( stores the (capped, seeded-subsampled) context rather than running gradient descent. + ``offload_mode``/``batch_size``/``disk_offload_dir`` are passed + straight through to ``TabICLClassifier`` -- see that class's own + docstring for what each controls (in short: where the column-wise + embedding tensor, the module's documented memory bottleneck, is + materialized, and how many ensemble members are processed per + forward pass). Recorded in the fitted model's ``params`` for + provenance, the same way ``n_context_rows``/``n_estimators`` already + are. + ``cache``, if given, is checked per horizon before fitting and written to immediately after -- see :mod:`odyssey.inference.fit_cache`. ``_load_tabicl_classifier`` is deferred until the first horizon that @@ -306,12 +348,17 @@ def _fit_one_tabicl( n_estimators=n_estimators, device=device, random_state=seed, + offload_mode=offload_mode, + batch_size=batch_size, + disk_offload_dir=disk_offload_dir, ) check_inference_cost( len(keep), int(x_all.shape[1]), n_estimators, context=f"{event_name}@{h:g}h ({feature_set} features)", + offload_mode=offload_mode, + disk_offload_dir=disk_offload_dir, ) clf.fit(x_fit, y_fit) out[h] = TabICLBaselineModel( @@ -321,6 +368,9 @@ def _fit_one_tabicl( params={ "n_context_rows": float(len(keep)), "n_estimators": float(n_estimators), + "offload_mode": offload_mode, + "batch_size": batch_size, + "disk_offload_dir": disk_offload_dir, }, all_nan_cols=all_nan_cols if all_nan_cols.any() else None, ) @@ -350,6 +400,9 @@ def fit_tabicl_baselines( device: str | None = None, cache: Optional[FitCache] = None, features: Optional[dict[str, np.ndarray]] = None, + offload_mode: str = "auto", + batch_size: Optional[int] = 8, + disk_offload_dir: Optional[str] = None, ) -> dict[tuple[str, float], TabICLBaselineModel]: """One TabICLv2 context per (event, horizon), on the same features as the GBM. @@ -372,6 +425,12 @@ def fit_tabicl_baselines( ``feature_set`` (see :func:`odyssey.inference.baseline_prep.prepare_baseline_data`); ``train_events_binned`` is then unused and may be empty. + + ``offload_mode``/``batch_size``/``disk_offload_dir`` pass straight + through to every ``TabICLClassifier`` this call fits -- see + :func:`_fit_one_tabicl` and ``TabICLClassifier``'s own docstring. + Defaults match ``TabICLClassifier``'s own defaults (``"auto"``/``8``/ + ``None``), so omitting them reproduces prior behavior exactly. """ models: dict[tuple[str, float], TabICLBaselineModel] = {} if features is None: @@ -392,6 +451,9 @@ def fit_tabicl_baselines( n_estimators=n_estimators, device=device, cache=cache, + offload_mode=offload_mode, + batch_size=batch_size, + disk_offload_dir=disk_offload_dir, ) for h, model in per_horizon.items(): models[(name, h)] = model diff --git a/scripts/tabicl_strong_compare.py b/scripts/tabicl_strong_compare.py new file mode 100644 index 0000000..f1e4f92 --- /dev/null +++ b/scripts/tabicl_strong_compare.py @@ -0,0 +1,188 @@ +"""Full-capability TabICL(strong) vs GBM(strong) on MIMIC, with bootstrap CIs. + +Reuses the existing GBM(strong) scores already dumped in alerts_rows_v3.parquet +(fit on 30 train shards, scored on 4 held-out shards, protocol v3) instead of +refitting the GBM. Fits TabICL at its real default capability (n_estimators=8, +TABICL_MAX_ROWS=50,000, both fit_tabicl_baselines/tabicl_baseline.py defaults) +on the strong 609-feature panel, offload_mode="cpu" (this host has 165GB free +RAM, no need for disk offload). One model fit+scored+dropped at a time, same +discipline as scripts/rescore_extra_baselines.py. +""" + +import argparse +import dataclasses +import gc +import json +import logging +import time +from pathlib import Path + +import numpy as np +import polars as pl + +from odyssey.data.alert_events import ALERT_EVENTS +from odyssey.data.value_binning import QuantileBinner +from odyssey.inference.baseline_prep import prepare_baseline_data +from odyssey.inference.tabicl_baseline import fit_tabicl_baselines +from odyssey.inference.uncertainty import bootstrap_auroc +from odyssey.training.shard_stream import make_preparer, shard_paths +from odyssey.training.train import TrainingConfig + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("tabicl_full_compare") + +HORIZONS = (8.0, 24.0, 72.0) +CORE_EVENTS = ("acute_kidney_injury", "death", "icu_admission", "vasopressor_start") + + +def main() -> None: # noqa: PLR0915 + """Fit TabICL(strong, full capability) per core event, score vs. the GBM dump.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", required=True, type=Path) + parser.add_argument("--train-shard-dir", required=True, type=Path) + parser.add_argument("--held-out-shard-dir", required=True, type=Path) + parser.add_argument("--existing-dump", required=True, type=Path) + parser.add_argument("--max-train-shards", type=int, default=8) + parser.add_argument("--max-held-out-shards", type=int, default=4) + parser.add_argument("--landmark-hours", type=float, default=4.0) + parser.add_argument( + "--only-event", default=None, help="restrict to one event, for validation" + ) + parser.add_argument("--offload-mode", default="cpu") + parser.add_argument("--output-json", required=True, type=Path) + args = parser.parse_args() + + raw_config = json.loads((args.run_dir / "config.json").read_text()) + known_fields = {f.name for f in dataclasses.fields(TrainingConfig)} + dropped = sorted(set(raw_config) - known_fields) + if dropped: + logger.info( + "config.json has fields no longer on TrainingConfig, dropping: %s", dropped + ) + config = TrainingConfig( + **{k: v for k, v in raw_config.items() if k in known_fields} + ) + binner = QuantileBinner.load(args.run_dir / "quantile_binner.json") + source = getattr(config, "source", "mimic_iv") + prepare = make_preparer( + normalize_medications=getattr(config, "normalize_medications", False), + history_recap=getattr(config, "history_recap", False), + source=source, + ) + alerts = [a for a in ALERT_EVENTS if a.name in CORE_EVENTS] + if args.only_event: + alerts = [a for a in alerts if a.name == args.only_event] + + t0 = time.time() + logger.info("preparing %d train shard(s)", args.max_train_shards) + train = prepare_baseline_data( + shard_paths(args.train_shard_dir, max_shards=args.max_train_shards), + prepare, + binner, + alerts=alerts, + feature_sets=("strong",), + source=source, + landmark_hours=args.landmark_hours, + ) + logger.info("train prep done in %.0fs", time.time() - t0) + for name, rows in train.rows.items(): + logger.info(" train candidate rows %s: %d", name, len(rows)) + + t0 = time.time() + logger.info("preparing %d held-out shard(s)", args.max_held_out_shards) + held = prepare_baseline_data( + shard_paths(args.held_out_shard_dir, max_shards=args.max_held_out_shards), + prepare, + binner, + alerts=alerts, + feature_sets=("strong",), + source=source, + landmark_hours=args.landmark_hours, + ) + logger.info("held-out prep done in %.0fs", time.time() - t0) + + existing = pl.read_parquet(args.existing_dump) + + results = {} + for alert in alerts: + event = alert.name + rows = train.rows.get(event, []) + if not rows: + continue + t0 = time.time() + models = fit_tabicl_baselines( + pl.DataFrame(), + {event: rows}, + {event: train.times[event]}, + horizons=HORIZONS, + source=source, + feature_set="strong", + features={event: train.features["strong"][event]}, + offload_mode=args.offload_mode, + ) + fit_s = time.time() - t0 + logger.info("fit %s: %d models in %.0fs", event, len(models), fit_s) + + held_rows = held.rows.get(event, []) + held_feats = held.features["strong"][event] + existing_ev = existing.filter(pl.col("event") == event) + + for h in HORIZONS: + model = models.pop((event, h), None) + if model is None: + continue + t0 = time.time() + proba = model.predict_proba(held_feats) + predict_s = time.time() - t0 + del model + gc.collect() + + new_cols = pl.DataFrame( + { + "subject_id": [float(r.subject_id) for r in held_rows], + "visit_id": [float(r.visit_id) for r in held_rows], + "time_hours": [r.time_hours for r in held_rows], + f"tabicl@{h:g}h": [float(v) for v in proba], + } + ) + joined = existing_ev.join( + new_cols, on=["subject_id", "visit_id", "time_hours"], how="inner" + ) + y = joined[f"y@{h:g}h"].to_numpy() + gbm_p = joined[f"gbm@{h:g}h"].to_numpy() + tabicl_p = joined[f"tabicl@{h:g}h"].to_numpy() + sid = joined["subject_id"].to_numpy().astype(int) + mask = ~np.isnan(y) + y, gbm_p, tabicl_p, sid = y[mask], gbm_p[mask], tabicl_p[mask], sid[mask] + + gbm_ci = bootstrap_auroc(y, gbm_p, sid) + tabicl_ci = bootstrap_auroc(y, tabicl_p, sid) + + cell = { + "n": int(mask.sum()), + "fit_s": fit_s, + "predict_s": predict_s, + "gbm": None if gbm_ci is None else vars(gbm_ci), + "tabicl": None if tabicl_ci is None else vars(tabicl_ci), + } + results[f"{event}@{h:g}h"] = cell + logger.info( + "%s@%gh n=%d gbm=%.4f tabicl=%.4f (predict %.0fs)", + event, + h, + cell["n"], + gbm_ci.point_estimate if gbm_ci else float("nan"), + tabicl_ci.point_estimate if tabicl_ci else float("nan"), + predict_s, + ) + args.output_json.write_text(json.dumps(results, indent=2)) + del models + gc.collect() + + args.output_json.write_text(json.dumps(results, indent=2)) + logger.info("wrote %s", args.output_json) + + +if __name__ == "__main__": + main() diff --git a/tests/odyssey/inference/test_tabicl_baseline.py b/tests/odyssey/inference/test_tabicl_baseline.py index 1794bb2..2ad4853 100644 --- a/tests/odyssey/inference/test_tabicl_baseline.py +++ b/tests/odyssey/inference/test_tabicl_baseline.py @@ -452,6 +452,98 @@ def test_inference_cost_guard_matches_the_measured_configurations() -> None: check_inference_cost(5_000, 609, 8, context="strong, small context") +# --------------------------------------------------------------------------- +# offload_mode / batch_size / disk_offload_dir: threaded through to the +# classifier and recorded for provenance (Track: strong-feature TabICL) +# --------------------------------------------------------------------------- + + +def test_fit_tabicl_baselines_passes_offload_params_through( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(tabicl_module, "TABICL_MIN_ROWS", 10) + events = _events(24) + binned = add_value_tokens(events) + times = all_event_times(binned, ALERT_EVENTS, "mimic_iv") + rows = _index_rows_from_events(binned, ALERT_EVENTS, landmark_hours=4.0) + + fit_tabicl_baselines( + binned, + rows, + times, + horizons=(8.0,), + feature_set="strong", + offload_mode="disk", + batch_size=1, + disk_offload_dir="/tmp/tabicl_offload", + ) + assert _RecordingFakeClassifier.instances + fit_kwargs = _RecordingFakeClassifier.instances[0].kwargs + assert fit_kwargs["offload_mode"] == "disk" + assert fit_kwargs["batch_size"] == 1 + assert fit_kwargs["disk_offload_dir"] == "/tmp/tabicl_offload" + + model = list( + fit_tabicl_baselines( + binned, + rows, + times, + horizons=(8.0,), + feature_set="strong", + offload_mode="disk", + batch_size=1, + disk_offload_dir="/tmp/tabicl_offload", + ).values() + )[0] + assert model.params["offload_mode"] == "disk" + assert model.params["batch_size"] == 1 + assert model.params["disk_offload_dir"] == "/tmp/tabicl_offload" + + +def test_fit_tabicl_baselines_defaults_reproduce_prior_offload_behavior( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Omitting the new kwargs must fit exactly as before their addition.""" + monkeypatch.setattr(tabicl_module, "TABICL_MIN_ROWS", 10) + events = _events(24) + binned = add_value_tokens(events) + times = all_event_times(binned, ALERT_EVENTS, "mimic_iv") + rows = _index_rows_from_events(binned, ALERT_EVENTS, landmark_hours=4.0) + + fit_tabicl_baselines(binned, rows, times, horizons=(8.0,), feature_set="basic") + fit_kwargs = _RecordingFakeClassifier.instances[0].kwargs + assert fit_kwargs["offload_mode"] == "auto" + assert fit_kwargs["batch_size"] == 8 + assert fit_kwargs["disk_offload_dir"] is None + + +def test_check_inference_cost_skips_the_ram_budget_for_disk_offload() -> None: + """Disk offload bypasses the RAM-budget gate; the bare flag alone does not.""" + from odyssey.inference.tabicl_baseline import check_inference_cost # noqa: PLC0415 + + with pytest.raises(ValueError, match="per predict_proba call"): + check_inference_cost(50_000, 609, 8, context="strong, no offload") + with pytest.raises(ValueError, match="per predict_proba call"): + # offload_mode alone, with no disk_offload_dir, does not bypass the + # gate -- "disk" without a directory cannot actually offload. + check_inference_cost( + 50_000, + 609, + 8, + context="strong, disk requested but no dir", + offload_mode="disk", + ) + # disk offload WITH a directory bypasses the RAM-budget gate entirely. + check_inference_cost( + 50_000, + 609, + 8, + context="strong, disk offload configured", + offload_mode="disk", + disk_offload_dir="/tmp/tabicl_offload", + ) + + def test_fit_cache_keys_include_the_feature_set() -> None: """A fit is only reusable for the feature matrix it was fit on.""" import inspect # noqa: PLC0415